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

要約

コンテキスト ウィンドウ サイズを拡張すると、大規模言語モデル (LLM) でより長いシーケンスを処理し、より複雑なタスクを処理できるようになります。
Rotary Positional Embedding (RoPE) は、ロングコンテキストのトレーニングに役立つ相対位置エンコーディング プロパティにより、事実上の標準になっています。
ただし、RoPE を BFloat16 形式で使用すると数値的な問題が発生し、特にロングコンテキストのシナリオで、意図された相対位置エンコーディングから逸脱することが観察されています。
この問題は BFloat16 の精度の制限から発生し、コンテキストの長さが増加するにつれて累積し、最初のトークンがこの問題に大きく寄与します。
これに対処するために、BFloat16 によって引き起こされる数値的な問題を軽減し、ロングコンテキスト機能を改善し、トレーニングを高速化するプラグアンドプレイのアテンションメソッドである AnchorAttendance を開発しました。
AnchorAttention は、最初のトークンを一貫した位置 ID を持つ共有アンカーとして扱い、トレーニング コンテキスト内のすべてのドキュメントに表示できるようにすることで、不必要なアテンションの計算を削減し、セマンティックな一貫性を維持し、計算効率を高めます。
3 種類の LLM の実験では、AnchorAttendant が、一般的なタスクに関する元の LLM の機能を維持しながら、標準のフル アテンション メカニズムと比較して、ロングコンテキストのパフォーマンスを大幅に向上させ、トレーニング時間を 50\% 以上短縮することが実証されました。
私たちのコードは https://github.com/haonan3/AnchorContext で入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google