【他山之石】NIPS 2024 | 注意力掩码和LayerNorm在Transformer中的作用

科技   2025-01-14 18:04   北京  

“他山之石,可以攻玉”,站在巨人的肩膀才能看得更高,走得更远。在科研的道路上,更需借助东风才能更快前行。为此,我们特别搜集整理了一些实用的代码链接,数据集,软件,编程技巧等,开辟“他山之石”专栏,助你乘风破浪,一路奋勇向前,敬请关注!

前言本文与之前将自注意力视为连续时间动力系统的研究不同,作者采用了离散时间建模,更贴近实际Transformer架构。

论文信息

题目:On the Role of Attention Masks and LayerNorm in Transformers

注意力掩码和LayerNorm在Transformer中的作用

作者:Xinyi Wu,Amir Ajorlou,Yifei Wang,Stefanie Jegelka,Ali Jadbabaie

论文创新点

  1. 注意力掩码对秩崩溃的影响分析:作者首次系统性地分析了注意力掩码对Transformer中秩崩溃现象的影响。通过引入图论方法,作者证明了在准强连通图的情况下,即使使用稀疏或局部注意力掩码,令牌的秩崩溃仍然会发生,但速率会减缓。这一发现为设计更高效的注意力机制提供了理论基础。
  2. LayerNorm对秩崩溃的缓解作用:作者通过构建非平凡的反例,证明了LayerNorm在某些情况下可以有效缓解令牌的秩崩溃问题。在适当选择值矩阵的情况下,带有LayerNorm的自注意力动态可以同时拥有从1到满秩的任意秩的平衡点。

摘要

自注意力机制是Transformer的关键机制,而Transformer是现代基础模型的核心构建块。最近的研究表明,纯自注意力机制在深度增加时会经历秩崩溃,限制了模型的表达能力和进一步利用模型深度的能力。然而,现有的关于秩崩溃的文献大多忽略了Transformer中的其他关键组件,这些组件可能缓解秩崩溃问题。在本文中,作者对自注意力机制下的秩崩溃进行了综合分析,考虑了注意力掩码和层归一化(LayerNorm)的影响。具体来说,作者发现尽管纯掩码注意力仍然会指数级地崩溃到一个秩为1的子空间,但稀疏或局部掩码注意力可以证明减缓崩溃速率。在LayerNorm的情况下,作者首先展示了对于某些类别的值矩阵,秩为1的子空间崩溃仍然会指数级发生。然而,通过构建非平凡的反例,作者证明了在适当选择值矩阵的情况下,一类通用的序列可能不会收敛到秩为1的子空间,并且带有LayerNorm的自注意力动态可以同时拥有从1到满秩的任意秩的平衡点。作者的结果反驳了之前关于LayerNorm在自注意力秩崩溃中不起作用的假设,并表明带有LayerNorm的自注意力构成了一个比最初认为的更具表达力和多功能的非线性动力系统。

4. 主要结果:带掩码和LayerNorm的注意力

为了研究令牌表示在自注意力动态下的演变及其长期行为,作者通过来衡量令牌相似性:
该度量在数学上等价于[12]中使用的度量,但(4)的形式在一般分析中更容易处理,并且更直接计算。作者的公式的另一个优点是它清楚地表明定理1和定理2不依赖于的具体选择:这些结果适用于任何Lipschitz连续的,其Lipschitz常数使得当且仅当,因为作者可以直接推导出
最后,作者在分析中采用以下假设:
  • A1 包含自环,即对于每个令牌
  • A2 存在常数,使得
A1确保每个令牌在每一层都有一个邻居,从而使得掩码注意力的计算对于每个令牌在每一层都是良定义的,而A2假设键和查询权重矩阵是有界的,这在实践中对于高效的注意力计算至关重要。

掩码注意力

作者首先分析不带LayerNorm的情况,并关注注意力掩码的影响。为了确保在没有LayerNorm的情况下,令牌轨迹对于所有都是有界的,作者进一步假设:
  • A3 序列是有界的。

然后,对于一般的注意力掩码,令牌之间仍然存在强烈的连接,并且令牌表示会指数级地崩溃到秩为1。定理1。考虑定义在(2)中的不带LayerNorm的自注意力动态。在A1-A3下,如果是准强连通图,则存在,使得对于所有
因此,令牌的秩崩溃会指数级发生,即存在,使得
其中的半径,这意味着令牌会指数级收敛到一个共同向量。上述结果表明,在纯自注意力下,只要序列中存在一个令牌,所有其他令牌都可以在固定层数内直接或间接参与,那么令牌的秩崩溃就会指数级发生。特别是,它将[12]中的主要结果从是完全图推广到更一般的注意力模式类别:注意力模式只需要是准强连通的,这意味着结果适用于实践中使用的各种注意力掩码,包括GPT系列中使用的因果掩码,或许多高效Transformer模型中部署的稀疏注意力模式。作者讨论了以下几个有趣的含义。局部 vs. 全局注意力 指数速率在图半径上单调。这意味着对于半径较大的图,秩崩溃应该较慢。作者的结果间接支持使用局部注意力模式,这不仅使注意力计算更高效(这是这些工作的最初动机),而且隐式地缓解了秩崩溃问题。聚焦 vs. 均匀注意力 此外,指数速率在上单调递减,这意味着越小,秩崩溃越慢。可以将解释为注意力在可达令牌之间的“聚焦”程度,因为在注意力均匀分布在可达令牌时达到最大值。除了应用注意力掩码和限制可达令牌的数量外,控制注意力聚焦程度的另一种方法是通过温度项。较大的值会使可达令牌之间的注意力分配更加均匀,从而使秩崩溃在各层之间更快发生。秩崩溃与通用逼近能力的权衡 最后,对于强连通图,上述结果还揭示了通用函数逼近能力与秩崩溃速率之间的权衡。Yun等人表明,带有强连通图掩码的Transformer是序列到序列函数通用逼近器,然而,对于掩码,它们需要至少的直径层才能实现完整的序列到序列函数逼近属性。这意味着直径较小的掩码(因此半径较小,因为半径直径半径)在函数逼近能力方面更高效,但它们更容易发生秩崩溃。

