Table of Contents

Why Transformer

在 Transformer 之前，序列数据的处理主要依赖 RNN（循环神经网络）和其改进版本 LSTM（长短时记忆网络）。这些模型通过循环结构逐步处理输入序列，每一步依赖于上一步的输出，因此擅长捕捉序列中的时间依赖关系。然而，这种递归结构存在一些的缺陷：

难以并行化，训练速度慢；
难以捕捉长距离依赖；
梯度消失/爆炸问题。

Attention Is All You Need 这篇 2017 年的论文提出了 Transformer 这一模型架构，其中运用的 attention 机制能有效缓解 RNN 的问题，但是标题提及的 attention 机制早在 Transformer 提出之前就存在了，并且已被运用在了自然语言处理与计算机视觉任务当中。

IMPORTANT

Attention 机制本质上是一种加权求和方法，根据输入序列中不同位置的信息重要性分配权重，从而动态调整模型的关注点。在普通的神经网络语言模型中，对于一个确定的单词，它的向量是固定的，但是引入 attention 机制后，对于同一个单词，在不同上下文语境下它的向量表达是不一样的。

在传统的 Seq2Seq 模型¹中，Attention 最早是作为一种辅助机制引入到 Encoder-Decoder 框架中，用以缓解固定维度的上下文向量无法有效表达长序列信息的问题。

其实，Attention is all you need 这篇论文的重要性不是提出了 attention 这一概念，甚至 Attention is Not All You Need，更重要的是提出了 Transformer 这一完全基于 attention 的结构。Transformer 的革命性在于彻底摆脱了循环神经网络和卷积神经网络，转而完全依赖注意力机制来建模序列中的依赖关系

Transformer 这一新的模型架构，不仅克服了 RNN 处理序列数据的缺陷，而且改变了自然语言处理（NLP）的范式，更奠定了如今大规模语言模型（LLM）如 GPT 的基础。Transformer 之所以能够成为大语言模型的基础架构，主要有以下几个原因：

可扩展性：Transformer 的并行计算特性使其能够高效地处理大规模数据和参数。这一点对于训练如 GPT 这样拥有数千亿参数的模型至关重要。
长距离依赖：Transformer 的自注意力机制能够有效捕捉序列中的长距离依赖关系，这对于理解和生成连贯的长文本至关重要。
迁移学习能力：以 Transformer 为基础的预训练模型（如 BERT 和 GPT）展现出了强大的迁移学习能力，可以在各种下游任务上通过微调快速适应。
多模态潜力：虽然最初为NLP设计，Transformer 已被证明在计算机视觉等其他领域也有出色表现，为多模态AI的发展奠定了基础。

How to Build A Transformer

研究生时学习深度学习的各种算法和模型结构，纯看公式和图解我都还是觉得很难懂，直到我自己动手用 python 学着实现梯度下降，还有用 PyTorch 实现 CNN，慢慢的才豁然开朗起来，尽管实现过后间隔一段时间一些细节还是会忘，但是算法与模型大体的框架与原理会深刻地刻在脑海里，动手实现一些算法与模型，对学习者真的很重要。

配合 Attention Is All You Need 这篇论文与 The Annotated Transformer 这篇文章，我们将使用 PyTorch，像搭积木一样一步步实现 Transformer 这个模型，论文解析+代码实践，希望能让你我更加深刻的理解 Transformer 的架构以及原理。

Model Architecture

论文：Most competitive neural sequence transduction models have an encoder-decoder structure......

当时大部分比较好的神经网络序列模型都采用了编码器（Encoder）和解码器（Decoder）的架构。Encoder 将输入序列 $x_1, x_2, ..., x_n$ 编码成更抽象且更具语义的连续型向量表示 $z = (z_1, z_2, ..., z_n)$ ，Decoder 则根据编码后的 $z$ 序列，结合目标序列的已生成部分，将其解码逐步生成最终的输出序列 $y_1, y_2, ..., y_m$ 。注意输入序列 $x$ 的长度 $n$ 与输出序列 $y$ 的长度 $m$ 是不一定相等，好比一句话英文翻译成中文时，词的字数不一定是相等的，一开始 Transformer 的设计也是为了翻译任务，所以 Transformer 的模型结构也是采用 Encoder-Decoder 的结构。

论文：At each step the model is auto-regressive, consuming the previously generated symbols as additional input when generating the next.......

