位置编码是Transformer中很重要的一环,在《让研究人员绞尽脑汁的Transformer位置编码》中我们就总结了一些常见的位置编码设计。大体上,我们将Transformer的位置编码分为“绝对位置编码”和“相对位置编码”两类,其中“相对位置编码”在众多NLP/CV的实验表现相对来说更加好些。

然而,我们可以发现,目前相对位置编码几乎都是在Softmax之前的Attention矩阵上进行操作的,这种施加方式实际上都存在一个理论上的缺陷,使得Transformer无法成为“万能拟合器”。本文就来分析这个问题,并探讨一些解决方案。

简单探针 #

顾名思义,位置编码就是用来给模型补充上位置信息的。那么,如何判断一个模型有没有足够的识别位置的能力呢?笔者之前曾构思过一个简单的探针实验:

对于一个有识别位置能力的模型,应该有能力准确实现如下映射 \begin{equation}\begin{array}{lc} \text{输入:} & [0, 0, \cdots, 0, 0] \\ & \downarrow\\ \text{输出:} & [1, 2, \cdots, n-1, n] \end{array}\end{equation}

也就是说,输入$n$个0,能有序地输出位置编号$1\sim n$。这个探针实验的思想很简单,即模型如果有能力做到这一点,说明识别位置是模型自身具备的能力,跟外部输入无关,这正是我们想要的。不难发现,绝对位置由于是直接施加在输入上的,所以它很容易能够完成探针测试。

无法胜任 #

然而,当笔者带着这个简单的探针实验去思考带有相对位置编码的Transformer模型时,却发现它们几乎都不能完成上述任务。

具体来说,除了《Self-Attention with Relative Position Representations》所提出的设计外,其余所有相对位置编码(包括笔者所提的RoPE)都只修改了Softmax前的Attention矩阵,那么带有相对位置信息的Attention矩阵依然是一个概率矩阵(即每一行求和等于1)。

另一方面,对于Transformer模型来说,Token之间的交互的唯一来源是Self Attention的$\boldsymbol{A}\boldsymbol{V}$这一步,或者写成$\boldsymbol{o}_i = \sum\limits_j a_{i,j}\boldsymbol{v}_j$。相同的输入意味着每个$\boldsymbol{v}_j$都是相同的,所以
\begin{equation}\boldsymbol{o}_i = \sum_j a_{i,j}\boldsymbol{v}_j = \sum_j a_{i,j}\boldsymbol{v} = \left(\sum_j a_{i,j}\right)\boldsymbol{v} = \boldsymbol{v}\end{equation}
这意味着每个$\boldsymbol{o}_i$也是相同的。换句话说,模型的每个位置自始至终都输出相同的结果,所以模型根本不可能输出各不相同的$[1, 2, \cdots, n-1, n]$。

类似的发现也出现在最近的论文《Your Transformer May Not be as Powerful as You Expect》中,作者构建了略有不同的例子来演示相对位置编码Transformer的拟合能力缺陷问题,两者异曲同工、不谋而合了。此外,本文开头说的是“万能拟合”,那解决了这个反例是不是就能做到“万能拟合”了呢?该论文也有相应的理论分析来肯定这一事实,这里就不详述了。

初步方案 #

稍加思考就可以发现,其实问题主要出在Attention矩阵的每一行求和等于1,要解决这个问题,想办法打破这个约束就行了。为此,《Your Transformer May Not be as Powerful as You Expect》在其发现之上进一步提出了如下设计
\begin{equation}\boldsymbol{O} = (\boldsymbol{A}\odot \boldsymbol{C})\boldsymbol{V}\quad \text{或者等价地}\quad\boldsymbol{o}_i = \sum_j a_{i,j}c_{i,j}\boldsymbol{v}_j\end{equation}
其中$\boldsymbol{C}$是一个可训练的参数矩阵,$\odot$是逐位相乘(Hadamard积)。为了使得整个模型依然只包含相对位置信息(因为本文就是讨论相对位置编码Transfomrer的缺陷),我们要约束$\boldsymbol{C}$为Toeplitz矩阵,即$c_{i,j}=g(i-j)$。

有了$\boldsymbol{C}$的加入,$\boldsymbol{A}\odot \boldsymbol{C}$作为一个整体,每一行的和显然不一定为1,从而打破了这个限制,因此是可以解决问题的(更多的实验结果请自行看原论文)。但这样一来,引入了新的参数矩阵不说,由于$\boldsymbol{C}$本身是有限大小的,所以它就不能很好地支持变长输入(或者矩阵$\boldsymbol{C}$相应地要做一些截断,即$c_{i,j}=g(\text{clip}(i-j, p_{\min}, p_{\max}))$的形式),总的来说显得不够简洁优雅。

去掉分母 #

再次回到问题所在:Attention矩阵的每一行求和等于1。是什么操作导致了这一现象呢?答案很显然,是Softmax:
\begin{equation}a_{i,j} = \frac{e^{b_{i,j}}}{\sum\limits_j e^{b_{i,j}}}\end{equation}
这里的$\boldsymbol{B}=(b_{i,j})$是Softmax前的矩阵。很明显,就是“除以$\sum\limits_j e^{b_{i,j}}$”这一步导致了$\sum\limits_j a_{i,j}=1$,那么一个很直接的想法就是:

