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

現代の大規模言語モデル（LLM）は主にTransformer (Vaswani et al., 2017) アーキテクチャに基づいている。その中核的な構成要素の一つがアテンション機構であり、以下のように表現できる：

ここで、アテンションロジット行列 $\mathbf{A}\triangleq(A_{ij})\in R^{T\times T}$、クエリとキーベクトルはそれぞれ $\mathbf{q}_{i}=\mathbf{W}_{Q}\mathbf{x}_{i}$ と $\mathbf{k}_{j}=\mathbf{W}_{K}\mathbf{x}_{j}$ の形をとる。$\mathbf{W}_{Q},\mathbf{W}_{K}\in\mathbb{R}^{d\times d}$ はクエリとキー行列であり、$\mathbf{x}_{i}$、$\mathbf{x}_{j}$ は $\mathbf{X}\in R^{T\times d}$ の $i$ 番目と $j$ 番目のトークンを表す。アテンションスコア行列 $\text{ATTN}(\mathbf{X})$ において、我々は表記を簡略化するために Vaswani et al. (2017) によって導入されたスケーリング係数 $1/\sqrt{d}$ を無視している。 因果マスク $\mathbf{M}^{-\infty}\in\mathbb{R}^{T\times T}$（ここで $\mathbf{M}^{-\infty}_{ij}=-\infty$ は $i<j$ の場合、そうでなければ $\mathbf{M}^{-\infty}_{ij}=0$）は、$\mathbf{A}$ の上三角部分を負の無限大に設定することで、各トークンが将来のトークンにアテンションを向けることを防いでいる。

Transformerアーキテクチャは、順序情報をエンコードするために位置埋め込みを使用しており、これは順序データを処理する上で不可欠である。RoPE (Su et al., 2021) は現在、特にLLMにおいて最も広く採用されているエンコーディングである。 RoPEはクエリとキーに作用し、それらを2次元のチャンクに分割し、各チャンクを異なる周波数で回転させる。これにより、式 1 のアテンションロジットの計算に回転行列を適用し、以下のように表現できる：

ここで、$m=j-i$ は $i$ と $j$ の相対距離である。回転行列 $R$ の詳細は付録 A に記載する。 RoPEは、絶対位置エンコーディングとして実装することで、効率的な相対位置エンコーディングを実現している。これは、キーとクエリベクトル、具体的には $R_{i,\theta}q_{i}$ と $R_{j,\theta}k_{j}$ に直接回転行列を適用することで、絶対位置エンコーディングの操作を模倣している。 なお、回転行列はFlashAttention2 (Dao, 2024) モジュールの外部でFloat32精度を使用して適用される。しかし、キーとクエリの内積はFlashAttention2モジュール内で行われ、FlashAttention2の要件に従ってBFloat16にキャストされなければならない。 また、回転角の選択は $\theta_{i}=base^{-2i/d}$ を満たし、現在のLLMにおける典型的な $base$ の値は10,000である。

RoPEは理論上、トークン間の相対的な位置距離$m=j-i$のみに依存するように設計されている。 重要な含意は、すべての位置インデックスに定数の位置シフト$\Delta$を加えても、注意計算が変化しないはずだということである。形式的には、以下のように表現できる：

ここで、$A_{ij}$は位置$i$と$j$の間の注意ロジットであり、$R_{m,\theta}$は相対距離$m$に対応する回転行列である（導出は付録A参照）。 しかし、重要な疑問は、この相対的性質が実際に成り立つのかということである。

我々は、位置シフトが注意計算に与える影響を経験的に調査する。 分析を明確にするために、注意計算を$\text{ATTN}^{l,h}(\mathbf{X},\Delta)$と再定義する。これは、位置シフト$\Delta$が適用された場合の$l$番目のトランスフォーマー層の$h$番目のヘッドによって生成される注意行列を表す。具体的には、元の位置インデックスが$[0,1,\dots,L-1]$である場合、シフト適用後は$[\Delta,\Delta+1,\dots,\Delta+L-1]$となる。 我々は、2つの異なる位置シフト$\Delta_{1}$と$\Delta_{2}$による注意の差を以下の指標を用いて測定する：

