arXiv:2411.13476cs.CL2024-11被引 17

BFloat16精度不足导致RoPE在长文本训练中失效,新方法AnchorAttention解决此问题。

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

  • 将首词设为共享锚点,用固定位置ID避免数值误差
  • 长上下文任务性能提升,训练速度加快超50%
  • 无需修改模型结构,可直接替换标准注意力机制

扩展上下文窗口使大语言模型能处理更长序列和复杂任务。旋转位置编码(RoPE)因具备相对位置编码特性,已成为长上下文训练的主流方法。然而我们发现,在BFloat16格式下使用RoPE会引发数值问题,导致其偏离预期的相对位置编码效果,尤其在长上下文场景中更为显著。该问题源于BFloat16精度有限,随上下文长度增加而累积,首词贡献尤为突出。为此,我们提出AnchorAttention,一种即插即用的注意力机制,可缓解BFloat16带来的数值误差,提升长上下文能力并加速训练。该方法通过将首词作为共享锚点并赋予固定位置ID,使其对所有文档中的内容可见,减少冗余注意力计算,保持语义连贯性,提升计算效率。在三种不同类型的大语言模型上实验表明,AnchorAttention显著改善长上下文表现,相比标准全注意力机制训练时间减少超过50%,同时保持原模型在通用任务上的性能。代码已公开于https://github.com/haonan3/AnchorContext。

原文摘要 · Abstract (English)

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.

长上下文RoPEBFloat16注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。