经典结构Transformer解读

谷歌公司在2017年发布了革命性的Transformer结构。Transformer最初是为机器翻译任务而设计的，然而由于其出色的特征学习能力，特别是在NLP领域上长距离的建模能力使其继卷积神经网络CNN和RNN之后，成为了语言建模最流行的结构。因此，现在的大语言模型大多数以Transformer结构为基础来进行训练。
基于 Transformer 结构的编码器（Encoder，图左侧）和解码器（Decoder，图右侧）结构如图所示。编码器和解码器均由若干个基本的 Transformer Block组成（图中的灰色框）。N× 表示按顺序 N 次计算。Transformer主要有以下几个模块组成：
输入模块：Transformer Encoder的开始部分，并对每个词Token添加位置编码便于模型在处理文本时能够Token时序。
注意力模块：Transformer的核心模块，用来捕获输入序列中的长距离依赖关系。
前馈模块：由全连接层—激活函数—全连接层的结构组成，目的是在模型中引入更多非线性。
残差连接模块：解决深层网络训练中出现梯度消失或梯度爆炸问题的经典方法。
层归一化模块：对各层参数的标准化处理，使模型的参数保持在一个合理数值范围。

图1 Transformer结构

输入表示与位置编码

Transformer Encoder的开始部分，输入句子中的每个词都表示为一个独立的Token Embedding。然而这样的Embedding仅包含了句子的语义信息。为了使模型在处理文本能够考虑到Token的顺序和位置信息，Transformer 模型对Token Embedding添加了不同频率的正余弦函数，如下所示：
$PE(pos,2i)=sin(\frac{pos}{{10000}^{2i/d}})$
$PE(pos,2i+1)=cos(\frac{pos}{{10000}^{2i/d}})$
其中，pos表示单词所在的位置，和表示位置编码向量中的对应维度，则对应位置编码的总维度（与词特征维度相同）。
通过将词汇特征（Token Embedding）与位置特征（Position Embedding）相加，得到每个Token的最终输入特征（Input Embedding），如下图所示。

多头注意力机制及其改进

1.多头注意力机制

自注意力机制是Transformer的核心模块，用来计算输入序列中每个Token与其他Token的相关性，并将输入序列的加权和作为该元素的向量表示。因此，自注意力机制使得模型具备长距离的建模能力，这也是为什么大模型大多数用Transformer为基础架构的原因。具体来说，自注意力机制中涉及到的三个关键元素：查询$q_i$（Query），键$k_i$（Key），值$v_i$（Value），这三个元素用于计算上下文单词所对应的权重得分。对于输入序列中的每一个单词表示，通过三个不同的权重矩阵$W^Q\in R^{d\times d_q},W^K\in R^{d\times d_k},W^Q\in R^{d\times d_v}$ 转换为其对应的向量，用于后续的注意力机制计算。