この指標は、すべての層とヘッドにわたる注意スコアの累積差を測定する。 $t_{\text{row}}$と$t_{\text{col}}$は注意行列の行と列を反復する。 そして、正規化ベクトル$\mathbf{n}=\left[\frac{1}{L},\frac{1}{L-1},\dots,1\right]$が要素ごとに適用される（$\odot$で表記）。これは、因果マスキングにより下三角行列の要素数が変化することを考慮するためである。

実験設定。 図1（左）では、$\Delta_{2}=16$を設定し、$\Delta_{1}$をリスト$[0,2,4,6,8,10,12,14,15,17,18,20,22,50,100,200,500,1000,2000]$にわたって変化させる。 青線はBFloat16精度のLLaMA-2-7Bの事前学習済みパラメータを使用した結果を表す。黄線は同じ事前学習済みパラメータをFloat32精度に変換した結果を示す。緑線はBFloat16精度でランダムに初期化されたパラメータ（PyTorchのデフォルトであるkaiming_uniform_で初期化）を使用した結果に対応する。各ケースにおいて、長さ$T=4096$の50のテキストシーケンスにわたって平均をとり、差$D$を計算する。 なお、可視化を改善するために$\Delta_{1}=16$を切り取っている。予想通り、$\Delta_{1}$と$\Delta_{2}$が等しい場合、差はゼロとなる。より詳細な可視化と議論は付録Bに提供されている。

結果の考察。 我々は、BFloat16精度を使用する場合、位置シフト$\Delta$が事前学習済みのLLaMA-2-7Bの注意計算に影響を与えることを観察した。対照的に、すべての計算をFloat32精度で行う場合、この効果は消失する（青線と緑線の比較）。さらに、事前学習済みパラメータとランダムに初期化されたパラメータを比較すると、どちらもBFloat16精度下で（青線と黄線の比較）、事前学習が不一致を増幅させることがわかった。 これらの発見は、BFloat16精度下では、RoPEが理論的に主張されている相対的位置エンコーディングから（わずかに）逸脱し、モデルが事前学習された後にこの性質が崩壊することを示唆している。

さらに、我々は各個別のトークンに対する$\Delta_{1}=0$と$\Delta_{2}=16$の間の注意の差の詳細を調査する（つまり、トークンごとの研究のために我々の指標で$t_{\text{row}}$にわたる合計を行わない）。これらのトークンごとの差を図1（中央）に可視化する。結果は明確に、最初のトークンが注意の差に最も大きく寄与していることを示している。最初のトークンを除外すると、残りのトークンは相対的位置の性質を保持している。

前の実験では、シーケンス長を$T=4,096$に固定した。シーケンス長が注意ロジットの差にどのように影響するかを調査するために、我々の研究を$[64,128,256,512,1024,2048,4096,8192]$のシーケンス長に拡張した。この試行では、異なる長さにわたるソフトマックス操作が最初のトークンの結果を比較不可能にするため、注意スコアではなく注意ロジットを測定した。明確にするために、注意ロジットの差を以下のように定義する：

ここで、$A_{ij}^{l,h}(\Delta)$は位置シフト$\Delta$を持つ$l$番目の層の$h$番目のヘッドの注意ロジット行列である。 我々は$\Delta_{1}=0$と$\Delta_{2}=16$を保持し、異なる文脈長を持つ50のシーケンスにわたって平均をとった指標$D_{\text{logit}}$を測定する。 図1（右）にこれらの結果をまとめる。シーケンス長が増加するにつれて、最初のトークンの注意ロジットの差も増加することが観察される。これは、より長いシーケンスが最初のトークンの不一致を悪化させることを示している。

