精度と位置が合致するとき：BFloat16が長いコンテキストのRoPEを解析する

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

November 20, 2024

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

cs.AI

要旨

コンテキストウィンドウサイズを拡張することで、大規模言語モデル（LLM）はより長いシーケンスを処理し、より複雑なタスクを処理できるようになります。Rotary Positional Embedding（RoPE）は、長いコンテキストのトレーニングに利益をもたらす相対位置エンコーディング特性を持つため、事実上の標準となっています。ただし、RoPEをBFloat16形式で使用すると、数値的な問題が発生し、特に長いコンテキストのシナリオでは、意図した相対位置エンコーディングから逸脱することが観察されます。この問題は、BFloat16の精度が限られていることから生じ、コンテキストの長さが増すにつれて蓄積され、最初のトークンがこの問題に大きく影響を与えています。この問題に対処するために、BFloat16によって引き起こされる数値的な問題を緩和し、長いコンテキストの能力を向上させ、トレーニングを高速化するAnchorAttentionというプラグアンドプレイのアテンション手法を開発しました。AnchorAttentionは不要なアテンション計算を削減し、意味的な一貫性を維持し、最初のトークンを一貫した位置IDを持つ共有アンカーとして扱うことで、トレーニングコンテキスト内のすべてのドキュメントに可視化します。3種類のLLMでの実験では、AnchorAttentionが長いコンテキストのパフォーマンスを大幅に向上させ、標準の完全なアテンションメカニズムと比較してトレーニング時間を50％以上短縮し、一般的なタスクにおける元のLLMの機能を維持します。当該コードはhttps://github.com/haonan3/AnchorContextで入手可能です。

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.