Transformer详解

发表于 分类于
较为详细地介绍了Transformer模型架构，包括其内部的Multi Head Attention模块，Position-wise Feed Forward模块和Position Embedding模块。Transformer可以说是LLM的基石，需要重点掌握。

原文是Attention Is All You Need 文章链接

经典的Transformer架构图，梦开始的地方~

Transformer模型架构的特点：

  • 基于Encoder-Decoder架构
  • 完全基于注意力机制（输入与输入，输入与输出，输出与输出，对应三种注意力机制），不使用递归或者卷积操作
  • 支持并行运算，训练速度更快

Transformer模型架构相较于之前模型的优势：

  • 诸如LSTM这样的模型计算当前时刻的隐藏状态时，依赖于前一时刻的隐藏状态，这种序列特征妨碍了并行训练
  • 使用CNN这样的模型试图并行计算隐藏层状态时，难以捕获远距离位置之间的依赖关系
  • 对于长序列来说，随着长度增加，对记忆模块的要求会越来越高，而Transformer每个位置独立地关注其它位置

模型架构

整体还是Encoder-Decoder架构，其中Encoder模块将输入序列映射为中间表示，该中间表示作为Decoder模块输入的一部分。Decoder模块的另一部分输入是已经解码的序列。整体的生成过程是自回归（auto-regressive）的，也就是每一步生成一个token。

Encoder包含6层，每层有两个子层：

  • 多头注意力层
  • position-wise的全连接层

每个子层都有残差连接 + Layer Normalization模块，即每个子层的输出是： 。模型中所有的embedding维度都是512。

Decoder也包含6层，每层有三个子层。相较于Encoder部分，额外多的一个子层是对Encoder模块的输出进行多头注意力操作。

此外，Decoder部分的注意力子层和Encoder部分的略有不同，它将当前位置及后续位置mask掉，从而只关注已解码部分。

注意力模块

自注意力机制

self-attention，也叫intra-attention，内部注意力。该注意力机制将单个序列的不同位置联系起来，以计算该序列的表示（其实每个位置的token都会有相应的隐藏层表示，只不过原始论文中的任务是机器翻译，所以需要获取序列的表示）。

注意力函数的本质是，将一个query和一组key-value对映射为一个输出。对于自然语言处理任务来说，当前词就是query，其它位置的词就是key，至于value通常和key保持一致。在计算时，不管是query还是key，value，都是向量。
论文中使用的是Scaled Dot-Product Attention：

代码实现： 

1
2
3
4
5
6
7
8
9
10
11
12
import torch
import math

def attention(query, key, value, mask=None, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

一般来说，，和都是输入经过线性变换得到，而且权重矩阵是不同的。如果输入不经过线性变换，那么每个词对应的，，都是完全一样的，那么在经过softmax操作之后的加权平均后，该词本身的比重将会是最大的，使得其它词的比重很少，无法有效利用上下文信息来增强当前词的语义表示。此外，线性变换的权重矩阵不同可以保证在不同空间进行投影，增强了表达能力，提高了泛化能力。

使用Scaled Dot-Product Attention的原因：

最常使用的两个注意力函数是additive attention和dot-product (multiplicative) attention。

其中additive attention使用前馈神经网络计算相似度，该神经网络仅有一层隐藏层（涉及线性变换、求和与非线性激活函数如tanh、权重计算和加权求和等步骤）。

虽然两者的理论复杂度差不多，但是使用优化的矩阵乘法时，dot-product attention的计算速度会快很多，而且空间效率也更高。在key对应的向量维度比较小的时候，两者的表现差不多，但是随着向量维度增大，additive attention的表现会超过dot product attention。为了克服这个问题，论文中增加了scaling factor以提升表现。

对此，论文作者猜测，随着向量维度增大，点积的结果在数值上也会增加，这导致进行softmax计算时这些区域的梯度会变得非常小。

至于scaling factor为什么设定为，有说法是：假设每个维度的分布都是均值为0，方差为1的正态分布，那么缩放后依旧保持了均值为0，方差为1的正态分布。

多头注意力机制

多头注意力机制分别将queries，keys和values用不同的、可学习的线性映射层，线性映射多次，每个映射结果可以看作子queries，子keys和子values。对每一组queries，keys和values，并行进行自注意力操作，之后将每一组的结果拼接，最后再对拼接的结果做映射，得到最终的结果。

代码实现： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import torch.nn as nn

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class MultiHeadedAttention(nn.Module):
    def __init__(self, num_head, d_model, dropout=0.1):
        super(MultiHeadedAttention, self).__init__()
        assert d_model % num_head == 0
        self.d_k = d_model // num_head
        slef.num_head = num_head
        # clones函数用于拷贝nn.Module，这里没写
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        if mask is not None:
            # same mask applied to all num_head heads.
            mask = mask.unsqueeze(1)
        n_batches = query.size(0)

        # 1) Do all the linear projections in batch, d_model -> num_head * d_k
        query, key, value = [
            lin(x).view(n_batches, -1, self.num_head, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        ]

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout
        )

        # 3) "Concat" using a view and apply a final linear.
        x = (
            x.transpose(1, 2)
            .contiguous()
            .view(n_batches, -1, self.num_head * self.d_k)
        )

        del query
        del key
        del value
        return self.linears[-1](x)

多头注意力机制的优势

多头注意力机制使得每个头可以关注输入序列中不同位置的信息，并从不同的表示子空间中提取有用的特征。举个例子，使用多头注意力机制找出句子中所有重要的名词时，一个头可能专注于句子中的主语，另一个头可能关注宾语，还有一个头关注其它类型的名词短语。每个头都会独立地计算权重，然后这些权重会被组合起来，以产生一个更加全面的表示，这个表示会同时考虑多个子空间的信息。

三种使用方式：

  • Encoder-Decoder：queries来自之前的decoder层，keys和values是encoder的输出
  • Encoder：queries、keys和values都来源于同一个地方，即encoder中上一层的输出
  • Decoder：和Encoder部分类似，不过不同于Encoder部分能关注所有位置，Decoder部分只能关注当前位置之前的元素

逐位置的前馈神经网络

每一个位置分别有对应的全连接的前馈神经网络。

代码实现： 

1
2
3
4
5
6
7
8
9
10
11
import torch.nn as nn

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))

