One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States

この論文は、LLM の隠れ状態から軽量な投影ヘッドを用いて直接検索埋め込みを生成する手法を提案し、従来の「生成後エンコード」パイプラインを不要にしながらも、QReCC ベンチマークで基線モデルの検索品質の 97% を維持できることを示しています。

要約

この論文は、**「AI が検索をするとき、なぜわざわざ『翻訳』と『辞書』を 2 回も使う必要があるのか？」**という疑問から始まります。

通常、AI（大規模言語モデル）が何かを検索するときは、以下の 2 つのステップを踏みます。

1. 思考: AI が「何を探せばいいか」を文章（クエリ）として考え出す。
2. 翻訳: その文章を、別の専門の AI（埋め込みモデル）に渡して、「検索用の数字のリスト（ベクトル）」に変換する。

この論文の著者は、**「待てよ、AI はすでに『何を探すべきか』を頭の中で完璧に理解しているはずだ。わざわざ文章に変えてから、また別の AI に翻訳させるなんて、無駄な手間じゃないか？」**と考えました。

そこで提案されたのが、**「AI 自身に検索能力を内蔵させる」**というアイデアです。

🌟 簡単な比喩で説明

この仕組みを、**「天才的な料理人（AI）」と「レシピ翻訳者（埋め込みモデル）」**の例で考えてみましょう。

🍳 従来の方法（2 つのモデルを使う）

1. 料理人（AI）が「今日はパスタを作りたいな」と考えます。
2. 料理人はその考えを「パスタのレシピ」という文章に書き出します。
3. その文章を、別の専門家の「レシピ翻訳者（埋め込みモデル）」に渡します。
4. 翻訳者は「パスタのレシピ」を、倉庫の棚番号（検索用ベクトル）に変換して、必要な食材を探し出します。

問題点: 料理人はすでに「パスタを作りたい」という完璧なイメージを持っています。それを一度文章にして、また別の人が翻訳するのは、**「料理人が考えたことを、一度メモに書き写し、それを別の人が読み取って翻訳する」**ようなもので、非常に時間と手間がかかります。

🚀 新しい方法（この論文の提案）

1. 料理人（AI）が「今日はパスタを作りたいな」と考えます。
2. 料理人の頭の中（隠れ層）には、すでに「パスタのイメージ」が完璧に詰まっています。
3. そこで、料理人の頭に**「小さな変換器（投影ヘッド）」**という小さな装置を取り付けます。
4. この装置は、料理人の「頭の中のイメージ」を直接、倉庫の棚番号（検索用ベクトル）に変換します。

メリット:

• 不要な工程の削除: 「文章に書き出す」→「翻訳する」という手間がなくなります。
• 超高速: 翻訳者が不要になったので、検索が劇的に速くなります（実験では約 22 倍速くなりました！）。
• 性能はほぼ同じ: 翻訳者がいなくても、料理人の頭の中のイメージを直接変換すれば、97% の精度で同じ食材が見つかります。

🔑 重要なポイント

• 「3 つの魔法の薬」: この「小さな変換器」を教えるために、3 つの異なる学習方法（アライメント損失、コントラスティブ損失、ランク蒸留）を組み合わせました。
  • 例えるなら、**「先生（既存の翻訳モデル）の答えに合わせる」「似たものは近くに、違うものは遠くにする」「先生が選んだ順位を真似する」**という 3 つのルールで訓練しています。
• 結果: 従来の方法と比べて、検索の精度はわずかに（3% 程度）落ちましたが、**「翻訳者（追加の AI モデル）を完全に不要にした」**という大きな勝利を収めました。

💡 まとめ

この研究は、**「AI には、自分の考えを直接検索に使える能力が最初から備わっている」**ことを証明しました。

これまでは、AI が検索するたびに「文章に変換して、別の AI に渡す」という**「二重の作業」を行っていましたが、これからは「AI 自身が、頭の中のイメージを直接検索に使える」**ようになります。

これは、**「通訳を介さず、直接現地の言葉で交渉できる」**ようなもので、AI システムをよりシンプルで、高速で、効率的にする大きな一歩です。

技術概要

以下は、提示された論文「One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States」の技術的サマリーです。

1. 問題設定と背景

従来のリトリエーショナル・ジェネレーション（RAG）システムにおける LLM エージェントの検索プロセスは、以下の 2 段階のモデルパイプラインが主流でした。

1. 検索クエリの生成: LLM が文脈に基づいて自然言語の検索クエリを生成する。
2. エンベディング化: 生成されたテキストを、別途用意された埋め込みモデル（Embedding Model）に入力し、ベクトル化してドキュメント検索を行う。

課題点:
この 2 モデル構成は、LLM がすでに生成プロセスの中で「ユーザーの意図、対話履歴、タスク要件」を内部の隠れ状態（Hidden States）にエンコードしているにもかかわらず、その情報を一度テキストとして出力し、それを別のモデルで再度ゼロから処理（再エンコード）させるという根本的な冗長性を含んでいます。これにより、インフラの複雑化と推論時のレイテンシ（遅延）が増大しています。

2. 提案手法：ネイティブ検索埋め込み

