基于PyTorch的Transformer组件实现

最近看了不少介绍LLM工作原理的文章，发现每一篇都会试图跟读者讲明白作为baseline的Transformer架构到底长啥样。但是好像比较少有代码实现的示例和具体的例子帮助理解。于是自己也想尝试着写一篇含有代码实现和具体例子解释的文章，希望能够给喜欢编程朋友带来一点点对于Transformer结构的启发😁。

整体架构和关键组件

注意力机制的提出极大提高了神经网络对序列相关信息的学习能力。Transformer结构正是完全基于注意力机制来完成对源语言序列和目标语言序列之间的全局依赖建模。它通过并行注意力计算来提取文本之间隐含的上下文信息,而不需要依赖递归结构或者时间步展开等方法。这给序列学习带来了巨大的改变。

Transformer的总体Architecture如下：

可以发现这里头主要有这么几个东西：

• Position Encoding部分
• Attention部分，包括Muiti-Head Attention和Masked Muiti-Head Attention
• 前馈层部分，对应图中的Feed-Forward
• 残差链接和归一化部分，对应图中的“Add&Norm”

接下来我尝试着去解释每个组件大致的功能，以及如何用基于Pytorch的代码实现这些模块

Position Encoding

原理

对于输入的文本序列，首先通过输入嵌入层将每个词语转换为其对应的向量表示。通常直接为每个词语创建一个向量表示。

由于Transfomer 模型不再使用基于循环的方式处理文本输入，序列中没有任何信息能够告诉模型词语之间的相对位置关系。在送入编码器端分析其上下文语义之前，一个非常关键的操作是在词嵌入中添加位置编码（Positional Encoding）这一特征。

具体来说，序列中每一个词语所在的位置都有一个向量。这一向量会与词语表示相应相加并送入到后续模块中做进一步处理。在训练的过程当中，模型会自动地学习到如何利用这部分位置信息。为了计算不同位置对应的编码，Transformer 模型使用不同频率的正余弦函数如下所示：

$PE(pos,2i)=sin(pos100002i/d)$

$PE(pos,2i+1)=cos(pos100002i/d)$

其中，pos 表示词语所在的位置，2i 和2i+1 表示位置编码向量中的相应维度，d 则对应位置编码的总维度。

通过上面这种方式得到位置编码有这样几个优点：

• 首先，正余弦函数的范围是在[-1,+1]，导出的位置编码与原词嵌入相加不会使得结果偏离过远而损害原有词语的语义信息。
• 其次，根据三角函数的基本性质，可以知道第pos+k 个位置的编码是第pos 个位置的编码的线性组合，这就意味着位置编码中包含着词语之间的距离信息(可以通过三角函数的变换得到)。
  对于输入的文本序列，首先通过输入嵌入层将每个词语转换为其对应的向量表示。通常直接为每个词语创建一个向量表示。

一个例子

举个例子说明这个：

假设我们有一个文本序列，包含四个词语：[我，爱，你，吗]。我们可以为每个词语分配一个编号：[0, 1, 2, 3]。然后我们可以根据上面的公式计算出每个词语所在位置的编码向量。假设位置编码的总维度是4，那么我们可以得到如下的矩阵：

$PE=\left[\begin{array}{cccc}sin(0)& cos(0)& sin(0)& cos(0)\\ sin(1/10000)& cos(1/10000)& sin(1/10000^2)& cos(1/10000^2)\\ sin(2/10000)& cos(2/10000)& sin(2/10000^2)& cos(2/10000^2)\\ sin(3/10000)& cos(3/10000)& sin(3/10000^2)& cos(3/10000^2)\end{array}\right]$

接下来，我们可以将每个词语的向量表示与其对应的位置编码向量相加，得到新的词嵌入矩阵：

$E=\left[\begin{array}{c}{e}_0+P{E}_0\\ e_1+P{E}_1\\ e_2+P{E}_2\\ e_3+P{E}_3\end{array}\right]$

