大语言模型架构中的注意力机制

在深入探讨大型语言模型中使用的复杂注意力机制之前,理解一个简化的自注意力(Simplified Self-Attention)版本非常有帮助。这个版本通常被用作教学工具,它抓住了自注意力的核心思想,但省略了标准自注意力中用于投影的可学习权重矩阵(Wq​,Wk​,Wv​)。

其核心思想是:基于输入序列元素自身的初始表示(例如,词嵌入向量),直接计算它们之间的相关性。模型关注的是元素间比较和加权求和这一机制本身,而不是通过学习参数来调整这些比较。在这种简化形式下,序列中的每个元素都会关注序列中的所有其他元素(包括自身),并根据某种相似性度量来确定关注的强度。

值得注意的是,一些研究确实探索了不依赖复杂QKV(Query, Key, Value)投影的注意力变体。例如,有研究表明,在某些特定领域(如量子物理波函数参数化),仅依赖位置信息的注意力机制也能取得有竞争力的结果,这引发了关于QKV投影在所有场景下是否必需的讨论 。此外,从某种角度看,卷积神经网络(CNN)也可以被视为一种简化的自注意力,其感受野是固定的,只关注局部信息,而自注意力则可以看作是具有可学习或全局感受野的CNN 。

简化的自注意力计算过程通常如下:

  1. 输入表示:从输入序列开始,每个元素(如单词)都有一个初始的向量表示,通常是词嵌入向量。假设我们有一个包含 T 个词元的序列,其嵌入向量为 x(1),x(2),…,x(T)。
  2. 计算相似度得分:计算序列中每对元素 (i,j) 之间的相似度得分。常用的方法是使用点积(dot product)或余弦相似度(cosine similarity)。例如,使用点积,得分 scoreij​=x(i)⋅x(j)。
  3. 归一化得分(获取权重):为了将得分转换为表示关注度分布的权重,通常使用Softmax函数对每个元素 i 的所有得分 scoreij​(j=1 到 T)进行归一化。这样可以得到注意力权重 αij​,其中 ∑j=1T​αij​=1。 αij​=∑k=1T​exp(scoreik​)exp(scoreij​)​
  4. 计算输出表示:序列中每个元素 i 的最终输出表示 z(i) 是通过将所有元素的输入表示 x(j) 按照注意力权重 αij​ 进行加权求和得到的。 z(i)=j=1∑T​αij​x(j) 这个输出向量 z(i) 就是元素 i 融合了序列中其他元素信息后的上下文表示。

“自”(Self)这个词在自注意力中具有特殊含义,它强调该机制是在单个序列内部的不同位置之间建立联系,以计算该序列自身的表示。这与早期主要用于连接编码器和解码器的注意力机制(例如,在基于RNN的机器翻译中,解码器关注编码器的输出)形成了对比。

即使在简化的形式下,自注意力的这一特性也得以体现:

  1. 输入只有一个序列的嵌入。
  2. 相关性计算发生在序列内部元素之间(x(i) vs x(j))。
  3. 输出是同一序列的新表示(z(1),…,z(T)),其中每个元素的表示都受到了序列内其他元素的影响。

这种序列内部的关注机制是自注意力的定义性特征。它使得模型能够直接从输入文本本身构建丰富的上下文理解,这对于超越简单的序列到序列转换、执行更复杂的语言理解(如BERT模型)或基于先前上下文生成文本(如GPT模型)的任务来说至关重要。

标准的自注意力机制(通常简称为自注意力)是在简化版本的基础上,引入了可学习的参数,使其更加灵活和强大。这是构成大型语言模型(如Transformer)的基础模块。

标准自注意力的核心是为输入序列中的每个元素(由其嵌入向量 x(i) 表示)生成三个不同的向量:查询向量(Query, q(i))、键向量(Key, k(i))和值向量(Value, v(i))。这些向量是通过将输入嵌入向量 x(i) 分别乘以三个独立的可学习权重矩阵 Wq​、Wk​ 和 Wv​ 得到的 :