在 Encoder-Decoder 的结构中，过去时刻的输出会最作为当前时刻的输入。也就是预测 $y_{t}$ 时，除了输入对应 $t$ 时刻的 $x_t$ ，也会同时将 $y_1$ 到 $y_{t-1}$ 作为 $t$ 时刻的输入，这就是常说的自回归（auto-regressive）。

在 Transformer 之前，基于 RNN 的 Encoder-Decoder 是常见的构造方式。Transformer 继承了传统的 Encoder-Decoder 架构，但不同的是通过使用自注意力（Self-Attention）和逐点前馈网络（Point-wise Feed-Forward Network）²来构造编码器和解码器。Transformer 将多个相同的 Encoder 和 Decoder 分别堆叠 $N$ 次，如下图，左边黄色块为 Encoder，最左边有一个"Nx"就是表示有 $N$ 个相同的 Encoder 块堆叠到一起，右边紫色的 Decoder 也是如此，然后将堆叠后的 Encoder 块与堆叠后的Decoder 块进行连接，这就是 Transformer 的基本结构：

在实现 Transformer 前，以防代码看的晕头转向，可以先看下面的代码组织结构图，这个图展示了以下关键点：

EncoderDecoder 是主要的类，它包含了 Generator、Encoder 和 Decoder。
Encoder 和 Decoder 都由多个层组成，分别是 EncoderLayer 和 DecoderLayer。
EncoderLayer 和 DecoderLayer 都使用了 SublayerConnection 来实现残差连接和层归一化。
LayerNorm 被用在 EncoderLayer 和 DecoderLayer 中进行层归一化。
Generator 用于最终的输出生成。

Encoder-Decoder 对应的代码实现可以参考如下：

class EncoderDecoder(nn.Module):
	"""
	   标准的编码器-解码器架构。该架构是许多模型的基础。
	   - encoder：编码器模块，用于处理输入（源）序列。
	   - decoder：解码器模块，用于生成输出（目标）序列。
	   - src_embed 和 tgt_embed：分别用于对源序列和目标序列进行嵌入（embedding）的模块，将输入序列转化为向量表示。
	   - generator：通常是一个全连接层，用于将解码器的输出转化为最终的预测结果，例如词汇表中的概率分布。
	"""

	def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
	   super(EncoderDecoder, self).__init__()
	   self.encoder = encoder
	   self.decoder = decoder
	   self.src_embed = src_embed
	   self.tgt_embed = tgt_embed
	   self.generator = generator

	def forward(self, src, tgt, src_mask, tgt_mask):
	   # 接收并处理加了掩码的 src 和目标序列。
	   return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

	def encode(self, src, src_mask):
	   # 将 src 序列进行嵌入编码
	   return self.encoder(self.src_embed(src), src_mask)

	def decode(self, memory, src_mask, tgt, tgt_mask):
	   # 将 memory 和 tgt 序列进行解码
	   return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

class Generator(nn.Module):
	"定义标准线性softmax生成步骤。"

	def __init__(self, d_model, vocab):
	   super(Generator, self).__init__()
	   # 这个线性层将输入从 d_model 维度投影到 vocab 维度
	   self.proj = nn.Linear(d_model, vocab)

	def forward(self, x):
	   return log_softmax(self.proj(x), dim=-1)

上述这两段代码定义了 Transformer 模型的核心架构组件：EncoderDecoder 和 Generator。

EncoderDecoder 类的 __init__ 定义好了模型的骨架，forward 定义了输入数据在模型中的流动顺序与处理方式，之后逐一实现其中的模块即可:
- 这是 Transformer 的主要架构，包含编码器、解码器、嵌入层和生成器。
- forward 方法定义了数据在模型中的流动过程：先编码，后解码。
- encode 方法将输入序列编码成中间表示。
- decode 方法使用编码后的信息和目标序列生成输出。
Generator 类:
- 这是 Transformer 的最后一层，用于将解码器的输出转换为词汇表上的概率分布。
- 使用一个线性层将输入投影到词汇表大小的空间。
- 应用 log_softmax 函数，得到每个词的对数概率。

这些组件共同工作，实现了 Transformer 的序列到序列转换功能：

源输入序列经过嵌入层和编码器处理。
目标序列经过嵌入层和解码器处理，同时利用编码器的输出。
解码器的输出通过生成器转换为最终的词概率分布。

定义好 EncoderDecoder 和 Generator 这两个 Transforms 的主要框架作为蓝图后，下面我们将解释并实现 EncoderDecoder 类中初始化的 encoder 与 decoder。

