try ai
科普
编辑
分享
反馈
  • 自注意力机制

自注意力机制

SciencePedia玻尔百科
核心要点
  • 自注意力机制通过计算一个词的“查询”(Query)向量与序列中所有其他词的“键”(Key)向量之间的相似度得分,使模型能够权衡序列中所有其他元素的重要性。
  • 多头注意力通过允许模型在并行的“头”中同时学习和关注不同类型的关系,从而增强了模型的能力。
  • 通过在所有词元对之间建立直接连接,自注意力克服了困扰循环架构(如 RNN)的长程依赖问题。
  • 自注意力的多功能性使其在语言之外的领域取得了突破,为基因组学、医学和物理学等领域的交互建模提供了一个通用框架。

引言

近年来,机器处理序列信息的方式发生了深刻的范式转变,从刻板的、逐步的分析转向了更整体、更具上下文感知能力的理解。这场革命的核心是自注意力机制,这是一个强大的概念,它使模型能够针对任何给定任务动态地权衡输入序列中不同部分的重要性。这一创新直接解决了先前架构(如循环神经网络 RNN)留下的一个关键知识空白,这些架构难以捕捉序列中相距遥远的元素之间的有意义关系。本文将对这一开创性机制进行全面探索。首先,在“原理与机制”一章中,我们将剖析自注意力精妙的机制,从其基本构建模块——查询(Query)、键(Key)和值(Value)——到 Transformer 模块的完整架构。随后,“应用与跨学科联系”一章将展示这一思想非凡的广度,揭示它如何成为一种通用语言,用于为基因组学、医学乃至理论物理学等不同领域的复杂交互建模。

原理与机制

想象一下,你置身于一个巨大的图书馆,正在寻找一个非常具体问题的答案。理论上,你可以把每一本书都从头到尾读一遍,但这会非常低效。相反,你有一个更聪明的策略。你的问题构成了一个​​查询​​(query)。你浏览书架上书籍的标题和章节标题——也就是​​键​​(keys)。当某个键与你的查询匹配时,你就会“注意”到它,并取下那本书来阅读其内容——即​​值​​(value)。你最终的理解是你最关注的那些书的值的综合。

这个简单的类比正是​​自注意力机制​​的核心。这个概念如此强大,以至于它彻底改变了机器处理信息的方式,从我们所说的语言到我们自身生物学中的复杂序列。它允许模型在生成序列表示时,权衡输入序列中不同部分的重要性。但与图书馆里的人不同,它可以同时为序列中的每一个词(或像素、或基因)做到这一点,让每个元素都能审视所有其他元素,并决定哪些元素与其自身意义最相关。

注意力的剖析:查询、键和值

让我们把图书馆的类比变得更精确。假设我们有一个句子,并且已经将每个词转换成一个数字向量,即一个捕获其初始意义的嵌入向量。对于某个特定的词,比如说“it”,我们想弄清楚“it”指的是什么。模型通过从“it”的初始嵌入向量生成三个不同的向量来做到这一点:

  1. 一个​​查询​​(Query)向量 (QQQ):该向量提出一个问题,封装了“it”正在寻找什么。例如,“我是一个代词;我正在寻找我的先行名词。”
  2. 一个​​键​​(Key)向量 (KKK):该向量如同一个标签,标示着词语自身的身份。像“robot”这样的词会有一个键,表示“我是一个名词,一个潜在的先行词。”
  3. 一个​​值​​(Value)向量 (VVV):该向量包含词语的实际内容或意义。对于“robot”来说,这将是其丰富的语义嵌入。

为了计算“it”这个词应该在多大程度上关注“robot”这个词,模型会计算一个​​相似度得分​​。这其实就是“it”的查询向量和“robot”的键向量的点积。高的点积意味着强匹配——问题找到了它的答案。这个过程是针对“it”和序列中每一个其他的词进行的。

这些原始得分随后被送入一个 ​​softmax​​ 函数,该函数做两件事:它使所有得分都变为正数,并强制它们的总和为一。结果是一个优美的概率分布,一组​​注意力权重​​,它精确地告诉模型如何分配其注意力。如果“robot”的键与“it”的查询非常匹配,它就会得到一个很高的注意力权重,比如 0.80.80.8,而其他词则得到较小的权重。

最后,模型通过对句子中所有​​值​​(Value)向量进行加权求和,来计算“it”的一个新的、具有上下文感知能力的表示。来自“robot”的值向量乘以 0.80.80.8,而其他的值向量则按其较小的权重进行缩放。结果是,“it”的新意义现在深深地融入了“robot”的意义。从得分到加权求和的整个过程被称为​​缩放点积注意力​​(scaled dot-product attention)。

