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

TL;DR

ETS introduces a training-free, energy-guided inference method for RL alignment in language models, improving generation quality efficiently without costly retraining.

Contribution

It proposes a novel test-time scaling algorithm that estimates energy terms via online Monte Carlo, enabling effective RL alignment without additional training.

Findings

ETS improves generation quality across multiple benchmarks.

The method reduces inference latency significantly.

Experiments validate the effectiveness and efficiency of ETS.

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.