为了使编码单词$x_i$与其他单词建立起注意力关系，Transformer通过位置查询向量与其他位置的键向量进行点积得到匹配分数$q_i\cdot k_1,q_i\cdot k_2,\dots ,q_i\cdot k_t$。如上图所示，计算过程可以被形式化地表述如下：
$Z=\text{Attention}(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{K}})V$
其中$Q\in R^{L\times d_q},K\in R^{L\times d_k},V\in R^{L\times d_v}$分别有输入序列不同单词去$q,k,v$向量所拼接组成的矩阵，$L$为输入序列的长度，$d$为缩放因子，大小为特征维度。
为了捕捉到输入序列中单词之间的各种关系，Transformer引入了“多头”机制。在多头自注意力机制中，模型有多个独立的注意力头，每个头都有自己独立的$Q,K,V$矩阵，其目的是通过不同的注意力头来捕捉输入序列中的不同关系。然后把所有注意力头的输出拼接起来，再通过一个线性变换$W^0$来综合不同的注意力头(输入序列中的不同关系)，上述过程可以被形式化表述如下：
$multihead(Q,K,V)=Concat(hea{d}_1,\dots ,hea{d}_h)W^0$
$hea{d}_i=Attention(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$
其中$Q{W}_i^Q,K{W}_i^K,V{W}_i^V$为不同头$hea{d}_i$对于的权重矩阵。
使用 Pytorch 实现的多头注意力参考代码如下：

class MultiHeadAttention(nn.Module):
    # heads多头注意力的数量
    # hid_dim 每个词输入的向量维度
    def __init__(self, heads, hid_dim, dropout):
        super().__init__()
        self.hid_dim = hid_dim
        self.heads = heads
        self.d_k = hid_dim // heads
        self.w_q = nn.Linear(hid_dim, hid_dim)
        self.w_v = nn.Linear(hid_dim, hid_dim)
        self.w_k = nn.Linear(hid_dim, hid_dim)
        self.fc = nn.Linear(hid_dim, hid_dim)
        self.dropout = nn.Dropout(dropout)
    def attention(q, k, v, scale, mask=None, dropout=None):
        attention = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(scale)
        if mask is not None:
            mask = mask.unsqueeze(1)
        # 把mask为0的位置的attention分数设置为-1e10
        attention = attention.masked_fill(mask == 0, -1e10)
        attention = F.softmax(attention, dim=-1)
        if dropout is not None:
            attention = dropout(attention)
        # attention结果与v相乘，得到多头注意力的结果
        output = torch.matmul(attention, v)
        return output
    def forward(self, query, key, value, mask=None):
        bs = query.shape[0]
        # 把K Q V矩阵拆分为多头注意力，头数为self.heads
        Q = self.w_q(query).view(bs, -1, self.heads, self.d_k)
        K = self.w_k(key).view(bs, -1, self.heads, self.d_k)
        V = self.w_v(value).view(bs, -1, self.heads, self.d_k)
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        # 注意力计算
        attention_score = self.attention(Q, K, V, self.d_k, mask, self.dropout)
        # 把多个头拼接在一起来作为全连接层的输入
        output = attention_score.transpose(1, 2).contiguous().view(bs, -1, self.hid_dim)
        output = self.fc(output)
        return output

2.注意力机制的改进

在Transformer的自注意力机制中，每一个token都会与所有的token算注意力，所以自注意力机制的时间和存储空间的复杂度与序列的长度呈平方关系，即$O(n^2)$，其中$n$为输入序列的长度，因此如何优化自注意力机制的时间和空间复杂度是人们越来越关心的问题。一些研究从近似注意力的角度研究来减少时间和空间的需求，提出了稀疏注意力和局部注意力等方法。一些常见的稀疏注意力如下图所示。

（1）全局注意力：Transformer的注意力是所有token进行全局注意力，这里选择了3个点token作为全局节点来进行全局注意力。
（2）带状注意力：由于大部分数据具有局部性，即对于大多数的token只与附近的token有强相关性，因此带状注意力将现在每个token只与邻近的节点进行注意力交互。
（3）随机注意力：通过随机采样来对token进行注意力交互，其目的是为了提高非局部的交互能力。
（4）局部注意力：对输入序列划分不重叠的局部区域块（Block），仅在该区域块内进行全局注意力。图中选择了Block的大小为$2\times 2$。

前馈神经网络

前馈层将自注意力的输出作为输入，并通过一个带有激活函数的两层全连接网络来进行复杂的非线性变换，其目的是在模型引入更多的非线性，以此来增强模型的表达能力。激活函数通常采用ReLU(Rectified Linear Unit)或GeLU(Gaussian Error Linear Unit)。形式化表示如下。
$FFN(x)=ReLU(x{W}_1+b_1)W_2+b_2$
与上一小节多头注意力比较，多头注意力机制是为了使模型获得输入序列中长距离的依赖关系，FNN则是为了使模型能够学习到这些依赖关系的复杂模式。
使用 Pytorch 实现的前馈神经网络参考代码如下：

class FeedForwardNetwork(nn.Module):
    def __init__(self, hid_dim, d_ff=2048, dropout=0.1):
        super().__init__()
        self.linear_1 = nn.Linear(hid_dim, d_ff)
        self.dropout_1 = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, hid_dim)
        self.dropout_2 = nn.Dropout(dropout)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.dropout_1(self.relu(self.linear_1(x)))
        x = self.dropout_2(self.linear_2(x))
    return x

残差连接与层归一化

残差连接与层归一化技术的引入是为了进一步提升Transformer在训练中的稳定性。残差连接与层归一化分别对应图1中的Add模块与Norm模块。

残差连接示意图
具体来说，Add模块实现了残差连接，该方法已经被证明用来有效解决网络退化的问题，也就是深层网络的效果反而比浅层网络效果差的问题。如图所示，其原理将某一层的输入$x_i$加上它的输出$f(x_i)$作为这一层的输出来保证深层网络的效果不差于浅层网络，方法形式化表示如下：
$x_{i+1}=f(x_i)+x_i$
其中$x_i$表示第$i$层的输入，$f(\cdot )$表示一个映射函数。
对于深度神经网络的每一层都可以看作是一层相对独立的分类器，即对上一层输出数据进行分类。然而，由于每一次输出的数据都有可能不同，这可能使得网络内部发生数据分布偏移，随着网络层数的增加最终导致网络性能的下降。为了解决这一问题，标准化方法通过及时将数据拉回到正态分布来保持神经网络的稳定性。
根据标准化维度的不同，可以将其分为批标准化（Batch Normalization）和层标准化（Layer Normalization）。批标准化是通过对批量大小（Batch size）维度的标准化来稳定数据的分布，而层标准化是通过隐藏层（hidden size）维度的标准化来稳定某一层的分布。
Transformer模型的编码器和解码器都是6层的神经网络，为了解决网络层级越多时出现的梯度消失和梯度爆炸的问题，Transformer模型在Norm模块中实现了层归一化来解决这一问题。
在层标准化中，对于给定的样本$x\in R^{B\times D}$，其中$B$是批量大小，$D$是特征维度大小。每个样本的均值$\mu$和方差${\sigma }^2$计算如下：
$\mu =\frac{{\sum }_{i=1}^D{x}_i}{D}$
${\sigma }^2=\frac{{\sum }_{i=1}^D{(x_i-\mu )}^2}{D}$
之后通过均值$\mu$和方差${\sigma }^2$来对输入进行标准化：
$\widehat{x_i}=\frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}$
其中$ϵ$是一个非常小的正数来避免父母为的情况。最后通过重新缩放和偏移后的输出。
$LN(x_i)=\gamma \widehat{x_i}+\beta$
其中$\gamma$和$\beta$是可学习的参数。
使用 Pytorch 实现的层归一化参考代码如下：