一点稳定性:缩放因子

你可能已经注意到上文中的“缩放”一词。这不仅仅是一个细节;它是一个关键的、体现数学精妙之处的设计。当我们计算两个向量的点积时,结果的大小取决于它们的维度 dkd_kdk​。随着维度的增长,点积往往会变得更大。如果将这些大值输入 softmax 函数,函数可能会“饱和”——产生极其尖锐的分布,其中一个权重几乎为 111,而所有其他权重几乎为 000。这使得模型很难学习,因为梯度变得小到可以忽略不计。

最初的 Transformer 论文中提出的解决方案惊人地简单:将点积除以维度的平方根,即 1dk\frac{1}{\sqrt{d_k}}dk​​1​。为什么是这个特定值?这源于对点积统计特性的考量。如果查询和键向量的分量都来自标准分布,那么它们的点积的方差将为 dkd_kdk​。通过除以 1dk\frac{1}{\sqrt{d_k}}dk​​1​,可以优雅地将方差归一化回 111,从而使 softmax 函数的输入保持在一个稳定、“健康”的范围内,而与嵌入维度无关。这是一个微小的改动,却使得训练深度强大的模型成为可能。

顺序问题:我在哪里?

我们所描述的机制有一个有趣的特性:它本质上是非时序的、与顺序无关的。它将输入视为一个无序的元素​​集合​​。如果你打乱一个句子中词语的顺序,注意力机制将产生完全相同的输出向量集合,只是顺序也被打乱了。这个属性被称为​​置换等变性​​(permutation equivariance)。对于某些任务,比如分析来自对撞机事件的一组粒子(其中顺序无关紧要),这是一个特性。

但对于语言,或病人血压的时间序列数据,顺序就是一切。“病人没有表现出恢复的迹象”(The patient showed no signs of recovery)与“迹象没有表现出恢复的病人”(The signs showed no patient of recovery)的意思大相径庭。为了解决这个问题,我们必须打破这种对称性。我们通过为每个输入嵌入添加一个唯一的“地址”来赋予模型位置感。这个地址是一个称为​​位置编码​​(Positional Encoding)的向量。通过添加一个仅取决于序列中位置 iii 的固定向量,我们确保了位于位置 1 的词“the”的总输入与位于位置 5 的“the”的输入是不同的。这个看似简单的补充为模型提供了学习依赖顺序的模式所需的信息,使其能够解决纯粹对称的注意力机制无法完成的任务,比如判断一个数字序列是否严格递增。

集体智慧:多头注意力

单个注意力机制可能会学习关注某一种关系,例如句法依赖。但语言是层次丰富且复杂的。一个词通过句法、语义、否定、共指等多种方式与其他词相关联。当你可以拥有多个视角时,为什么只满足于一个呢?

这就是​​多头注意力​​(multi-head attention)背后的洞见。我们不再只有一组查询、键和值的投影矩阵,而是创建几组——比如 8 组或 12 组。这些“头”中的每一个都并行运作,拥有自己学习到的参数。因此,每个头都可以学习专门化。一个头可能专注于追踪哪个动词支配哪个名词。另一个可能学会将一个医学发现与否定它的词(如“no”或“denies”)联系起来。第三个可能追踪蛋白质序列中对应其三维折叠模式的长程依赖关系。

在每个头都生成了其各自的、具有上下文感知能力的输出向量之后,我们只需将它们的结果连接起来,并通过一个最终的线性投影层将它们混合回一个单一、统一的表示。这使得模型能够同时关注来自不同位置的不同表示子空间的信息,从而对输入序列产生极其丰富和细致的理解。

完整机制:一个完整的 Transformer 模块

自注意力虽然强大,但它只是一个完整的 ​​Transformer 模块​​中的一个组件。一个标准的模块有两个主要的子层:

  1. 一个​​多头自注意力​​层。正如我们所见,该层作为通信中心,从整个序列中收集信息。
  2. 一个​​逐位置前馈网络 (FFN)​​。这是一个小型的、两层的神经网络,它独立地应用于注意力层在每个单一位置的输出。

这里有一个漂亮的分工。注意力层在​​时间或序列维度上​​混合信息,而 FFN 则在​​每个时间步内部​​对特征执行复杂的非线性变换。FFN 可以被认为是“思考”或“处理”部分,它接收注意力收集到的信息,并计算出一个更抽象的表示。

这两个子层通过另外两个关键组件粘合在一起:​​残差连接​​和​​层归一化​​。每个子层的输出都与其输入相加(一个残差连接),然后对结果进行归一化。这个看似简单的技巧对于训练由许多 Transformer 模块组成的非常深的网络至关重要。残差连接为梯度在网络中反向传播创建了一条“高速公路”,极大地缓解了困扰旧架构的梯度消失问题。

