无位置编码 (NoPE) 也有长度泛化问题？首个针对NoPE的长度外推方法

2024-05-16 05:05

©作者 | FudanNLP

单位 | 复旦大学NLP实验室

基于 Transformer 的大型语言模型（LLM）展示了强大的文本理解能力，颠覆了整个 NLP 领域的应用范式。然而，在长度有限文本上预训练的语言模型却无法像人类一样泛化到任意长度文本，如何解决长度泛化问题成为了 LLM 的一项主要挑战。

目前主流的长度泛化工作认为：Transformer 模型中显式的位置编码（例如旋转位置编码 RoPE）带来了长度泛化问题，并试图 1）提出新型可泛化的位置编码（e.g., Alibi[1], XPOS[2]）; 2）二次微调 LLM 到更长的文本上（e.g., LLaMA2Long [3]）; 3）以内插、外推等方式修改 RoPE 位置编码（e.g., PI[4], YaRN[5]）。

这三类工作现阶段都取得了不错的进展，其中一些方法已经成功将 LLM 的长度扩展到 100K 以上 [6]。然而，位置编码是否是导致长度泛化问题的唯一因素？或者说，如果没有位置编码问题，长度泛化问题会消失吗？

无位置编码（NoPE）的 Transformer 已经被证明在自回归语言模型任务上和 Transformer+RoPE 效果相当 [7,8,9]，但是 NoPE 的长度泛化问题并没有改善，和 RoPE 一样严重。华师、复旦、上海 AI Lab 联合团队基于 NoPE，在排除位置编码影响下，研究长度泛化失败的表现和原因，并首次提出适用于 NoPE 的长度泛化方法。

论文标题：

Length Generalization of Causal Transformers without Position Encoding

论文链接：

https://arxiv.org/pdf/2404.12224.pdf

代码链接：

https://github.com/AntNLP/nope_head_scale

论文速看

本文在 NoPE 场景中探究 LLM 的长度泛化问题，发现即使没有位置编码，LLM 仍然无法直接外推。经过注意力熵的可视化，我们观察到无论 RoPE 还是 NoPE，泛化失败的长度上都表现出注意力严重分散的问题。通过简单的注意力缩放（引入 SoftMax 温度超参让注意力强制集中）就能显著提升 NoPE 的长度泛化能力，但是对 RoPE 没有显著效果。

此外，本文发现不同的注意力头表现出不一致的注意力分散模式，例如 20% 的头完全不需要强制集中，剩余的头所需的集中力度也不尽相同。因此本文提出了一种基于梯度下降的高效超参数搜索，为每个注意力头搜索一个最佳温度超参数。对于 1B 模型，仅需在 30M Token（预训练的 0.03%）的长文上搜索 704 个超参，就能从 2k 扩展到 16k。

本文的核心贡献是“抛砖引玉”提出在 NoPE 场景下研究长度泛化问题的原因和解决方案。因为即使传入显式 PE（如 RoPE），Transformer 也不可避免的会以某种形式自行学习位置信息。未来的外推方法可能是“PE 自身的外推”+“NoPE 的外推”，甚至是只有“NoPE 的外推”。

现象：注意力分散

NoPE 模型中还有什么因素会制约外推能力？在 Transformer-NoPE 架构中，Embedding 层和 FFN 层均与位置无关，只有注意力模块与位置相关。那么，注意力模块是否会在不同位置上表现出不同的行为？本文发现，Transformer 模型外推失败的时候，总是伴随着注意力分散。

为了定量描述注意力模块的聚焦程度，作者可视化并分析了注意力分布的熵。点积注意力可以表示为：

其中，温度系数默认取：

对注意力矩阵的每一行分别计算熵，就得到了每个位置的注意力分布情况。熵越小，注意力越集中；反之熵越大，注意力越分散。

