ETS: Energy-Guided Test-Time Scaling for Training-Free RL Alignment
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.
Lay Summary
Training large language models to be helpful, safe, and aligned is difficult. The standard method, reinforcement learning fine-tuning, is effective but complex, costly, and unstable. We propose a training-free alternative that guides text generation at inference time with an energy score measuring alignment, without updating model parameters. This energy is estimated by efficient Monte Carlo sampling and accelerated via modern hardware and importance sampling, preserving output quality. Experiments on reasoning, coding, and science benchmarks show our Energy-Guided Test-Time Scaling (ETS) consistently improves generation. This offers a simpler, more stable path to better AI systems.