技術的サマリー：Give it Space！エンコーダにおける位置情報と意味表現の明示的解離

問題定義

位置エンコーディング（PE）は、順列不変なトランスフォーマーがシーケンスの順序を表現する上で不可欠であるが、位置情報がどのように処理・保存されるかの内部メカニズムは未だ十分に理解されていない。現代のアプローチ、特にロータリー位置エンコーディング（RoPE）などの相対的位置エンコーディング（RPE）手法は、長文脈の理解や検索タスクにおいてしばしば困難に直面する。さらに、既存の RPE 手法は長期減衰などのヒューリスティックに依存しており、これが位置バイアスを生み出し、性能のボトルネックとなっている可能性がある。訓練されたトランスフォーマーにおいて位置信号と意味信号がほぼ直交する部分空間を占めるという証拠がある一方で、標準的なアーキテクチャはこれらの信号を絡み合わせ、それぞれの役割を特定したり独立して最適化したりすることを困難にしている。著者らは、これらのストリームを明示的に分離することで、メカニズムに関する洞察が得られ、より優れた位置表現へとつながると仮定している。

手法

著者らは、意味、絶対位置（AP）、**相対位置（RP）**の 3 つの情報ストリームを明示的に解離する新しいアーキテクチャ、DSTG-NeoBERTを提案する。

アーキテクチャの修正

NeoBERT アーキテクチャに基づき、モデルは以下の個別の修正を通じてこれらのストリームを処理する：

1. 入力表現：各トークンは、$d_{AP}$ 次元の AP 埋め込みと $d_{sem}$ 次元の意味埋め込みという、2 つの独立した埋め込みによって表現される。隠れ層の総サイズは $d_{model} = d_{AP} + d_{sem}$ である。
2. アテンションメカニズム：
   • クエリとキーは、各アテンションヘッド内で各ストリームに対して独立して計算される。
   • アテンションログイットは、相対位置（RP）バイアス、意味アテンション重み、および AP アテンション重みの和である。
   • バリューは各ストリームに対して別々に計算され、それぞれの空間内に留まる。
   • 重要な設計選択として、分散の一致を確保し、高次元空間が支配的になるのを防ぐため、両ストリームをヘッドあたりの同じサイズに射影する。
3. レイヤー正規化：相互依存を避けるため、AP 埋め込みと意味埋め込みに対して、アテンションの前後に個別の RMSNorm レイヤーを適用する。
4. フィードフォワードネットワーク（SwiGLU）：意味埋め込みと位置埋め込みは連結され、隠れ空間に射影される。決定的な点として、ダウン射影は情報を個別の AP ストリームと意味ストリームに再分離する。ストリームが直接相互作用するのはこのコンポーネントのみであり、これにより意味条件付きの構造的更新が可能となる。
5. 訓練目的（MLM デコーディング）：マスク言語モデル（MLM）デコーディングヘッドは、最終隠れ状態の意味部分のみに適用される。最終層の AP 成分は MLM 目的からの勾配を受けないため、位置部分空間は予測圧力から実質的に解放される。
6. ベースライン：著者らは、DSTG-NeoBERT を、同一データ（FineWeb）で同一ハイパーパラメータで訓練された 3 つのベースラインと比較する：
   • RoPE-NeoBERT：ロータリー位置エンコーディングを使用。
   • AP-NeoBERT：学習された絶対位置埋め込み（絡み合った状態）を使用。
   • RP-NeoBERT：学習された相対位置バイアスを使用。

主要な発見と貢献

1. 低周波構造多様体の出現

解離モデルにおける孤立した AP 部分空間は、自発的に低周波の 2 次元多様体へと収束する。

• 次元削減：学習された AP 埋め込み行列に対する主成分分析（PCA）により、最初の 2 つの主成分が全空間分散の**94.1%**を説明することが明らかになった。
• 周波数分析：タイプ II 離散コサイン変換により、これらの成分が低周波で動作し、スペクトルパワーのほぼすべてを最初の 4 つのビンに集中させていることが確認された。
• 対比：これに対し、絡み合った AP ベースライン（AP-NeoBERT）は高次元のまま留まり、高周波信号に支配されており、最初の 2 つの主成分で説明される分散は 23.4% に過ぎない。

2. アテンションヘッドの専門化

解離アーキテクチャは、アテンションヘッドがほぼ専ら 2 つのグループに特化する明確な機能的分類を示す：

• 構造指向ヘッド：AP 情報に強く影響を受ける。
• 意味指向ヘッド：意味内容によって駆動される。
• RP の役割：著者らは純粋な RP 指向ヘッドは見つからなかった。代わりに、RP バイアスは意味指向ヘッドのための局所的な支援メカニズムとして機能する（例：自己アテンションが遠隔文脈に焦点を当てないようにする）一方、AP は明確な構造情報ストリームとして機能する。

3. 巨視的構造の保持

標準的な位置エンコーディングは、巨視的文書構造を頑健に保持することに失敗する：

• 絡み合った AP：AP-NeoBERT は中間層で構造情報をエンコードするが、MLM 予測目的からの干渉（位置次元が意味予測に転用される）により、最終層ではこの情報の大部分を失う。
• 相対的手法（RoPE/RP）：これらの手法は、MLM 事前訓練下では隠れ状態において巨視的文書構造をエンコードすることに大きく失敗する。
• 解離された解決策：位置を孤立した部分空間にオフロードし、MLM ヘッドから除外することにより、DSTG-NeoBERT はネットワーク全体、最終層を含めて構造情報（文や段落の境界）を保持する。

4. プロビング結果

著者らは、3 つの階層レベルと Flash-Holmes ベンチマークに対して線形プロブを用いてモデルを評価した：

• トークンレベルおよびセグメントレベルの AP：DSTG-NeoBERT は最終層において絶対位置およびセグメントレベルの位置に対して高い $R^2$ スコアを維持するのに対し、絡み合った AP および相対的手法のベースラインは著しい劣化を示す。
• セグメント内位置：興味深いことに、セグメント内の位置情報は、すべてのモデルにおいて意味部分空間から取得可能であり、この特定の局所情報は AP ストリームを必要としないことを示唆している。
• 言語的忠実度（Flash-Holmes）：解離アーキテクチャは、5 つの分野（形態論、統語論、意味論、談話、推論）にわたる 65 の言語現象のうち49においてプロビング結果を改善する。特に、統語論、意味論、談話においてすべてのベースラインを上回る。

意義と主張

本論文は、位置ストリームと意味ストリームを明示的に解離することで、トランスフォーマー内部の「クリーンなメカニズム的研究」を提供すると主張している。主な意義は以下の点にある：

1. 位置情報は本質的に低次元かつ構造的であり、意味予測と競合させない場合、単純な多様体を形成する。
2. 標準モデルにおける構造情報の喪失は、最終層において MLM 目的が位置次元を上書きすることに直接起因する。
3. これらのストリームを結合解除することで、巨視的文書構造を保持するより豊かな埋め込みが可能となり、標準的な下流タスクのパフォーマンスを低下させることなく（GLUE、MTEB、SQuAD の結果はベースラインと同等）、言語表現が向上する。

著者らは、このアプローチが意味表現と干渉しない位置エンコーディングの設計に対する新たな方向性を提供し、RAG における位置非依存の意味表現のキャッシングや、低次元 AP 多様体の操作による長さ外挿などの実用的な応用を可能にする可能性があると示唆している。ただし、これらの知見は現在、エンコーダのみのモデルと、比較的小規模なスケール（6 レイヤー、220 億トークン）に限定されている点に留意が必要である。