突破:攻克长程依赖

要理解这个架构的真正天才之处,我们必须将其与它的前身——​​循环神经网络 (RNN)​​——进行比较。RNN 一步一步地处理序列,维持一个从一个时间步传递到下一个时间步的隐藏状态,就像一个传话游戏。要让一个长段落开头和结尾的两个词联系起来,信息必须穿过所有中间的词。在长距离上,这个信号可能会衰减为零(​​梯度消失问题​​)或爆炸成无意义的值(​​梯度爆炸问题​​)。

自注意力完全回避了这个问题。它在序列中的每一对词元之间创建了直接的、并行的连接。信息在任意两点之间传播的路径长度始终恰好为一。这为梯度和信息提供了一个“虫洞”,使得模型可以轻而易举地学习相距甚远的元素之间的依赖关系。这种毫不费力地为​​长程依赖​​建模的能力是 Transformer 成功的主要原因。

然而,这种能力是有代价的。因为每个元素都必须与所有其他元素进行比较,所以自注意力的计算和内存成本随序列长度 NNN 呈二次方增长。其复杂度大约是 O(N2d)\mathcal{O}(N^2 d)O(N2d),其中 ddd 是模型的隐藏维度。这种二次方扩展是该架构的阿喀琉斯之踵。将一份临床记录的长度加倍,计算量不止是加倍,而是翻了两番。处理高分辨率的三维医学图像(可以看作是极长的像素序列)变得成本高昂,难以承受。这催生了实际的工程解决方案,如将序列分块并独立处理这些块,同时也推动了寻找更高效注意力机制的大量研究工作。

最后一点警示:注意力不等于可解释性

看着注意力权重并将其解释为模型行为的合理解释,是件极具诱惑力的事情。我们看到一个模型正确地翻译了一个句子,并注意到它将一个代词的高度注意力放在了它的先行词上,我们便会想:“啊哈!原来是这样才做对的。”

但我们必须谨慎。注意力图向我们展示了模型收集了什么信息,但它并不一定告诉我们这些信息是如何被使用的。从值的加权和到最终输出的路径还要经过更多的非线性层(FFN、残差连接)。一个词元可能获得了很高的注意力,但它的值向量可能在后续计算中被投射掉或被忽略。模型是一个复杂的动态系统,而注意力权重只是其中的一个中间部分。

更微妙的是,模型的目标是最小化预测误差,它会利用数据中的任何统计规律性来达到这个目的。这意味着它将学习相关性,而不一定是因果性。一个复杂的实验可能会表明,模型在两个神经元 iii 和 jjj 之间产生了注意力,不是因为 jjj 直接导致了 iii,而是因为它们都由第三个未观察到的神经元 ccc 驱动。注意力反映的是由混杂因素引起的相关性,而不是像 Granger 因果关系这类方法试图识别的直接因果联系。因此,虽然注意力图是窥探这些模型内部的一个迷人而有用的工具,但我们必须抵制将其视为简单、直接的“解释”的冲动。它们是线索,而非结论。

应用与跨学科联系

在了解了自注意力的原理和机制之后,你可能会留下这样的印象:我们一直在讨论的是一个由计算机科学家为处理人类语言而设计的、虽然巧妙但可能应用范围狭窄的工具。事实远非如此。我们真正一直在探索的是一种描述交互的、全新的、深刻的语言。自注意力机制以其优雅的简洁性,为模拟任何复杂系统中各部分如何相互关联提供了一种通用句法。它是一个框架,用于在系统的每个点上提问:“我应该关注哪些其他部分?关注多少?才能理解我自己的角色?”

事实证明,这个问题不仅适用于句子和语法。它也是一条DNA链、一个折叠的蛋白质、一位诊断病人的医生以及一位为宇宙建模的物理学家所提出的基本问题。让我们开启一段跨越现代科学和工程领域的旅程,看看这一个优美的思想是如何提供意想不到的答案并建立起令人惊讶的联系的。

生命的语言:从DNA到蛋白质

生物学的最核心是一种由四个字母组成的字母表书写的语言:A,C,G,TA, C, G, TA,C,G,T。基因组,我们的生命之书,是一个长度惊人的序列,在其中找到有意义的“短语”——如基因、启动子、增强子——是一项艰巨的任务。传统上,科学家们寻找固定的模式或基序(motifs)。但如果我们能教一台机器阅读DNA呢?这正是配备了自注意力的 Transformer 所能做到的。通过将DNA片段视为一个词元序列,可以训练模型区分功能性区域(如启动子)与周围的基因组文本。模型学习了基因组复杂、长程的“语法”,其中数百个碱基之外的信号可以影响一个基因的活性。