Encoder

Transformer 编码器由 6 层堆叠而成，每一层包括 多头自注意力（multi-head self-attention） 和 前馈网络（position-wise fully connected feed-forward network） 两个子层。
每个子层的输出都与输入进行相加（残差连接），并通过 层归一化 处理，以提高训练稳定性。
每个子层的输出维度保持一致，方便堆叠使用。

按照原文，接下来实现 Encoder 、LayerNorm、SublayerConnection 和 EncoderLayer 这 4 个模块，Encoder 由 N 个 EncoderLayer 组成，SublayerConnection 则是实现了模型中的残差连接（residual connections）³以及 dropout⁴，每个 Encoder 的输出数据都会由 LayerNorm 进行归一化（记得结合上文图中的左半部分一起看代码）。

先定义一个工具函数 clone，用于复制 $N$ 个相同的 layer，方便我们堆叠：

def clones(module, N):
	"生成 N 个相同的 layer。"
	return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

接着实现 EncoderDecoder 类中初始化传入的 encoder，如前文所述，Transformer 的编码器 Encoder 是由 $N$ 个相同的子 encoder 层（对应代码中的 EncoderLayer）堆叠而成的：

class Encoder(nn.Module):
	"编码器由 N 个相同的层堆叠组成。"

	def __init__(self, layer, N):
	    super(Encoder, self).__init__()
	    # 将 encoder layer 复制 N 次进行堆叠
	    self.layers = clones(layer, N)
	    # 层归一化
	    self.norm = LayerNorm(layer.size)

	def forward(self, x, mask):
	    """
	    依次通过每个层传递输入 (和掩码)。
	    通过每个层运行 x，然后进行归一化。
	    """
	    for layer in self.layers:
	        x = layer(x, mask)
	    return self.norm(x)

在 Transformer 中，SublayerConnection 被用在编码器和解码器的每个子层中。它确保了信息可以有效地在网络中流动，同时保持了网络的稳定性和正则化。这个结构是 Transformer 能够训练深层网络并取得卓越性能的关键因素之一。

对应论文：That is, the output of each sub-layer is $LayerNorm (x + Sublayer (x))$ , where $Sublayer(x)$ is the function implemented by the sub-layer itself. To facilitate these residual connections, all sub-layers in the model, as well as the embedding layers, produce outputs of dimension $dmodel = 512$ .

class SublayerConnection(nn.Module):
	"""
	A residual connection followed by a layer norm.
	Note for code simplicity the norm is first as opposed to last.
	"""

	def __init__(self, size, dropout):
	    super(SublayerConnection, self).__init__()
	    self.norm = LayerNorm(size)
	    self.dropout = nn.Dropout(dropout)

	def forward(self, x, sublayer):
	    "Apply residual connection to any sublayer with the same size."
	    return x + self.dropout(sublayer(self.norm(x)))

定义好 Encoder 的结构与残差连接后，我们需要实现 Encoder 中传入的 layer 模块，也就是对应上图 Transformer 图解中左边橙色框圈出来的 Encoding layer。

对应论文：Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position-wise fully connected feed-forward network.

class EncoderLayer(nn.Module):
	"编码器由自注意力机制和前馈网络 (定义如下) 组成"

	def __init__(self, size, self_attn, feed_forward, dropout):
	    super(EncoderLayer, self).__init__()
	    self.self_attn = self_attn # self_attn将在下面的代码实现
	    self.feed_forward = feed_forward
	    self.sublayer = clones(SublayerConnection(size, dropout), 2)
	    self.size = size

	def forward(self, x, mask):
	    "Follow Figure 1 (left) for connections."
	    x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
	    return self.sublayer[1](x, self.feed_forward)

最后实现很多地方都用到的 LayerNorm 层，对输入数据进行归一化（normalization），对应论文：We employ a residual connection around each of the two sub-layers, followed by layer normalization.

$y = \gamma \cdot \frac{x - \mu}{\sigma + \epsilon} + \beta$

$y$ 是归一化后的输出
$x$ 是输入
$\mu$ 是输入的平均值（对应代码中的 mean）
$\sigma$ 是输入的标准差（对应代码中的 std）
$\epsilon$ 是为了数值稳定性添加的小常数（对应代码中的 self.eps）
$\gamma$ 是可学习的缩放参数（对应代码中的 self.a_2）
$\beta$ 是可学习的偏移参数（对应代码中的 self.b_2）

