技術要約: PJ-RoPE: 相対アテンションのためのフーリエ・ジェット・アフィン位置空間

問題提起

相対的な位置表現は、Transformerにおける長文脈アテンションの中心的な要素であるが、既存の手法は統一された構造の構成要素としてではなく、競合する設計として扱われることが多い。

• RoPE は、フーリエ特性（回転）を介して相対的なラグをエンコードする。
• ALiBi は、アテンション・ロジットに単調なアフィン再帰項を直接加える。
• Jordan-RoPE は、半単純な回転周波数を有限のジョルダンブロックに置き換え、微分のような「有限ジェット」特徴を生成する。

本論文は、これらの手法が定数係数差分加群に根ざした共通の代数的構造を共有していると主張する。具体的には、単純な根はフーリエ/RoPE特性に対応し、重複する非ゼロの根はジョルダン/フーリエ・ジェットを生成し、重複する単位根はALiBiのようなアフィン再帰を生み出す。核心となる問題は、これらを単一の学習可能な空間内の選択可能なセクターとして扱う統一フレームワークの欠如と、特定のタスクがどのセクターを利用しているかを判断するための診断法の不在である。さらに、高次のジェット特徴（相対距離の累乗を含む）は、長文脈において安定性の課題を突きつけ、変換された特徴量やロジットを劇的に増大させる。

手法

著者らは、これらのプリミティブを学習可能なフーリエ・ジェット・アフィン相対位置空間へと統合するフレームワークである PJ-RoPE (Poincaré-type RoPE) を提案する。

1. 理論的定式化

このフレームワークは、定数係数差分加群の代数に基づいている。位置空間 $P_{PJ}$ を以下の3つのセクターの直和として定義する：

• 斉次フーリエ・ジェット・セクター ($H_{FJ}$): RoPE（次数-0）およびJordan-RoPE（次数-$r > 0$）を含む。代数的には、これらは重複する非ゼロの根を持つ方程式の解であり、$(d/L)^r e^{-cd/L} e^{i\omega d}$ のような項をもたらす。
• アフィン再帰セクター ($A_{rec}$): 定数項および線形項（例：$-sd/L$）を含み、これはローレンツ群のアフィン拡張に対応する「並進のような」方向を提供する（ゆえに「ポアンカレ型」と呼ぶ）。
• ライトコーン (LC) セクター ($H_{LC}$): 高次挙動に対処するためのコンパクト化された座標チャート。生の距離 $d$ をラピディティ（速力）に似た位相 $\phi_L(d) = L \text{asinh}(d/L)$ と、飽和振幅 $\beta_L(d) = d/\sqrt{d^2 + L^2}$ に置き換える。

フレームワークは2つの実装レジームを区別する：

• PJ-bias: フーリエ特性、減衰ジェット、アフィン項、およびLC変種を含むスカラー加算アテンション・カーネル。これは学習可能であるが、厳密な特徴変換ではない。
• PJ-rotary: クエリおよびキーの特徴に対して相対的な作用を適用する厳密な特徴変換実装であり、RoPEおよびJordan-RoPEの閉包性を検証するために使用される。

2. 適応的診断

タスクが空間のどの領域を選択するかを決定するために、著者らは適応的セクターゲート ($g_{FJ}, g_{aff}, g_{LC}$) と3つの診断法を導入する：

• 有効質量 (Effective Mass): ゲート化されたパラメータの大きさが、各セクターおよびジェット次数にどのように割り当てられているかを測定する。
• 関数エネルギー (Functional Energy): 減衰や位相の打ち消しを考慮に入れつつ、評価ウィンドウ全体におけるカーネル成分の実効的なスケールを測定する。
• Leave-One-Order-Out (LOO): 特定のジェット次数を除去することで適合関数が劣化するかどうかを評価し、パラメータの選択と実際の機能的寄与を分離する。

3. ライトコーンによる安定化

長文脈における高次ジェットの指数関数的な増大を防ぐため、LCブランチは座標変換を適用する。これにより、変換された特徴量およびロジットの大きさが制限されるが、長距離における位相分解能が圧縮される。これは、安定性と分解能の間のトレードオフを明示的に定式化している。

主な貢献

