技術要約: PaTH Attention

1. 問題提起

アテンション機構は現代の大型言語モデル（LLM）の基盤であるが、本質的に置換不変（permutation-invariant）である。言語のような構造化されたドメインをモデリングするには、位置エンコーディングが必要となる。Rotary Position Embedding (RoPE) がデファクトスタンダードとなっているが、これには決定的な限界がある。それは、キーとクエリに適用される変換が**データに依存しない（data-independent）**ことである。RoPEにおいて、2つのトークン間の相互作用は、トークンの内容ではなく、単にそれらの相対的な距離のみに依存する。

このデータ非依存性は、RoPEベースのトランスフォーマーの表現力を制限する。理論的分析によれば、RoPEベースのモデルは$TC^0$ 計算複雑性クラス（通常の絶対位置埋め込みを持つトランスフォーマーと同等）に留まり、シーケンシャルな推論や状態追跡（例：フリップフロップ言語モデリング、状態追跡問題）を必要とするタスクを解く能力が制限される。その結果、これらのモデルは入力シーケンスへの動的な適応を要求される合成タスクにおいて苦戦することになる。

2. 手法: PaTH Attention

著者らは、柔軟でデータ依存的な位置エンコーディングスキームであるPaTH (Position Encoding via Accumulating Householder Transformations) を提案している。

コアメカニズム

固定された回転行列を適用するRoPEとは異なり、PaTHは、クエリとキーの位置間のパスに沿ったデータ依存の累積的なハウスホルダー変換の累積積を用いてアテンション・ロジットを計算する。

• アテンション・ロジットの定式化: クエリ $q_i$ とキー $k_j$ の間のアテンション・スコアは次のように定義される：
  $$A_{ij} \propto \exp\left(k_j^\top \left(\prod_{s=j+1}^{i} H_s\right) q_i\right)$$
• データ依存の遷移行列 ($H_s$): シーケンス内の各行列 $H_s$ は、ハウスホルダー型の単位行列プラスランク1の行列である：
  $$H_s = I - \beta_s w_s w_s^\top$$
  ここで、$w_s$ と $\beta_s$ は入力トークン $x_s$ の関数である。具体的には、$\beta_s = 2 \times \text{sigmoid}(u^\top x_s + b)$ であり、これにより遷移行列が負の固有値を持つことを可能にしている。これは状態追跡性能において極めて重要である。

理論的な表現力

本論文は、PaTHがトランスフォーマーの理論的な能力を拡張することを確立している。

• 複雑性クラス: 2つのアテンションヘッドを持つ1層のPaTHトランスフォーマーは、$AC^0$ 還元下で$NC^1$ 完全問題を解くことができる。これは、RoPEベースのモデルを制約している $TC^0$ の壁をPaTHが打破できることを示している。
• 線形RNNとの関係: PaTHは、表現力の高い線形RNN（DeltaNetなど）と密接に関連しており、ソフトマックス・アテンションの連想想起（associative recall）の利点を保持しつつ、状態を追跡する能力を継承していることが示されている。

拡張: PaTH-FoX

著者らは、PaTHが他のアテンション修正手法と互換性があることを示している。彼らは、PaTHの乗法的位置エンコーディングと、Forgetting Transformer (FoX) の加法的でデータ依存の忘却ゲートを組み合わせたPaTH-FoXを提案している。このハイブリッド・アプローチは、特に長さの外挿（length extrapolation）において性能をさらに向上させる。

3. 効率的な学習と推論

PaTHにおける大きな課題は、フルサイズの $L \times L$ アテンション行列を実体化することなく、ハウスホルダー行列の累積積を効率的に計算することである。

• UT変換による表現: 著者らは、ハウスホルダー行列の積をコンパクトに表現するためにUT変換を利用している。$L$ 個の変換の積は、$P = I - W^\top T^{-1} W$（ここで $W$ はベクトル $w_s$ を含み、$T$ は $\beta$ 値から導出される三角行列である）として表現できる。
• FlashAttentionスタイルのアルゴリズム: ハードウェア効率の高い学習を可能にするため、著者らはブロック単位のアルゴリズムを開発した：
  • 各ブロックに対して「境界調整済み」のクエリおよびキー行列を事前計算する。
  • キー/バリューのブロックに対して右から左へのスキャンを行い、ストリーミング形式で累積的なハウスホルダー積を蓄積する。
  • このアプローチは、動的な位置エンコーディングを組み込みながら、FlashAttentionのI/O効率（SRAM内での動作）を維持する。
