A Theoretical Lens for RL-Tuned Language Models via Energy-Based Models
Zhiquan Tan, Yinrong Hong
TL;DR
The paper addresses the limited theoretical understanding of KL-regularized RL for large language models by casting the KL-regularized objective as a conditional energy-based formulation. It derives a multiplicative detailed-balance structure with a canonical potential for instruction-tuned policies, enabling KL convergence to a stationary distribution and quantifying mixing via the spectral gap. For reasoning with verifiable rewards, it shows an exact equivalence to minimizing KL to an optimal reasoning distribution and reveals a Bernoulli KL relationship between accuracy and the KL gap under natural gradient flow. This energy-based perspective unifies RLHF-style tuning and reasoning under a common framework, explains entropy-accuracy trade-offs, and provides principled guidance for designing and analyzing RL-alignment methods.
Abstract
Large language models (LLMs) trained via KL-regularized reinforcement learning demonstrate strong instruction following, self-correction, and reasoning abilities. Yet their theoretical underpinnings remain limited. We exploit the closed-form energy-based model (EBM) structure of the optimal KL-regularized policy to provide a unified variational analysis of LLMs. For instruction-tuned models, under natural assumptions on reward potentials and pretraining symmetry, we prove that the transition kernel satisfies detailed balance with respect to a scalar potential encoding response quality. This yields monotonic KL convergence to a high-quality stationary distribution, bounded hitting times to superior states, and exponential mixing governed by the spectral gap. For reasoning models trained with verifiable rewards (RLVR), we show the objective is equivalent to expected KL minimization toward an optimal reasoning distribution, with the suboptimality gap reducing to the Bernoulli KL between target and current accuracies along the natural gradient flow. This helps explain empirical entropy-accuracy trade-offs.