class LayerNorm(nn.Module):
    "构造一个层归一化（LayerNorm）。"

    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        # 初始化3个可训练参数，用于缩放和平移
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

Decoder

解码器结构在编码器的基础上增加了一个对编码器输出进行注意力的子层。
使用了掩码机制，确保每个位置的预测只能依赖前面已生成的输出，从而实现自回归的生成过程。

Decoder 由 $N$ 个相同的层（DecoderLayer）堆叠而成，每个层都包含自注意力机制、源注意力机制和前馈神经网络。Decoder 的输出通常会传递给 Generator 生成器（通常是线性层加 softmax），用于预测下一个词。

class Decoder(nn.Module):
	"生成具有掩码的通用N层解码器。"
	def __init__(self, layer, N):
	       super(Decoder, self).__init__()
	       self.layers = clones(layer, N)
	       self.norm = LayerNorm(layer.size)

	def forward(self, x, memory, src_mask, tgt_mask):
	       """
	       x：目标序列的嵌入表示。
	       memory：来自编码器的输出，包含源序列的信息。
	       src_mask：源序列的掩码，用于源注意力机制。
	       tgt_mask：目标序列的掩码，用于自注意力机制，确保在生成每个词时，只能看到之前生成的词，而不能看到未来的词。
	       """
	       for layer in self.layers:
	              x = layer(x, memory, src_mask, tgt_mask)
	       return self.norm(x)

同样定义好 Decoder 后，我们需要实现其中的 decoder layer，即下面的 DecoderLayer。其中涉及到的注意力机制 self_attn 和 src_attn 下一节会讲。

对应论文：In addition to the two sub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head attention over the output of the encoder stack. Similar to the encoder, we employ residual connections around each of the sub-layers, followed by layer normalization.

class DecoderLayer(nn.Module):
	"""
	解码器层，由三个主要部分组成：自注意力、源注意力和前馈神经网络。
	"""

	def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
	   """
	   初始化解码器层。

	   参数:
	   - size: 模型的维度
	   - self_attn: 自注意力机制
	   - src_attn: 源注意力机制（也称为交叉注意力）
	   - feed_forward: 前馈神经网络
	   - dropout: Dropout 比率
	   """
	   super(DecoderLayer, self).__init__()
	   self.size = size
	   self.self_attn = self_attn  # 自注意力模块
	   self.src_attn = src_attn    # 源注意力模块
	   self.feed_forward = feed_forward  # 前馈神经网络
	   # 创建三个子层连接，用于实现残差连接和层归一化
	   self.sublayer = clones(SublayerConnection(size, dropout), 3)

	def forward(self, x, memory, src_mask, tgt_mask):
		 """
		 前向传播函数。

		 参数:
		 - x: 解码器的输入
		 - memory: 编码器的输出
		 - src_mask: 源序列的掩码
		 - tgt_mask: 目标序列的掩码

		 返回:
		 - 经过完整解码器层处理后的张量
		 """
		 m = memory
		 # 1. 自注意力子层
		 # 这里的 x 为解码器的输入，x 被用作 query、key 和 value。
		 x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
		 # 2. 源注意力子层（交叉注意力）
		 # 这里 x 是来自解码器的查询（query），而 m（即 memory）是来自编码器的键（key）和值（value）
		 x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
		 # 3. 前馈神经网络子层
		 return self.sublayer[2](x, self.feed_forward)

定义一个工具函数 subsequent_mask，这个函数创建一个掩码，用于确保在解码过程中，每个位置只能注意到它之前的位置，而不能看到未来的位置，即预测 t 时刻的数据时，只能看到 t 时刻之前的数据，因此需要屏蔽掉 t 时刻之后的序列（推荐观看🔗 3Blue1Brown视频11:08-Masking ）。这是实现自回归（autoregressive）行为的关键，也就是 next token prediction 的关键。

对应论文：We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position ii can depend only on the known outputs at positions less than $i$ .

def subsequent_mask(size):
	"创建一个掩码来屏蔽后续位置。"
	# 创建一个形状为 (1, size, size) 的三维张量
	attn_shape = (1, size, size)

	# 使用 torch.triu 创建一个上三角矩阵，对角线上移一位
	# diagonal=1 表示主对角线上方的第一条对角线
	subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)

	# 将上三角矩阵取反，得到最终的掩码
	return subsequent_mask == 0

masking 掩码效果如下图：