如果我不想$\sum\limits_j a_{i,j}=1$,那么干脆别除以$\sum\limits_j e^{b_{i,j}}$就行了?

事实上确实可以!实验结果显示,不除以该分母的Transformer确实能成功地完成前述探针测试。此时就不得不感概一下GAU的“先见之明”了,它提出的新式Attention直接是$\text{relu}^2$激活然后简单除以$n$来归一化,避免了$\sum\limits_j a_{i,j}=1$,从而增强了模型的理论能力(当然也许作者根本没想那么多,是笔者想象的成分居多)。

新归一化 #

然而,我们在《听说Attention与Softmax更配哦~》发现像GAU里的不进行概率归一化的Attention设计可能存在外推能力欠佳的问题。也就是说,进行概率归一化导致了前面说的理论缺陷,简单地除以$n$来归一化则外推能力可能欠佳,有没有同时能兼顾两者的方案呢?

让我们再发散一下脑洞。从范数的角度来看,$\sum\limits_j e^{b_{i,j}}$实际上是向量$e^{b_{i,:}}$的$l_1$范数,所以Softmax实际上就是向量的$e^{b_{i,:}}$的$l_1$归一化操作,那么要避免$\sum\limits_j a_{i,j}=1$,又有保留归一化,换成其他的归一化操作是否可以呢?比如$l_2$归一化:
\begin{equation}a_{i,j} = \frac{e^{b_{i,j}}}{\sqrt{\sum\limits_j e^{2b_{i,j}}}}\end{equation}

经过笔者测试,这种$l_2$归一化的Attention,确实能成功完成探针实验。那么,这个改动对我们更关心的NLP预训练场景有没有帮助呢?笔者也做了相应的对比实验,结果是分两部分:

1、对于标准的Attention + FFN组合,应用$l_2$归一化Attention之前要缩小一下Attention的$\boldsymbol{W}_V,\boldsymbol{W}_O$的初始方差,实验结果则是略差于常规的$l_1$归一化Attention;

2、对于全GAU的架构,可以直接应用$l_2$归一化Attention,不需要改动初始化,实验结果则是略优于常规的$l_1$归一化Attention。

两者的差别大概是源于它们本身的初始化方式不同,在标准的Attention + FFN组合中,初始Attention矩阵接近一个均匀矩阵(每个数都相同),而在《门控注意力单元(GAU)还需要Warmup吗?》我们则分析过,GAU的初始Attention矩阵更接近一个单位阵(的若干倍)。

峰回路转 #

再次纵观前文,我们发现是因为“每个$\boldsymbol{v}_j$都是相同的”,所以“$\sum\limits_j a_{i,j}=1$的模型无法完成探针实验”。但如果每个$\boldsymbol{v}_j$不全相同呢?

我们知道,从BERT开始,主流的Transformer模型都是像“[CLS] SENT [SEP]”设计输入的,也就是在输入前后会附加一些标记性的Token,如果我们将这些标记Token当作模型的一部分而不是输入(也就是说输入“[CLS] 0 0 ⋯ 0 0 [SEP]”而不是全0),那么是否有可能完成探针呢?

笔者也对此做了实验,发现对输入补充上标记行Token后,不需要对相对位置编码Transformer的其他部分做修改,确实也能够完成探针实验。这结果就有点啼笑皆非了,原来BERT的作者们也很有“先见之明”啊,所添加的特殊Token [CLS][SEP]还有辅助定位的作用,我们分析那么久的理论缺陷,居然就这样被两个特殊Token解决了。这不禁让人想起《How Much Position Information Do Convolutional Neural Networks Encode?》所提到的“CNN是通过padding来识别绝对位置的”这一结论,两者有一定的相通之处。

当然,这也不意味着我们前面的思考全无意义。比如对GAU模型来说,Attention换用$l_2$归一化确确实实有加快收敛、轻微提升效果的作用。此外,既然可以接受$l_2$归一化,那么$e^{b_{i,j}}$是不是还可以换成一般的激活函数(比如去掉非负性约束)呢?笔者也简单做了“$\text{swish}(b_{i,j})$ + $l_2$归一化”的实验,发现有一定的可行性。从这个角度来看,$l_2$归一化下的Attention实际上有更多的拓展空间。

曲终人散 #

本文分析了相对位置编码Transformer的一个隐含缺陷,并探讨了相应的对策,从中引申出关于Attention矩阵的非负性、归一化方式的思考。

转载到请包括本文地址:https://kexue.fm/archives/9105

更详细的转载事宜请参考:《科学空间FAQ》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Jun. 07, 2022). 《相对位置编码Transformer的一个理论缺陷与对策 》[Blog post]. Retrieved from https://kexue.fm/archives/9105

@online{kexuefm-9105,
        title={相对位置编码Transformer的一个理论缺陷与对策},
        author={苏剑林},
        year={2022},
        month={Jun},
        url={\url{https://kexue.fm/archives/9105}},
}