ETS: Energy-Guided Test-Time Scaling for Training-Free RL Alignment

AI summary

Overview

Research area: Machine learning — reinforcement learning (RL) alignment of language models, test-time scaling (TTS), and masked language modeling (MLM) theory.

Technical level: Advanced. The paper is built around a closed-form characterization of the KL-regularized RL optimum, an explicit backward transition kernel, Monte Carlo estimators, and total-variation convergence bounds, alongside systems-level acceleration.

Scope (one sentence): The paper derives a closed-form decomposition of the optimal RL-aligned transition kernel for general masked language models and uses it to build ETS, a training-free inference-time sampler that estimates an "energy" term online and accelerates it with importance sampling.

What This Paper Is About

RL post-training for language model alignment works well but is expensive and unstable: it needs reward modeling, large-scale human preference data, and reruns whenever the reward changes. The authors observe that the standard KL-regularized RLHF objective already has a closed-form optimal solution (Proposition 1), so instead of approximating it with iterative gradient training, one could sample from it directly at inference time. The goal is therefore to construct a training-free sampling procedure that converges to that optimal RL distribution, and to make it fast enough to be practical for both autoregressive models (ARMs) and diffusion language models (DLMs).

Key Contributions

1. A closed-form energy-reweighted backward transition for general MLMs. Proposition 2 shows the optimal RL transition kernel factorizes into (a) the reference model's transition kernel p_ref(x_s | x_t, y) and (b) an energy term E(y, x_s) = E_{p_ref(x_0 | y, x_s)}[exp(r(y, x_0)/λ)], a conditional expectation of exponentiated rewards. The paper notes this differs from prior continuous-time diffusion control-theory results because it applies to general discrete MLM and is derived directly from MLM properties (Appendix A.2).

2. The ETS algorithm with a proven convergence rate. Algorithm 1 samples candidate continuations at guidance steps, estimates their energy by online Monte Carlo with K independent completions (Equation 8), and resamples candidates via self-normalizing importance sampling. Proposition 3 bounds the total variation distance between the induced distribution q and the target p by O~(I/√M + I·ε), where I is the number of guidance steps, M the number of candidates, and ε the energy estimation error; the bound vanishes as M → ∞ and ε → 0.

3. Provably corrected acceleration (ETS-IS). Algorithm 2 replaces reference-model samples in the energy estimator with samples from a cheaper proposal model p_small, corrected by the importance weight p_ref(x_0|y,x_s)/p_small(x_0|y,x_s). Proposition 4 gives a high-probability error bound L·e^{D/λ}·√(log(2/δ)/(2K)), and Theorem 1 combines this with Proposition 3 to give O~(I/√M + I/√K) — the same asymptotic accuracy as Algorithm 1 when K is chosen sufficiently large.

4. A training-free proxy reward and a cross-paradigm empirical study. Because ground-truth rewards are unavailable at inference, the paper designs a self-consistency proxy reward (sample K completions, reward 1 for the majority-vote answer, 0 otherwise) and compares it against token-level confidence (logits), predictive entropy, and self-certainty, reporting that self-consistency yields a reward distribution closest to the ground-truth reward (Appendix B.4 and D.2). Experiments span ARMs and DLMs on mathematics, coding, and STEM benchmarks.

Main Findings

• Sampling from the RL optimum without training is possible. Under the unified MLM framework (which subsumes ARMs with fixed left-to-right masking and DLMs with flexible, top-K non-sequential unmasking where K = d_x/T), the optimal RL policy's backward transition is exactly the reference transition reweighted by an energy term (Proposition 2), so no gradient-based post-training is required.

• ETS consistently improves generation quality. The authors report that across reasoning, coding, and science benchmarks, ETS improves over standard inference and test-time scaling baselines, and — per their claim in Section 6.2 — matches and even surpasses the post-trained RL policy (GRPO) without any parameter updates. Note: the specific per-benchmark accuracy and latency values live in Table 1, which is not included in the truncated content provided here, so no numeric results can be quoted.

• Acceleration trades a small amount of quality for speed. The importance-sampling acceleration speeds up sampling "with slight performance degradation," and ETS latency is reported as comparable to standard TTS methods such as Best-of-N while achieving significantly better accuracy.

• Best-of-N is a special case of ETS. Setting I = 1 and taking λ → 0 makes the normalized resampling weights degenerate to the single largest energy within the finite candidate set, recovering Best-of-N.

• Parallelism is a structural advantage. For ARMs, batching and early-stop mechanisms decode M × K sequences in parallel and often terminate early, yielding lower actual per-token cost and fewer generated tokens than sequential approaches like Power Sampling (Section 6.2, Appendices B.1 and B.2).

• Error grows linearly in guidance steps. The TV bound scales with I, similar to error accumulation results in continuous diffusion models. The authors caution this does not imply I = 1 is optimal, because they assume the guidance error ε is identical at each step, which may not hold in practice.

• Scale-aligned models work as proposals. For ARMs, Qwen3 lightweight models sharing the reference tokenizer are used as p_small. For DLMs, no aligned small model exists for LLaDA-8B-Instruct, so Fast-dLLM (block-wise approximate KV cache plus confidence-aware parallel decoding) is treated as the proposal distribution instead.

Methodology in Plain English