Attention

Scaled Dot-Product Attention

Scaled Dot-Product Attention 是 Transformer 模型采用的核心注意力机制，用于捕捉输入序列中不同位置的依赖关系。其主要思想是根据输入中各个元素之间的相关性来动态调整信息的权重。具体来说，它计算序列中每个元素之间的相似性（相关性）并使用该相似性来加权其他元素。

400

IMPORTANT

注意力函数是 query 和 key-value 键值对到输出的一个映射⁵，注意力函数的输出就是 value 的加权和⁶，而这里每个 value 的权重就是根据每个 value 的 key 与输入的一条 query 的相似度得来的⁷。

那么 query，key 和 value 到底是什么？

Query (查询)：
- 表示当前我们关注的内容或者想要查找的信息在序列处理任务中，通常是指我们当前正在处理的元素，比如单词。这里可以理解为一个"问题"或"搜索词"。
Key (键)：
- 用来和 Query 进行匹配的对象，可以理解为可能包含答案的"索引"，用于计算与 Query 的相关性得分。
Value (值)：
- 包含实际的内容信息。当 Key 与 Query 匹配度高时，对应的 Value 会被选中，可以理解为 Key 对应的"答案"或"内容"，被选中的 value 则会使用 Query 与 Key 的相似度（向量点积）作为权重进行加权和。

NOTE

在 Transformer 中，Q、K、V 通常是通过对输入进行线性变换得到的，也就是说一开始大家都是从输入的序列元素的向量初始化而来的，最初大家都一模一样，后面在模型训练中才变的不一样。 Q（查询）、K（键）和 V（值）并不是模型直接输入的原始数据，而是通过线性变换从输入生成的。具体来说，Q、K、V 是通过输入向量乘以三个不同的可学习参数矩阵得到的：

$Q = XW_Q$
$K = XW_K$
$V = XW_V$

其中：

$X$ 是输入序列的表示（通常是词嵌入或者前一层的输出），
$W_Q$ 、 $W_K$ 、 $W_V$ 是学习的权重矩阵，它们是模型的参数。

这些权重矩阵在训练过程中是模型需要学习的参数。通过这些权重矩阵，输入被映射成 $Q$ 、 $K$ 和 $V$ ，然后这些向量会被用来计算注意力分数（即 $QK^T$ ）以及最终的加权和（通过 $V$ ）。

Transformer 的优势之一就是可并行计算，其中支持可并行计算的一点就是矩阵乘法。用 $Q$ ， $K$ 和 $V$ 表示上述包含 $n$ 个元素对应 Query 矩阵， Key 矩阵和 Value 矩阵之后，可以将 attention 机制写成以下公式：

\begin{equation} \text{Attention}(Q, K, V) = \text{softmax}\left (\frac{QK^T}{\sqrt{d_k}}\right) V \end{equation}

$QK^T$ ：点积操作，用于计算 Query 和 Key 之间的相似度。Query 是当前元素，而 Key 代表其余元素。通过点积，衡量当前元素和其他元素之间的相似性。
$\sqrt{d_k}$ ：缩放因子。由于点积的值会随着维度增大而增大，为了避免相似度值过大，导致梯度消失问题，缩放因子 $\sqrt{d_k}$ 被引入进行归一化⁸。这是所谓的"Scaled"部分。

softmax：对相似性分数进行归一化，得到概率分布。softmax 将分数转换为权重，使得权重的总和为 1。

由于 $Q$ 与转置后的 $K$ 要进行点积运算，所以这两个矩阵的维度要求是一样的，比如说 $d_k$ ，但是 $d_k$ 不一定要与 $d_v$ 的维度一致。

Attention 部分的讲解推荐李沐老师的视频 🔗34:00，当然 3Blue1Brown 的视频 🔗13:10 也很好很丝滑。

NOTE

直观理解：在语言模型的上下文中，每个单词都对应一个 Query，它会通过与其他单词的 Key 进行匹配，找到那些与自己相关的单词，然后根据这些相关单词的 Value 来更新自己的表示。通过这样一个过程，模型能够灵活地根据上下文关系，重点关注某些单词或信息，而忽略其他不相关的信息。

最后 Attention 对应的实现代码如下：

def attention(query, key, value, mask=None, dropout=None):
	"Compute 'Scaled Dot Product Attention'"
	d_k = query.size(-1)
	scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
	if mask is not None:
		# 可以使用subsequent_mask生成的掩码矩阵
	    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