class NormLayer(nn.Module):
    def __init__(self, size, eps=1e-6):
        super().__init__()
        self.size = size
        # 层归一化包含两个可以学习的参数
        self.alpha = nn.Parameter(torch.ones(self.size))
        self.bias = nn.Parameter(torch.zeros(self.size))
        self.eps = eps

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

解码器

交互自注意力机制

上图给出了一个Transformer的Encoder-Decoder交互结构，Encoder通过将输入序列转化为对应的编码信息，之后和Decoder与这些编码信息进行充分交互。正如图10.2.1所示，Encoder与Decoder的结构类似，两者都拥有Add模块、Feed Forward模块等常规模块，解码器端还额外增加了掩码多头自注意力机制和交互自注意力机制。
交互自注意力机制，即Encoder-Decoder架构中的多头自注意力（Multi-head Self Attention），是现代序列到序列（Seq2Seq）模型的关键组成部分，如图10.2.5.1所示，用于解决自然语言处理任务，如机器翻译、文本摘要等。这一机制允许解码器（Decoder）在生成输出序列的过程中，动态地参考编码器（Encoder）所捕获的整个输入序列的语境信息。具体而言，当解码器在生成序列的某个时刻t的输出时，它会利用该时刻的隐藏状态作为查询（Query），与编码器所有时刻的隐藏状态——这里扮演着键（Key）和值（Value）的角色——进行交互。这个过程通过计算查询向量与键向量之间的相似度（通常采用点积注意力机制），得到一个权重分布，随后使用这个权重分布对值向量进行加权求和，从而获得一个上下文相关的向量表示。
引入多头的概念进一步增强了这种机制的效果。多头自注意力允许模型同时关注输入的不同位置，从不同的子空间中捕捉信息，这样可以更全面地理解输入序列中的各种依赖关系。每个注意力头独立工作，处理输入的线性变换，并行地执行注意力计算，最后将所有头的结果拼接起来并通过另一个线性层转换，形成最终的注意力输出。
因此，交互自注意力机制确保了解码器在生成每个输出词时，都能充分考虑到输入序列的全局信息，而不是仅仅依赖于局部或最近的输入词。这使得模型能够更好地处理长距离依赖问题，提高生成序列的质量和连贯性，尤其是在处理复杂且信息丰富的文本数据时表现尤为突出。

掩码多头自注意力机制

掩码矩阵(Mask矩阵)作用在自注意力矩阵上的示意图。
掩码多头自注意力机制（Masked Multi-head Self Attention），作为Transformer模型中解码器 （Decoder）的一个核心组件，其设计目的是为了在生成序列时维持序列的因果关系，即确保在生成序列中的每一个位置的输出时，模型只能访问到该位置之前的序列信息，而不能“预见”未来的信息。这一点对于诸如机器翻译、文本生成等任务至关重要，因为在实际应用中，我们不可能在生成一个单词前就预知后续的所有单词。
在标准的多头自注意力机制中，模型通过计算查询（Query）、键（Key）和值（Value）三者之间的相互关系来生成输出。然而，在解码器内部，如果不对这些计算加以限制，那么在生成某个位置的输出时，模型可能会不恰当地利用到序列中尚未生成的部分，破坏了生成过程的顺序性和逻辑性。为了避免这种情况，掩码多头自注意力机制在计算注意力权重时，会应用一个特殊的掩码矩阵。
这个掩码矩阵是一个下三角矩阵，如图10.2.5.2，其中对角线以上的元素被设为非常低的数值（例如负无穷大），而对角线及其以下的元素则保留为正常值。这样的设计使得在注意力权重的计算过程中，对于任何给定的位置，模型只能关注到位置至的信息，而位置及之后的信息会被忽略。这是因为当使用点积注意力机制时，掩码矩阵中的负无穷大值会导致相应的注意力权重被归一化为零，从而实现对未生成部分的屏蔽。
通过这种方式，掩码多头自注意力不仅保留了多头自注意力机制并行处理和捕捉多种依赖关系的能力，还严格遵循了序列生成的因果性原则，确保了在任何时刻，模型的决策都只基于它已经生成的信息，而非未来的输入,这种机制对于训练模型以正确顺序生成高质量的序列输出是极其关键的。