長文脈モデルの学習には通常、長い配列の処理が必要である。しかし、この直接的な方法はBFloat16下での精度誤差を悪化させる。図1（右）が示すように、処理されるトークン数が増加するにつれて、不一致は悪化する。 配列長を短くすることで精度の問題を軽減できるが、短いウィンドウサイズで長文脈モデルを学習することは矛盾しているように思える。 この矛盾を解消するために、我々は文献の技術を再検討し、文書間の注意を遮断する文書内注意(Zhao et al., 2024b)（図2左に示す）が、完全な注意と比較してより短い配列長で長文脈モデルを学習するのに使用できることを発見した。さらに、経験的に、文書内注意はLLaMA-3シリーズのようなオープンソースLLMで成功裏に適用されている。Gao et al. (2024)による最近の研究も、文書境界を越えた注意を遮断することが短文脈と長文脈の両方のパフォーマンスを向上させることを示し、文書内注意の有効性を確認している。

以前、我々はBFloat16を使用する際のRoPEの相対位置エンコーディングにおいて、注意スコアの差異として現れる偏差を特定した。この観察は次の疑問を提起する：この不一致は長文脈性能に重大な影響を与えるのか？特に、文書内注意が注意スコアの不一致を緩和するように見えることを考慮すると、文書内注意で訓練されたモデルにおいては、BFloat16の影響は長文脈能力に対して無視できるかもしれない。これを調査するために、我々は2種類の位置インデックスを比較する。1つ目は、文書内注意の一般的な実践Zhao et al. (2024b); Gao et al. (2024)に従い、シーケンスの始めから終わりまで連続的に位置IDを割り当てる方法である（図2左）。2つ目は、各文書内で位置インデックスを1にリセットする方法である（図2中央）。 理論的には、両方の位置付け方式は同一の結果をもたらすはずである。 この比較は、相対位置エンコーディングの偏差がさらに長文脈性能に測定可能な差異をもたらすかどうかを評価するために使用される。

結果の考察。 結果は図3に示されている。我々は、位置IDをリセットすることが一貫してモデルの長文脈性能を向上させることを観察した。対照的に、位置IDをリセットしない従来のアプローチ（図2左）は、劣った性能をもたらす。この性能差は、各文書の最初のトークンに異なる位置IDを割り当てることで生じる可能性があり、これがモデルを混乱させる可能性のある不整合を引き起こす。位置IDをリセットすることで、一貫した位置関係（最初の位置IDが最初のトークンにある）を維持する。 なお、この分析はRoPEがBFloat16精度下で何らかの形で絶対位置エンコーディングメカニズムとして機能するという仮説に基づいている。この仮定を厳密に検証し、その意味を完全に理解するためには、さらなる研究が必要である。

位置IDのリセットはモデルのパフォーマンスを向上させますが、重大な欠点も伴います：モデルは、コンテキスト長に達するか超える配列を処理する場合にのみ、回転角の全スペクトルを学習できるのです。この制限により、コンテキストウィンドウサイズを増やすにつれて、全コンテキストウィンドウを埋める十分な長い配列を収集することが、そのような長大なデータの希少性のため現実的でなくなるため、モデルがより長いコンテキスト長のシナリオに一般化する能力が妨げられます。

この課題に対処するため、我々はAnchorAttentionを提案します。これは、長いコンテキストウィンドウ内のすべての文書に可視的な共有共通アンカートークンを導入する注意機構です。さらに、文書間の冗長な注意接続は無視され、モデルが一貫性のある情報に集中できるようになります。AnchorAttentionの図解は図2（右）に示されています。

AnchorAttentionでは、アンカートークン$\mathscr{A}$がすべての文書の共通の開始点として機能します。この設計は、最初のトークンがRoPEにおける相対位置エンコーディングからの偏差の原因であるという我々の観察に動機づけられています。最初のトークンを除いて、後続のトークンは望ましい相対位置エンコーディングの特性を保持できます。固定位置ID（開始ID）を持つ共有アンカートークンを導入することで、モデルは文書の冒頭と開始位置IDの関係について混乱しなくなります。 具体的には、我々は文頭トークン<bos>をこの共通アンカーとして指定します。<bos>トークンは通常、明示的な意味を持たないため、このトークンに集中するBFloat16による精度誤差は、モデルの全体的なパフォーマンスに最小限の影響しか与えません。

