🤖 machine learning

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

The paper proposes ETS, a training-free inference method that leverages energy-guided test-time scaling with online Monte Carlo estimation and importance sampling to efficiently sample from optimal RL policies, thereby improving generation quality across reasoning, coding, and science benchmarks without the cost and instability of traditional RL post-training.

Original authors: Xiuyu Li, Jinkai Zhang, Mingyang Yi, Yu Li, Longqiang Wang, Yue Wang, Ju Fan

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xiuyu Li, Jinkai Zhang, Mingyang Yi, Yu Li, Longqiang Wang, Yue Wang, Ju Fan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Teaching AI is Expensive and Messy

Imagine you have a very talented student (a Large Language Model) who knows a lot but sometimes gives answers that are technically correct but not very helpful or logical.

To fix this, researchers usually use Reinforcement Learning (RL). Think of this as hiring a strict tutor to grade the student's homework over and over again. The tutor gives points for good answers and subtracts points for bad ones. The student then studies hard to get more points next time.

The problem? This "tutoring" process is:

Expensive: It requires massive computing power.
Unstable: Sometimes the student gets confused and starts giving weird answers.
Slow: You have to re-teach the student every time you want them to learn a new skill.

The Solution: "ETS" (Energy-Guided Test-Time Scaling)

The authors of this paper propose a clever trick. Instead of re-teaching the student (training), they change how the student takes the test (inference).

They call their method ETS. Here is how it works, broken down into three simple concepts:

1. The "What If" Game (Sampling from the Best)

Usually, when an AI answers a question, it picks the first answer that seems okay and moves on.
ETS says: "Wait, let's not just pick one answer. Let's imagine many different versions of the answer at the same time."

Think of it like a detective solving a mystery. Instead of guessing one suspect immediately, the detective writes down 10 different theories about who did it. Then, they check which theory holds up best. ETS does this by generating multiple "candidate" sentences simultaneously.

2. The "Energy" Score (The Magic Compass)

How does the AI know which of those 10 theories is the best without a human teacher?
The paper introduces an "Energy" term. Imagine every possible answer has an invisible "energy score."

