理解 Transformer 中的多头注意力

自注意力机制简介

注意力机制，或者更具体地说，自注意力机制，最早由 Vaswani 等人在 2017 年的著名论文  “Attention is All You Need“ 中提出。

这一机制为Transformer 模型提供了动力，而 Transformer 又支撑了我们如今一些最强大的模型，如 LLM 和 VLM。

所以，是的，它非常重要！

但问题在于：很多人（我曾经也是其中之一！）认为注意力是 2017 年才被发明的。

事实并非如此。

最初的注意力机制于2015 年由Bahdanau 等人在机器翻译的背景下提出。他们的想法是让解码器在每一步解码时关注输入句子的不同部分，而不是依赖于一个固定的上下文向量。

这一小小的转变带来了巨大的差异——并为后来发展成为自注意力，以及最终的多头注意力埋下了种子。

但此时您可能会有一个问题：

“注意力到底是什么意思？”

我用两个简单的句子来引入这个概念：

• “The cat sat on the mat.” → 我们的模型可能会把注意力从 sat 聚焦到 cat 和 mat 上。
• “She drank her tea because she was thirsty.” → 从 drank 指向 tea 和 thirsty 的注意力，连接了原因与动作。

从这些例子中，我们可以开始看到注意力试图做什么。它的核心是聚焦：帮助模型决定与当前处理的词相关的其他词，通常基于语义意义或语法关系。

一个重要的点是，在计算注意力权重和分数时，会使用所有的标记（稍后详述），而不只是某些特定的标记，如上图所示。这也意味着某些词之间的注意力会更强，如箭头的粗细所示。

这正是 Transformer 擅长理解语言的原因；它们利用注意力动态地将每个词与其他所有词关联起来。

中间步骤

我们已经弄清楚了什么是注意力，但仍然不知道内部是如何处理的。

这就是查询（Queries）、键（Keys）和值（Values）（通常缩写为Q、K 和 V）发挥作用的地方。接下来我们将一直沿用示例句子——“the cat sat on the mat”。

首先，当一个句子输入 Transformer 后，会立即被分词器进行分词，其中一个著名的分词方法是字节对编码（Byte-Pair Encoding）。

可以将分词理解为把文本拆分成更小的单元。由于我们句子里的词都很短，在这个示例中，我们可以假设每个单元都等于一个单词。从视觉上看，句子经过分词器后的样子应如下所示：

请注意，连标点符号（句号）也被转换成了一个 Token。同样为简化起见，我们假设句号是一个独立的 token。

分词后，我们使用分词器词表将每个 token 转换为相应的 Token ID。可以将分词器词表理解为一个数据库/表，存放着模型可识别的所有 token 及其对应的 Token ID。

随后，每个唯一的 Token ID 会通过嵌入矩阵映射到相应的 token 向量（embedding）。需要注意的是，所有这些 token 向量的维度相同；在本例中我们假设为 512。

最后一步是加入位置嵌入，以便模型知道每个 token 的位置信息。常见的两种方法是：

1. 正弦位置编码（用于最初的 Transformer）
2. 可学习的位置嵌入（用于 BERT、GPT）

这些位置嵌入的底层机制超出了本教程的范围。目前我们只需要知道，我们通过以下公式得到了输入向量：

输入向量 = Token 向量 + 位置向量

查询、键和值

现在我们终于可以引入 Q、K、V 向量了。首先要注意的是，对于输入中的每个词，我们会创建三个向量：

• 查询（Q）——表示该词在寻找什么
• 键（K）——表示该词能为他人提供什么
• 值（V）——该词实际承载的信息

这些是可学习的变换，同时也是各自的矩阵。从数学上，这些变换可以写成：

由此，每个 token 都可以用三个新向量表示。下面这张图有助于可视化这个过程：

计算缩放点积

在 “Attention is all you need” 一文中，给出了一个数学公式，展示如何使用这些新向量计算注意力权重。让我们一步步走过这一过程，假设我们正在计算位于 “sat” 时的注意力权重：

首先，我们使用“sat” 的查询向量与每个词（包括它自己）的键向量进行比较。得到一组相似度分数（标量），表示每个词与“sat”的相关程度。这些更正式地称为注意力分数。

然后这些分数会通过 sqrt(d_k) 进行缩放，以防止分数过大，并有助于在反向传播训练时稳定梯度。接着将缩放后的分数传入 softmax，将其转换为注意力权重。

softmax 会给出概率，告诉模型每个词应被赋予多少权重，因此某个特定 token 的注意力权重之和将等于 1。

最后，我们使用这些注意力权重对值向量进行加权求和。

这样就得到了一个新的向量，表示在上下文中的 “sat”，由所有 token 的信息丰富而成，但主要来自 cat、mat 以及其他相关 token。

