A Theoretical Lens for RL-Tuned Language Models via Energy-Based Models
This work addresses a foundational problem in AI by offering theoretical insights into RL-tuned language models, which is incremental as it builds on existing methods to explain empirical behaviors.
The paper tackles the limited theoretical understanding of large language models trained via KL-regularized reinforcement learning by exploiting their energy-based model structure to provide a unified variational analysis, proving results such as monotonic KL convergence and bounded hitting times for instruction-tuned models, and showing equivalence to expected KL minimization for reasoning models.
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.