虽然不同位置的线性变换是一样的，但是每一层使用的参数是不同的。

Feed Forward层的作用
在Multi-Head Attention中，主要是进行矩阵乘法，即都是线性变换，而Feed-forward模块在每次线性变换都引入了非线性激活函数ReLU，它的作用是将数据映射到高维度的空间然后再映射到低维度的空间，提取到更深层次的特征。此外，Self-Attention输出的并不是上下文语义嵌入，而是带有权重的原始上下文本身，所以需要FFN来进行特征提取。

残差连接

每个子层都有残差连接 + Layer Normalization模块，即每个子层的输出是： 。

残差连接主要解决了两个问题：

  1. 保留知识问题
  2. 梯度消失问题

随着网络层数的增加，越深的网络层其保留的原始信息越少，通过在网络中跨越多个层级直接传递信息，可以使得后续的网络可以提取有效的高层次信息，避免信息的丢失。

此外，梯度信息在反向传播的过程中很容易出现消失的情况。这是由于梯度在多层网络中传递时会受到多次权重矩阵的连续乘法操作，导致梯度逐渐变小，使得网络参数无法得到有效更新。通过将前一层的特征直接添加到后续层级，避免了梯度消失问题，同时也有助于加速收敛和优化过程。

位置编码

由于自注意力机制对待每个位置都是一样的，为了利用序列信息，模型引入了Positional Embedding，从而注入相对位置信息和绝对位置信息。具体来说，直接将positional embeddings和input embeddings相加，维度也保持一致。

其中是位置，是维度。

此外，可以改写为:

代码实现： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import torch.nn as nn
import troch
import math

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout, max_len=50000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # compute the positional encodings
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, x):
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)

关于Positional Embeddings的选择

论文作者也使用了可学习的positional embeddings，效果和余弦版本差不多。不过作者最终还是选择了余弦版本的positional embeddings，对此，作者的解释是，这可能允许模型外推到更长的序列（相比于训练过程）。

训练过程

Attention Is All You Need这篇文章本身是为了解决机器翻译任务提出的，因此也是按照机器翻译来从头开始训练的，不存在预训练。

后续的BERT，GPT等预训练语言模型则是采用了Transformer架构，在大量的无标注语料上进行自监督学习，即通过某种策略为无标注语料构造标签，然后再以监督学习的方式进行训练。

后续发展的不同的预训练模型的主要区别在于预训练任务、模型结构和训练语料。

文章小结

本文介绍了Transformer的模型结构，它有着相比于LSTM和CNN的巨大优势，这也是Transformer能够广泛应用于NLP领域的重要原因。但Transformer也不是没有缺点，它有如下缺点：

  1. 计算复杂度：Transformer模型在训练时可能需要大量计算，尤其是对于大型数据集和长序列，这使得在实时应用程序或资源受限设备上使用Transformer变得具有挑战性
  2. 并行化困难：Transformer模型的顺序性质可能导致难以并行化训练过程，从而减慢训练时间
  3. 缺乏可解释性：Transformer模型难以解释，因为它们不像其它一些机器学习模型那样具有清晰的输入到输出的映射
  4. 对超参数的敏感性：Transformer模型对超参数的值比较敏感，这使得想要通过调整超参数获得最佳性能变得更具挑战性
  5. 有限的输入长度：Transformer模型通常会受限于它们可以处理的输入序列的长度，这对于需要更长上下文的任务来说是个问题。

后续有些工作就是为了解决或缓解这些问题，同时LLM时代也对原始的Transformer结构做了一些调整，此处不展开介绍。

【参考文献】

  1. The Annotated Transformer
  2. 手把手教你了解Transformer——Part 1
  3. 手把手教你了解Transformer —Part 2
  4. 线性Transformer应该不是你要等的那个模型