著者は、追加の埋め込みモデルを不要にし、LLM エージェント自体にネイティブな検索能力を持たせる手法を提案しました。

核心となるアプローチ:
LLM の生成プロセスで既に計算されている「隠れ状態（Hidden States）」を、軽量な**投影ヘッド（Projection Head）**を介して直接、既存の埋め込み空間へマッピングします。これにより、生成されたテキストをエンコードするステップを完全に排除します。

技術的詳細:

• 隠れ状態の抽出: LLM のオートレグレッシブ生成過程において、各トークンの最終層の隠れ状態 $h_i$ を収集します（特別なトークンを除く）。これらは生成時に既に計算されているため、追加の計算コストはほぼゼロです。
• 投影ヘッドのアーキテクチャ:
  • 入力投影: 線形層で LLM の隠れ次元を内部次元へ変換。
  • Transformer エンコーダ: 位置埋め込みを加え、複数層の Transformer エンコーダでトークン間の依存関係を統合。
  • プーリング: 有効な位置に対して平均プーリングを行い、固定次元のベクトルへ圧縮。
  • 出力投影と正規化: 最終的な埋め込み次元へ変換し、L2 正規化を適用（ドット積による検索がコサイン類似度と等価になるように）。
• 学習目的（3 つの損失関数の組み合わせ）:
  既存の埋め込みモデル（教師モデル）からの知識蒸留を用いて、以下の 3 つの損失を同時に最適化します。
  1. アライメント損失 (Alignment Loss): 予測ベクトルと教師モデルの埋め込みベクトル間の角度距離を最小化（直接的な位置合わせ）。
  2. コントラスト損失 (Contrastive Loss): バッチ内の異なるクエリ間の相対的な構造を保持するよう促す（InfoNCE 損失）。
  3. ランク蒸留損失 (Rank Distillation Loss): 教師モデルが持つドキュメントのランキング順序（相対的なスコア）を KL 発散を用いて転送。

3. 実験設定

• データセット: 会話型検索ベンチマーク「QReCC」を使用（346 の会話、2,189 の検索トリガー）。
• モデル構成:
  • LLM エージェント: Qwen3-8B
  • 教師埋め込みモデル: Qwen3-Embedding-8B
  • 同一ファミリー（Same-family）の設定を採用。
• ベースライン: 従来の「生成→エンコード」パイプライン（LLM 生成テキストを埋め込みモデルへ投入）。

4. 主要な結果

• 検索品質: ベースラインと比較して、Recall@10 で 97%、MRR@10 で 96.4% の性能を維持しました（ベースラインの 97% 相当）。
• レイテンシ: 埋め込みモデルの推論パスを削除したことで、推論遅延が 43.5ms から 2.0ms へ劇的に減少し、21.8 倍の高速化を達成しました。
• 統計的有意性: McNemar 検定により、品質の差は統計的に有意（p=0.0005）ですが、実用上は極めて近い性能であることが示されました。

5. 重要な知見（アブレーション研究）

• 損失関数の相乗効果:
  • アライメント損失単体では最も高い性能を示しますが、コントラスト損失やランク蒸留損失を組み合わせることでさらに向上します。
  • ランク蒸留損失単体では性能が崩壊（Recall@10 ≈ 0）しますが、アライメント損失で空間の幾何学的基盤を確立した上で組み合わせることで、ランキング構造の微調整に寄与します。
• 学習レシピの重要性:
  • 学習率とエポック数の組み合わせが極めて重要です。学習率が高すぎると学習が崩壊し、80 エポックという長期学習と低い学習率（2×10⁻⁴）の組み合わせが最良の結果をもたらしました。これは、LLM の隠れ空間と埋め込み空間の間の微妙な幾何学的関係を学習するには、慎重な最適化が必要であることを示唆しています。

6. 貢献と意義

1. 冗長性の解消: LLM 検索パイプラインにおける「生成→再エンコード」という冗長なプロセスを理論的に指摘し、隠れ状態の直接投影による代替を提案しました。
2. インフラの簡素化: 推論時に埋め込みモデルを不要にすることで、システム構成を単純化し、レイテンシを大幅に削減しました。
3. 実用性の証明: 12 種類のアブレーション設定と統計的検証を通じて、追加モデルなしでも基線モデルと同等に近い検索品質を達成可能であることを実証しました。

7. 限界と今後の課題

• データセットの限定: 現在は QReCC のみでの評価であり、他のドメインやオープンドメイン検索への汎化性は未検証です。
• 学習フェーズの依存: 学習時には教師モデル（埋め込みモデル）が必要であり、推論時のみ不要になります。
• クロスファミリー課題: 異なるモデルファミリー間（例：LLM と全く異なる系列の埋め込みモデル）での投影は、共有事前学習表現がないためより困難である可能性があります。

結論:
この研究は、LLM エージェントが自身の内部表現を活用することで、外部の埋め込みモデルに依存せず高速かつ高精度な検索を実現できる可能性を示しました。特に、推論コストの削減とシステム設計の簡素化において、RAG システムの新たなパラダイムを提示する重要な貢献と言えます。