Apple researchers derive a theoretical ceiling for KL-regularized model alignment

Apple researchers derived an exact expression for the reward improvement available to an optimally solved, KL-regularized language-model alignment objective. The result gives researchers a theoretical frontier for measuring how much reward current alignment methods leave unrealized under a fixed divergence budget.

In two tested configurations—safety alignment with BeaverTails and summarization with Reddit TL;DR—the authors found that best-of-N sampling stayed on or near their estimated reward-versus-KL frontier. PPO and GRPO remained substantially below it across the KL budgets tested, according to the paper. The comparison does not establish that best-of-N is superior across other models, tasks, reward functions or alignment objectives.

KL regularization penalizes an aligned model for moving too far from its base model. A parameter called lambda controls that penalty: a smaller lambda permits more aggressive movement, while a larger one constrains it more strongly. The paper’s first theorem states that, for a fixed query and the optimal policy under its objective, expected reward gain equals lambda multiplied by the Jeffreys divergence between the base and aligned policies. Jeffreys divergence is the sum of KL divergence in both directions.

The result is narrower than a universal limit on model alignment. It characterizes the optimal solution to the paper’s particular KL-regularized reward objective. The authors contrast the equality with earlier analyses based on a square-root-of-KL upper bound.

Because the equality includes the optimally aligned policy, it cannot by itself predict the frontier before alignment. The authors address that problem by rewriting reward gain as a covariance under the base policy, then constructing Monte Carlo estimators for reward gain and KL divergence from base-model generations and reward scores. They used 10,000 samples per query to estimate the experimental frontier and swept lambda from 0.01 to 10. The paper says estimation converges more slowly at smaller lambda because exponential weighting increases variance.

The proxy-reward analysis addresses a separate problem: systems normally optimize a learned reward model rather than an unavailable “gold” measure of human preference. The paper compares policies optimally aligned to gold and proxy rewards and defines their difference as a preference residual. Its Proposition 3 makes the true-reward gap depend on an exponential of that residual divided by lambda. Under that model, larger reward error can widen the gap, and reducing lambda can magnify the effect.

That mechanism gives a mathematical account of reward hacking, in which optimization raises the proxy score while performance under the intended measure deteriorates. An earlier ICML study independently found this pattern in a synthetic gold-versus-proxy setup using reinforcement learning and best-of-N sampling.

The Apple authors also prove a conditional result for reward-model ensembles. If independently sampled proxy rewards are bounded and unbiased for a bounded gold reward, they say the expected gold reward of the ensemble-aligned policy converges toward ideal-reward alignment at a rate of O(n^-1/2). The guarantee does not cover systematic bias, which the paper leaves for future work.

For the experiments, the safety setup paired a Zephyr-7B base model with BeaverTails, while the summarization setup used a Pythia-1B-based model with Reddit TL;DR. The researchers varied N from 1 to 10,000 for best-of-N and reported the best PPO and GRPO runs from hyperparameter sweeps. They excluded DPO because their comparison required an explicit reward model.

Apple has also examined automated evaluation in a separate study of LLM judge quality, but that work covered personalized writing rather than reward-model alignment.

The alignment results were measured with learned reward-model scores rather than direct human evaluation of aligned outputs. The paper has not supplied evidence that the same ordering holds at other model scales or with other tasks and reward models. Its theory also does not yet extend to alternatives such as f-divergence regularization, win-rate maximization or per-token KL divergence.