• 推論: 推論においては、現在のタイムステップの遷移行列を使用して履歴のキーをインプレースで更新するため、変換ベクトルの個別のキャッシュを保持する必要がない。これにより、既存の推論カーネル（FlashDecodingなど）との互換性が確保される。

4. 実験結果

合成ベンチマーク

PaTHは、シーケンシャルな推論と状態追跡を調査するために設計されたタスクで評価された：

• フリップフロップ言語モデリング (FFLM): PaTHベースのモデル（1層、2ヘッド）はこのタスクでほぼ完璧な精度を達成したが、RoPEベースのモデルは著しく苦戦した。
• 数理問題 ($A_5$ 群): PaTHは、ベースライン（RoPE, SBA, FoX）よりも少ないレイヤー数で、$NC^1$ 完全な数理問題を解いた。
• マルチクエリ反復連想想起 (MQRAR-N): Paプトは、$N$-back想起（$N < 4$）において変数値を正常に追跡し、順序付けられた履歴の追跡に失敗した最近のベースラインを上回った。

言語モデリング (760M パラメータ)

Fineweb-Eduコーパス（50B トークン）を用いて実験が行われた：

• 標準ベンチマーク: PaTHは、ゼロショットの常識推論タスク（LAMBADA, PIQA, HellaSwagなど）においてRoPEおよびFoXを一貫して上回り、Wikitextにおいてより低いパープレキシティを達成した。
• 長さの外挿: 4Kコンテキストで訓練されたモデルは、より長いシーケンスに対しても良好に汎化された。
  • RoPE: 4Kを超えると急激に失敗した。
  • PaTH: 32Kまで安定した性能を維持した。
  • PaTH-FoX: 状態追跡が重要なコード生成（CodeParrot）において特に強みを示し、64Kトークンまで堅牢に汎化された。
• 長文脈ベンチマーク: RULER, BABILONG, PhoneBookにおいて、PaTH-FoXは最高の検索および状態追跡スコアを達成し、Multi-Needle-In-A-Haystackおよび変数追跡タスクにおいてRoPEおよびFoXを大幅に上回った。

モデル変換

著者らは、学習済みのRoPEモデルをPaTHへ変換する検討を行った：

• SmolLM2-1.7B: 高品質な数学/コードデータにさらされる前にモデルを変換し、その後に継続学習を行った場合、PaTHは元のRoopeモデルよりも数学（GSM-8K）およびコーディング（HumanEval, MBPP）のベンチマークで優れた性能を示した。
• Qwen2.5-7B: すでに過学習（overtrained）されたRoPEモデルからの蒸留では、一部の性能向上は見られたものの、特定のベンチマークで低下も見られ、混合した結果となった。これは、変換はモデルが「硬化（ossify）」する前に行うのが最も効果的であることを示唆している。

5. 意義と主張

本論文は、データ依存の乗法的変換を導入することで、現在の位置エンコーディングスキームの理論的および経験的な限界に対処すると主張している。

• 表現力: トランスフォーマーを $TC^0$ 計算複雑性クラスを超えて進展させ、$NC^1$ 完全問題を解くための原理的な方法を提供する。
• 実用性: UT変換とブロック単位のアルゴリズムを通じて、PaTHは標準的なアテンションメカニズムに匹敵する効率性を達成しており、大規模な学習への適用が可能である。
• 状態追跡: ハウスホルダー変換のデータ依存性により、モデルは入力シーケンスに動的に適応でき、RoPEベースのモデルでは困難な状態追跡タスクを解決できる。

6. 限界

著者らは以下のいくつかの注意点を認めている：

• 数値的安定性: 学習は精度に対して敏感になる可能性があり、特にBF16では累積積において不安定になることがあり、スケーリング係数 $\beta$ のクリッピングが必要となる。
• 計算オーバーヘッド: 効率的ではあるものの、ハウスホルダー積の計算におけるヘッド次元の増大は、標準的なアテンションと比較して計算およびメモリのオーバーヘッドを増加させる。
• 幾何学的バイアス: PaTHはネイティブには回転をモデル化しない（ハウスホルダー行列は反射であるため）。これは、RoPEが持つような、相対位置に対する構造的な依存性といった特定の幾何学的帰納バイアスを制限する可能性がある。ただし、著者らは将来の研究において、反射の合成を通じて回転を近似できる可能性を示唆している。