The authors start from a fact already known in the RLHF literature: the objective "maximize reward while staying close to a reference model" has an exact optimal answer, proportional to the reference model's probability times an exponential of the reward. Training-based methods such as PPO, DPO, and GRPO approximate this answer with gradient updates. The paper instead asks a different question: can we just sample from that answer?

To do so they work in the masked-language-model view of generation, where text is produced by repeatedly unmasking tokens, a process that covers both ordinary left-to-right language models and diffusion language models. They show that the optimal RL version of this unmasking process is the ordinary process multiplied by a per-state "energy" — an expectation, under the reference model, of the exponentiated reward of future completions. States that tend to lead to high-reward completions get up-weighted.

Since that expectation cannot be computed exactly, ETS approximates it with Monte Carlo: for each candidate continuation, draw several independent completions, score them with a reward function, average the exponentiated scores, and then resample candidates in proportion to these estimated energies. Doing this at several points during decoding re-aligns the generation toward the optimum repeatedly. The authors prove how far the resulting distribution can drift from the target, in terms of the number of candidates and the energy estimation error.

To cut latency, they exploit the fact that the Monte Carlo estimator only needs samples from some model resembling the reference. They substitute a fast, smaller, tokenizer-aligned model (for ARMs) or an accelerated decoding scheme (Fast-dLLM, for DLMs), and correct the resulting bias with importance weights. They prove this corrected estimator retains the same asymptotic convergence order. Finally, because the true reward is unavailable at test time, they use majority-vote self-consistency across sampled completions as a stand-in reward.

Why This Matters

Impact on research. The paper reframes RL alignment as an inference-time sampling problem rather than a training problem, and provides a discrete-MLM counterpart to existing continuous-time diffusion control-theory derivations. It also supplies a convergence-rate guarantee that the authors point out prior related training-free methods (Dang et al., Uehara et al., Singhal et al.) lack, and situates ETS as theoretically more general (covering ARMs, not just diffusion models), more efficient, and usable without a verifiable reward function.

Real-world applications:

• Deploying aligned behavior without retraining. A system can change its reward or preference criterion at inference time without rerunning an RL pipeline — relevant whenever reward designs change frequently.
• Cost-sensitive LLM serving. Operators who cannot afford reward-model training or large-scale preference data collection can apply energy-guided decoding on top of an existing reference checkpoint.
• Domains requiring verifiable or self-consistency-checkable answers. Mathematics, code generation, and multiple-choice science reasoning (the paper's own evaluation areas) are natural fits, since self-consistency provides a usable proxy reward.
• Diffusion language models. ETS is one of the few approaches explicitly designed to align diffusion LMs at inference time, where the authors note TTS methods currently incur substantially higher latency than for ARMs.

Industry relevance. The methods are implemented on standard serving stacks — vLLM is referenced for ARMs, Fast-dLLM for DLMs, and lm-eval for benchmarking — and experiments run on 4× H20-3e-141GB GPUs. The paper compares against Verl-based RL and LLaDA-1.5, indicating a practical framing aimed at production alignment workflows where retraining cost and hyperparameter sensitivity are real bottlenecks.

Future Directions

• Tighter error bounds that account for varying guidance error. The authors explicitly state that their bound assumes identical estimation error ε at every guidance step, whereas in practice ε may vary across states, and they leave a more delicate bound to future work.
• Integrating additional acceleration techniques. The paper notes that quantization offers no operator-level speed advantage at their computational scale, and that speculative decoding via EAGLE-3 does not support their batching techniques because of tree attention and vLLM integration; several further acceleration approaches are left as future work.
• Extending to better reward proxies. The reward study compares self-consistency against token-level confidence, predictive entropy, and self-certainty; improving the training-free proxy — and closing the gap to true reward supervision — remains open.
• Broadening the applicability and theory relative to diffusion-only methods. The authors position their unified framework as generalizing prior diffusion-specific approaches, suggesting further work unifying ARMs and DLMs under one sampling theory.

Target Audience

Researchers and engineers working on LLM post-training and alignment, test-time scaling, and inference-time algorithms; practitioners interested in diffusion language models; and readers with a background in RL, Monte Carlo estimation, and probability who want a training-free alternative to RLHF/GRPO-style pipelines. The paper assumes familiarity with KL-regularized RL objectives, Markov chains, importance sampling, and total variation distance, so it is best suited to readers at an intermediate-to-advanced level.

Authors’ abstract

Reinforcement Learning (RL) post-training alignment for language models is effective, but also costly and unstable in practice, owing to its complicated training process. To address this, we propose a training-free inference method to sample directly from the optimal RL policy. The transition probability applied to Masked Language Modeling (MLM) consists of a reference policy model and an energy term. Based on this, our algorithm, Energy-Guided Test-Time Scaling (ETS), estimates the key energy term via online Monte Carlo, with a provable convergence rate. Moreover, to ensure practical efficiency, ETS leverages modern acceleration frameworks alongside tailored importance sampling estimators, substantially reducing inference latency while provably preserving sampling quality. Experiments on MLM (including autoregressive models and diffusion language models) across reasoning, coding, and science benchmarks show that our ETS consistently improves generation quality, validating its effectiveness and design. The code is available at https://github.com/sheriyuo/ETS.