Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings

技術要約：極座標位置エンベディングによる「何（What）」と「どこ（Where）」のデカップリング

問題提起
Transformerアーキテクチャにおいて、アテンション機構は「内容（何）」と「シーケンス内の位置（どこ）」という2つの要素に基づいて、キーをクエリにマッチングさせます。著者らは、Llama 3やDeepSeekといった最先端モデルで使用されている広く採用された回転位置エンベディング（RoPE）が、これら2つの要素を絡み合わせていると主張しています。RoPEでは、キーおよびクエリベクトルに適用される回転が、それらの内容の大きさ（マグニチュード）と相対的な位置の両方に依存します。この絡み合いは「what-where」の混同を生み出し、内容によるマッチングまたは位置によるマッチングのいずれか一方のみに独立して依存するルールをモデルが学習することを妨げます。具体的には、RoPEのアテンションスコアにおける相互作用項により、回転の位相がコンテンツの特徴によって影響を受けるため、純粋に位置インデックスやコンテンツのマッチングに依存するルールを学習することが困難になります。

手法：極座標位置エンベディング（PoPE）
著者らは、キーとクエリの表現を極座標形式に再定式化することで、コンテンツと位置を分離するRoPEの修正案であるPoPEを提案しています。

• 再定式化: RoPEでは、コンポーネントは位置に基づいて回転しますが、初期位相はコンテンツによって決定されます。PoPEでは、著者らは実数値のキーおよびクエリベクトルを以下の性質を持つ複素ベクトルへと変換します：
  • 大きさ ($\mu$): 非負性を保証するためにソフトプラス活性化関数 ($\sigma(x) = \ln(1 + e^x)$) を用いて、コンテンツ特徴からのみ導出されます。
  • 位相 ($\phi$): シーケンス位置 ($s$ または $t$) とコンポーネント固有の周波数 ($\theta_c$) によって厳密に決定されます。
• アテンションスコア: アテンションスコアは、クエリの共役転置とキーとの内積の実部として計算されます。数学的には、以下のように記述されます：
  $$a^{PoPE}_{ts} = \sum_{c=1}^{d} \mu_{\tilde{q}_{tc}} \mu_{\tilde{k}_{sc}} \cos((s - t)\theta_c + \delta_c)$$
  ここで、$\delta_c$ は学習可能なバイアス項です。
• RoPEとの主な相違点:
  1. デカップリング（分離）: RoPEにおいてコンテンツと位置を混合させる相互作用項 ($\phi_{k_{sc}} - \phi_{q_{tc}}$) が排除されています。大きさ（コンテンツ）と位相（位置）は厳密に分離されています。
  2. 周波数解像度: PoPEは $d/2$ 組のペアではなく、$d$ 個のコンポーネント（個々の要素）に対して動作するため、事実上、周波数チャネルの数が倍増します。
  3. 学習可能なバイアス: 学習可能なバイアス $\delta_c$ を導入することで、コンテンツと位置の混同を再導入することなく、各周波数の最適な相対オフセットを調整できます。

主な貢献

1. 診断タスク（間接インデックス作成）: 著者らは、ソース文字からの相対的なオフセットに基づいてターゲット文字を特定する必要があるタスク（例：「'z'の3つ左の位置にある文字を見つけよ」）を導入しました。このタスクは、コンテンツと位置の独立した操作を必要とします。
2. 理論的洞察: 論文では、RoPEがいかにコンテンツと位置を混同させるかを示す代数的な導出を提供すると同時に、PoPEがいかにアテンション機構が独立した「what」と「where」のマッチングの結合としてルールを実装できるかを示しています。
3. 効率的な実装: 標準的なアテンションと比較して、フル複素行列を実体化することなく、複素内積の実部を計算する、TritonとFlash Attention 2を用いた効率的な実装が提供されています。これは、単一の乗算オーバーヘッドのみを追加します。

結果

• 間接インデックス作成: 診断タスクにおいて、RoPEは平均精度がわずか11.16%にとどまり、必要な独立したマッチングの学習に苦戦しました。対照的に、PoPEは94.82%の精度を達成し、これら2つの要素を分離する能力を実証しました。
• シーケンスモデリング（音楽およびゲノミクス）:
  • 音楽 (JSB および MAESTRO): PoPEは両方のデータセットにおいて、負の対数尤度（NLL）においてRoPEを上回り、より低いパープレキシティを達成しました。
  • ゲノミクス (ヒト参照ゲノム): PoPEはRoPEのベースラインと比較して、NLLの大幅な低下を示しました。
• 言語モデリング (OpenWebText): モデルスケール（124M, 253M, 774M パラメータ）全体にわたり、PoPEは一貫してRoPEよりも低いパープレキシティを達成しました。性能差はモデルサイズとともに安定、あるいは拡大しました。
• ダウンストリームタスク: 6つのダウンストリームタスク（LAMBADA, BLiMP, CBT, HellaSwag, PIQA, ARC-E）におけるゼロショット評価において、PoPEベースのモデルは、すべてのモデルサイズにおいてRoPEベースのモデルよりも高い平均精度を達成しました。
• 長さの外挿（Length Extrapolation）:
  • ゼロショット: PoPEは、微調整や周波数補間を行うことなく、訓練コンテキスト（1024トークン）の最大10倍の長さ（10240トークン）のシーケンスに対して、強力なゼロショット長外挿能力を示しました。
  • YaRNとの比較: 外挿のために設計された手法であるYaRNは、微調整の限界（4096トークン）まではRoPEの劣化を緩和しますが、それを超えると大幅に劣化します。PoPEは、YaRNが必要とする補間技術や広範な微調整を必要とすることなく、長いシーケンス上での性能を維持しました。
  • 周波数使用分析: 特徴量ノルムの分析により、RoPEモデルはノイズを避けるために高周波チャネルを抑制する傾向があるのに対し、PoPEモデルは層を横断して高周波を含む、より分散された範囲の周波数を利用していることが明らかになりました。

意義と主張
本論文は、PoPEがRoPEに内在する「what-where」の絡み合いに対する原理的な解決策を提供すると主張しています。コンテンツと位置を分離することで、PoPEはデータの効率性、漸近的な精度、そして最も重要なことに、コンテキスト長の汎化性能を向上させます。著者らは、これらの利点が既存のRoPEアーキテクチャへの軽微な修正によって達成され、複雑なアーキテクチャ変更を必要としないことを強調しています。

本研究は、シーケンスモデリングに関する基礎的な研究として提示されています。著者らは、この手法が音楽、ゲノミクス、自然言語などの領域で能力を向上させる（タンパク質折り畳みや創薬への応用を助ける可能性がある）一方で、有害なコンテンツの生成や高度なディープフェイクの作成といった、基盤モデルに共通するデュアルユース（両義的利用）のリスクを伴うことを述べています。主要な貢献は、コンテンツ情報と位置情報を分離することが、より堅牢で汎用性の高いTransformerモデルにつながることを示した点にあります。