When Precision Meets Position: BFloat16 Breaks Down RoPE in Long-Context Training

📄 arXiv: 2411.13476v2 📥 PDF

作者: Haonan Wang, Qian Liu, Chao Du, Tongyao Zhu, Cunxiao Du, Kenji Kawaguchi, Tianyu Pang

分类: cs.CL

发布日期: 2024-11-20 (更新: 2024-11-26)

🔗 代码/项目: GITHUB

💡 一句话要点

提出AnchorAttention，解决BFloat16下RoPE长文本训练的数值问题。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本建模 旋转位置编码 BFloat16 数值稳定性 注意力机制 大型语言模型 AnchorAttention

📋 核心要点

  1. RoPE在长文本建模中被广泛使用，但BFloat16精度下存在数值问题，导致性能下降。
  2. AnchorAttention将首个token作为共享锚点，减少不必要的注意力计算，缓解数值问题。
  3. 实验表明，AnchorAttention在长文本任务上显著提升性能，并减少了超过50%的训练时间。

📝 摘要（中文）

扩展上下文窗口大小使得大型语言模型(LLMs)能够处理更长的序列和更复杂的任务。旋转位置嵌入(RoPE)因其相对位置编码特性而成为长上下文训练的事实标准。然而，我们观察到，将RoPE与BFloat16格式一起使用会导致数值问题，使其偏离其预期的相对位置编码，尤其是在长上下文场景中。这个问题源于BFloat16的有限精度，并随着上下文长度的增加而累积，其中第一个token对此问题贡献显著。为了解决这个问题，我们开发了AnchorAttention，这是一种即插即用的注意力方法，可以缓解BFloat16引起的数值问题，提高长上下文能力，并加快训练速度。AnchorAttention通过将第一个token视为具有一致位置ID的共享锚点，使其对训练上下文中的所有文档可见，从而减少不必要的注意力计算，保持语义连贯性，并提高计算效率。在三种LLM上的实验表明，与标准全注意力机制相比，AnchorAttention显著提高了长上下文性能，并将训练时间减少了50%以上，同时保留了原始LLM在通用任务上的能力。我们的代码可在https://github.com/haonan3/AnchorContext获得。

🔬 方法详解

问题定义：论文旨在解决在使用BFloat16格式进行长文本训练时，RoPE（旋转位置编码）出现的数值精度问题。现有方法在长文本场景下，由于BFloat16的精度限制，RoPE无法准确地进行相对位置编码，导致模型性能下降。尤其是在长文本序列中，误差会随着序列长度的增加而累积，严重影响模型的训练效果。

核心思路：论文的核心思路是将序列中的第一个token视为一个“锚点”（Anchor），并赋予其一个固定的位置ID。通过这种方式，所有其他的token都相对于这个锚点进行位置编码，从而减少了由于BFloat16精度限制而导致的误差累积。这种方法可以有效地缓解长文本训练中的数值问题，提高模型的性能。

技术框架：AnchorAttention是一个即插即用的注意力机制，可以很容易地集成到现有的Transformer模型中。其主要流程如下：1) 将输入序列的第一个token作为锚点；2) 为锚点分配一个固定的位置ID；3) 在计算注意力权重时，所有token都相对于该锚点进行位置编码；4) 使用修改后的注意力权重进行加权求和，得到最终的输出。

关键创新：AnchorAttention的关键创新在于引入了“锚点”的概念，并将其应用于位置编码中。与传统的RoPE方法相比，AnchorAttention可以有效地减少由于BFloat16精度限制而导致的误差累积，从而提高模型在长文本场景下的性能。此外，AnchorAttention还减少了不必要的注意力计算，提高了计算效率。

关键设计：AnchorAttention的关键设计在于如何选择锚点以及如何为锚点分配位置ID。论文中选择序列的第一个token作为锚点，并为其分配一个固定的位置ID。这种设计简单有效，并且易于实现。此外，AnchorAttention还可以与其他注意力机制结合使用，以进一步提高模型的性能。没有提及损失函数或网络结构的修改。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AnchorAttention在三种不同的LLM上都取得了显著的性能提升。与标准的全注意力机制相比，AnchorAttention在长文本任务上将训练时间减少了50%以上，同时保持了模型在通用任务上的性能。这些结果表明，AnchorAttention是一种有效且高效的长文本建模方法。

🎯 应用场景

AnchorAttention可应用于各种需要处理长文本序列的自然语言处理任务，如文档摘要、机器翻译、问答系统等。该方法能够有效提升模型在长文本场景下的性能，并降低训练成本，具有重要的实际应用价值和潜力。未来可进一步探索其在其他领域的应用，例如长序列视频理解等。

📄 摘要（原文）

Extending context window sizes allows large language models (LLMs) to process longer sequences and handle more complex tasks. Rotary Positional Embedding (RoPE) has become the de facto standard due to its relative positional encoding properties that benefit long-context training. However, we observe that using RoPE with BFloat16 format results in numerical issues, causing it to deviate from its intended relative positional encoding, especially in long-context scenarios. This issue arises from BFloat16's limited precision and accumulates as context length increases, with the first token contributing significantly to this problem. To address this, we develop AnchorAttention, a plug-and-play attention method that alleviates numerical issues caused by BFloat16, improves long-context capabilities, and speeds up training. AnchorAttention reduces unnecessary attention computations, maintains semantic coherence, and boosts computational efficiency by treating the first token as a shared anchor with a consistent position ID, making it visible to all documents within the training context. Experiments on three types of LLMs demonstrate that AnchorAttention significantly improves long-context performance and reduces training time by over 50\% compared to standard full attention mechanisms, while preserving the original LLM's capabilities on general tasks. Our code is available at https://github.com/haonan3/AnchorContext.