[随笔] nn.Embedding的前向传播与反向传播

nn.Embedding的前向传播与反向传播

nn.Embedding的前向计算过程

embedding module 的前向过程其实是一个索引（查表）的过程
表的形式是一个 matrix（embedding.weight, learnable parameters）

matrix.shape: (v, h)
v：vocabulary size=num_embedding
h：hidden dimension=embedding_dim

仅从数学的角度来说(方便推导模型),具体索引的过程，可以通过 one hot + 矩阵乘法的形式实现的

input.shape: (b, s)
> b：batch size
> s：seq len当执行下行代码时,会进行如下计算 embed = embedding(input) > input.shape(b,s) 	e.g [[0, 2, 2,1]]
> 最终的维度变化情况:(b, s) ==> (b, s, h)1.(b, s) 经过 one hot => (b, s, v)
inputs: [[0, 2, 2, 1 , 1]] 
inputs One-Hot: 数值分类(0-4 => 五分类) 0:[1,0,0,0,0][[[1,0,0,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,1,0,0,0]]]
matrix(embedding.weight):[[ 1.0934,  1.7521, -1.9529, -1.0145,  0.5770],[-0.4371, -0.4270, -0.4908, -0.3988,  0.9695],[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],[ 0.7268, -0.4491, -0.8089,  0.7516,  1.2716],[ 0.7785, -0.4336, -0.7542, -0.1953,  0.9711]]

2.(b, s, v) @ (v, h) ==> (b, s, h)

x(b, s, h):[[[ 1.0934,  1.7521, -1.9529, -1.0145,  0.5770],[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],[-0.4371, -0.4270, -0.4908, -0.3988,  0.9695],[-0.4371, -0.4270, -0.4908, -0.3988,  0.9695]]]

但本质上,embedding(input) 是一个内存寻址的过程:

假设 inputs 是一个包含词索引的张量,i是数值化文本 inputs上的Token，weight 是嵌入矩阵。对于每个索引 i，嵌入向量 v_i 对应的计算过程是：v_i=embedding.weight[i]

nn.Embedding的反向传播过程

只有前向传播中用到的索引会接收梯度。

假如反向传播过来的梯度是 [0.1,0.1,0.3] ,原始的embedding矩阵= [[1. ,1. ,1.],[1. ,1. ,1.]] , lr=0.1

那么 反向传播以后embedding的参数就为 [[1. ,1. ,1.],[1. ,1. ,1.]] - 1 * [[0.1,0.1,0.3],[0.,0.,0.]]

即 [[0.99. ,0.99 ,0.97],[1. ,1. ,1.]]