带LayerNorm的掩码注意力:秩崩溃

到目前为止,作者已经考虑了不带LayerNorm的纯自注意力动态,并关注了注意力掩码的作用。如果在(3)中添加LayerNorm并考虑注意力动态会发生什么?在本节中,作者首先展示一个负面结果,表明对于某些类别的值矩阵,令牌的指数级崩溃到一个共同向量仍然会发生。定理2。考虑定义在(3)中的带LayerNorm的自注意力动态。设是强连通图。假设A1-A2,并且对于所有是正交的,并且初始输入满足:
  • () 是满秩的。
则存在,使得,并且
其中的半径,这意味着令牌会指数级收敛到上的一个共同点。该结果可以看作是[15]中定理4.1的广义离散版本。值得注意的是,作者的分析纯粹基于高级线性代数工具:非负矩阵的无限乘积及其遍历性,并且可以考虑随时间变化的权重和一般的注意力掩码,而[15]中则是固定是完全图。满足初始输入条件的一种方法是要求,并以均匀随机的方式初始化令牌在上,那么条件几乎必然成立。这是[15]中处理条件的方式。注意,条件()意味着存在,使得对于所有,这可以通过超平面分离定理或Farkas引理得出(见附录C中的引理6)。如果初始令牌几何满足比上述更强的条件,则()不再必要,定理2甚至可以直接推广到准强连通图。作者定义,表示令牌之间的最小余弦相似度。如果初始时所有令牌对的余弦相似度都是非负的,那么只要是准强连通的,秩崩溃就会指数级发生。推论1。考虑定义在(3)中的带LayerNorm的自注意力动态。设是准强连通图。在A1-A2下,如果对于所有是正交的,并且,则存在,使得,并且
其中的半径,这意味着令牌会指数级收敛到上的一个共同点。全掩码 vs. 因果掩码 作者可以通过指定掩码中的中心节点数量来细化推论1,那么指数速率的上界将是,这意味着秩崩溃的速率会受到掩码中中心节点数量的负面影响。在全注意力的情况下,是完全图,掩码将有个中心节点,与定理2中的上界匹配。在因果注意力的情况下,因果图,掩码只有一个中心节点,上界会更宽松,表明因果掩码在缓解秩崩溃速率方面相对于全掩码具有优势。后LN vs. 前LN LayerNorm的定义在(3)中遵循原始Transformer论文,如今称为后LN。许多LLMs中的另一种LayerNorm用法是前LN,其中LayerNorm位于自注意力之前,可以写为
注意,定理2和推论1直接适用于前LN的情况,证明类似。

带LayerNorm的掩码注意力:反例

前几节的主要结果似乎有些悲观:无论是带LayerNorm还是不带LayerNorm,自注意力动态似乎注定会在长期内崩溃到秩为1的子空间。然而,在本节中,作者首先构建了一个非平凡的反例,其中仅使用LayerNorm,对于一类通用的输入序列,令牌收敛到一个平衡点,其中秩崩溃不会发生。请注意,对于Transformer模型来说,重要的是它能够防止一类通用输入序列的秩崩溃,而不仅仅是一个特定的输入序列。然后,作者展示了一个一般结果,表明在LayerNorm和适当选择值矩阵的情况下,自注意力动态可以同时拥有从1到满秩的任意秩的平衡点。此外,对于一类通用的输入序列,令牌在结果动态下不会收敛到秩为1的子空间。

4.3.1 说明性反例

为简单起见,作者考虑,并且是因果掩码。然后设,这导致注意力矩阵
作者进一步设
对于。不失一般性,设。然后经过仔细分析,作者发现,根据其初始位置,第一个令牌将收敛到。假设第一个令牌收敛到。那么第二个令牌的收敛情况如图1所示,其中。详细的证明可以在附录E中找到。注意,由于LayerNorm的缩放效应,的任何缩放版本在这里同样有效。备注2。对于任何正交矩阵在这个例子中同样有效,并且结果令牌轨迹为这个非平凡的反例表明,在LayerNorm动态下,存在适当选择的矩阵,可以防止令牌崩溃到秩为1的子空间,对于非零测度的输入序列。

5. 数值实验

击呗!

人工智能前沿讲习
领先的人工智能知识平台
 最新文章