Multi-Head Attention

Multi-Head Attention 是 Transformer 模型中用于增强模型表达能力的关键机制，它在处理序列数据时能从不同的子空间中捕获多种不同的模式⁹。相比于单头注意力机制（上一节所讲的完整一个 Scaled Dot-Product Attention 计算过程），Multi-Head Attention 能够在多个子空间上同时应用注意力，从而让模型在不同的角度下理解输入数据¹⁰。

在 Multi-Head Attention 中，输入数据会通过线性变换被投影成多个不同的子空间，每个子空间计算一个独立的Scaled Dot-Product Attention¹¹。然后，所有这些独立的注意力结果会被拼接在一起，并通过一个线性变换进行整合，形成最终的输出（如下图）。

表示式可以这样写：

\begin{aligned} MultiHead(Q, K, V) &= Concat(head_1, ..., head_h)W^O \\ where\ head_i &= Attention(QW^Q_i, KW^K_i, VW^V_i) \end{aligned}

这里的 $W^O, W^Q_i, W^K_i, W^V_i$ 是线性投影的矩阵，权重都是在模型训练中得出。

经过每个头的独立注意力计算后，将所有头的结果拼接起来，形成一个整体矩阵。假设有 $h$ 个头，每个头输出一个大小为 $d_v$ 的矩阵，拼接后的矩阵大小为 $h \times d_v$ 。拼接后的矩阵再通过一个线性变换 $W^O$ 映射回原来的维度，得到最终的 Multi-Head Attention 输出。

NOTE

头数和维度的平衡：在实践中，每个注意力头的维度通常是将输入的维度 $d_{\text{model}}$ 均分给 $h$ 个头。例如，如果输入维度为 512，头数为 8，则每个头的维度为 $d_k = d_{\text{model}} / h = 64$ 。这样保证每个头的计算量不会过大。
线性变换的权重共享：每个头都有独立的线性变换矩阵 $W^Q, W^K, W^V$ ，但是在最后输出的线性变换 $W^O$ 是共享的，用来将多头的输出重新映射到输出空间。

最后ulti-Head Attention 对应的实现代码如下：

class MultiHeadedAttention(nn.Module):
	def __init__(self, h, d_model, dropout=0.1):
	    "Take in model size and number of heads."
	    super(MultiHeadedAttention, self).__init__()
	    assert d_model % h == 0
	    # 我们假设d_v总是等于d_k，但也可以不一样
	    self.d_k = d_model // h
	    self.h = h
	    # 4个用于可学习线性转换权重矩阵 W^Q, W^K, W^V, W^O
	    self.linears = clones(nn.Linear(d_model, d_model), 4)
	    self.attn = None
	    self.dropout = nn.Dropout(p=dropout)

	def forward(self, query, key, value, mask=None):
	    "Implements Figure 2"
	    if mask is not None:
	        # 将 mask 扩展到每个头（head）
	        mask = mask.unsqueeze(1)
	    nbatches = query.size(0)

	    # 1. 对query, key, value分别进行线性变换
	    # 2. 将变换后的结果重塑为 4 维张量：(nbatches, seq_len, self.h, self.d_k)
	    # 3. 交换第二和第三维，得到(nbatches, self.h, seq_len, self.d_k)
	    #
	    # 具体步骤：
	    # - zip(self.linears, (query, key, value))将前 3 个线性层和输入配对
	    # - lin(x)对输入进行线性变换
	    # - .view(nbatches, -1, self.h, self.d_k)重塑张量
	    # - .transpose(1, 2)交换维度
	    #
	    # 这样处理后，每个头都有自己的query, key, value表示
	    query, key, value = [
	        lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
	        for lin, x in zip(self.linears, (query, key, value))
	    ]

	    # 将注意力应用在批处理的所有投影向量上。
	    x, self.attn = attention(
	        query, key, value, mask=mask, dropout=self.dropout
	    )

	    # "Concat" using a view and apply a final linear.
	    x = (
	        x.transpose(1, 2)
	        .contiguous()
	        .view(nbatches, -1, self.h * self.d_k)
	    )
	    # 删除不再需要的中间变量，以节省内存
	    del query
	    del key
	    del value
	    # 最后，通过最后一个线性层进行最终的线性变换，得到最终的输出
	    return self.linears[-1](x)

Attention in Transformer

在 Transformer 模型中，自注意力机制被应用于三个不同的场景：

