注意力机制详解
在一切开始之前,我们应当先强调,深度学习这个领域至今仍然没有太多的可解释性。了解神经网络的人都应当知道,我们在网络结构上所做的创新,大多是因为在某种启发性的思想下对网络的结构进行一些改变,剩下的全部交给梯度下降。这样改变出来的结构,它可能会训练得更好,也有可能更差。如果更好,我们往往会尝试使用之前对其进行改造时所受的启发,对其进行解释。但一个网络经过改进以后变得更好,其真实原因是否真的就是我们所解释的那样,实际上很多时候仍然不得而知。
当然,这并不是说神经网络中的每一个设计都没有原因。有些设计存在相当明确的数学或者工程动机,而另一些设计,我们更多只能给出一种启发性的理解。因此,在大多数教程和详解当中,对于注意力机制的解释,基本上都会直接给出注意力层的计算流程,而不是从原理上一步一步推出为什么这个架构是这样的。这种做法在工程上当然没有问题,但它会直接给出大量陌生的术语和概念,对学习和理解造成一定困难。
因此本文尝试从一种可解释的角度去讲解注意力架构。但是很多部分可能会显得略有牵强,甚至和实际工程的思考方式有所不同。在讲解过程当中,我们会尽可能从可解释的角度出发,同时也会指出:哪些地方存在比较明确的数学原因,哪些地方更多只是一种帮助理解的解释,而实际训练得到的神经网络是否真的按照这种方式工作,我们并不知道。
注意力机制在当今 LLM 以及许多其他神经网络应用当中都起着十分重要的作用。本文希望从注意力的原理讲起,一步一步讲通注意力原理与注意力架构。
零、生物注意力简介
在正式开始之前,我们首先需要简单讲一下人类所说的"注意力"。
我们每时每刻都会接收到大量的信息,但是显然不会以完全相同的程度去处理这些信息。例如,当你正在一间十分嘈杂的房间里面和某个人交谈时,耳朵实际上接收到了房间里面大量不同的声音,但是你的注意力可能主要集中在与你交谈的人的声音上。又比如,当你在桌面上寻找自己的手机时,你的视觉系统实际上看到了桌面上的大量物体,但是和"找到手机"这个目标更加相关的信息,会更容易进入你的关注。
因此,如果只从一种十分抽象的角度来看,所谓注意力,可以理解为:面对大量输入信息时,根据当前的需求,有选择地更加关注其中的一部分信息。
机器学习当中的注意力机制,正是受到了这种思想的一定启发。当然,这里必须再次强调,我们并不是说神经网络当中的注意力层真的模拟了人脑注意力的工作过程。它更多只是借用了这样一种启发:既然不同的信息对当前任务的重要程度可能不一样,那么我们是否可以让神经网络自己学习"应该更加参考哪些信息",然后按照这种程度去汇总信息?
接下来的整篇文章,其实都在尝试回答这个问题。
一、RAG 简介
我们希望把这种"根据当前需求,从很多信息当中寻找更加相关的信息"的特点转移到神经网络当中。为了便于理解,我们要先从 RAG 当中常见的向量检索这一很多人都听说过的东西开始讲起。
当然,RAG 和注意力机制只是有一部分思想上的相通性,它们俩在具体实现细节上仍然有诸多不同。因此这里从 RAG 讲起更多是为了便于理解和深入浅出的目的,并不是说注意力机制来自 RAG,或者注意力机制本质上就是 RAG。后续我们也会逐渐看到两者的区别。
如果你已经对 RAG 背后的向量检索原理有了非常深刻的理解,那么可以跳过此段。如果你只是听说或者使用过 RAG,但对其背后原理没有那么理解,建议阅读。
RAG,也就是 Retrieval-Augmented Generation,中文一般称为检索增强生成。它并不等同于向量数据库或者向量搜索,但是在很多 RAG 系统当中,向量检索是其中非常常见的一步。
传统关键词搜索存在一个很明显的问题,也就是它对于语义相关性并不一定能够很好地识别。例如,我说我"脖子很痛"和我"颈部不舒服"。这两者在语义上是强相关的,甚至表达的意思几乎一样,但是如果仅仅使用关键词匹配,就不一定能够很好地把它们关联到一起。
在这种背景下,我们可以尝试使用一个模型。这个模型输入一段文字,输出一个固定维度的向量,即 A⟶f(A)=a。
我们希望,当两段语料 A 和 B 在我们关心的语义上越相似,它们得到的向量 a 和 b 也越相近。在这种目的下训练出来的模型通常被称为 Embedding 模型或者嵌入模型。
为了方便计算,我们还可以对向量做归一化,即 a^=∣a∣a。这样两个向量的余弦相似度就是 sim(A,B)=cosθ=a^⋅b^。因此,在向量已经归一化以后,我们只需要计算点积,就可以得到余弦相似度。
需要注意,点积越接近 1,表示两个归一化向量方向越接近;越接近 −1,表示两个向量方向越相反。但是这里不能简单地把 −1 理解成"毫无关联",因为一个 Embedding 模型的向量空间究竟如何组织语义,本身就是训练得到的。
RAG 中常见的向量查询原理是,假设我们有 S 段语料 D1,D2,...,DS。首先将它们全部转化为向量并储存,即 Di⟶di。接下来假设我们有一个查询文本 Q,同样将其转换为向量 q,然后依次计算 si=q⋅di。最后按照 si 的大小寻找 Top K,就得到了在这种向量相似度定义下,与查询最相关的若干段语料。
注意,这里描述的只是最基本的原理。在实际实现上会存在很多差别。例如,在语料数量非常大的时候,我们通常不会真的把所有结果完整排序,而会使用专门的 Top K 或近似最近邻搜索算法。
为什么要先讲这一部分呢?
因为我们发现,这种向量检索过程提供了一个注意力机制当中非常需要的思想,也就是寻找关联。
但是从这里开始,我们需要逐渐摆脱"语料"这个限制。前面之所以一直使用语料,是因为 RAG 本身就是一个很容易理解的例子。实际上,注意力机制并不要求我们处理的东西一定是自然语言。
接下来,我们可以把一段语料更加一般地看作一个"信息单元"。这个信息单元可以是一段文字,也可以是图片当中的一部分、声音中的某一段,甚至可以只是神经网络上一层产生的某一个向量。
注意力机制真正关心的不是"这是不是语言",而是一系列任意同类信息的关联度如何寻找。
二、交叉注意力:获取相关性权重
交叉注意力的目标是,我们存在一系列内容和一个查询,我们需要去识别出每段内容和查询之间的相关性,并且以相关性作为一个权重来参考这些内容,即相关性越高,对该内容参考得越多,最后综合生成一个结果。在这方面,它和前面讲到的向量检索十分相似。
我们更形式化地描述这个问题。假设现在有 S 个信息单元 X1,X2,...,XS,以及一个查询 XQ。
首先,为了可计算,我们需要使用神经网络前面的某种表示模块,将这些内容转化为向量,即 Xi⟶xi,查询同样得到 XQ⟶xQ。
注意,这里和前面的 RAG 已经开始产生区别。RAG 中我们往往希望训练出一个专门用于语义检索的 Embedding 模型,而注意力层接收到的通常只是网络当前已有的向量表示。它们可能来自输入的某种嵌入,也可能直接来自网络上一层的输出,并不要求是一个专门训练的语义 Embedding 模型。
接下来,我们希望获得一种专门用于"匹配"的表示。因此对于被查询的内容,我们令 ki=xiWK;而对于查询,我们令 q=xQWQ。我们将 ki 称为键,也就是 Key;将 q 称为查询,也就是 Query。
为什么需要再乘一次 WK 和 WQ,而不是和RAG一样,直接拿原来的 xi 和 xQ 做点积?这里可以先给出一种工程上的理解:原来的向量需要承担网络当中的很多任务,而现在我们希望从中学习出一种专门适合当前注意力层进行"查询"和"被查询"的表示,因此给它们增加两个可学习的线性变换。
有了 Query 和 Key 以后,我们可以计算 si=q⋅ki,并把它作为一个匹配程度的分数。
但是在实际注意力中,我们通常还会除以 Key 向量维度 dk 的平方根,也就是:
si=dk q⋅ki
为什么要这么做?
假设向量不同维度上的数值尺度大致相近,那么随着维度 dk 增大,点积实际上是越来越多项的累加,即 q⋅ki=∑j=1dkqjkij,其典型数值尺度也会随维度增大。如果直接把越来越大的数送入 Softmax,Softmax 很容易变得非常尖锐,从而使训练变得不稳定。因此除以 dk ,可以使不同维度下点积分数的尺度更加稳定。
有了所有分数 s1,s2,...,sS,我们仍然不能直接把它们作为权重,因为它们可能是正数,也可能是负数,而且总和也不一定为 1。
因此,我们对所有分数进行一次 Softmax:
αi=∑j=1Sesjesi
于是 αi>0,并且 ∑i=1Sαi=1。这样我们就得到了一个归一化后的权重向量 α=(α1,α2,...,αS),第 i 位表示:在这一轮注意力计算当中,第 i 个信息单元应当被参考多少。
三、获取相关性权重以后:生成结果
现在我们已经获得了相关性权重。当然这也是一个向量,第 i 位表示查询和第 i 个 Key 之间经过 Softmax 以后得到的权重。
有了这个相关性,我们应当如何生成结果呢?
为了按照权重生成结果,我们必须要把每一个信息单元转换为一种能够参与加权求和的表示。我们很容易想到,又是向量。
因此,我们需要另一种方法,想办法把每个信息单元转为一种新的向量。这种向量在理论解释上应当具有一定的"可加可乘性":如果一个信息单元占 70 的权重,另一个占 30 的权重,那么 0.7v1+0.3v2 理论上应当能够被解释为某种"更多地包含第一个信息单元的信息,同时也包含第二个信息单元的信息"的混合表示。
我们把这种向量称为 Value,也就是值。
实际计算中,我们仍然从原来的向量表示 xi 出发,再乘上另外一个可学习矩阵,即 vi=xiWV,得到所有 Value v1,v2,...,vS。
然后按照前面得到的权重做加权求和:
o=i=1∑Sαivi
这就是这一轮注意力计算得到的输出。
把前面所有过程合在一起,如果我们暂时忽略一系列向量同时计算时的矩阵形式,那么交叉注意力所做的事情其实就是:
Query 与 Key 计算匹配程度 → Softmax 得到权重 → 按照权重对 Value 做加权求和。
注意,这里所谓 Value 在"语义上可加可乘",仍然只是我们为了理解这个架构而提出的一种理论解释。
数学上,向量当然可以做加法和数乘。但是数学上可以计算,并不代表这个运算在人类可以理解的语义上一定有意义。
我们希望 0.7v1+0.3v2 能够解释成"混合了两部分信息",这样前面的整套逻辑才十分自然。但是实际训练的时候,没有任何一个 loss 在显式要求 Value 必须满足我们心中所想的这种语义代数。
网络真正需要做到的事情只是:经过这一系列计算之后,最终的 loss 能够降低。
因此,训练结束以后,Query 和 Key 的点积究竟对应着一种什么样的人类可解释"相关性",Value 又是否真的满足我们想象中的语义可加可乘性,我们并不知道。这些解释更多是我们为了理解为什么会设计出这样一个架构,而建立起来的一套逻辑模型。
四、自注意力
仍然强调前面的观点,我们可以从一种相当自然的逻辑去理解注意力架构为什么可能这样设计,但是这并不代表训练出来的网络内部一定真的按照我们赋予它的语言进行工作。尤其到了自注意力,这种"理论解释"和实际神经网络内部表示之间的距离会更加明显。因此这一部分我们既尝试给出一个直观理解,也要时刻记住:我们真正确定的是它做了什么计算,而不是训练以后每一个向量在人类语义中到底代表什么。
交叉注意力所应对的是:存在一系列内容和另外一个查询内容,我们要根据查询,有选择地从这些内容当中获取信息。
那么接下来可以提出一个很自然的问题:如果现在不存在一个单独放在外面的查询,而是希望这一系列内容当中的每一个内容,都能够从其他内容当中获取和自己有关的信息,应该怎么办?这就是自注意力要解决的问题。
假设现在我们有 S 个信息单元,并且它们当前的向量表示为 x1,x2,...,xS。把它们写成一个矩阵:
X=x1 x2 ⋮ xS
接下来,我们让每一个信息单元同时生成自己的 Query、Key 和 Value,也就是 Q=XWQ, K=XWK, V=XWV。因此对于 S 个信息单元,我们会得到 S 个 Query、 S 个 Key 和 S 个 Value,一共 3S 个向量。
对于第 i 个信息单元,我们把它产生的 qi 作为查询,然后和所有信息单元产生的 Key 计算点积:
sij=dk qi⋅kj
这里 sij 表示:当我们以第 i 个信息单元作为查询者的时候,第 j 个信息单元得到的匹配分数。
对固定的 i,我们再沿着所有 j 做 Softmax,即 αij=∑m=1Sesimesij。于是,第 i 个位置新的表示就是 oi=∑j=1Sαijvj。如果一次性使用矩阵来表示,这整套过程实际上非常简单:
A=Softmax(dk QKT)
然后:
O=AV
其中 A 是一个 S×S 的矩阵。矩阵中的第 i 行表示:以第 i 个信息单元作为 Query 时,它对全部信息单元分配出来的注意力权重。
于是经过一次这样的计算以后,第 i 个位置新的向量 oi,实际上混合了所有信息单元的 Value,并且不同信息的混合程度由它自己的 Query 与各个 Key 的匹配程度决定。
因此,如果从前面的交叉注意力推广过来,自注意力实际上可以理解成:
原来是"一个查询去查询一系列内容",现在变成"一系列内容中的每一个,都把自己当成查询,再去查询这一整组内容"。
至于这样进行一次信息混合以后,为什么一定能够产生对任务有用的表示,我们并不能单纯依靠上述逻辑证明。我们能够确定的是,这种结构给了网络一种根据当前内容动态决定"应该从哪些位置获取多少信息"的能力,而这些信息究竟被训练成什么含义,最终仍然交给梯度下降决定。
五、注意力层在神经网络中的位置
讲到这里还需要补充一个容易被忽略的问题:注意力层本身通常并不是一个完整的神经网络,也不负责直接生成最终答案。
前面我们一直写 X⟶Q,K,V⟶O,但是这里的 X 本身,通常就是网络前面某些模块产生出来的向量表示;而注意力计算得到的 O,也通常只是新的中间表示,接下来还会继续交给后续网络进行计算。也就是说,一个很粗略的神经网络流程可以理解为:
输入 ⟶ 前序表示 ⟶ 注意力层 ⟶ 新的表示 ⟶ 后续网络 ⟶ 最终输出。
因此,注意力机制本身并不要求输入一定是一段语言,也不要求输出一定是一段语言。它真正规定的只是这样一种计算结构:
对一组向量产生 Query、Key 和 Value,利用 Query 和 Key 得到权重,再利用这个权重对 Value 进行混合。
至于这些向量原来代表文字、图像、声音还是其他特征,是网络其他部分需要决定的事情。
实际工程当中还存在多头注意力等大量变体,但是它们并不影响我们这一篇文章想解释的最基本原理,因此这里暂时不展开。
六、总结
最后的最后,还是要回到开头的核心观点。
神经网络仍然是一个可解释性比较低的领域。我们可以给注意力机制建立一套相当自然的解释,但是一定要注意区分:什么是这个结构确实在做的数学运算,什么是设计它时可能采用的工程和数学考虑,什么又只是我们为了方便理解而赋予这些向量的语义解释。
也因此,在真正工程上实践有效的东西,和我们理论思考当中的理想解释之间,会存在各种不同。
在本文当中就有不少,我们可以做一些梳理。
既然我们是利用点积计算两段内容的匹配程度,为什么不直接使用原来的向量生成 K 和 Q,反而在生成以后还要分别乘上不同的权重矩阵?
从一种比较自然的工程视角来看,原来的向量可能同时要承担很多任务,而 Query 和 Key 承担的是两种不同的角色:"我想寻找什么"和"我能够通过什么特征被别人找到"。给它们分别乘上 WQ 和 WK,相当于允许网络自己学习两种不同的表示方式,即 Q=XWQ, K=XWK。当然,这同时也增加了可学习参数和整个网络的表达能力。
使用和生成 K、Q 同源的向量,只不过乘上了另外一个不同的权重矩阵 WV,为什么这样得到的向量就可能在语义上可加可乘?
我们不知道。
从理论解释上,为了让 ∑iαivi 能够被理解成"按照权重混合不同的信息",我们自然希望 Value 具有某种语义上的线性混合性质;但是训练过程本身没有直接要求这一点。它只要求通过 V=XWV 以及后续所有网络结构,最终得到一个更低的 loss。因此最终学出来的 Value 究竟是不是我们想象中的这种"可加可乘语义空间",并不能因为前面的解释听起来合理就认为它一定如此。
K、Q、V 当中为什么只差一个线性变换?凭什么它们之间是线性的?
这里同样不需要认为存在什么深刻的理论要求,规定 Query、Key 和 Value 必须通过线性关系得到。线性变换计算简单,而且可以让网络为三个不同角色学习不同的表示,即 Q=XWQ, K=XWK, V=XWV。而至于更加复杂的非线性表达能力,可以由网络的其他非线性模块以及多层网络共同提供。因此这里采用线性变换,更多是一种结构简单、计算方便而且实践有效的工程设计。
为什么自注意力机制把整组信息当中和某一个位置相关的内容进行加权混合,它就会有效?
严格来说,我们不能从前面的直觉直接推出"它一定有效"。我们能够说的是,自注意力提供了一种十分灵活的信息交换方式:对于不同的输入,每个位置得到的权重 αij 都可以不同,因此网络可以根据当前输入动态决定从哪些地方获取更多信息。至于这种结构在某个任务上究竟能不能学习到有用的东西,最终仍然是一个需要由训练和实验回答的问题。
我们可以发现,很多我们在理论上难以完全解释的东西,如果从工程视角看,它实际上就会自然很多。但有些地方也存在明确的数学动机,例如除以 dk 是为了控制点积进入 Softmax 时的数值尺度。
而另一些地方,例如为什么恰好选择这样的 Q、K、V 结构、为什么某种信息混合方式最终会学到十分优秀的表示,就很难像物理定律一样从一个基本理论严格推导出来。
很多时候,可能只是我们受到某种思想启发,设计出一种具有足够表达能力、能够方便训练的网络结构,然后让梯度下降去寻找其中的参数。如果它工作得很好,我们当然可以尝试解释它为什么有效。但是必须时刻记住:
能够帮助我们理解一个结构的解释,并不自动等于模型内部真实采用了这种人类能够理解的语义机制。
而这一点从纯理论向理论与工程混合视角的转变,也是神经网络初学者必须要迈出的一步。
我们仍然应该追问"为什么",因为很多地方确实存在漂亮的数学原因;但是与此同时,我们也必须接受,神经网络中并不是每一个成功的工程结构,都能够像经典物理理论一样从几个基本原则严格推导出来。