其中，$e_i$ 表示第 i 个词语的向量表示。这样，我们就为每个词语添加了位置信息，并且保证了不同位置的词语有不同的编码。

最后，我们将新的词嵌入矩阵送入编码器端进行后续处理。

代码

使用Pytorch 实现的位置编码参考代码如下所示：

from kiwisolver import Variable
import torch
import torch.nn as nn
import math

class PositionEncorder(nn.Module):
    def __init__(self, d_model, max_seq_len):
        super().__init__()
        self.d_model = d_model # 模型维度
        self.max_seq_len = max_seq_len # 序列的最大长度

        pe = torch.zeros(max_seq_len, d_model) # 储存位置编码的Tensor
        for pos in range(max_seq_len): # 遍历序列每个位置
            for i in range(0, d_model, 2): # 遍历每一个维度，步长为2（两个维度一组进行正余弦计算）
                pe[pos,i] = math.sin(pos / (10000 ** ((2 * i)/d_model)))
                pe[pos,i+1] = math.cos(pos / (10000 ** ((2 * (i+1))/d_model)))

        pe = pe.unsqueeze(0) # 添加了一个额外的维度，使其形状从(max_seq_len, d_model)变为(1, max_seq_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x * math.sqrt(self.d_model) # 使得单词嵌入表示相对大一些
        seq_len = x.size(1)
        x  = x + Variable(self.pe[:,:seq_len], requires_grad=False).cuda() # 加上位置信息

MutiHead Self Attention

原理

自注意力

自注意力操作是基于Transformer 的机器翻译模型的核心操作，它能够在源语言的编码和目标语言的生成中有效地利用源语言、目标语言任意两个单词之间的依赖关系。给定由单词语义嵌入和其位置编码相加得到的输入表示。为了实现对上下文语义依赖的捕捉，进一步引入在自注意力机制中涉及到的三个要素：查询$q_i$（Query），键$k_i$（Key），值$vi$（Value）。

在编码输入序列中每一个单词的表示的过程中，这三个要素用于计算上下文单词所对应的权重得分。直观地说，这些权重反映了在编码当前单词的表示时，对于上下文不同部分所需要的关注程度。通过三个线性变换$W^Q\in R^{d\times d_q}$，$W^K\in R^{d\times d_k}$，$W^V\in R^{d\times d_v}$将输入序列中的每一个单词表示$x_i$ 转换为其对应的$q_i\in R^{d_k}$，$k_i\in R^{d_k}$，$vi\in R^{d_v}$ 向量。

为了得到编码单词$x_i$ 时所需要关注的上下文信息，通过位置$i$查询向量与其他位置的键向量做点积得到匹配分数$q_i\cdot k_1,q_i\cdot k_2,...,q_i\cdot k_t$。

为了防止过大的匹配分数在后续Softmax 计算过程中导致的梯度爆炸以及收敛效率差的问题，这些得分会除放缩因子$\surd d$ 以稳定优化。放缩后的得分经过Softmax 归一化为概率之后，与其他位置的值向量相乘来聚合希望关注的上下文信息，并最小化不相关信息的干扰。上述计算过程可以被形式化地表述如下：

$Z=Attention(Q,K,V)=Softmax(Q{K}^T/\surd d)V$

其中$Q\in R^{L\times d_q}$，$K\in R^{L\times d_k}$，$V\in R^{d\times d_v}$ 分别表示输入序列中的不同单词的q, k, v 向量拼接组
成的矩阵，L 表示序列长度，$Z\in R^{L\times d_v}$表示自注意力操作的输出

多头注意力

为了更好地利用自注意力机制来捕捉上下文信息的多样性，提出了多头自注意力（Multi-head Attention）的方法，它可以同时关注上下文的不同方面。其实就是每一个单词跟多组$W^Q$，$W^K$，$W^V$ 做乘法，映射到多个表征空间，其中每一个表征空间可能偏向于关注不同的上下文内容

一个例子

举个例子：

假设我们有一个句子：“I love to eat pizza.”