q(i)=Wq​x(i)
k(i)=Wk​x(i)
v(i)=Wv​x(i)
这三个向量在概念上扮演着不同的角色,可以通过一些类比来帮助理解:

  • 数据库/搜索引擎类比
    • Query (Q):如同你在搜索引擎中输入的查询词,代表当前词元想要查找的信息或它关注的方面。
    • Key (K):如同数据库中每条记录的索引或网页的标题,代表序列中每个词元“广告”出的自身信息或特征,用于与查询进行匹配。
    • Value (V):如同数据库记录或网页的实际内容,代表序列中每个词元实际携带的信息。一旦查询与某个键匹配成功,对应的值就会被提取出来。
  • 图书馆类比
    • Query (Q):你的检索主题
    • Key (K):每本书的索引条目关键词
    • Value (V):书本的具体内容
  • 舞台剧演员类比
    • Query (Q):主角演员提出的问题(“我该怎么演?”)。
    • Key (K):其他关键角色所掌握的相关信息片段(提示)。
    • Value (V):其他角色的观点、动机或行动,这些信息共同塑造了主角的决策。

总的来说,Query 代表当前元素正在“寻找”什么;Key 代表每个元素能够“提供”什么以供查找;Value 代表每个元素在被关注时实际“贡献”什么信息。

与简化自注意力不同,标准自注意力引入了三个权重矩阵 Wq​,Wk​,Wv​。这些矩阵是模型的参数,通过训练数据学习得到。它们的作用至关重要:

  • 投影到不同子空间:这些矩阵将原始的输入嵌入向量投影到三个不同的表示空间(查询空间、键空间、值空间)。这使得模型能够学习到对于计算注意力最有效的表示方式。例如,模型可以学习到应该关注输入嵌入的哪些方面来形成一个好的“查询”,哪些方面适合作为“键”来匹配查询,以及哪些方面应该包含在“值”中以供后续聚合。
  • 增加模型容量和灵活性:通过引入这些可学习的参数,模型的表达能力大大增强。模型不再局限于使用原始嵌入进行比较和加权,而是可以学习到更复杂的、任务相关的交互方式。
  • 维度控制:这些矩阵还可以控制Q, K, V向量的维度。通常要求查询向量和键向量具有相同的维度 dk​,因为它们需要进行点积运算。而值向量的维度 dv​ 可以不同,它决定了最终输出上下文向量的维度。

值得一提的是,有研究探索了共享 Wq​,Wk​,Wv​ 矩阵的可能性,以减少参数量和计算复杂度,并在某些任务上取得了不错的效果,这表明在某些情况下,这三个投影可能学习到了相似的特征。

缩放点积注意力是实际中最常用的一种自注意力计算方式。其计算步骤如下:

  1. 计算相似度得分:对于当前正在处理的词元 i(由其查询向量 q(i) 代表),计算它与序列中所有词元 j(由其键向量 k(j) 代表)的点积。这个点积值 scoreij​=q(i)⋅k(j) (或者写作 q(i)Tk(j))衡量了查询 i 和键 j 之间的相似度或对齐程度 。
  2. 缩放(Scaling):将计算得到的点积得分除以键向量维度 dk​ 的平方根 dk​​。即 scaled_scoreij​=dk​​q(i)⋅k(j)​。这一步非常重要,因为当 dk​ 较大时,点积的结果可能会变得非常大,导致Softmax函数进入梯度饱和区(梯度接近于0),从而使得模型训练困难。通过缩放,可以保持点积结果的方差在一个合理的范围内,有助于稳定训练过程。
  3. 归一化(Softmax):对缩放后的得分应用Softmax函数。Softmax函数作用于给定查询 i 的所有 j(j=1 到 T)的得分上,将它们转换为一组非负且总和为1的注意力权重 αij​。 $$ \alpha_{ij} = \text{softmax}\left(\frac{score_{ij}}{\sqrt{d_k}}\right) = \frac{\exp(scaled_score_{ij})}{\sum_{p=1}^{T} \exp(scaled_score_{ip})} $$ αij​ 表示词元 i 对词元 j 的关注程度。
  4. 加权求和:计算词元 i 的最终输出上下文向量 z(i)。这是通过将序列中所有词元的值向量 v(j) 按照对应的注意力权重 αij​ 进行加权求和得到的。 z(i)=j=1∑T​αij​v(j) 这个 z(i) 向量融合了来自整个序列的信息,其中与词元 i 相关性更高的词元(具有更高的 αij​ 权重)贡献更大。