编码器 - 解码器自注意力（Encoder-Decoder Attention）：在解码器的每一层中，查询（Query）来自于前一层的解码器输出，而键（Key）和值（Value）来自于编码器的输出。这种机制允许解码器在生成输出序列的每一步骤中，都能关注到输入序列的所有部分。
编码器自注意力（Self-Attention in the Encoder）：在编码器的每一层中，查询、键和值都来自于同一层的前一步骤的输出。这意味着每个位置在编码器的输入序列中都能关注到序列中的其他位置，从而使得模型能够捕捉到序列内部的长距离依赖关系。
解码器自注意力（Self-Attention in the Decoder）：解码器的自注意力层与编码器的自注意力层类似，但是需要进行掩码处理以确保在生成输出序列时，每个位置只能关注到当前位置及其之前的位置，这样可以保持自注意力机制的自回归特性。这种掩码机制防止了信息在序列中向后流传，确保了预测的依赖性只能是单向的。

Position-wise Feed-Forward Networks

FFN(x) = max(0, xW_1 + b1)W_2 + b_2

class PositionwiseFeedForward(nn.Module):
	"Implements FFN equation."

	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()))

Embeddings and Softmax

class Embeddings(nn.Module):
	def __init__(self, d_model, vocab):
	    super(Embeddings, self).__init__()
	    self.lut = nn.Embedding(vocab, d_model)
	    self.d_model = d_model

	def forward(self, x):
	    return self.lut(x) * math.sqrt(self.d_model)

Positional Encoding

不同位置的编码向量在不同的嵌入维度上有不同的变化。位置编码用于为输入序列中的每个位置提供唯一的表示，从而帮助模型在没有显式顺序信息的情况下感知输入序列中的相对或绝对位置。

在 Transformer 中，位置编码通过以下公式生成：

偶数位置维度：使用正弦函数。 $PE (pos, 2 i) = \sin\left (\frac{pos}{10000^{\frac{2 i}{d_{\text{model}}}}}\right)$
奇数位置维度：使用余弦函数。 $PE (pos, 2 i+1) = \cos\left (\frac{pos}{10000^{\frac{2 i}{d_{\text{model}}}}}\right)$ 其中 pos 是位置， $i$ 是维度索引， $d_{\text{model}}$ 是嵌入的维度。

位置 0 的向量：图中用红色框标出了位置 0 的向量。我们可以看到这个向量的值在不同维度上有周期性变化，符合正弦和余弦的模式。

随着位置增加的变化：随着位置从 0 到 9 递增，每个位置的编码向量也在不同维度上呈现出周期性变化。通过这种方式，Transformer 能够为每个 token 分配一个独特的位置信息，使得即便没有循环结构，模型依然能够感知输入序列中的顺序。

class PositionalEncoding(nn.Module):
	"Implement the PE function."

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

	    # Compute the positional encodings once in log space.
	    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)

Full Model

把前面所有实现的每个模块的代码搭建起最终完整的 Transformer：

def make_model(
	    src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
	):
	    "Helper: Construct a model from hyperparameters."
	    c = copy.deepcopy
	    attn = MultiHeadedAttention(h, d_model)
	    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
	    position = PositionalEncoding(d_model, dropout)
	    model = EncoderDecoder(
	        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
	        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
	        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
	        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
	        Generator(d_model, tgt_vocab),
	    )

	    # This was important from their code.
	    # Initialize parameters with Glorot / fan_avg.
	    for p in model.parameters():
	        if p.dim() > 1:
	            nn.init.xavier_uniform_(p)
	    return model

Why Self-Attention

Layer Type	Complexity per Layer	Sequential Operations	Maximum Path Length
Self-Attention	$O(n² · d)$	$O(1)$	$O(1)$
Recurrent	$O(n · d²)$	$O(n)$	$O(n)$
Convolutional	$O(k · n · d²)$	$O(1)$	$O(log_k(n))$
Self-Attention (restricted)	$O(r · n · d)$	$O(1)$	$O(n/r)$

$n$ 是序列长度， $d$ 是模型的隐藏维度或特征维度， $r$ 是在受限自注意力（restricted self-attention）中，代表注意力的范围或窗口大小， $k$ 是在卷积层中，通常代表卷积核的大小。