这个过程会在句子中的每个 token 上发生。每个 token 都会关注其他所有 token，甚至是它自己。这就是我们称之为自注意力的原因。

另外请注意，我们计算的是单个 token 的新的富集向量，但在实际中，我们会一次性对许多（如果不是全部的话）token 进行计算，因为矩阵乘法非常高效。

之所以可以这样做，是因为该架构完全可并行化。与循环神经网络（RNN）按序处理不同，这一机制让 Transformer 能一次性处理所有 token，使之快速而强大。

为便于理解，请看上图。可以看到输入矩阵 X 被转换为相应的 Query、Key 和 Value 矩阵，然后应用缩放点积注意力公式，从而获得输入中所有 token 的完整富集嵌入向量 A。

在 PyTorch 中实现缩放点积注意力

现在让我们用 PyTorch 编写刚刚讨论的内容。出乎意料地，这很简单：

import torchimport torch.nn.functional as Fdef scaled_dot_product_attention(Q, K, V):    d_k = Q.size(-1)  # the dimension of the key vectors    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))    attn_weights = F.softmax(scores, dim=-1)    output = torch.matmul(attn_weights, V)    return output, attn_weights技术支持

通过实现一个简单的函数，我们就能复现上述过程！以下是要点：

• Q、K 和 V 的张量形状为 [batch_size, seq_len, dim]
• 输出是一组新向量，其中每个 token 都在整个序列的上下文中被表示
• 注意力分数是原始的未归一化相似度
• 注意力权重是经 softmax 归一化后的结果，用于聚焦，且总和为 1

不同类型的注意力

不过，还有一个我尚未提及的重要细节。Transformer 中并不只有一种注意力——实际上有三种：

1. 编码器自注意力
2. 解码器自注意力
3. 编码器—解码器（交叉）注意力

它们各自承担略有不同的角色，我们快速过一遍。

编码器自注意力

这正是我们目前为止所探讨的：在编码器中，每个 token 会关注同一输入序列中的所有其他 token，例如 “the cat sat on the mat.” 这个例子。

这里要注意的是没有限制，因为整个输入是已知的，使模型能够为每个 token 构建深层的上下文表示，既理解其前文也理解其后文。

还请注意，这发生在编码器中。

解码器自注意力

这里会稍微复杂一些。由于我们现在位于解码器中，我们不希望每个 token 看到未来，因为那会在文本生成时作弊。

再举一个例子：“The dog barked very loudly.” 当模型生成 “barked”（为简化假设每个词都是一个 token）时，它不应该知道下一个 token 是 “very”，否则模型将无法学会生成新的 token。

为解决这个问题，我们应用因果掩码（causal mask）（有时也称为三角掩码），它确保每个 token 只能关注自身及其之前的 token，而不会关注之后的 token。

从上图可以看到，通过掩蔽去除了未来的注意力分数，使模型能够学习生成新的、准确且合理的 token。

编码器—解码器（交叉）注意力

一旦解码器开始生成，它仍然需要来自原始输入的信息（即它要响应的对象）。这就是交叉注意力发挥作用的地方。

此时，解码器会关注编码器的输出。因此，虽然解码器自注意力是有掩码的，交叉注意力则没有，解码器可以一次性查看所有编码器 token。

在我学习这个概念时，非常喜欢翻译的例子。假设您要把 “Bonjour tout le monde” 翻译成英文，解码器在决定每个英文单词时，应该能够查看所有法语单词，对吗？

这正是交叉注意力所做的。不过我们不会在此进一步展开，因为这超出了本文的范围。

为何需要多头注意力

现在我们已经理解了自注意力如何工作。但您可能会问：

既然已经有注意力了，为什么还需要多个头？

为此，我想用手电筒的类比来说明，这个类比在我最初学习时帮助我建立了直觉。

把一个注意力头想象成一支手电筒，试图照亮句子中重要的部分。但一句话往往有多种不同的关系需要理解：主谓、宾语—代词、形容词—名词、长程依赖等。

如果我们每次不仅能捕捉到一种关系，而是能捕捉到更多，不是更好吗？

因此，解决方案就是给模型多支手电筒，每支手电筒关注句子的不同部分。这就是多头注意力。

与其计算一组 Q、K、V 并运行一次注意力计算，不如将嵌入向量拆分为更小的部分，并并行运行多个独立的注意力头。

如上图所示，输入嵌入（为简化，这里以单个 token 为例）被平均拆分到n 个不同的头上。

在我们的例子中，假设维度是 512，头数是 8，那么每个注意力头将并行处理 (512/8=64) 个维度，并行地尝试捕捉不同的关系。

为了进一步建立直觉，了解这些头可能（我用“可能”，因为我们仍然不完全清楚每个头具体关注哪种关系）处理哪些关系，我们来看一个简单的例子：

考虑这个句子：