这个过程通常使用矩阵运算来实现,以提高效率。如果将整个序列的查询、键、值向量分别堆叠成矩阵 Q,K,V,则整个序列的输出 Z 可以一次性计算:

Attention(Q,K,V)=softmax(dk​​QKT​)V

标准自注意力机制通过计算出的注意力权重 αij​ 来显式地权衡输入序列中不同部分(词元 j)对于理解或表示当前部分(词元 i)的重要性。权重 αij​ 越大,表示词元 j 的信息(体现在其值向量 v(j) 中)在构建词元 i 的新表示 z(i) 时所占的比重越大。

由于这些权重是基于 Q 和 K 向量之间的交互(点积相似度)动态计算出来的,并且 Q,K 本身是通过可学习的 Wq​,Wk​ 矩阵从输入 x 导出的,因此这种重要性权衡是上下文相关的并且是可学习的。模型在训练过程中会学习如何调整 Wq​,Wk​,Wv​ 矩阵,以便为特定任务生成最有用的注意力模式和上下文表示。

自注意力机制,特别是作为Transformer架构的核心组件,构成了绝大多数现代大型语言模型(LLM)的基础。它的成功源于其并行计算能力和有效捕捉长距离依赖关系的能力,这使得训练规模空前庞大的模型成为可能。无论是像BERT这样用于理解任务的模型,还是像GPT这样用于生成任务的模型,其核心都是基于自注意力机制来处理和理解输入的文本序列。

虽然注意力权重 αij​ 直观地显示了模型在计算 z(i) 时对 v(j) 的关注程度,但将这些权重直接等同于模型做出最终预测的“解释”需要谨慎。

一方面,注意力权重确实反映了模型在特定层中信息聚合的方式。可视化注意力权重(例如使用BertViz或热力图 )可以提供关于模型内部工作机制的线索,例如模型是否学习到了语法结构或指代关系。

然而,另一方面,有研究指出,注意力权重与模型预测之间的关系可能并不直接。在某些任务(如文本分类)中,即使显著改变注意力权重分布(例如,使其均匀分布),模型的最终预测结果也可能变化不大,这表明注意力权重可能不是影响预测的唯一或决定性因素。模型的最终输出是经过多层非线性变换、残差连接和层归一化等复杂操作的结果,单一层的注意力权重很难完全解释最终决策。此外,不同的注意力头可能学习到冗余或难以解释的模式。

因此,虽然注意力权重是分析模型行为的有价值的工具,但不应将其视为模型决策的完整或绝对可靠的解释。对于需要高可信度解释的场景,可能需要结合其他可解释性方法(如梯度方法、显著性图)进行综合分析。

在许多自然语言处理任务中,特别是文本生成,模型需要按照时间顺序逐个生成词元。这种类型的模型被称为自回归模型(Autoregressive Models)。例如,在生成句子“我喜欢吃苹果”时,模型需要先生成“我”,然后基于“我”生成“喜欢”,再基于“我喜欢”生成“吃”,以此类推。