High Energy = A bad, confusing, or wrong answer (like a heavy rock you don't want to carry).
Low Energy = A good, clear, and correct answer (like a light feather).

The AI's goal is to find the path with the lowest energy. The paper proves mathematically that if you can calculate this "energy" correctly, you can find the perfect answer without ever having to retrain the model.

3. The "Speed Trick" (Importance Sampling)

Calculating this "energy" for every single candidate is slow. It's like checking the weight of 100 rocks one by one; it takes forever.

To fix this, the authors use a Speed Trick (called Importance Sampling):

They use a smaller, faster AI (a "lightweight" model) to quickly guess which candidates look promising.
They then use the big, smart AI only to double-check the most promising ones.
It's like having a quick assistant scan a room to find the heavy rocks, so you don't have to lift every single object yourself. This makes the process fast enough to be practical.

The Result: Better Answers, No Extra Training

The paper tested this method on two types of AI models:

Autoregressive Models (ARMs): The standard "read one word, then write the next" models (like most chatbots).
Diffusion Language Models (DLMs): A newer type that builds answers by gradually filling in blanks (like a painter filling in a sketch).

The findings were surprising:

Better than Training: The ETS method produced better answers on math, coding, and science tests than models that had actually been "trained" with the expensive RL method.
No Training Needed: It works with the model exactly as it is, right out of the box.
Efficient: Even though it checks many possibilities, the "speed trick" keeps it fast.

Summary Analogy

Imagine you are trying to find the highest peak in a foggy mountain range.

Old Way (RL Training): You hire a guide to climb the mountain, map it, and teach you the route. This takes weeks and costs a fortune.
ETS Way: You stand at the base. Instead of climbing one path, you send out 20 drones (candidates) to fly up different paths. You use a special sensor (Energy) to see which path leads to the highest peak. You use a fast, cheap drone to scout the easy paths and a high-tech drone only for the tricky ones. You pick the best path instantly.

The paper claims: This "drone scouting" method (ETS) finds the highest peak (the best answer) faster and more accurately than the expensive "guide training" method, without needing to change the map (the model) beforehand.

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

Problem Statement

Reinforcement Learning (RL) post-training has become a central paradigm for aligning Large Language Models (LLMs) with human intent, particularly for enhancing reasoning and instruction following. However, this approach suffers from significant practical limitations: it requires costly reward modeling, relies on large-scale human preference data, exhibits unstable training dynamics, and is sensitive to hyperparameters. Furthermore, any change in reward design necessitates retraining, leading to substantial computational and human costs.

While the optimal RL distribution admits a closed-form solution theoretically, existing methods rely on iterative gradient-based optimization to approximate this solution rather than exploiting its structure directly. This raises a fundamental question: Can we sample directly from the optimal RL distribution at inference time without any additional training?

Methodology

The authors propose Energy-Guided Test-Time Scaling (ETS), a training-free inference method designed to sample directly from the optimal RL policy. The method is built upon a unified Masked Language Modeling (MLM) framework that encompasses both Autoregressive Models (ARMs) and Diffusion Language Models (DLMs).

1. Theoretical Foundation: Energy-Reweighted Backward Transition

The core theoretical insight is that the optimal backward transition kernel for the MLM framework decomposes into two factors:

Reference Transition: The transition kernel provided by the reference model ( $p_{ref}$ ).
Energy Term: A conditional expectation of exponentiated rewards, acting as an energy term.

Specifically, the target backward transition $p(s | t, y)$ is proportional to:
$p_{ref}(s | t, y) \cdot \mathbb{E}_{p_{ref}(x_0|y,s)} \left[ \exp\left(\frac{r(y, x_0)}{\lambda}\right) \right]$
where $r(y, x_0)$ is the reward function and $\lambda$ is a regularization constant. This formulation allows for training-free sampling by estimating the energy term via Monte Carlo methods.

2. Algorithm: Energy-Guided Test-Time Scaling (ETS)

The algorithm operates as an iterative guidance process:

Sampling: At each guidance step, the algorithm samples $M$ candidate continuations from the reference model.
Energy Estimation: For each candidate, the energy term is estimated using Monte Carlo sampling. Specifically, $K$ independent completions are generated for each candidate to compute a self-normalized importance sampling estimator of the expected reward.
Selection: Candidates are reweighted based on their estimated energies, and one is selected to proceed to the next step.
Reward Design: To operate without ground-truth rewards during inference, ETS employs a self-consistency reward. It extracts final answers from the $K$ completions, identifies the majority vote as the consensus, and assigns a reward of 1 to completions matching the consensus and 0 otherwise.

3. Acceleration via Importance Sampling (ETS-IS)

To address the latency bottleneck caused by repeated model queries for energy estimation, the authors introduce ETS-IS.

Proposal Model: Instead of sampling from the large reference model ( $p_{ref}$ ) for energy estimation, a computationally cheaper, lightweight proposal model ( $p_{small}$ ) is used.
Correction: Importance sampling weights are applied to correct the bias introduced by the proposal model, ensuring the estimator remains unbiased and theoretically consistent with the target distribution.
Implementation:
- For ARMs: A smaller, scale-aligned model (e.g., Qwen3-1.7B for Qwen3-8B) serves as the proposal.
- For DLMs: Fast-dLLM is utilized, leveraging block-wise KV caching and parallel decoding to accelerate inference.

Key Contributions

Training-Free Alignment: The paper presents a method to sample directly from the optimal RL policy without iterative gradient updates or fine-tuning, resolving the instability and cost issues of traditional RL post-training.
Unified Framework: The approach is applicable to both ARMs and DLMs under a unified MLM formulation, providing a general solution for training-free scaling.
Theoretical Guarantees: The authors provide rigorous proofs regarding the convergence rate of the sampling distribution to the target optimal distribution in terms of Total Variation Distance. The error bound is shown to scale linearly with the number of guidance steps and inversely with the square root of the sample size.
Efficient Acceleration: By integrating importance sampling with lightweight proposal models and modern acceleration frameworks (like Fast-dLLM), the method significantly reduces inference latency while preserving sampling quality.

Experimental Results

The authors evaluated ETS on reasoning (MATH500, GSM8K), coding (HumanEval), and science (GPQA) benchmarks using both ARMs (Qwen3-1.7B, Qwen3-8B) and DLMs (LLaDA-8B-Instruct).

Performance: ETS consistently improves generation quality over standard inference and other test-time scaling baselines (e.g., Beam Search, Best-of-N, Power Sampling).
Comparison to RL: Notably, ETS matches or surpasses the performance of post-trained RL policies (such as GRPO) without any parameter updates.
Efficiency: The acceleration method (ETS-IS) achieves a favorable trade-off between latency and accuracy. For instance, on Qwen3-8B, ETS-IS achieves accuracy comparable to standard ETS but with significantly reduced latency.
Scaling: Experiments show that increasing the number of candidates ( $M$ ) is more effective for improving accuracy than increasing the number of completions ( $K$ ), and that intermediate guidance steps ( $I$ ) provide proportional improvements in accuracy.

Significance and Claims

The paper claims that ETS establishes a practical, training-free pathway toward reward-aware generation. By demonstrating that high-quality alignment can be achieved through inference-time scaling rather than costly post-training, the work offers a cost-effective and stable alternative for resource-constrained settings. The authors assert that their method democratizes alignment techniques and reduces the environmental footprint of LLM deployment by eliminating the need for extensive retraining cycles. The framework is presented as a foundation for future research into efficient inference and test-time scaling in generative models.