Technical Summary: Energy-Based Decoding (EBD)

1. Problem Statement

The paper addresses a critical mismatch in the evaluation of pre-trained Large Language Models (LLMs). While standard benchmarks assume models will interpret prompts as tasks and produce well-formed answers, frozen pre-trained models are optimized primarily for next-token prediction (continuation). Consequently, under standard prompting and direct decoding, these models often generate continuation-style text rather than task-oriented responses. This leads to unreliable benchmark scores that conflate a model's actual capabilities with its inability to adopt a "task mode" without parameter updates.

Existing solutions face significant limitations:

• Post-training (SFT/RLHF): Effective but costly and obscures the distinction between pre-trained knowledge and learned behaviors.
• Likelihood-only decoding (e.g., temperature reduction, beam search): These methods reshape the output distribution without external utility signals, often amplifying spurious high-likelihood continuation modes rather than activating task-solving behavior.
• Sample-heavy search (e.g., Best-of-N, reranking): While they use reward signals, they are computationally expensive and inefficient when the base distribution is heavily biased toward non-task continuations.

The core challenge is to activate task-oriented, high-utility behaviors from frozen pre-trained models at inference time without parameter updates, while maintaining efficiency and faithfulness to the original model distribution.

2. Methodology: Energy-Based Decoding (EBD)

The authors propose Energy-Based Decoding (EBD), a training-free, reward-guided framework that steers generation toward high-utility responses while anchoring them to the pre-trained prior.

Theoretical Formulation

EBD frames decoding as sampling from a reward-tilted posterior distribution:
$$\pi^*_\beta(y | x) \propto p_\theta(y | x) \exp(\beta S_x(y))$$
Where:

• $p_\theta(y | x)$ is the frozen pre-trained model's prior.
• $S_x(y)$ is a prompt-normalized reward score derived from an external lightweight reward model.
• $\beta$ is an inverse temperature controlling the trade-off between the prior and the reward.

This formulation is equivalent to an energy-based distribution where the energy $E_x(y) = -\log p_\theta(y | x) - \beta S_x(y)$. The goal is to find low-energy responses that are both plausible under the pre-trained model and high-utility according to the reward.

Algorithm: Block-wise Metropolis-Hastings

To sample from this target distribution efficiently, EBD implements a two-stage pipeline using a Block-wise Metropolis-Hastings (MH) sampler:

1. Initialization (Stage I):

   • Sample a small pool of candidates ($n_{init}$) from the frozen prior.
   • Score them with the reward model and compute prompt-level normalization statistics (mean $\mu_R$ and standard deviation $\sigma_R$).
   • Define a standardized advantage score: $A_x(y) = (R(y, x) - \mu_R) / \sigma_R$.
   • Select the best candidate from the pool as the initial state $y^{(0)}$.

2. Refinement (Stage II):

   • Perform $K$ refinement steps. In each step:
     • Block Selection: Partition the current response into $M$ contiguous blocks and uniformly sample a cut position $c$.
     • Proposal: Keep the prefix ($y_{<c}$) fixed and regenerate the suffix ($y'_{\ge c}$) from the matched conditional prior $p_\theta(\cdot | x, y_{<c})$.
     • Acceptance: Compute the MH acceptance ratio. Crucially, because the proposal suffix is sampled from the exact same conditional distribution as the prior term, the proposal-prior cancellation occurs:
       $$\frac{p_\theta(y' | x) Q(c, y | y')}{p_\theta(y | x) Q(c, y' | y)} = 1$$
     • The acceptance probability simplifies to a reward comparison:
       $$\alpha = \min(1, \exp(\beta [A_x(y') - A_x(y)]))$$
   • If accepted, update the state; otherwise, retain the current state.

Key Efficiency Mechanism

The Proposal-Prior Cancellation is the central innovation. It eliminates the need to re-evaluate the full-sequence likelihood under the base model during the acceptance step. The computational cost per step reduces to one autoregressive suffix generation and one reward model evaluation, scaling linearly with the number of refinement steps.

3. Key Contributions

1. Formulation: The paper formulates reward-guided decoding as sampling from a reward-tilted posterior, enabling the activation of task-oriented behavior under a frozen prior.
2. Algorithmic Innovation: It introduces EBD with Proposal-Prior Cancellation, which allows for efficient updates based solely on reward differences without explicit full-sequence likelihood rescoring.
3. Behavioral Shift: The authors demonstrate that EBD does not merely improve aggregate scores but shifts the behavioral distribution of base models to closely resemble their post-trained (instruct-tuned) counterparts.
4. Empirical Validation: Extensive evaluation across five models and six benchmarks shows consistent improvements over direct decoding and prior decoding-based methods (e.g., Power Sampling).

4. Experimental Results

The authors evaluated EBD on five open-source models (Llama-3-8B, Mistral-7B, Qwen2.5-7B, Qwen3-8B-Base, Olmo-7B) across objective (GPQA, Math500, HumanEval) and subjective (AlpacaEval2.0, MT-Bench, WritingBench) benchmarks.

• Performance Gains:
  • AlpacaEval2.0: Qwen3-8B-Base improved from 8.8 to 44.5 (+405%).
  • Objective Tasks: Average objective scores increased significantly (e.g., Qwen2.5-7B from 0.398 to 0.557).
  • Behavioral Alignment: The Pearson correlation of binary correctness with post-trained instruct models increased from 0.256 (Direct) to 0.385 (EBD) on average.
• Efficiency:
  • EBD is substantially faster than sample-heavy baselines like Power Sampling. On Mistral-7B Math500, EBD reduced latency by 18.9× relative to prior decoding work.
  • It remains robust across different reward model sizes (0.6B to 8B), with smaller models often matching the performance of larger ones.
• Robustness:
  • EBD improves performance on both zero-shot and cold-start SFT settings.
  • Unlike SFT, which suppresses certain continuation modes, EBD applies a reward tilt without altering the underlying distribution, making it robust across different task types (e.g., improving math/reasoning while maintaining code generation capabilities).

5. Significance and Claims

The paper claims that decoding is a critical, often overlooked component of model evaluation. Standard direct decoding underestimates the capabilities of frozen pre-trained models in task-oriented settings.

• Fair Evaluation: EBD enables a fairer inference-time evaluation of pre-trained models by activating task-oriented behaviors that are otherwise inaccessible via standard prompting.
• Complementarity to Post-training: EBD serves as a complementary approach to parameter updates. It can close a substantial portion of the performance gap between pre-trained and post-trained models without any training, acting as an effective inference-time boost or a diagnostic for checkpoint sensitivity.
• Efficiency vs. Effectiveness: The method achieves a superior quality-efficiency trade-off compared to existing decoding-only approaches, proving that external reward signals combined with efficient MCMC sampling can unlock latent capabilities in frozen models.

The authors conclude that reward-guided inference can significantly improve alignment-relevant behavior without a single parameter update, suggesting that the "missing piece" in pre-trained model evaluation is often the decoding strategy rather than the model weights themselves.