さらに、共有アンカーを使用することで、モデルが文の冒頭と開始位置IDの関係を正しく解釈することを保証するために位置IDをリセットする必要がなくなります。これにより、ウィンドウ長の最初から最後まで連続した位置ID（図2の右に示すノンリセット方式）を使用できます。結果として、各イテレーションで回転角の全スペクトルを学習できます。文書内注意アプローチと比較して、本稿の方法は、全学習ウィンドウ長を埋める豊富な長い配列に大きく依存せずに、全回転スパンを学習できます。これにより、長い配列データの収集とアップサンプリングへの依存が軽減されます。

ベースモデル。 我々の実験では、主にLlaMA-2-7Bモデルをベースモデルとして使用する。提案手法の有効性を異なる事前学習モデルにわたって評価するため、Llama-3-8B (Dubey et al., 2024)、Qwen-1.5 (1.8B) (Yang et al., 2024)、およびMistral-7B-v0.3 (Jiang et al., 2023)についても第5.3節で評価する。これらのモデルは、アーキテクチャとベースモデルの事前学習パラダイムの範囲を代表しており、我々のアプローチの一般性をテストすることを可能にする。

データセット。 我々は、長文脈訓練のためにSlimPajamaデータセット(Soboleva et al., 2023)を使用する。これはLLaMAの事前学習データ混合(Touvron et al., 2023)のオープンソースレプリケーションである。このデータセットには、Common Crawl (CC)、C4、GitHub、ArXiv、Books、Wikipedia (Wiki)、StackExchangeなど、多様なテキストソースが含まれている。オリジナルのSlimPajamaに加えて、Fu et al. (2024)のデータ混合手法を適用し、元のデータセットの全体的な分布を維持しながら、各ソース内の長いシーケンスをアップサンプリングする。 このアップサンプリングは、訓練中にモデルをより長文脈シナリオに露出させることを目的としている。

我々の実験では、オリジナルとアップサンプリングされたSlimPajamaデータセットの両方から20億トークンをサンプリングする。64Kおよび128Kトークンのシーケンス長を持つデータセットをそれぞれSlimPajama-64KおよびSlimPajama-128Kと呼ぶ。アップサンプリングされたデータセットはUpSampledMix-128Kと表記する。詳細な統計は付録Cに記載されている。

訓練構成。 我々の訓練ハイパーパラメータは主に(Zhang, 2023)に基づいている。すべてのモデルは8台のNVIDIA A100 GPUで訓練される。学習率を$2\times 10^{-5}$に設定し、重み減衰0.1、$\beta_{1}=0.9$、および$\beta_{2}=0.95$のAdamWオプティマイザを使用する。各モデルは2000ステップ訓練され、これは20億トークンのデータセットに対して約1エポックに相当する。バッチサイズは8に設定され、64Kコンテキストでは1バッチあたり50万トークン、128Kコンテキスト長では100万トークンに相当する。 表1に訓練構成をまとめる。

適切な指標を用いた長文脈能力の測定。 我々が提案する注意機構がベースモデルの長文脈能力を向上させる効果を決定的に評価するためには、堅牢で適切な評価指標を使用することが重要である。

パープレキシティ（PPL）は長文脈言語モデルの評価によく用いられるが、最近の研究ではその長文理解能力の評価における信頼性に疑問が投げかけられている。Hu et al. (2024)は、PPLが主に局所的な情報の捕捉を測定するため、モデルの長距離依存関係の理解能力との相関が低いことを示している。さらに、Fang et al. (2024)は、PPLが長文脈入力の理解に不可欠な重要なトークンを見落とし、信頼性の低い評価につながるという経験的証拠を提供している。また、Gao et al. (2024)は、長文データの量を増やすとPPLは改善するものの、長文データのみを使用すると下流の長文脈タスクのパフォーマンスが大幅に低下することを発見している。これらの観察と一致して、我々も初期の訓練ステップ後にPPLが変化しない一方で、RULER ベンチマーク(Hsieh et al., 2024)のパフォーマンスは向上し続けることを発見した（図4参照）。これはさらに、PPLが長文脈パフォーマンスの向上を適切に反映していない可能性を示している。