在多头自注意力机制中，我们可以通过使用多组独立的线性变换来创建多个注意力头。每个注意力头都有自己的权重矩阵 $W^Q$，$W^K$ 和 $W^V$，用于将输入的单词映射到不同的表征空间。

假设我们使用两个注意力头，即 heads = 2。在这种情况下，我们将有两个独立的注意力头，每个头都有自己的权重矩阵。

第一个注意力头可能更关注动词和名词之间的关系，例如 “love” 和 “pizza” 之间的关系。在这个注意力头中，权重矩阵 $W_1^Q$ 可以将 “love” 映射到一个表征空间，$W_1^K$ 可以将 “pizza” 映射到另一个表征空间，$W_1^V$ 可以将 “pizza” 映射到另一个表征空间。这样，第一个注意力头可以关注句子中动词和名词之间的语义关联。

第二个注意力头可能更关注形容词和名词之间的关系，例如 “I” 和 “pizza” 之间的关系。在这个注意力头中，权重矩阵 $W_2^Q$ 可以将 “I” 映射到一个表征空间，$W_2^K$ 可以将 “pizza” 映射到另一个表征空间，$W_2^V$ 可以将 “pizza” 映射到另一个表征空间。这样，第二个注意力头可以关注句子中形容词和名词之间的语义关联。

所以说，通过使用多头自注意力机制，我们可以同时关注句子中不同方面的语义关系，从而捕捉到更丰富的上下文信息。每个注意力头都可以偏向于关注句子中不同的语义关联，从而提升模型的表征能力和语义理解能力。

代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super().__init__()

        self.d_model = d_model
        self.d_k = d_model // heads  # 确保每个头的维度加起来等于模型的总维度
        self.h = heads

        self.q_linear = nn.Linear(d_model, d_model)  # 线性变换，将输入维度调整为d_model
        self.v_linear = nn.Linear(d_model, d_model)  # 线性变换，将输入维度调整为d_model
        self.k_linear = nn.Linear(d_model, d_model)  # 线性变换，将输入维度调整为d_model
        self.dropout = nn.Dropout(dropout)  # Dropout层，用于随机置零一部分输入
        self.out = nn.Linear(d_model, d_model)  # 线性变换，将输入维度调整为d_model

    def attention(q, k, v, d_k, mask=None, dropout=None):
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)  # q 和 k 点乘后除以根号d_k, 得到注意力分数

        # 掩盖掉那些为了填补长度增加的单元，使其通过softmax 计算后为0
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)

        scores = F.softmax(scores, dim=-1)  # 对注意力分数进行softmax操作，得到注意力权重

        if dropout is not None:
            scores = dropout(scores)  # 对注意力权重进行dropout操作

        output = torch.matmul(scores, v)  # 注意力权重与v相乘得到输出
        return output

    def forward(self, q, k, v, mask=None):

        bs = q.size(0)  # 获取batch size

        # 进行线性操作划分成h个头
        k = self.k_linear(k).view(bs, -1, self.h, self.d_k)  # 对k进行线性变换并调整维度形状
        q = self.q_linear(q).view(bs, -1, self.h, self.d_k)  # 对q进行线性变换并调整维度形状
        v = self.v_linear(v).view(bs, -1, self.h, self.d_k)  # 对v进行线性变换并调整维度形状

        # 调换维度顺序使得最后一个维度为单元长度
        k = k.transpose(1, 2)  # 转置k的维度
        q = q.transpose(1, 2)  # 转置q的维度
        v = v.transpose(1, 2)  # 转置v的维度

        # 计算attention
        scores = attention(q, k, v, self.d_k, mask, self.dropout)

        # Concat多个头
        concat = scores.transpose(1, 2).contiguous().view(bs, -1, self.d_model)  # 将多个头的结果拼接起来
        output = self.out(concat)  # 进行线性变换得到最终输出

        return output
        

前馈层

原理

