A paper published on Apple's Machine Learning Research site examines how much reward improvement language model alignment can theoretically achieve under a fixed divergence budget. The work, titled "Theoretical Limits of Language Model Alignment", is authored by Lucas Monteiro Paes, Natalie Mackraz, Barry-John Theobald and Federico Danieli, with Paes and Mackraz credited as equal contributors.
Alignment methods aim to shape model outputs so they reflect human preferences while retaining the capabilities of the base model. The authors note that two families of techniques dominate practice: reinforcement learning, which maximizes expected reward subject to a KL-divergence constraint, and best-of-N alignment, which picks the highest-scoring output out of N independently drawn samples. Despite how widely both are used, the fundamental ceiling on reward gains under a KL budget has remained poorly understood.
The paper characterizes the information-theoretic limits of KL-regularized alignment by deriving the largest expected reward gain possible for a given KL-divergence budget. Its first result is a closed-form expression for the optimal improvement, which is governed by a Jeffreys divergence term rather than the square root of the KL divergence used in earlier analyses. The researchers also restate this expression as a covariance under the base model, which yields a practical estimator able to predict achievable alignment gains from base model samples alone.
The analysis is then extended to the proxy reward setting. The authors show that the gap between ideal and proxy alignment, commonly described as reward hacking, widens as the magnitude of reward error grows and as the KL penalty factor decreases. They additionally prove that reward ensembling mitigates reward hacking, giving a theoretical basis for a technique already used in practice.
On the empirical side, the team computed the KL-reward Pareto frontier for two language model tasks, safety and summarization. Best-of-N sampling was found to closely approach the theoretical limit, while PPO and GRPO remained substantially suboptimal. According to the authors, the results help explain several phenomena previously observed empirically in the alignment literature and indicate that algorithmic improvements will be needed to reach optimal alignment without incurring high inference costs.
The paper was accepted at the Fine-Tuning in Modern Machine Learning (FITML) Workshop at NeurIPS 2024.