LLM：ALiBi - 给注意力加上线性偏置

论文：https://arxiv.org/pdf/2108.12409.pdf

代码：https://github.com/ofirpress/attention_with_linear_biases

发表：2021

长度外推

参考：https://spaces.ac.cn/archives/9431#ALIBI

长度外推性是一个训练和预测的长度不一致的问题。具体来说，不一致的地方有两点：

1、预测的时候用到了没训练过的位置编码（不管绝对还是相对）；

2、预测的时候注意力机制所处理的token数量远超训练时的数量。

第1点：可能大家都容易理解，没训练过的就没法保证能处理好，这是DL中很现实的现象，哪怕是Sinusoidal或RoPE这种函数式位置编码也是如此。

第2点：可能读者会有些迷惑，Attention理论上不就是可以处理任意长度的序列吗？训练和预测长度不一致影响什么呢？答案是熵，我们在《从熵不变性看Attention的Scale操作》也已经分析过这个问题，越多的token去平均注意力，意味着最后的分布相对来说越“均匀”（熵更大），即注意力越分散；而训练长度短，则意味着注意力的熵更低，注意力越集中，这也是一种训练和预测的差异性，也会影响效果。

摘要

如何让模型在推理时实现对训练期间未见过的序列的外推？我们引入了一种更简单、更有效的位置表示方法，即带有线性偏置的注意力（ALiBi）。

ALiBi核心思想：没有将位置嵌入添加到单词嵌入中，而是用一个和query, key之间的距离成比例的一个“惩罚项”来偏置query-key的注意力得分。训练速度提高了 11%，内存使用量减少了 11%。

1 简介 

我们发现：使用Sinusoidal PE嵌入的Transformer语言模型（LMs）的外推能力非常弱。

我们证明：这种外推失败是由位置嵌入方法引起的。如下图所示，对于Sinusoidal PE（黄色）方法的最新替代方案 RoPE（橙色） 已经改进了外推。但最好的是T5 bias（紫色），其比正弦方法慢得多，但会引入额外的内存和参数。

图1：左边的是模型在512数据集训练的模型，右边是在1024数据集上训练的模型。横坐标是推理是输入的句子长度，纵坐标的困惑度。我们的目标是，困惑度越小越好。通过观察上图我们可以看到，Sinusoidal编码，Rotary编码，T5 Bias这三种位置编码，当输入特别长的时候，困惑度就会飙升，但是ALiBi编码会保持一个平稳的水平（即使超过训练的token长度），说明ALiBi编码有很好的长度外推能力。 

2 目前的方法无法有效地进行外推

从技术上讲应该能够外推的正弦位置方法，在实际应用中其实际的外推能力非常有限。尽管旋转位置方法RoPE比正弦方法有所改进，但它仍然没有取得令人满意的结果。T5 bias方法比这两种方法中的任何一种都能带来更好的外推效果，因此我们得出结论，外推能力很大程度上取决于位置嵌入。遗憾的是，T5 bias 在计算上成本很高。 

图2：批处理训练、推理速度和内存使用的比较。ALiBi基本上还是不错的。 

1：Sinusoidal PE 是不用学习的常数向量，

2：Sinusoidal PE直接加到transformer第一层的输入token embeddings上

1：不同于在Transformer底部添加正弦嵌入，RoPE将每个注意力层的键和查询都乘以正弦嵌入。

2：与正弦嵌入或学习的位置嵌入方法不同，RoPE每个层都注入位置信息，而不仅仅是初始层。

3：RoPE没有向自注意力层（self-attention）的值添加位置信息。自注意力层的输出是一个对输入值向量进行线性变换、加权和的结果；因此，通过不将位置信息插入到值中，每个Transformer层的输出都不包含任何显式位置信息。我们怀疑这种位置信息的隔离可能对外推有益。

1：尽管大多数模型使用训练过的或正弦的位置嵌入，T5模型使用一种相对位置方法，这种方法不向词嵌入添加位置信息（与之前的方法不同）。相反，它修改了注意力值的计算方式。我们将此称为“T5 bias”方法。

2：与旋转方法一样，T5 bias将位置信息注入到每个层的模型中，并没有将任何显式位置信息整合到自注意值向量中。
3：T5 bias 确实允许语言模型进行外推。但这种令人印象深刻的性能需要付出代价：训练速度至少是使用正弦模型的两倍。

3 ALIBI

1：ALIBI 不会在网络中添加positon embeddings。唯一的修改就是在query-key内积后，添加一个静态的、非学习的bias。 换个说法：ALiBi 直接作用在attention score中，给 attention score 加上一个预设好的偏置矩阵。

2：公式第一项是注意力的分数。第二项是一个相对距离的矩阵。

        例如、，他们直接的距离为1-1=0，所以对应的位置是0.

        例如、，他们直接的距离为1-2=-1，所以对应的位置是-1.

为什么是一个下三角矩阵，因为我们研究的是，autoregressive language modeling，所以这就是上三角被MASK的原因，我们不关注未来，只关注过去，这个编码只应用在Query和Key中，不会应用到Value中。

 

3：给定的预设矩阵中还会乘上m的调节因子，m的设置与attention的头数有关。论文中也做了尝试把m作为学习参数，但是并没有获得更好的效果。初始数值和间隔相同，计算公式为：

如果head = 8，m的数值分别为：

，

,

,

,

...

如果head = 16，m的数值分别为：

，

,

,

,

...

我们观察到：这组斜率在各种文本领域和模型大小上都有效。在新数据集上训练新模型时不需要调整这些斜率值。这使得我们的方法与正弦方法相似，正弦方法中的超参数一次性设置的，然后在不同大小的不同数据集上的不同模型中重复使用。ALiBi具有对新颖性的归纳偏差；它会惩罚查询-键对之间的注意力分数，随着键与查询之间的距离增长，惩罚会增加。不同的头会以不同的速率增加他们的惩罚，这取决于斜率的幅度。

我们最初尝试让斜率可训练，但并没有得到很好的外推结果。大约十个斜率集的简短手动探索让我们发现了最终选择的斜率集。

ALiBi是一种相对位置方法，我们在每个层的键和查询中添加位置信息，但不在值中添加，这与T5偏差和旋转方法中的做法不同。

代码实现（研究明白了再补）