前馈层将自注意力子层的输出作为输入，并通过一个包含Relu 激活函数的双层全连接网络对输入进行更加高级的非线性变换。这一非线性变换对模型最终的效果有着非常重要的作用。

代码

import torch
import torch.nn as nn
import torch.nn.functional as F

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff=2048, dropout=0.1):
        super(FeedForward, self).__init__()

        # d_ff默认为2048
        self.linear_1 = nn.Linear(d_model, d_ff)  # 线性变换，将输入维度调整为d_ff
        self.dropout = nn.Dropout(dropout)  # Dropout层，用于随机置零一部分输入
        self.linear_2 = nn.Linear(d_ff, d_model)  # 线性变换，将输入维度调整为d_model

    def forward(self, x):
        x = self.dropout(F.relu(self.linear_1(x)))  # 使用ReLU激活函数并应用dropout
        x = self.linear_2(x)  # 进行线性变换
        return x

残差连接和归一化

原理

由Transformer 架构构成的网络模型通常都是非常巨大。编码器和解码器都由许多层基本的Transformer 单元构成，每一层里面都包含复杂的非线性变换，这就导致模型的训练比较难。因此在Transformer 单元中进一步引入了残差连接与层归一化技术以进一步提升训练的稳定性。

• 残差链接：残差链接是一种在每个子层（如多头自注意力层和前馈层）后添加输入与输出之和的方法。在Transformer模块，残差连接主要是指使用一条直连通道直接将对应子层的输入连接到输出上去，从而避免由于网络过深在优化过程中潜在的梯度消失问题，公式如下：

  $x_{l+1}=f(x_l)+x_l$

• 层归一化：一种对每个子层的输出进行规范化处理的方法，用于将数据平移缩放到均值为0，方差为1 的标准分布。这样能让训练变得更加稳定

例子

层归一化：

假设一个Transformer模型的隐藏层维度为d=4，输入为一个三词的序列，即x=[x1,x2,x3]，其中x1,x2,x3都是4维的向量。那么隐藏层的输出为h=[h1,h2,h3]，也是一个三词的序列，其中h1,h2,h3也都是4维的向量。我们可以用矩阵形式表示h：

$h=[[h_{11},h_{12},h_{13},h_{14}],[h_{21},h_{22},h_{23},h_{24}],[h_{31},h_{32},h_{33},h_{34}]]$

我们需要分别计算每个词向量的均值和方差，即：

$mea{n}_i=\frac{1}{4}{\sum }_{j=1}^4{h}_{ij}$

$va{r}_i=\frac{1}{4}{\sum }_{j=1}^4(h_{ij}-mea{n}_i{)}^2$

其中 $i=1,2,3$ 分别表示第一个，第二个和第三个词。然后，我们用这些均值和方差对每个词向量进行归一化，即：

$z_i=\frac{h_i-mea{n}_i}{\sqrt{va{r}_i+eps}}$

其中 $eps$ 是一个很小的正数，用来避免除零错误。这样，我们就得到了归一化后的词向量矩阵：

$z=[[z_1],[z_2],[z_3]]$

这种归一化方法可以使每个词向量在不同的样本和层之间具有相同的尺度，从而减少了梯度消失或爆炸的风险，提高了模型的稳定性和效率。

残差连接

假设我们有一个由两层Transformer 单元组成的编码器，每个单元包含一个多头自注意力层和一个前馈层。输入序列为$x_0,x_1,...,x_n$，其中$x_i$ 是一个$d$ 维的向量。那么第一层Transformer 单元的输出为：

$$\begin{array}{rl}{z}_0& =\text{MultiHead}(x_0,x_0,x_0)+x_0\\ z_1& =\text{MultiHead}(x_1,x_1,x_1)+x_1\\ & ⋮\\ z_n& =\text{MultiHead}(x_n,x_n,x_n)+x_n\\ y_0& =\text{FFN}(z_0)+z_0\\ y_1& =\text{FFN}(z_1)+z_1\\ & ⋮\\ y_n& =\text{FFN}(z_n)+z_n\end{array}$$