但故事还在深入。我们能做的不仅仅是得到最终答案;我们还能窃听模型的内部思考过程。通过检查注意力权重,我们可以问模型在做决定时“看”了DNA的哪些部分。在一个与生物学家思考方式惊人平行的发现中,研究人员发现不同的注意力头可以学习专门化,有效地成为特定基序的检测器,比如转录因子 (TFs) 的结合位点 [@problem-id:2373335]。更令人兴奋的是,通过观察哪些头关注序列的哪些其他部分,我们可以提出关于不同转录因子如何以组合之舞的方式协同调控基因的新假说。这就像有了一位不知疲倦的助手,他阅读了数十亿条DNA序列,现在能指出我们可能错过的那些微妙的交互模式。

当然,这引出了一个诱人但危险的类比:我们能说从位点 AAA 到位点 BBB 的高注意力权重意味着 AAA 导致了在 BBB 处的效果吗?答案通常是否定的。注意力权重反映的是相关性,而不必然是因果性。一个大的权重是一条线索,一个值得研究的提示,但它不是影响力的直接度量。要做出因果声明,必须小心翼翼,例如,通过在来自精心设计的干预实验的数据上训练模型,这是一个要求高得多的门槛。

从序列到生命的旅程从DNA延续到蛋白质。蛋白质是细胞的劳动力,是由氨基酸链折叠成复杂三维形状以执行其功能的。分子生物学的中心法则告诉我们,序列决定结构,结构又决定功能。研究人员通过在庞大的蛋白质序列数据库上训练 Transformer,构建了大规模的“蛋白质语言模型” (PLMs)。他们使用了我们在前一章看到的相同的掩码语言模型任务:隐藏一个氨基酸,并要求模型根据其上下文来预测它。

为什么这能行得通?这里有一个优美的统计学论证。蛋白质的三维结构和功能可以被看作是一种潜在的或隐藏的属性,它约束了进化所选择的序列。在1D序列中相距遥远但在3D折叠结构中接触的两个氨基酸是高度相互依赖的。为了从一个正确预测另一个,模型别无选择,只能学习连接它们的底层三维结构。通过这样做,训练过程含蓄地将关于结构和功能的信息打包到模型的嵌入向量中,使它们在下游任务(如预测药物如何与靶蛋白结合)中变得异常强大。一些注意力头甚至在没有任何显式监督的情况下,学会了生成与蛋白质接触图谱(其折叠形状的直接可视化)极为相似的注意力图。

这一思路在 AlphaFold 这一科学领域的里程碑式成就中达到了顶峰。其架构中的一个关键创新是一种称为“三角自注意力”的机制。想象一下,模型正试图完善它对两个氨基酸 iii 和 jjj 之间关系的信念。它通过与蛋白质中所有其他氨基酸 kkk 进行通信来做到这一点。模型实际上是在问,对于每一个 kkk,“鉴于我所知道的 iii 和 kkk 之间,以及 kkk 和 jjj 之间的关系,这告诉我关于 iii 和 jjj 之间的关系是什么?”这个过程是强制执行几何一致性的有力方式。它在计算上等同于三角不等式:如果你知道从 AAA 到 BBB 和从 BBB 到 CCC 的距离,你对从 AAA 到 CCC 的距离就有了强有力的约束。通过反复应用这种三角更新,模型“推理”出全局一致的三维结构。

见所未见:科学与医学中的注意力

自注意力的力量远远超出了生物学的线性序列。让我们考虑一张三维医学图像,比如病人肺部的CT扫描。像间质性肺病这样的疾病可能表现为弥漫性、广泛分布的模式,如果计算机只看小块区域,就很难识别。在这里,一种混合方法已被证明非常强大。卷积神经网络 (CNN) 作为前端,用于处理高分辨率图像,它非常擅长高效地提取纹理和边缘等局部特征。CNN 逐步对图像进行下采样,创建一个更小、更抽象的特征图。在这个阶段,自注意力接管了工作。将特征图视为一组词元,一个 Transformer 层可以应用其全局视野,连接来自肺部上下叶的微妙信号,以识别弥漫性疾病的典型特征。这种架构的结合,既利用了CNN进行局部感知的效率,又发挥了 Transformer 的全局推理能力,同时在计算上也是可行的 [@problem_-id:4534202]。

