技術要約: StateBridge

問題提起

大規模言語モデル（LLM）ベースのマルチエージェント（MA）システムは、現在、エージェント間の通信に自然言語（離散的なトークン）に依存しています。このアプローチは「離散的なボトルネック」を導入します。すなわち、送信側は自身の連続的な内部隠れ状態を離散的なトークンへと圧縮しなければならず、その際、トークンの識別子だけでは捉えきれない情報を破棄してしまいます。受信側がこれらのトークンを埋め込み層を介してベクトルへと再マッピングする際、受信できるのはトークンの識別子のみであり、送信側の元の連続的な隠れ状態ではありません。これは情報の損失と調整能力の低下を招きます。

生の隠れ状態を直接転送することは潜在的な代替案ですが、幾何学的なミスマッチに直面します。学習済みのLLMは入力層でのトークン埋め込みを処理するように訓練されていますが、デコーダーの隠れ状態は表現空間の異なる領域を占めています。隠れ状態を直接渡すと、受信側はそのベクトルが持つ特定の幾何学的整合性を解釈できないため、意味的なミスマッチが生じます。既存の解決策は、内部状態をレイヤーごとに注入するか（移植性を制限する）、あるいは特定のモデルへの汎用性を制限する学習済みのプロジェクターを必要とします。

手法: StateBridge

著者らは、閉形式の直交変換を用いて、送信側の最終層の隠れ状態を受信側の入力埋め込み空間に整合させる、トレーニングフリーの通信インターフェースであるStateBridgeを提案しています。この手法は以下の3つのステージで動作します。

1. メッセージ抽出 (Message Extraction):
   送信エージェントがメッセージを生成します。全シーケンスを転送する代わりに、システムは生成されたトークンの最終 $K$ 個の隠れ状態 ($S$) を抽出します。また、共有語彙埋め込み行列から、対応する参照トークン埋め込み ($R$) を取得します。これらの参照埋め込みは、整合のターゲットとして機能します。

2. 整合インターフェース (Alignment Interface):
   StateBridgeの核心は、$S$ を受信側と互換性のある整列済みプレフィックス $\bar{S}$ に変換するための3ステップの変換です。

   • 直交プロクラスティス整合 (Orthogonal Procrustes Alignment): この手法は、$S$ と $R$ を $\mathbb{R}^d$ における点群として扱います。まず、グローバルなオフセットを除去し、主方向を等しくするために、両方の集合をセンタリングおよびホワイトニング（白濁化）します。次に、ペアごとの距離と角度を維持しながら、ホワイトニングされた送信側の状態をホワイトニングされた参照埋め込みへと写像する最適な回転行列 $Q^*$ を求めるため、直交プロクラスティス問題を解きます。これにより、送信側の表現の幾何学的構造（意味的類似性）が維持されます。
   • ノルム校正 (Norm Calibration): 最終層の隠れ状態は、通常、入力トークンの埋め込みよりも大幅に大きなノルムを持ちます。StateBridgeは、整列された各ベクトルを語彙埋め込みの平均 $\ell_2$ ノルムに一致するように再スケールします。これにより、プレフィックスが単に大きさによって受信側の注意スコアを支配してしまうことを防ぎます。
   • 語彙アンカリング (Vocabulary Anchoring): 整列されたベクトルが、事前学習中に遭遇した埋め込み空間の領域内に確実に存在するように、各校正済みベクトルを（コサイン類似度に基づく）最も近い語彙埋め込みに向かってわずかに移動させます。このステップは、連続性を維持しながら、表現を受信側が期待する入力分布に固定します。

3. プレフィックス注入 (Prefix Injection):
   得られた整列済みプレフィックス $\bar{S}$ は、受信側のプロンプト埋め込みの先頭に付加されます。受信側は、アーキテクチャの変更やトークン化を行うことなく、標準的な言語モデリングレイヤーを使用してこの連続的なプレフィックスを処理します。

主な貢献

1. 表現ミスマッチの解決: 本論文は、閉形式の整合のみで、表現の幾何学的ミスマッチを解決し、学習やアーキテクチャの変更なしに効果的な潜在通信を可能にできることを示しています。
2. StateBridgeフレームワーク: 著者らは、プロクラスティス整合、ノルム校正、および語彙アンカリングを組み合わせた、パラメータフリーの具体的なインターフェースを提案しています。このアプローチは、送信側の状態のペアごとの幾何学的関係を保持しながら、受信側の入力空間との互換性を確保します。
3. 実証的な優位性: 2つのファミリー（Qwen3およびOLMo3）に属する4つのモデルと8つのベンチマーク（数学、QA、コード）において、StateBridgeはテキストベースの通信およびKVキャッシュ転送のベースラインを一貫して上回りました。26のモデル・タスクのペアのうち、22個で最高スコアまたは同等の最高スコアを達成しました。

結果

• パフォーマンス: StateBridgeは、すべてのモデル設定において、最高のベースラインに対して平均2.4〜2.9パーセントポイントの改善を達成しました。
• 汎用性: 同手法は、異なるモデルアーキテクチャ間で堅牢な性能を示しました。特に、OLMo3-7B-Thinkモデルにおいて、KVキャッシュ転送（LatentMAS）はテキスト（73.9%）と比較して低い性能（平均55.1%）でしたが、StateBridgeは76.7%を達成しました。これは、StateBridgeの入力層のみの操作が、全トランスフォーマーレイヤーに状態を注入する手法よりも、アーキテクチャの違いに対して堅牢であることを示唆しています。
• タスク感度: 利得は、深い推論を必要とする困難なベンチマーク（例：MedQA、GPQA、AIME、コード生成）で最も顕著でした。これらのタスクでは、テキスト通信によって有用な連続情報が失われやすいためです。
• アブレーション研究:
  • プロクラスティス整合をリッジ回帰（非直交線形写像）に置き換えると、パフォーマンスが大幅に低下（平均82.4%から74.9%へ）し、ペアごとの幾何学的構造を保持することが極めて重要であることが確認されました。
  • ノルム校正または語彙アンカリングを除去すると、パフォーマンスも低下しました。これは、幾何学的整合だけでは、スケールおよび分布の適合性なしには不十分であることを示しています。
• ケーススタディ: 再構成タスクにより、整列されたプレフィックスには、可視的なサフィックス（接尾辞）トークンには存在しない意味的な情報（例：「koilonychia」などの特定の医学用語）が含まれていることが実証され、プレフィックスが離散的なテキスト出力を超えた情報を運んでいることが証明されました。

重要性と主張

本論文は、潜在通信における主要な課題は、単に豊かな状態を転送することではなく、それらの状態を次のエージェントにとって「読める（理解可能な）」ものにすることであると主張しています。StateBridgeは、転送される状態が情報量豊かであり、かつ受信側の入力空間と幾何学的に互換性があることを保証することで、これに対処します。

著者らは、彼らのアプローチがトレーニングフリーであり、アーキテクチャの変更を必要としないことを強調しており、市販のLLMに適用可能であることを述べています。彼らは、通信インターフェースのデザインが、同質なマルチエージェントシステムにおいて、重要でありながらしばしば見落とされている要素であると主張しています。現在の研究は共有されたモデルの重みと固定されたプレフィックス長を前提としていますが、結果は、整合ベースのインターフェースが、特に離散的なトークン化によって実質的な情報損失が生じるタスクにおいて、マルチエージェントの協調を大幅に強化できることを示唆しています。