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

TL;DR

This work introduces Energy-Guided Test-Time Scaling (ETS), a training-free method to sample directly from the Reinforcement Learning optimum for language-model alignment. By deriving a closed-form target distribution $p^*(\boldsymbol{x}_0|\boldsymbol{y})$ and decomposing the backward transition with an energy term $\mathcal{E}(\boldsymbol{y},\boldsymbol{x}_s)$, ETS performs Monte Carlo energy estimation and uses lightweight proposal models with importance sampling to maintain efficiency. The authors prove total-variation convergence guarantees and demonstrate empirically that ETS improves generation quality on autoregressive and diffusion language models across reasoning and coding benchmarks, with latency comparable to or better than existing test-time scaling baselines. Additionally, ETS-IS extends the approach to accelerate energy estimation via cheaper proposals for ARMs and DLMs, providing a practical, training-free pathway to reward-aware generation in NLP.

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.

Figures (10)

• Figure 1: Unified MLM framework. Generation is modeled as a backward Markov chain from $x_T$ to $x_0$. Case 1 shows the fixed left-to-right decoding path of ARMs; Case 2 illustrates the flexible, non-sequential unmasking of DLMs.
• Figure 2: Overview of Energy-Guided Test-Time Scaling (ETS). ETS performs iterative guidance on the unified MLM framework. At each guidance step (zoomed-in, right), the algorithm evaluates $M$ candidates. Their associated energy $\mathcal{E}$ is estimated via Monte Carlo method using $K$ independent completions of the corresponding candidate. We utilize an aligned lightweight model $p_{\text{small}}$ with Importance Sampling (IS) correction to accelerate energy estimation while maintaining theoretical consistency with the target optimal distribution.
• Figure 3: Effect of total samples on ETS. We ablate the total samples with Qwen3-8B and plot HumanEval accuracies (left) with corresponding latencies (right) for various sample counts.
• Figure 4: Effect of guidance steps on ETS. We evaluate Qwen3-8B on HumanEval (left) with corresponding latencies (right) are reported under various guidance steps.
• Figure 5: Comparisons between TTS methods. We ablate different latencies and plot corresponding Humaneval accuracies with Qwen3-8B, for various training-free TTS methods.

Theorems & Definitions (20)

• Proposition 1
• Proposition 2
• Remark 1
• Proposition 3
• Remark 2
• Proposition 4
• Theorem 1
• Proposition 4
• proof
• proof