1. 位置空間の定式化: 相対アテンションを、スカラーバイアスカーネルと厳密な回転特徴変換を分離した、学習可能なフーリエ・ジェット・アフィン空間として定式化した。
2. セクターの包含: RoPE、Jordan-RoPE（有限ジェット）、ALiBi（アフィン再帰）、およびLC座標を、単一の数学的構造内の異なるセクターとして特定した。
3. 適応的診断: セクターゲート、有効質量、関数エネルギー、およびLOOアブレーションの導入により、タスクレベルでのセクター選択を経験的に測定する方法を提供した。
4. エビデンスの連鎖: 制御されたカーネル、合成タスク、バイトレベル言語モデル、および記譜音楽トークンストリームを用いた、セクターの回復、学習可能な使用、および自然なタスクへの割り当てを評価した。

実験結果

実験では、短文脈で学習され、長文脈（最大32,768トークン）で評価される小規模な因果的Transformer（約0.15M–0.20Mパラメータ）を使用している。

• 固定カーネルの回復: スカラーPJ-biasは、純粋な位相、有限ジェット、アフィン、およびLCターゲットを正常に回復し、この空間が設計されたプリミティブ・ファミリーを網羅していることを確認した。
• 適応的セクターの回復: 「教師」信号（例：純粋なアフィンまたは特定のジェット次数）を用いて学習した際、学習されたゲートと診断法は、対応するセクターへと正しく質量をシフトさせた。
• 合成シーケンスタスク: 学習可能なアテンションモデルは、位置ブランチが信号を保持している場合にこれらのセクターを利用できる。アフィン教師は再帰ブランチを活性化させ、符号付きジェット教師はFJ変種を活性化させる。
• 自然言語 (バイトレベル): 言語タスク（Tiny Shakespeare, WikiText-2, War and Peace）において、モデルは主にアフィン/再帰の挙動を選択する。高次のFJ/LCの質量は、軽微な診断的補正としてのみ現れる。NTKスタイルのRoPEスケーリングとアフィン再帰の組み合わせが、最良の長文脈性能を示す。
• 音楽トークンストリーム: 記譜音楽タスク（MAESTRO, MusicNet）では、割り当てが異なる。LC-アフィン変体は依然として強力であるが、高次補正（小さくても非ゼロの質量）への一貫した、測定可能な割り当てが見られる。これは、モチーフの回帰やリズムのエンベロープをモデル化する必要性と一致している。
• LCの安定性 vs 分解能: LC変体は、32,768トークンにおいてQ/Kノルムおよびキャッシュ/ロジットのスケールを正常に制限する。しかし、これには代償が伴う。「遠距離の位相スパン」が大幅に縮小するため、生の座標と比較して、ハードネガティブ・プローブにおける検索分解能が低下する。

意義と主張

本論文は、PJ-RoPEを単なる新しいアーキテクチャとしてではなく、位置空間の診断フレームワークとして位置づけている。その主な意義は、相対的な位置メカニズムを、異なるタスクが異なる領域を選択する、統一された学習可能な空間として再定義したことにある：

• タスク依存の選択: 実証的な証拠は、言語タスクはアフィン/再帰構造を好む一方で、構造化された記譜ストリーム（音楽など）は、アフィンと高次ジェット補正の混合を利用することを示唆している。
• 安定性と分解能のトレードオフ: ライトコーン座標によって導入されるトレードオフを明示的に定量化している。これらは成長を制限することで長文脈への外挿を安定させるが、位相分解能を圧縮する。
• 関心の分離: 「表現の正確性」（厳密なPJ-rotaryの閉包性）と「長文脈の安定性」（スカラーLCパス）を区別し、厳密な特徴変換は安定化なしでは長距離で急激に劣化する可能性があることを指摘している。

著者らは、安定した長文脈の使用は、アフィン再帰、フーリエ・ジェット構造、およびLCコンパクト化の間のタスク依存の割り当てに依存すると結論付けている。彼らは、これらの知見が小規模なモデルと特定のトークナイゼーション・スキームに基づいていることを認め、これらの割り当てをデプロイメント・ルールとして扱う前には、より大規模な検証が必要であるとして、一般化については慎重な姿勢を保っている。