HELMET (Yen et al., 2024)、LongBench (Bai et al., 2023)、LongBench-Cite (Zhang et al., 2024a)、InfiniteBench (Zhang et al., 2024c)、NoCha (Karpinska et al., 2024)などの既存の長文脈言語モデル用ベンチマークは、指示追従能力に大きく依存しており、長文脈ベースモデルの評価には適していない。さらに、LongBench (Bai et al., 2023)のタスクは16Kトークンの文脈長で十分に対応できる。16Kトークンは中程度の長さの文脈を表すが、64Kあるいは128Kトークンを扱える我々のようなモデルを十分に挑戦させるものではない。 これらの制限にもかかわらず、我々はLongBenchで我々のモデルを評価した。なぜなら、それが実世界の長文脈タスクを含んでおり、そのfew-shot in-context learning (ICL)タスクがベースモデルの評価に適しているからである。ベースモデルの評価におけるfew-shot ICLの有効性は、Gao et al. (2024)の同時期の研究でも観察されている。

しかし、13のタスク全てにわたる全体的なパフォーマンスを測定する公式評価とは異なり、我々はLLaMA-2-7Bモデルに対して2つのタスク（NIAH-Single 3とNIAH-Multikey 3）を除外している。この除外は、長文脈訓練が主にモデルの既存の能力を元の文脈ウィンドウからより長いシーケンスに拡張するという仮定と、LLaMA-2-7Bが元の4K文脈ウィンドウ内でもこれら2つのタスクに苦戦していたという観察に基づいている。したがって、モデルが元のウィンドウ内で実行できなかったスキルを拡張された文脈で新たに獲得することを期待するのは不合理である。（LLaMA-2-7BとLLaMA-2-7B-Chatの13タスク全てにわたる詳細な結果と包括的な議論は付録Dに記載されている）。なお、クロスアーキテクチャ評価セクション（セクション5.3）の他の高度なモデル（LLaMA-3-8B、Mistral-7B-v0.3、Qwen-1.5-1.8B）については、13タスク全てを使用している。

さらに、単一のチェックポイントでRULERのパフォーマンスを報告する以前の研究とは異なり、我々は最後の50訓練ステップにわたって10ステップごとにチェックポイントを保存し、5つのチェックポイントにわたる平均RULERパフォーマンスを報告することを提案する。異なるランダムシードで5回モデルを再訓練することも、パフォーマンスを平均化する合理的な方法であるが、実験コストが大幅に増加する。複数のチェックポイントにわたる平均化の必要性を示すために、我々は16Kウィンドウサイズの元のSlimPajamaで標準的な全注意を用いて訓練されたLLaMA-2-7BのRULERパフォーマンスを評価した。図4に示すように、16K RULERベンチマークのパフォーマンスは訓練中に変動を示している。特定のチェックポイントを恣意的に選択することを防ぐために、我々は将来の研究が複数のチェックポイントにわたる平均パフォーマンスを報告することを提案する。

LLLMの一般的能力。 継続的な事前訓練がモデルの初期文脈長内での中核的な能力を保持していることを確認するために、我々はMMLU (Hendrycks et al., 2021)とHellaSwag (Zellers et al., 2019)ベンチマークを使用する。MMLUはSTEM、人文科学、社会科学、哲学、法律、医学を含む57の多様な科目にわたる知識と推論を評価し、一般的な理解の広範な尺度を提供する。HellaSwagは、モデルにシナリオの最も妥当な続きを選択させることで常識的推論をテストし、文脈的予測スキルに挑戦する。 これらのベンチマークは基礎的能力を評価するための標準的な選択肢であり、我々はベースモデルとチャットモデル間でパフォーマンスの差がほとんどないことを観察しており、これらがベースモデルの能力評価に適していることを支持している。