图中实线表示每个位置上的熵（所有注意力头的平均），虚线表示每个位置上的对数困惑度。可以发现，模型困惑度的崩溃总是伴随着熵的上升，即注意力分散。

因此作者修改了 SoftMax 的最佳温度超参数，使注意力重新聚焦，从而显著地增强了 NoPE 模型的上下文长度（左图）。然而，这一方法却无法直接应用于 RoPE 模型（右图）。

按注意力头分别控制

实验表明，简单地修改模型整体的最佳温度超参数，得到的上下文扩展的范围有限。作者进一步细化了熵的可视化，发现不同的注意力头表现出不一致的注意力模式，有些熵较小，有些则比较大。这暗示了每个注意力头的最佳温度超参数或许不同。

于是，作者为每个头分配不同的最佳温度超参数，对于 1B 大小的模型，共计 22 层*32 头=704 个参数。通过基于梯度下降的高效超参数搜索，找到每个头的最佳温度超参数 λ^(h)，使其外推效果达到最佳。

最佳温度超参数与熵的大小是否有关呢？作者可视化了两者的相关关系，发现注意力越聚焦（熵更小），则对应的 λ^(h) 越大。另外，这种相关性与层数有关，模型的不同层呈现出不同的行为模式。

实验结果

本文提出的方法能够显著提高 NoPE 模型在长序列语言建模、合成任务（passkey 检索）和真实世界的长上下文任务（LongBench）上的性能。

由于只有 704 个可调整的参数，本文的 NoPE 长度外推方法只需要极少的训练 token 即可收敛。在相同的训练量下，NoPE 与 RoPE 上的基线性能持平。然而，当增加训练量时，RoPE 基线的性能仍有提升空间，而 NoPE 上的超参数搜索已经达到上界。

另一方面，随着 NoPE 的外推长度不断变大，其在短上下文上的性能逐渐下降，这意味着限制 NoPE 模型专注于更少的 token 可能会对其性能产生不利影响。如何恢复短上下文的性能还需要进一步研究。

此外，NoPE 模型还呈现出远超 RoPE 的外推潜力。NoPE 模型不做任何操作，就能在 2 倍预训练长度上很好地完成 passkey 检索任务（图中黑色虚线为预训练长度或微调长度）。

总结与展望

本文讨论了无位置编码（NoPE）模型的长度泛化性问题。作者提出了基于梯度下降的高效超参数搜索方法，针对每个注意力头搜索最佳的温度系数。实验证明 NoPE 的长度泛化能力具有竞争力，为长上下文语言模型提供了一个新的选项。

本文为理解模型的位置信息提供了一个新的视角。NoPE 去除了显式位置编码的干扰，直击模型内部的位置信息表示。现有的长度外推工作大多聚焦在操作位置编码的表示上，本文揭示了长度泛化失败与注意力分散之间的相关性，为该领域的研究提供的新的方向。

参考文献

[1] Ofir Press, Noah A. Smith, Mike Lewis. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022

[2] Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, Furu Wei. A Length-Extrapolatable Transformer. ACL 2023

[3] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, Hao Ma. Effective Long-Context Scaling of Foundation Models.

[4] Shouyuan Chen, Sherman Wong, Liangjian Chen, Yuandong Tian. Extending Context Window of Large Language Models via Positional Interpolation.

[5] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context window extension of large language models. ICLR 2024

[6] Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, Lingpeng Kong. Training-Free Long-Context Scaling of Large Language Models.

[7] Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, Omer Levy. Transformer Language Models without Positional Encodings Still Learn Positional Information. Findings of EMNLP 2022

[8] Ta-Chung Chi, Ting-Han Fan, Li-Wei Chen, Alexander I. Rudnicky, Peter J. Ramadge. Latent Positional Information is in the Self-Attention Variance of Transformer Language Models Without Positional Embeddings. ACL 2023

[9] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, Siva Reddy. The Impact of Positional Encoding on Length Generalization in Transformers. NeurIPS 2023

更多阅读