其中$\text{MultiHead}$ 是多头自注意力函数，$\text{FFN}$ 是前馈网络函数。然后我们对每个$y_i$ 进行层归一化操作，得到第一层Transformer 单元的最终输出：

$$\begin{array}{rl}{h}_0& =\text{LayerNorm}(y_0)\\ h_1& =\text{LayerNorm}(y_1)\\ & ⋮\\ h_n& =\text{LayerNorm}(y_n)\end{array}$$

其中$\text{LayerNorm}$ 是层归一化函数。接着我们将$h_i$ 作为第二层Transformer 单元的输入，重复上述过程，得到编码器的最终输出。

代码

import torch
import torch.nn as nn

class NormLayer(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.size = d_model
        self.eps = eps

        # 两个可学习参数
        self.alpha = nn.Parameter(torch.ones(self.size))  # 可学习的缩放参数
        self.bias = nn.Parameter(torch.zeros(self.size))  # 可学习的偏置参数

    def forward(self, x):
        norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) \
        / (x.std(dim=-1, keepdim=True) + self.eps) + self.bias  # 归一化操作和残差连接
        return norm

开始搭建：Encorder & Deocrder

介绍完几个关键的组成部分，接下来就可以把他们搭建成Transformer的编码器和解码器了

编码器端的实现比较简单，就是用上面介绍的模块搭起来。解码器端就要复杂一些了。主要有两个不同的地方。

第一，解码器的每个Transformer 块里面，第一个自注意力子层要加一个注意力掩码，就是图中的掩码多头注意力（Masked Multi-Head Attention）部分。这是因为在翻译的时候，编码器端只是把源语言序列的信息编码好，这个序列是已经知道的，所以编码器只要考虑怎么把上下文语义信息融合好就行了。而解码器端要生成目标语言序列，这个过程是自回归的，也就是说，每生成一个单词，只能看到当前单词之前的目标语言序列，不能看到后面的。所以要加一个掩码，把后面的文本信息遮住，不让模型在训练的时候直接看到后面的文本序列，不然就没法训练好。

第二，解码器端还多了一个多头注意力模块，用交叉注意力（Cross-attention）方法，同时接收编码器端的输出和当前Transformer 块的前一个掩码注意力层的输出。查询是用解码器前一层的输出投影出来的，而键和值是用编码器的输出投影出来的。它的作用是在翻译的时候，为了生成合理的目标语言序列要看看源语言序列是什么。基于这样的编码器和解码器结构，源语言文本先经过编码器端的每个Transformer 块对它的上下文语义做一层层的提取，最后输出每个源语言单词相关的表示。解码器端按照自回归的方式生成目标语言文本，也就是说，在每个时间步t，根据编码器端输出的源语言文本表示，和前t − 1 个时刻生成的目标语言文本，生成当前时刻的目标语言单词。

代码如下：

• Encorder部分：

import torch.nn as nn
import torch.nn.functional as F

class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()

        self.norm_1 = Norm(d_model)  # 第一个归一化层
        self.norm_2 = Norm(d_model)  # 第二个归一化层

        self.attn = MultiHeadAttention(d_model, n_heads, dropout=dropout)  # 多头注意力机制
        self.ff = FeedForward(d_model, dropout=dropout)  # 前馈神经网络

        self.dropout_1 = nn.Dropout(dropout)  # 第一个丢弃层
        self.dropout_2 = nn.Dropout(dropout)  # 第二个丢弃层

    def forward(self, x, mask):
        x2 = self.norm_1(x)  # 对输入进行归一化
        x = x + self.dropout_1(self.attn(x2, x2, x2, mask))  # 应用多头注意力机制和残差连接，并使用丢弃层
        x2 = self.norm_2(x)  # 对上一步的输出进行归一化
        x = x + self.dropout_2(self.ff(x2))  # 应用前馈神经网络和残差连接，并使用丢弃层

        return x

