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

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

论文试图解决在使用BFloat16格式时，Rotary Positional Embedding (RoPE) 在长上下文场景中出现的数值问题，这些问题导致模型偏离其预期的相对位置编码效果。这是一个具体的技术问题，特别是在大规模语言模型中处理长序列时。

论文提出了一种名为AnchorAttention的新方法，通过将第一个token作为共享锚点并赋予一致的位置ID，减少不必要的注意力计算，从而缓解BFloat16带来的数值问题。这一方法不仅提高了长上下文处理能力，还显著加快了训练速度。

论文通过实验展示了AnchorAttention在三种不同类型的大规模语言模型上的有效性，表明该方法可以显著提高长上下文性能，并将训练时间减少超过50%，同时保持模型在一般任务上的原有能力。此外，作者提供了开源代码，方便其他研究者复现和进一步研究。未来的研究可以探索AnchorAttention在更多模型和任务中的应用。

近期在这个领域中，有几项相关研究值得关注：1.《Longformer: The Long-Document Transformer》提出了通过局部窗口注意力机制来处理长文档。2.《Big Bird: Transformers for Longer Sequences》引入了一种混合注意力机制，结合全局和局部注意力来处理更长的序列。3.《Performer: Generalized Attention with RFF Kernels》使用随机特征函数核来近似注意力机制，以提高效率。