“The quick brown fox jumps over the lazy dog.”

• 一个头可能关注形容词—名词对：“quick → fox”、“lazy → dog”
• 另一个头可能关注主语—动词：“fox → jumps”
• 再一个头可能捕捉到长距离关系：“jumps → over → dog”

要点在于，所有这些头是同时运行的，它们的输出会被拼接并投影回到原始嵌入维度（在我们的例子中为 512）。

通过代码可视化多个头

为了彻底巩固对多头注意力的理解，我们使用 HuggingFace 库来观察多头注意力机制如何工作。

!pip install transformers bertviz torch技术支持

首先运行上述代码单元，下载所有相关的包。

from transformers import BertTokenizer, BertModelfrom bertviz import head_viewimport torchmodel_name = 'bert-base-uncased' # We will be using the Bert modelmodel = BertModel.from_pretrained(model_name, output_attentions=True)tokenizer = BertTokenizer.from_pretrained(model_name)model.eval() # As we are not training, we have set the model to evaluation mode so no updates are made# This is our input sentencesentence = "The quick brown fox jumps over the lazy dog."# Using the tokenizer to convert the sentence into tokensinputs = tokenizer(sentence, return_tensors='pt')input_ids = inputs['input_ids']# Forward pass with attentionswith torch.no_grad():    outputs = model(**inputs)    attentions = outputs.attentions  # Tuple of (num_layers, batch, num_heads, seq_len, seq_len)# Here we decode the tokens tokens = tokenizer.convert_ids_to_tokens(input_ids[0])head_view(attentions, tokens) # Visualization技术支持

我们得到的输出如下图所示。注意有 12 种不同的颜色，这意味着该模型有 12 个头。连线的粗细也与 token 之间注意力分数的大小相关。

实现多头注意力：最佳实践

本节我们将编写多头注意力机制的伪代码。先假设我们的输入 X 的形状为 [batch, seq_len, d_model]。

然后我们可以按如下伪代码进行：

Q = X @ W_Q # Calculating the Query valuesK = X @ W_K # Calculating the Key valuesV = X @ W_V# Calculating the Value valuesQ_i, K_i, V_i = split_into_heads(Q, K, V, num_heads) # Now we split the Q,K,V into multiple different heads (by evenly splitting the dimensions)# Here we compute scaled dot-product attention for each headscores = (Q_i @ K_i.T) / sqrt(head_dim)scores += maskA_i = softmax(scores)A_i = dropout(A_i)Z_i = A_i @ V_iZ = concat_heads(Z_i) # We have to concatenate all of the Heads# This is our final linear projection matrixoutput = Z @ W_Ooutput = dropout(output)技术支持

以上内容您现在应该都很熟悉了，可能除了最后的线性投影矩阵。

我们快速过一下。

当所有单独的注意力头完成各自的工作后，它们的输出会被拼接回一个向量。但这个合并的输出仍需被重塑回原始的模型维度，这正是最终线性投影矩阵的作用所在。它只是将拼接后的输出映射回与输入嵌入相同的大小，以便传递给 Transformer 的下一层。

我们也来讨论一些优化与最佳实践：

1. 并行化：由于各个头相互独立，我们可以并行计算它们。
2. 矩阵乘法：通过线性代数操作（如 matmul）进行批处理，能获得高度优化的性能。
3. FlashAttention（进阶）：这是一种较新的技术，利用底层 CUDA 技巧来降低注意力的内存开销并加速计算，尤其适用于长序列。

在最佳实践方面，实现注意力层时有 3 个主要超参数需要关注：

1. num_heads：并行运行的注意力头数量，需确保它能被 d_model 整除。
2. head_dim：每个头的维度。
3. dropout_rate：用于防止过拟合。

结语

本文涵盖的内容很多。我们从一个简单的想法出发：注意力 关乎 聚焦。随后我们探讨了 Transformer 如何借助查询、键和值来计算每个词与其他所有词的关系，从而利用这种“聚焦”。

我们看到，自注意力（以及简要提到的因果注意力与交叉注意力）为每个 token 赋予了丰富的、具备上下文意识的表示。接着我们更进一步，引入多头注意力：不再只依赖一支“手电筒”，而是给模型多支手电筒，同时捕捉句子中的不同模式。

这正是 Transformer 的真正力量所在。它们不仅能“读”，还理解整个序列中的结构、含义和关系。

本文仅覆盖了 Transformer 的一个方面，仍有许多内容可学，欢迎查看这门完整、详尽的课程：Transformer Models with PyTorch 课程。

如果您想直接上手应用层面，我也推荐：Developing Multi-Input Models For OCR 项目，或Working with HuggingFace 课程。若您更偏好文章，请查看这篇直观的使用 PyTorch 构建 Transformer 模型的完整指南，以及我关于视觉 Transformer 的另一篇指南。