技術要約: Reason What Matters (ReWAM)

問題提起

Universal Multimodal Embeddings (UME) は、多様な検索タスクをサポートするために、モダリティを越えて統一された表現を学習することを目指しています。近年の手法では、複雑で構成的かつ微細な検索タスクを扱うために、Chain-of-Thought (CoT) 推論を UME に統合する試みが進んでいますが、以下の2つの決定的な制限が、コーパス規模での展開を阻害しています。

1. 粗いクレジット割り当て (Coarse Credit Assignment): 既存の手法は、多くの場合、検索ベースの報酬を用いた Group Relative Policy Optimization (GRPO) を採用しています。このアプローチは、軌跡内のすべての CoT トークンに対して単一のスカラー・アドバンテージを割り当てます。その結果、入力された証拠によって真に支持されている推論ステップと、幻覚（ハルシネーション）や無関係なステップを区別することができません。また、ポジティブなターゲットとハードネガティブを分かつ具体的な差異を明示的に特定することもできません。この粒度の欠如が、特に分布外 (OOD) タスクにおける性能向上を制限しています。
2. 推論レイテンシ (Inference Latency): 現在のアプローチは、部分的なトレースですでに検索に十分な証拠が得られている場合であっても、通常、埋め込みを生成する前に完全な CoT トレースを生成します。この自己回帰的なフルトレースの生成は、トレースの長さとコーパスのサイズの両方に比例して増大する大幅なレイテンシを導入し、明示的な推論を大規模システムにおいて非実用的なものにしています。

手法: Reason What Matters (ReWAM)

ReWAM は、推論の品質と推論効率を同時に最適化するように設計された、検索に基づいた推論フレームワークです。これは、固定された CoT 条件付き埋め込み器（embedder）が安定した検索信号を提供する、デカップルされた「推論器（reasoner）– 埋め込み器」アーキテクチャに基づいています。本フレームワークは、以下の2つのコアコンポーネントを導入しています。

1. 検索認識型自己蒸留 (Retrieval-aware Self-Distillation: RASD)

RASD は、トークンレベルの教師信号を精緻化するための「特権的なガイダンス (privileged guidance)」を構築することで、粗いクレジット割り当ての問題に対処します。

• 特権情報の構築: マルチモーダル解析器が、ポジティブなターゲットと検索されたハードネガティブを対照させ、2種類の特権情報を抽出します：(1) 根拠となる証拠 (Grounded Evidence)（マルチモーダル入力によって支持される事実）および (2) 決定境界 (Decision Boundaries)（ターゲットと混同しやすい候補を分かつ差異）。
• 主張の検証 (Claim Verification): 解析器は、サンプリングされた CoT 軌跡をこの特権情報と照合し、支持された主張と矛盾する記述を特定することで、「対照的証拠特権情報 (Contrastive Evidence Privileged Information: CEPI)」を形成します。
• トークンレベルの再重み付け: オンポリシーのセルフティーチャーは、この CEPI を使用してサンプリングされた軌跡の再スコアリングを行います。これは、特権ガイダンスがある場合とない場合での各トークンの対数尤度差を計算します。この差は、ストップグラディエント信号として機能し、トークン固有の重みを生成します。
• 目的: これらの重みは、検索報酬から導出される軌跡レベルのアドバンテージを調整します。これにより、方策の更新が、軌跡全体を一様に強化するのではなく、入力によって支持され、かつ識別的な推論ステップに集中することを保証します。

2. 検索適応型推論 (Retrieval-adaptive Inference: RAI)

RAI は、CoT の長さを適応的に制御し、トークン生成を加速させることで、推論レイテンシに対処します。

• 検索信頼度ヘッド (Retrieval Confidence Head): 軽量な MLP ヘッドが、部分的な Co-T の「残りの検索有用性 (remaining retrieval utility)」を予測します。これは、推論トレースを継続することによって得られる追加の検索ゲインと、停止することによるゲインを推定します。
• 早期打ち切り (Early Truncation): 推論時、予測される残りの有用性が閾値を下回った場合、CoT 生成は直ちに打ち切られ、不毛な推論ステップを回避します。
• 投機的デコーディング (Speculative Decoding): 有用な継続生成を加速するために、ReWAM は投機的デコーディングを採用しています。セミ自己回帰的なドラフトモデルが複数のトークンを並列に提案し、それを凍結されたターゲット推論器が検証します。これにより、明示的な CoT の品質を維持しながら、1トークンあたりのデコーディングコストを削減します。

主な貢献

1. ReWAM フレームワーク: 検索の品質と推論の効率性の間の溝を埋める、新しい検索に基づいた推論フレームワークであり、大規模コーパスに対するスケーラブルなパラダイムとして、明示的な推論を確立します。
2. 検索認識型自己蒸想 (RASD): 検索証拠の検証を利用して、入力に支持された推論と幻覚を区別する、トークンレベルのクレジット割り当て手法です。これにより、スカラーの軌跡報酬を超えた精緻な学習信号を実現します。
3. 検索適応型推論 (RAI): 予測された残りの有用性による軌跡レベルの打ち切りと、投機的デコーディングによるトークンレベルの加速を組み合わせ、検索性能を損なうことなくレイテンシを大幅に削減する適応型推論戦略です。

実験結果

MMEB-V2（汎用マルチモーダル検索）および MRMR（推論集中的な検索）ベンチマークにおいて広範な実験を実施しました。

• 検索性能: ReWAM は両方のベンチマークにおいて最先端（SOTA）の性能を達成しています。
  • MMEB-V2 において、ReWAM (Qwen2-VL-2B) は総合スコア 66.5 を達成し、従来の最良の明示的 CoT 手法である Embed-RL を上回り、潜在的推論手法である PLUME を 4.9 ポイント上回ると同時に、明示的な CoT の解釈可能性を維持しています。
  • MRMR において、ReWAM は推論集中的なタスクでトップスコアを記録し、顕著な差（例：RIME に対して +3.1 ポイント、Embed-RL に対して +6.5 ポイント）でリードしており、数学、物理学、工学などのドメイン固有のタスクにおいても強力な性能を示しています。
• 推論効率: ReWAM は大幅なスループット向上を実現しています。
  • 競合する明示的 CoT UME 手法（例：Embed-RL）と比較して、最大 5倍のスループット を達成しています。
  • 効率性を重視した潜在的推論手法である PLUME を、検索性能とスループットの両面で上回っています。
  • アブレーション研究により、検索信頼度ヘッドが Co-T の長さを 26–36% 短縮し、投機的デコーディング単独と比較してスループットをさらに 26–64% 向上させることが確認されました。

重要性

本論文は、明示的な推論はレイテンシと粗い最適化信号のために、大規模な展開にはコストがかかりすぎるとこれまで軽視されてきたと主張しています。ReWAM は、推論を検索フィードバックに根ざさせ（RASD を通じて）、推論を適応的に制御する（RAI を通じて）ことにより、明示的な CoT を極めて効果的かつ効率的にできることを証明しています。これらの結果は、推論強化型の UME が、明示的なテキストトレースによる解釈可能性と構成的推論能力を犠牲にすることなく、大規模なコーパス上で動作する実世界のシステムにとって実用的であることを示唆しています。