手撕 Transformer (2)：嵌入层和位置编码的实现

上篇文章讲过Transformer 可分为四个部分输入、输出、编码器、解码器。本文主要手撕输入部分。输入部分由嵌入层(Embedding)和位置编码(Positional Encoding)组成。本文进行嵌入层和位置编码的代码实现和原理讲解。1 嵌入层嵌入层的作用为了将文本中词汇的数字表示转换为向量表示语义向量这样后续神经网络就可以对其进行计算了。1.1 代码实现importtorchimporttorch.nnasnnimportmathfromtorch.autogradimportVariableclassEmbeddings(nn.Module):def__init__(self,d_model,vocab):# d_model: 词嵌入的维度# vocab: 词表的大小super(Embeddings,self).__init__()self.lutnn.Embedding(vocab,d_model)self.d_modeld_modeldefforward(self,x):# 前向传播# x 是输入进模型的文本通过映射后的数字张量returnself.lut(x)*math.sqrt(self.d_model)if__name____main__:d_model512vocab1000xVariable(torch.LongTensor([[123,233,510,998],[985,211,110,996]]))embEmbeddings(d_model,vocab)emb_resultemb(x)print(embedding result: ,emb_result)print(emb_result.shape)# torch.Size([2, 4, 512])在Embeddings类中self.lut nn.Embedding(vocab, d_model)会创建一个随机初始化的嵌入矩阵。PyTorch 的nn.Embedding模块默认使用均匀分布随机初始化权重。在模型训练过程中当计算损失函数并执行反向传播时嵌入层的权重会接收到梯度然后通过优化器如 Adam进行更新。这样模型会逐渐学习到更有意义的词向量表示这些表示会捕捉到词语之间的语义和语法关系。运行结果embedding result: tensor([[[-49.8672, -21.6785, 18.1069, ..., 0.2031, -28.3568, 5.5724], [-55.8387, -26.6077, 37.4205, ..., 7.8280, -5.1322, 8.1475], [ 21.9637, 9.6126, 53.4801, ..., 16.6295, 37.5978, 13.2768], [-19.0594, -13.2244, 16.7811, ..., 16.9383, -46.1544, -3.1326]], [[ 9.8451, 22.9543, 3.1216, ..., 18.1514, 24.2709, 31.3333], [ 30.5660, -9.3572, -5.8656, ..., 4.3933, 9.5235, 9.1021], [ 14.2475, 28.2354, 49.7318, ..., 9.2369, -23.4376, -7.1588], [ 15.4746, 40.1049, -19.8356, ..., -25.1046, 13.6735, -18.5525]]], grad_fnMulBackward0) torch.Size([2, 4, 512])为什么需要学习嵌入向量随机初始化的嵌入向量只是初始值不包含任何语义信息。通过训练模型会根据具体任务如机器翻译、文本分类等的目标调整嵌入向量使得相似含义的词在向量空间中距离更近不同含义的词距离更远。为什么计算完Embedding之后要乘以d m o d e l \sqrt{{d_{model}}}dmodel​​self.lut(x)*math.sqrt(self.d_model)放大信号词嵌入通常是随机初始化的其方差较小通过乘以d m o d e l \sqrt{{d_{model}}}dmodel​​来放大嵌入向量的幅度确保嵌入向量的尺度与位置编码通常使用正弦/余弦函数生成相当。注意此处不是注意力机制中的缩放点积那个是除以d m o d e l \sqrt{{d_{model}}}dmodel​​。我们可以单独把Embedding的作用拿出来看一下embeddingnn.Embedding(10,3)input1torch.LongTensor([[1,2,4,5],[4,3,2,9]])print(整数张量表示词ID,input1)print(input1转成向量表示,embedding(input1))从下面的运行结果可以明显看出Embedding的数字表示转向量表示的作用。整数张量表示词ID tensor([[1, 2, 4, 5], [4, 3, 2, 9]]) input1转成向量表示 tensor([[[-0.6656, 1.6754, -0.5841], [ 1.1583, 0.0122, 0.0297], [-1.5521, 1.9699, 0.0168], [ 0.9703, -0.0608, -0.6835]], [[-1.5521, 1.9699, 0.0168], [ 1.1763, 0.1059, -0.6196], [ 1.1583, 0.0122, 0.0297], [-0.7003, 0.6548, 0.0784]]], grad_fnEmbeddingBackward0)如果词ID中有数字0计算出的嵌入向量会是 0 吗从下面的例子中可以看到并不是。因为嵌入向量是随机初始化的并且在训练过程中不断更新。示例embeddingnn.Embedding(10,3)input2torch.LongTensor([[0,2,0,5]])print(整数张量表示词ID,input2)print(input2转成向量表示,embedding(input2))从下面的运行结果可以看出虽然词ID中有 0但是嵌入向量中并没有 0 。整数张量表示词ID tensor([[0, 2, 0, 5]]) input2转成向量表示 tensor([[[-2.0432, 0.4369, -0.4257], [-0.1574, 0.1013, -0.1821], [-2.0432, 0.4369, -0.4257], [ 0.0601, 0.9223, 0.3128]]], grad_fnEmbeddingBackward0)在实际应用中有的时候是需要嵌入向量中有 0 的让这些参数在训练的过程中不更新。当数据批量输入进模型时序列的长度可能不一致这时候就需要对短序列的特定维度进行补 0 使其与最长序列相等。初始化时对应的嵌入向量会初始化为 0 。训练时填充位置的嵌入向量不会被更新梯度为 0避免填充位置对模型训练产生干扰。推理时填充位置的嵌入向量保持为 0不影响模型对有效序列的处理。例如在机器翻译任务中输入句子[“I love you”, “He eats”]假设最长的序列是I love you长度为3短序列为He eats长度为2。填充后[[1, 2, 3], [4, 5, 0]]。具体代码只需要添加一个参数即可示例embeddingnn.Embedding(10,3,padding_idx0)input3torch.LongTensor([[0,2,0,5]])input4torch.LongTensor([[1,2,3],[4,5,0]])print(整数张量表示词ID,input3)print(input3转成向量表示,embedding(input3))print(input4转成向量表示,embedding(input4))运行结果整数张量表示词ID tensor([[0, 2, 0, 5]]) input3转成向量表示 tensor([[[ 0.0000, 0.0000, 0.0000], [-0.7443, 0.0692, 0.0825], [ 0.0000, 0.0000, 0.0000], [-0.1140, -0.5122, -0.4336]]], grad_fnEmbeddingBackward0) input4转成向量表示 tensor([[[-1.0667, -0.9710, -0.4726], [-0.7443, 0.0692, 0.0825], [-0.8729, 0.7102, -1.5695]], [[ 0.7366, 1.0636, 0.5947], [-0.1140, -0.5122, -0.4336], [ 0.0000, 0.0000, 0.0000]]], grad_fnEmbeddingBackward0)2 位置编码2.1 为什么需要位置编码RNN 和 LSTM 是一个词一个词按顺序进模型自然知道先后CNN 有卷积核能看到局部顺序Transformer 不含循环结构、也不含卷积操作的自注意力是并行的同时看所有词没有顺序概念。如果没有位置编码那么“我爱你”和“你爱我”的词嵌入完全一样注意力计算结果完全一样。所以需要有位置编码。位置编码的作用补上顺序信息。在嵌入向量进入编码器和解码器之前我们需要把位置信息加入嵌入向量。位置编码与嵌入向量具有相同的维度d model d_{\text{model}}dmodel​二者可以直接相加。我们在本文只讨论 Transformer 原论文中使用的位置编码——正余弦位置编码这是一种相对位置编码。这种位置编码是固定、不可训练的。但不是所有的位置编码都是不可训练的如 BERT 和 GPT-1/2 用的位置编码是可学习的位置嵌入我们在此处不展开。Transformer 使用不同频率的正余弦函数构造位置编码其公式如下P E ( p o s , 2 i ) sin ⁡ ( p o s 10000 2 i / d model ) PE_{(pos,2i)} \sin(\frac{pos}{10000^{2i/d_{\text{model}}}})PE(pos,2i)​sin(100002i/dmodel​pos​)P E ( p o s , 2 i 1 ) cos ⁡ ( p o s 10000 2 i / d model ) PE_{(pos,2i1)} \cos(\frac{pos}{10000^{2i/d_{\text{model}}}})PE(pos,2i1)​cos(100002i/dmodel​pos​)其中p o s pospos是词在句子中的位置比如第1个词第2个词d model d_{\text{model}}dmodel​是词向量的维度在原论文中是 512i ii是维度的索引。因为公式是把偶数维度 (2 i 2i2i) 给sin ⁡ \sinsin奇数维度 (2 i 1 2i12i1) 给cos ⁡ \coscos所以i ii的取值范围是0 , 1 , 2 , . . . , d model 2 − 1 0, 1, 2, ..., \frac{d_{\text{model}}}{2}-10,1,2,...,2dmodel​​−1。第一次看到这两个公式的时候很多人都是一头雾水。比如为什么是 10000 为什么要计算d model 2 \frac{d_{\text{model}}}{2}2dmodel​​其实这不是严格“数学推导出来”的而是根据设计目标推导出的一个合理形式。这里把公式记住就行本文不做展开详情见这篇文章浅谈正余弦位置编码的数学原理2.2 代码实现classPositionalEncoding(nn.Module):Implement the PE function.def__init__(self,d_model,dropout,max_len5000): d_model: 词嵌入的维度 dropout: 丢弃神经元的概率 max_len: 每个句子的最大长度 super(PositionalEncoding,self).__init__()self.dropoutnn.Dropout(pdropout)# 初始化位置编码矩阵petorch.zeros(max_len,d_model)# torch.arange(0, max_len)创建一维张量此时张量的形状为 torch.Size([max_len])# unsqueeze(dim) 表示在指定维度位置插入一个新维度# unsqueeze(0)在第 0 维插入形状变化 torch.Size([max_len]) → torch.Size([1, max_len])# unsqueeze(1)在第 1 维插入形状变化 torch.Size([max_len]) → torch.Size([max_len, 1])# 这样设计是为了后续与 div_term 进行广播运算时能正确计算出位置编码矩阵positiontorch.arange(0,max_len).unsqueeze(1)# 生成不同频率的缩放因子用于后续的正弦和余弦计算div_termtorch.exp(# torch.arange(0, d_model, 2)生成从0到d_model-1步长为2的序列如[0, 2, 4, ..., d_model-2]# math.log(10000.0)自然对数作为频率的基数torch.arange(0,d_model,2)*-(math.log(10000.0)/d_model))# 对位置编码张量的偶数维度从0开始步长为2应用正弦函数pe[:,0::2]torch.sin(position*div_term)# 对位置编码张量的奇数维度从1开始步长为2应用余弦函数pe[:,1::2]torch.cos(position*div_term)# 在第 0 维插入一个新维度作为后续的 batch_size 维度# [max_len, d_model] - [1, max_len, d_model]pepe.unsqueeze(0)# 将 pe 注册为模型的缓冲区使其成为模型的一部分自动保存和加载不参与梯度计算self.register_buffer(pe,pe)defforward(self,x):# x 的形状 [batch_size, seq_len, d_model]# self.pe[:, : x.size(1)]预计算位置编码张量 [1, max_len, d_model] 变为 [1, seq_len, d_model]# .requires_grad_(False)明确指定位置编码不参与梯度计算# x ...通过 pytorch 的广播机制位置编码会自动扩展为 [batch_size, seq_len, d_model]xxself.pe[:,:x.size(1)].requires_grad_(False)# dropout 随机将部分神经元的输出置为0防止过拟合returnself.dropout(x)2.3 代码和公式之间的关联步骤 1构造位置索引矩阵positionpositiontorch.arange(0,max_len).unsqueeze(1)torch.arange(0, max_len)生成一个一维张量[0, 1, 2, ..., max_len-1]形状为[max_len]。.unsqueeze(1)在第 1 维列方向插入一个维度得到形状[max_len, 1]的列向量。为什么需要列向量因为后面要与频率缩放因子div_term行向量进行广播乘法生成一个[max_len, d_model/2]的矩阵其中每个元素是pos * factor_i。步骤 2计算频率缩放因子div_termdiv_termtorch.exp(torch.arange(0,d_model,2)*-(math.log(10000.0)/d_model))torch.arange(0, d_model, 2)生成[0, 2, 4, ..., d_model-2]这些值正是公式中的2 i 2i2i。math.log(10000.0)是自然对数ln ⁡ ( 10000 ) \ln(10000)ln(10000)。将[2i]乘以-(ln(10000)/d_model)得到-(2i * ln(10000))/d_model。再取指数exp得到exp(-(2i * ln(10000))/d_model)。根据指数和对数性质exp ⁡ ( − 2 i ln ⁡ ( 10000 ) d model ) 10000 − 2 i / d model \exp\left(-\frac{2i \ln(10000)}{d_{\text{model}}}\right) 10000^{-2i/d_{\text{model}}}exp(−dmodel​2iln(10000)​)10000−2i/dmodel​这正是公式中分母部分的倒数。也就是说div_term实际上是一个向量其第i ii个元素为KaTeX parse error: Expected EOF, got _ at position 11: \text{div_̲term}[i] 1000…步骤 3计算position * div_termposition形状为torch.size([max_len, 1])div_term形状为torch.size([d_model/2])。通过广播机制两者相乘得到一个形状为[max_len, d_model/2]的矩阵矩阵的每个元素为KaTeX parse error: Expected EOF, got _ at position 39: …times \text{div_̲term}[i] pos …这正是正弦/余弦函数的自变量。步骤 4填充偶数和奇数维度pe[:, 0::2]选取所有行、从第 0 列开始每隔一列即偶数索引列赋值为sin(position * div_term)。这样就实现了P E ( p o s , 2 i ) sin ⁡ ( p o s ⋅ 10000 − 2 i / d model ) PE_{(pos, 2i)} \sin\left(pos \cdot 10000^{-2i/d_{\text{model}}}\right)PE(pos,2i)​sin(pos⋅10000−2i/dmodel​)pe[:, 1::2]选取所有行、从第 1 列开始每隔一列即奇数索引列赋值为cos(position * div_term)。这样就实现了P E ( p o s , 2 i 1 ) cos ⁡ ( p o s ⋅ 10000 − 2 i / d model ) PE_{(pos, 2i1)} \cos\left(pos \cdot 10000^{-2i/d_{\text{model}}}\right)PE(pos,2i1)​cos(pos⋅10000−2i/dmodel​)为什么代码中使用指数和对数变换直接计算10000 ** (-2i/d_model)也可以但存在两个问题幂运算在深度学习中可能不如指数对数稳定且高效。使用exp和log可以避免显式的除法更适合在 GPU 上并行计算。通过恒等变换10000 − 2 i / d model exp ⁡ ( − 2 i d model ln ⁡ ( 10000 ) ) 10000^{-2i/d_{\text{model}}} \exp\left(-\frac{2i}{d_{\text{model}}} \ln(10000)\right)10000−2i/dmodel​exp(−dmodel​2i​ln(10000))我们可以用一次exp和一次乘法完成所有频率的计算简洁高效。总结词嵌入Token Embedding负责语义这个词是什么意思位置编码Positional Encoding负责位置这个词排在第几位