class Encorder(nn.Module):
    def __init__(self, vocab_size, d_model, N, heads, dropout):
        super().__init__()
        self.N = N
        self.embed = Embedder(vocab_size, d_model)  # 嵌入层
        self.pe = PositionalEncoder(d_model, dropout=dropout)  # 位置编码
        self.layers = get_clones(EncoderLayer(d_model, heads, dropout), N)  # 创建 N 个编码器层
        self.norm = Norm(d_model)  # 归一化层

    def forward(self, src, mask):
        x = self.embed(src)  # 对输入源序列进行嵌入
        x = self.pe(x)  # 对嵌入序列应用位置编码
        for i in range(self.N):
            x = self.layers[i](x, mask)  # 将序列通过每个编码器层

        return self.norm(x)  # 对最终输出进行归一化

• Decorder部分：

import torch.nn as nn
import torch.nn.functional as F

class DecoderLayer(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super(DecoderLayer, self).__init__()

        self.norm_1 = Norm(d_model)  # 归一化层 1
        self.norm_2 = Norm(d_model)  # 归一化层 2
        self.norm_3 = Norm(d_model)  # 归一化层 3

        self.dropout_1 = nn.Dropout(dropout)  # 丢弃层 1
        self.dropout_2 = nn.Dropout(dropout)  # 丢弃层 2
        self.dropout_3 = nn.Dropout(dropout)  # 丢弃层 3

        self.attn1 = MultiHeadAttention(heads, d_model, dropout=dropout)  # 注意力机制 1
        self.attn2 = MultiHeadAttention(heads, d_model, dropout=dropout)  # 注意力机制 2
        self.ff = FeedForward(d_model, dropout=dropout)  # 前馈神经网络

    def forward(self, x, e_outputs, src_mask, trg_mask):
        x2 = self.norm_1(x)  # 对输入进行归一化
        x = x + self.dropout_1(self.attn1(x2, x2, x2, trg_mask))  # 应用注意力机制 1 和残差连接，并使用丢弃层
        x2 = self.norm_2(x)  # 对上一步的输出进行归一化
        x = x + self.dropout_2(self.attn2(x2, e_outputs, e_outputs, src_mask))  # 应用注意力机制 2 和残差连接，并使用丢弃层
        x2 = self.norm_3(x)  # 对上一步的输出进行归一化
        x = x + self.dropout_3(self.ff(x2))  # 应用前馈神经网络和残差连接，并使用丢弃层
        return x

class Decoder(nn.Module):
    def __init__(self, vocab_size, d_model, N, heads, dropout=0.1):
        super().__init__()
        self.N = N
        self.embed = Embedder(vocab_size, d_model)  # 嵌入层
        self.ps = PositionalEncoder(d_model, dropout=dropout)  # 位置编码
        self.layers = get_clones(DecoderLayer(d_model, heads, dropout), N)  # 创建 N 个解码器层
        self.norm = Norm(d_model)  # 归一化层

    def forward(self, trg, e_outputs, src_mask, trg_mask):
        x = self.embed(trg)  # 对目标序列进行嵌入
        x = self.ps(x)  # 对嵌入序列应用位置编码
        for i in range(self.N):
            x = self.layers[i](x, e_outputs, src_mask, trg_mask)  # 将序列通过每个解码器层
        return self.norm(x)  # 对最终输出进行归一化

• 最终的Transformer：

import torch.nn as nn

class Transformer(nn.Module):
    def __init__(self, src_vocab, trg_vocab, d_model, N, heads, dropout):
        super().__init__()
        self.encoder = Encoder(src_vocab, d_model, N, heads, dropout)
        self.decoder = Decoder(trg_vocab, d_model, N, heads, dropout)
        self.out = nn.Linear(d_model, trg_vocab)

    def forward(self, src, trg, src_mask, trg_mask):
        e_outputs = self.encoder(src, src_mask)
        d_output = self.decoder(trg, e_outputs, src_mask, trg_mask)
        output = self.out(d_output)
        return output