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

Technical Summary: Retrieval-Warmed Energy-Based Reasoning

Problem Statement
Iterative inference procedures, such as diffusion samplers and energy-based reasoning models, are increasingly "warm-started" using candidate solutions retrieved from a memory of past trajectories to reduce refinement steps. However, when such pipelines improve or fail, the specific source of the gain or loss is rarely clear. It is ambiguous whether performance changes stem from the retrieved content being genuinely task-relevant, from a shift in the initialization toward a better region regardless of content (bias shift), or simply from stochastic perturbations breaking a degenerate equilibrium. A single end-to-end accuracy metric cannot disentangle these confounded effects. This paper frames the issue as one of explainable failure attribution for structured-reasoning systems, aiming to localize which component of a reasoning pipeline drives an outcome.

Methodology
The authors study Retrieval-Warmed Energy-Based Reasoning (RW-EBR), which combines an IRED energy-based diffusion model with a Modern Hopfield trajectory memory. To isolate the effects of retrieval, they propose a five-arm ablation methodology that decomposes the pipeline into three distinct failure components:

1. Key Quality (K): Does the encoder map inputs to keys that retrieve useful values?
2. Warm-Start Mechanism (M): Given a retrieved value of fixed quality, does the inference loop refine it toward the true target?
3. Stored-Value Quality (V): Does the cold (unwarmed) model produce predictions useful as future warm-starts?

The five ablation arms are:

• Aligned Oracle: Uses ground-truth values retrieved by the learned key encoder.
• Shuffled Oracle: Uses ground-truth values but with $(key, value)$ pairings permuted across queries, destroying per-graph alignment while preserving the value distribution.
• Best-Constant: Uses a constant value for warm-starting, isolating the effect of bias shift.
• Per-Query Random: Uses random initialization (uniform or Gaussian) to control for stochastic warm-starting.
• Aligned Cold-Pred: The deployable pipeline using cold predictions as stored values.

This diagnostic logic is adapted from LLM-RAG evaluation frameworks but applied here to iterative diffusion inference.

Experimental Setup
The methodology is evaluated on two structured reasoning tasks:

1. Connectivity-2: Predicting all-pairs reachability on Erdős–Rényi graphs ($N=12, p=0.2$).
2. Sudoku: A 9x9 puzzle solving task.

The base model is IRED's GraphEBM. Retrieval keys are generated by a GIN encoder (for connectivity) or a ResNet (for Sudoku) trained with supervised contrastive loss. The memory stores trajectories populated during a warm-up phase. The primary metric is balanced accuracy, with a specific acceptance gate requiring a $\Delta$ balanced accuracy $\ge -2$ percentage points (pp) relative to the cold baseline to be considered a deployable warm-start.

Key Results

• Connectivity-2 Findings:

  • Dominant Lever is Alignment: The "aligned-vs-shuffled-oracle" swing reaches +35 pp in balanced accuracy. This indicates that per-graph value alignment is the dominant factor, far exceeding the contribution of bias shift (bounded at $\le +8$ pp by the best-constant arm) or stochasticity (random arms fell below cold performance).
  • Stored-Value Failure: While the oracle mechanism (using ground truth) achieves a massive gain (+22.39 pp), the deployable pipeline using cold predictions fails the acceptance gate ($\Delta = -4.09$ pp). The decomposition reveals that the failure is not in the mechanism or retrieval, but in Stored-Value Quality (V). The cold model's predictions are biased toward the positive class, and warm-starting with these biased values amplifies the error rather than correcting it.
  • $\beta$ Trend Inversion: The optimal retrieval temperature ($\beta$) depends on memory quality. High-quality (oracle) values benefit from smoothing (lower $\beta$), while low-quality (cold-pred) values prefer peaky retrieval (higher $\beta$) to focus on the "least bad" neighbor, yet neither clears the performance gate.

• Sudoku Findings:

  • Key-Quality Failure: Applying the same diagnostic logic to Sudoku, the pipeline fails at the Key Quality (K) component. The key encoder fails to clear the quality ratio gate ($\approx 0.42$ vs. required $\ge 0.85$).
  • Candidate Pool Ceiling: Even with a perfect encoder, the "best-in-batch" target similarity plateaus at $\approx 0.30$, indicating a fundamental limit in the candidate pool or target definition for this setup. Consequently, the warm-start mechanism and stored-value components were not evaluated for Sudoku, as the failure was isolated to the first component.

Significance and Claims
The paper claims three primary contributions:

1. Methodology: A five-arm ablation suite that separates bias shift, stochastic warm-starting, and aligned value reuse in retrieval-warmed iterative inference.
2. Alignment Effect: Empirical evidence on connectivity-2 that per-graph alignment is the dominant lever for performance, isolated via a +35 pp swing between aligned and shuffled oracle conditions.
3. Heterogeneous Failure Modes: A demonstration that the "first blocking component" varies by task. On connectivity-2, the bottleneck is stored-value quality (the cold model's predictions are insufficient), whereas on Sudoku (under the current setup), the bottleneck is key quality (the encoder cannot retrieve sufficiently similar candidates).

The authors emphasize that the decomposition names the specific component responsible for failure or success, providing a localized verdict that aggregate end-to-end scores cannot offer. They explicitly state that the decomposition, the partial-noise warm-start mechanism (previously seen in SDEdit), and the oracle-memory upper bound are not novel primitives; the novelty lies in the ablation methodology and its application to reveal these specific failure modes in structured reasoning.