自回归模型的核心要求是,在预测序列中的某个位置(例如,第 t 个词元)时,模型只能依赖于该位置之前已经生成或已知的词元(即位置 1 到 t−1 的词元),而不能“看到”或使用该位置之后的信息(即位置 t 及之后的词元)。如果模型在预测第 t 个词元时能够访问到第 t 个或其后的词元信息,就相当于在预测时“偷看”了答案,这在现实的、逐步生成的场景中是不可能的,并且会导致模型在训练时学到错误的依赖关系。

标准自注意力机制允许序列中的每个位置关注包括其自身和之后位置在内的所有位置。这对于需要理解整个句子上下文的任务(如文本分类或机器翻译的编码器)是合适的,但对于自回归生成任务则会破坏其时序约束。

因果注意力(Causal Attention),也常被称为掩码自注意力(Masked Self-Attention),正是为了解决这个问题而设计的。它的核心思想是在自注意力计算过程中引入一个掩码(mask),阻止信息从未来的位置流向当前或过去的位置。

具体实现方式通常是在计算完缩放点积得分(dk​​QKT​)之后,但在应用Softmax函数之前,将一个掩码矩阵加到得分矩阵上。这个掩码矩阵通常是一个上三角矩阵(或根据具体实现形式调整),其对角线及以下的元素为0,而对角线以上的元素(对应于未来位置)为负无穷大(−∞)。

$$ masked_score_{ij} = \begin{cases} scaled_score_{ij} & \text{if } j \le i \ -\infty & \text{if } j > i \end{cases} $$

当Softmax函数应用于这些修改后的得分时,那些被设置为−∞的位置的指数将趋近于0。因此,这些未来位置对应的注意力权重 αij​(其中 j>i)将变为0。这就有效地阻止了任何位置 i 去关注其后的位置 j。

这种机制确保了模型在处理位置 i 时,其生成的上下文向量 z(i) 只包含了来自位置 1 到 i 的信息,从而严格遵守了自回归的约束,防止了信息泄露(information leakage)。

因果注意力与标准自注意力的主要计算差异就在于Softmax之前的掩码操作。Q, K, V向量的生成、点积得分的计算、缩放以及最终的加权求和步骤都保持不变。这个看似简单的修改,却对模型的功能产生了根本性的影响,使其适用于自回归任务。

因果(掩码)自注意力是Transformer解码器(Decoder)架构中不可或缺的一部分。在解码器的每一层中,第一个自注意力模块通常就是掩码自注意力模块。它负责处理到目前为止已生成的输出序列,并为序列中的每个位置生成一个上下文表示,这个表示只依赖于该位置及其之前的位置。

这与解码器中的另一个注意力模块——编码器-解码器注意力(Encoder-Decoder Attention 或 Cross-Attention)形成对比。后者允许解码器的每个位置关注编码器输出的所有位置,以便从输入序列中获取信息,这里不需要因果掩码。同时,它也与编码器中的自注意力(通常是双向的,无因果掩码)不同。

因果注意力强制模型只能利用单向上下文(从左到右,或称过去的信息),这与自回归生成的本质需求完全吻合。它使得像GPT系列这样的模型能够有效地进行文本生成。

相比之下,掩码语言模型(Masked Language Modeling, MLM)的目标通常是理解语言,而不是生成。MLM的代表模型如BERT,在训练时会随机遮盖(mask)输入序列中的一部分词元,然后利用双向上下文(即被遮盖词元左右两边的信息)来预测被遮盖的词元。BERT中的自注意力机制是双向的,没有因果掩码。

这种根本性的差异导致了它们各自的优势领域:

  • 因果注意力(AR模型,如GPT)
    • 天然适合自回归生成任务。
    • 推理时通常更高效,可以利用KV缓存(Key-Value Cache)机制,避免重复计算先前词元的键和值向量,从而加速生成过程。
  • 双向注意力(MLM模型,如BERT)
    • 更擅长需要深度上下文理解的任务(如文本分类、问答、命名实体识别),因为它能同时考虑一个词左右两边的信息。
    • 不直接适用于从左到右的生成任务。虽然可以通过特殊方法(如FIM - Fill-in-the-Middle)进行填充任务,但这与纯粹的自回归生成不同。