需要分析的“序列”根本不必是空间上的;它也可以是时间上的。考虑一个病人的电子健康记录 (EHR),这是一个由临床事件——诊断、化验、用药——组成的序列,这些事件在多年间以不规则的时间间隔发生。我们如何预测病人未来发生不良事件的风险?像循环神经网络 (RNN) 这样的旧模型有一个很强的内置偏见:过去事件的影响往往随时间呈指数衰减。这是一个僵化的假设,在医学中可能不成立;童年时期的疾病可能在几十年后再次变得相关。

这正是自注意力灵活性的闪光之处。通过调整位置编码,我们可以让模型意识到分隔事件的实际物理量:时间差 Δt\Delta tΔt。模型不再仅仅被给予事件的顺序,而是它们之间精确的时间间隔。然后,自注意力机制可以直接从数据中学习一个“时间影响核”。它可能会发现某个特定化验的相关性在六个月后达到峰值然后减弱,或者另一个事件的影响遵循一个复杂的非单调模式——这种灵活性是那些带有固定指数衰减偏见的模型根本无法企及的。

创建具有物理意识的位置编码的原则是通用的。让我们从诊所转向天空。一颗高光谱成像卫星捕捉从地球上一个单点反射的光,但将其分成数百个狭窄的波长波段。由此产生的光谱包含了所存在物质的丰富特征,但这些波段通常间隔不规则,有些可能因大气吸收而缺失。简单地将这些波段输入一个带有基于整数的位置编码的标准 Transformer 在物理上是无意义的;它会把 400400400 nm 和 401401401 nm 之间的间隔与 150015001500 nm 和 160016001600 nm 之间的间隔等同对待。

解决方案是相同的:让模型意识到真正的物理“位置”——波长 λ\lambdaλ。通过提供作为实际波长值(或它们的差值 λi−λj\lambda_i - \lambda_jλi​−λj​)函数的位置编码,自注意力机制被赋予了学习光谱物理学中固有的、真正的长程相关性的能力。它可以学习到可见光谱中的一个窄吸收特征与红外光谱中的一个宽特征是耦合的,这是特定矿物的标志,从而能够对我们的星球进行更强大、更基于物理的分析。

学习自然法则

我们已经看到自注意力被用来分析由生物学和物理学定律支配的系统。但它能否更进一步,学习这些定律本身呢?在一个引人入胜的研究方向中,研究人员正在使用 Transformer 作为“神经算子”来学习由偏微分方程 (PDEs) 描述的物理系统的动力学。

考虑热方程,它描述了温度如何通过一种材料扩散。模拟这种情况的一个经典方法是有限差分法,其中一个点在下一个时间步的温度是根据其当前温度及紧邻点的温度的加权平均来计算的。这种计算“模板”(stencil)在某种程度上是一个微小的、固定的注意力机制,只关注其局部邻域。

如果我们将这个固定的模板替换为一个完整的自注意力层会发生什么?我们可以初始化一个系统(例如,一个二维的温度网格),并要求 Transformer 预测下一个时间步的状态,使用真实的 PDE 模拟作为基准真相。模型将每个网格点视为一个词元,学习一个近似物理定律的算子。因为它的注意力是全局的,它可以学会捕捉更复杂、非局部的物理现象,而这些是简单模板会错过的。它不仅仅是在解方程;它是在学习方程的本质。

这把我们带到了理论物理学前沿最深刻的联系之一。几十年来,研究一维量子多体系统的强大工具是矩阵乘积态 (MPS)。MPS 是一种绝妙的拟设(ansatz),或数学模板,它在表示“有能隙”系统(其中相关性随距离呈指数衰减)方面异常高效。然而,对于“临界”系统——那些处于相变点的系统——相关性以慢得多的幂律形式衰减,并且纠缠随系统大小呈对数增长。为了捕捉这一点,MPS 中的参数数量必须随系统大小呈多项式增长,很快变得难以处理。

在这里,Transformer 架构揭示了一个根本性的优势。临界系统是具有长程依赖系统的典型例子。一个带有相对位置编码的自注意力层,由于其本身的性质,可以用固定数量的参数学习一个幂律交互核,而与系统大小无关。对于一个足够大的临界系统,Transformer 在参数效率上变得远超专门定制的 MPS。这不仅仅是一个数值技巧;这是关于架构偏置的一个深刻陈述。诞生于语言世界的 Transformer,拥有一个惊人地适合描述临界现象的无标度、长程相关世界的内在结构。

从阅读基因组到破解量子世界,自注意力的旅程证明了一个伟大思想的统一力量。它不仅仅是一个工程工具;它是一个新的镜头,通过它我们可以审视复杂系统的相互联系,是一块罗塞塔石碑,帮助我们将自然界错综复杂的模式翻译成我们能开始理解的计算语言。