Retrieval-Warmed Energy-Based Reasoning: A Five-Arm Ablation Methodology for Diffusion-as-Inference on Structured Reasoning Tasks

技術要約：検索によるウォームスタートを用いたエネルギーベース推論

問題提起
拡散サンプラーやエネルギーベース推論モデルなどの反復的な推論手順は、洗練（refinement）のステップ数を削減するために、過去の軌跡のメモリから候補となる解を検索して「ウォームスタート」させる手法がますます活用されている。しかし、このようなパイプラインが性能を向上させた場合、あるいは失敗した場合、その増減の具体的な要因が不明確であることが多い。性能の変化が、検索された内容が真にタスクに関連しているためなのか、それとも内容に関わらず初期化がより良い領域へとシフトしたため（バイアスシフト）なのか、あるいは単に確率的な摂動が退化した平衡状態を打破したためなのか、その判別が困難である。単一のエンドツーエンドの精度指標では、これらの混同された効果を分離することはできない。本論文はこの問題を、構造化推論システムにおける「説明可能な失敗の帰属（explainable failure attribution）」の問題として定式化し、推論パイプラインのどのコンポーネントが結果を左右しているのかを特定することを目指している。

手法
著者らは、IREDエネルギーベース拡散モデルとModern Hopfield軌跡メモリを組み合わせた「検索によるウォームスタートを用いたエネルギーベース推論（RW-EBR）」を研究している。検索の効果を分離するため、彼らはパイプラインを以下の3つの異なる失敗コンポーネントに分解する5つのアームによるアブレーション手法を提案している。

1. キーの品質 (K): エンコーダが、有用な値を検索するためのキーを適切にマップできているか？
2. ウォームスタート・メカニズム (M): 固定された品質を持つ検索値が与えられたとき、推論ループはそれを真のターゲットへと洗練できるか？
3. 格納された値の品質 (V): コールド（ウォームアップ前）のモデルは、将来のウォームスタートとして有用な予測を生成できるか？

5つのアブレーション・アームは以下の通りである：

• Aligned Oracle (整列済みオラクル): 学習されたキーエンコーダによって検索された、グラウンドトゥルース（正解）の値を使用する。
• Shuffled Oracle (シャッフル済みオラクル): グラウンドトゥルースの値を使用するが、$(key, value)$ のペアリングをクエリ間で入れ替えることで、値の分布は維持しつつ、グラフごとの整合性を破壊する。
• Best-Constant (最良定数): 固定の値をウォームスタートに使用し、バイアスシフトの影響を分離する。
• Per-Query Random (クエリごとのランダム): 確率的なウォームスタートを制御するために、ランダムな初期化（一様分布またはガウス分布）を使用する。
• Aligned Cold-Pred (整列済みコールド予測): コールド予測を格納値として使用する、実際にデプロイ可能なパイプライン。

この診断ロジックは、LLM-RAGの評価フレームワークから適応されたものであるが、ここでは反復的な拡散推論に適用されている。

実験設定
この手法は、2つの構造化推論タスクで評価される。

1. Connectivity-2: Erdős–Rényiグラフ ($N=12, p=0.2$) における全対到達可能性の予測。
2. Sudoku: 9x9のパズル解決タスク。

ベースモデルはIREDのGraphEBMである。検索キーは、学習された対照学習損失（supervised contrastive loss）を用いて訓練されたGINエンコーダ（Connectivity用）またはResNet（Sudoku用）によって生成される。メモリには、ウォームアップフェーズ中に生成された軌跡が格納される。主要な指標は均衡精度（balanced accuracy）であり、デプロイ可能なウォームスタートと見なされるためには、コールド・ベースラインに対して $\Delta$ 均衡精度 $\ge -2$ パーセントポイント（pp）という特定の受理ゲートを通過する必要がある。

主な結果

• Connectivity-2 の知見:

  • 支配的なレバーは整列（Alignment）: 「aligned-vs-shuffled-oracle」の振れ幅は、均衡精度において +35 pp に達した。これは、グラフごとの値の整列が支配的な要因であることを示しており、バイアスシフト（Best-Constantアームにより $\le +8$ ppに限定）や確率性（Randomアームはコールドの性能を下回った）による寄与を大きく上回っている。
  • 格納された値の失敗: オラクル・メカニズム（グラウンドトゥルースを使用）は大幅な利得（+22.39 pp）を得る一方で、コールド予測を使用するデプロイ可能なパイプラインは受理ゲートに失敗した（$\Delta = -4.09$ pp）。この分解により、失敗の原因はメカニズムや検索にあるのではなく、格納された値の品質 (V) にあることが明らかになった。コールドモデルの予測は正例（positive class）に偏っており、これらの偏った値でウォームスタートを行うと、誤差を修正するどころか増幅させてしまう。
  • $\beta$ トレンドの逆転: 最適な検索温度 ($\beta$) はメモリの品質に依存する。高品質（オラクル）な値は平滑化（低い $\beta$）の恩恵を受けるが、低品質（コールド予測）な値は、「最悪ではない」近傍に焦点を当てるために鋭い（peaky）検索（高い $\beta$）を好む。しかし、いずれもパフォーマンスのゲートを突破することはできなかった。

• Sudoku の知見:

  • キー品質の失敗: Sudokuに同じ診断ロジックを適用したところ、パイプラインは キーの品質 (K) コンポーネントで失敗した。キーエンコーダは、品質比率ゲート（要求値 $\ge 0.85$ に対し $\approx 0.42$）をクリアできなかった。
  • 候補プールの天井: 完全なエンコーダを適用したとしても、「バッチ内の最良」のターゲット類似度は $\approx 0.30$ でプラトー（停滞）に達しており、これは現在のセットアップにおける候補プールまたはターゲット定義の根本的な限界を示している。その結果、Sudokuにおいては、最初のコンポーネントで失敗が特定されたため、ウォームスタート・メカニズムおよび格納された値のコンポーネントについては評価されなかった。

意義と主張
本論文は、主に3つの貢献を主張している。

1. 手法: バイアスシフト、確率的ウォームスタート、および整列された値の再利用を分離する、5つのアームによるアブレーション・スイート。
2. 整列効果: Connectivity-2において、グラフごとの整列が性能の支配的なレバーであることを、alignedとshuffledのオラクル条件間の+35 ppの振れ幅を通じて実証したこと。
3. 不均質な失敗モード: 「最初の阻止コンポーネント」がタスクによって異なることを示したこと。Connectivity-2ではボトルネックは格納された値の品質（コールドモデルの予測が不十分）であったが、Sudoku（現在のセットアップ下）では、ボトルネックはキーの品質（エンコーダが十分に類似した候補を検索できない）であった。

著者らは、この分解手法が失敗または成功の責任を負う特定のコンポーネントを命名し、集計されたエンドツーエンドのスコアでは提供できない局所的な判定を下すものであることを強調している。また、部分的なノイズを加えるウォームスタート・メカニズム（SDEditで既に見られるもの）や、オラクル・メモリの上限値は、新しいプリミティブではなく、新規性はアブレーション手法とその適用による、これら特定の失敗モードの解明にあると明言している。