首先，RNN 和 LSTM 是顺序处理数据，无法充分利用现代 GPU 的并行计算能力，而 Self-Attention 机制可以让每个位置的数据独立计算，极大提升了效率。其次，RNN 和 LSTM 在捕捉长距离依赖时表现不佳，而 Transformer 中每一层的自注意力机制可以直接关注输入序列中的任意位置，因此能更好地建模长距离依赖。

一个来自知乎回答的翻译例子： “I arrived at the bank after crossing the river” 这里面的 bank 指的是银行还是河岸呢，这就需要我们联系上下文，当我们看到 river 之后就应该知道这里 bank 很大概率指的是河岸。在 RNN 中我们就需要一步步的顺序处理从 bank 到 river 的所有词语，而当它们相距较远时 RNN 的效果常常较差，且由于其顺序性处理效率也较低。
Self-Attention 则利用了 Attention 机制，计算每个单词与其他所有单词之间的关联，在这句话里，当翻译 bank 一词时，river 一词就有较高的 Attention score。利用这些 Attention score 就可以得到一个加权的表示，然后再放到一个前馈神经网络中得到新的表示，这一表示很好的考虑到上下文的信息。如下图所示，encoder 读入输入数据，利用层层叠加的 Self-Attention 机制对每一个词得到新的考虑了上下文信息的表征。Decoder 也利用类似的 Self-Attention 机制，但它不仅仅看之前产生的输出的文字，而且还要 attend encoder 的输出。

Why LayerNorm

LayerNorm 是针对每个单一样本的特征向量进行归一化，而不是像 BatchNorm 那样在一个 batch 上进行归一化。LayerNorm 对每个输入样本的特征进行归一化，主要是为了稳定和加速训练过程。归一化后输入的数据具有零均值和单位方差，可以防止模型训练过程中由于不同特征的尺度差异导致反向传播时梯度更新不稳定的问题（梯度的尺度不变性：LayerNorm 使得梯度对输入的整体尺度变化不敏感）。

输入样本的均值和方差的导数对于模型的训练至关重要，它们通过重新中心化和重新缩放梯度来提高模型性能。重新中心化和重新缩放梯度可以减少内部协变量偏移（Internal Covariate Shift, ICS）¹²。有助于稳定每一层的输入分布，减少 ICS，从而使得模型能够使用更大的学习率，加速训练过程，并提高训练的稳定性。同时，重新中心化和重新缩放的梯度可以减少参数更新时的方差，使得参数更新更加稳定，有助于模型更快地收敛。

减少梯度消失或梯度爆炸问题：在深度网络中，梯度消失或梯度爆炸是常见的问题，这会导致模型训练困难。通过归一化，可以确保梯度在反向传播过程中保持在一个合理的范围内，从而避免这些问题。

促进更均匀的梯度流动：归一化操作有助于使得梯度在网络的不同层之间更均匀地分布，这有助于每一层的参数更新都能反映整个网络的梯度信息，从而提高训练效率。

这里更详细的讲解推荐李沐老师的视频🔗25:40。

参考与推荐阅读

欢迎关注微信公众号👏 code

欢迎微信扫码加入我的付费知识星球👏 code

Footnotes

Seq2Seq（Sequence to Sequence）模型是一种处理序列数据的神经网络架构，广泛应用于自然语言处理中的任务，如机器翻译、文本摘要和对话生成。该模型的目标是将输入序列（如一个句子）转换为目标序列（如翻译后的句子）。Seq2Seq 模型通常由两个部分组成：编码器（Encoder）和解码器（Decoder），这两部分通常由 RNN 或者 LSTM 构成。 ↩
简单讲就是 MLP，全连接的多层感知机。 ↩
残差连接允许信息直接从下层传递到上层，有助于解决深度网络中的梯度消失问题。 ↩
Dropout 通过在训练时随机丢弃一些神经元，来防止过拟合。 ↩
论文：An attention function can be described as mapping a query and a set of key-value pairs to an output. ↩
论文：The output is computed as a weighted sum of the values. ↩
论文：the weight assigned to each value is computed by a compatibility function of the query with the corresponding key，这里的 compatibility function 可以简单理解为相似度 ↩
论文：We suspect that for large values of $d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by $\sqrt{d_k}$ . ↩
论文：Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. ↩
一头雾水? 八头雾水???????? ↩
论文：On each of these projected versions of queries, keys and values we then perform the attention function in parallel, yielding $d_v-dimensional$ output values. ↩
在深度学习模型中，随着网络层数的增加，每一层的输入分布可能会发生变化，这种现象被称为内部协变量偏移 ↩