因此,选择因果注意力还是双向注意力(或无掩码注意力)是模型架构设计的核心决策之一,它直接决定了模型的主要能力和适用场景。这也推动了结合两者优势的研究方向,例如开发能够同时进行自回归生成和掩码填充的模型,或者探索在多模态场景下调整因果约束以允许跨模态信息流动的架构。

单头自注意力(无论是标准版本还是因果版本)虽然强大,但它迫使模型将所有关于词间关系的信息压缩到一组注意力权重中。多头注意力(Multi-Head Attention, MHA)机制通过并行运行多个注意力“头”(heads)来扩展了单一注意力的能力,允许模型同时关注来自不同表示子空间的信息。

MHA的核心思想是将计算拆分成多个部分并行执行。与其使用一组 Wq​,Wk​,Wv​ 矩阵计算一次注意力,MHA使用 h 组不同的权重矩阵(Wq,i​,Wk,i​,Wv,i​,其中 i 从1到 h),从而创建 h 个并行的注意力“头”。

对于每个头 i,输入的嵌入向量(或来自前一层的输出向量)X 会被独立的线性层(权重为 Wq,i​,Wk,i​,Wv,i​)投影,生成该头专属的查询、键和值矩阵 Qi​,Ki​,Vi​ 。

Qi​=XWq,i​
Ki​=XWk,i​
Vi​=XWv,i​
通常,为了保持计算成本与单头注意力相当(或在一个可控范围内),输入向量的维度 dmodel​ 会被分配给 h 个头。也就是说,每个头处理的查询和键的维度通常是 dk​=dmodel​/h,值的维度是 dv​=dmodel​/h。在实践中,这通常通过使用单个大的 Wq​,Wk​,Wv​ 矩阵,然后将其输出分割或重塑(reshape)成 h 个头所需的维度来实现,这样更高效。

每个注意力头 i 使用其自身的 Qi​,Ki​,Vi​ 矩阵,独立地执行缩放点积注意力计算(根据需要,可以是标准自注意力或因果自注意力)。

$$ \text{head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i $$

这里 dk​ 是每个头的键向量维度,即 dmodel​/h。

计算完所有 h 个头的输出 head1​,head2​,…,headh​ 后,它们的结果会被拼接(concatenate)在一起,形成一个维度为 h×dv​=dmodel​ 的大向量。

Concat(head1​,…,headh​)

最后,这个拼接后的向量会再通过一个最终的线性投影层(权重矩阵为 Wo​)进行转换,得到多头注意力层的最终输出 Z。这个输出通常也具有维度 dmodel​,以便能够输入到下一层。

MultiHead(Q,K,V)=Concat(head1​,…,headh​)Wo​

有些文献也从求和的角度来描述MHA。如果将输出投影矩阵 Wo​ 按行分解为 h 个矩阵 Wo,i​,那么MHA的输出可以看作是每个头的输出 headi​ 乘以其对应的 Wo,i​ 后的总和。

多头注意力的主要优势在于它赋予了模型同时从不同的表示子空间(representational subspaces)学习信息的能力。

由于每个头使用不同的、独立学习的投影矩阵(Wq,i​,Wk,i​,Wv,i​),每个头可以学习关注输入序列中关系的不同方面。例如:

  • 一个头可能学习关注句法依赖关系(如动词与其主语的关系)。
  • 另一个头可能学习关注语义相似性(如同义词或相关概念)。
  • 还有一个头可能学习关注相对位置信息。

通过并行地运行这些专门化的注意力头,并将它们的结果结合起来,模型可以获得对输入序列更丰富、更细致、更多角度的理解。这比单一注意力机制(它必须将所有类型的关系信息混合在单一的注意力权重集里)能捕捉到更复杂的模式。