長文脈学習に最適な位置埋め込みを見出すため、我々はLLaMA-2-7BをSlimPajama-64Kデータセットで訓練した。このデータセットは64Kトークンの系列に分割されている。我々は3つの手法を評価した：RoPE (Su et al., 2021)、NTK-aware (LocalLLaMA, 2023)、およびYaRN (Peng et al., 2023)である。(Men et al., 2024)に従い、埋め込み公式$\theta_{i}=\text{base}^{-2i/d}$（式6）のbase値を変化させた。これはTransformersのrope-thetaパラメータに相当する(Wolf et al., 2020)。 NTK-awareとYaRNは相対スケール$s=T_{\text{new}}/T_{\text{origin}}$に依存する。ここで$T_{\text{origin}}$は元の文脈サイズ（4K文脈長）、$T_{\text{new}}$は拡張サイズ（64K文脈長）であるため、我々は$s=16$と設定した。 (Fu et al., 2024)に従い、長文脈能力の学習には1Bトークンで十分であることが示唆されているため、我々は各訓練で1Bトークンを使用し、訓練時間と複数モデルの評価のバランスを取った。

表2は、RULERベンチマークにおける4Kから64Kトークンまでの性能を示している。適切なbase値を持つ通常のRoPEが訓練文脈長内で最も良い性能を示し、64K文脈長でNTK-awareとYaRN手法を上回っている。 訓練文脈長を超えた性能（汎化）を測定するため、我々はモデルを128K RULERベンチマークで評価した（表4）。YaRN埋め込みは未見のより長い文脈に対してより良く汎化する。一方、通常のRoPEは訓練された長さ内では良好な性能を示すが、より長い文脈では性能が低下する。 RoPEのbase値の重要性を認識し、我々は計算資源を節約するため32K文脈長でモデルを訓練し、rope-thetaの選択をさらに研究した。表3に示すように、性能は特定のtheta値でピークに達し、それ以上ではわずかな利益しか得られず、thetaをさらに増加させると性能がわずかに低下する。これは(Liu et al., 2023b; Men et al., 2024)の観察と一致している。 これらの知見に基づき、我々は文脈長に応じてbase thetaの値を適応させる：32Kに対して1M、64Kに対して5M、128Kに対して10M、256Kに対して50Mである。このアプローチにより、異なる文脈長にわたって最適な性能を確保できる。 表2: 通常のRoPEで訓練された長文脈モデルは、拡張された文脈ウィンドウサイズ内で評価した場合、NTKやYaRNを使用したモデルよりも優れた性能を示す。$\text{Free}^{\star}$は訓練なしで直接評価されたモデルを示す。 Base 64K 32K 16K 8K 4K $\text{Free}^{\star}$ RoPE 10K - - - - 85.29 YaRN 10K 0.18 1.09 6.02 18.38 29.06 RoPE 1M 66.89 70.79 77.02 84.64 90.42 5M 68.81 72.68 81.03 85.79 89.28 10M 69.43 71.01 81.01 85.28 88.13 YaRN 1M 64.22 65.69 73.28 79.98 85.02 5M 63.48 64.47 75.29 79.90 84.01 10M 58.91 62.03 70.68 75.62 83.04 NTK 10K 47.32 53.05 60.62 72.42 81.61 1M 62.47 70.49 75.87 83.48 86.41 5M 61.45 68.52 75.90 78.47 85.17 10M 58.75 65.95 72.27 77.43 83.99 表3: 適切な$base$は長文脈学習にとって重要である。大きすぎる値や小さすぎる値を使用すると性能が低下する可能性がある。 Base 32K 16K 8K 4K 平均 500 8.52 16.43 21.18 37.06 20.80 10K 15.21 28.79 43.61 83.03 42.66 200K 68.91 76.67 85.64 88.00 79.81 600K 74.61 81.79 85.03 89.36 82.70 900K 75.05 81.62 85.22 88.65 82.64 5M 76.09 77.15 85.66 87.00 81.48 1B 73.97 76.48 81.09 82.33 78.47 表4: 拡張された文脈ウィンドウを超えて評価すると、YaRNが優位性を示す。 Base 128K 64K RoPE 1M 35.93 (-30.95) 66.88 5M 53.98 (-14.83) 68.81 YaRN 10K 27.25 (-22.84) 50.09 1M 57.07 (-7.