Paper

Information Theoretic Guarantees For Policy Alignment In Large Language Models

Abstract

Policy alignment of large language models refers to constrained policy optimization, where the policy is optimized to maximize a reward while staying close to a reference policy based on an f-divergence like KL divergence. The best of n alignment policy selects the sample with the highest reward from n independent samples. Recent work shows that the reward improvement of the aligned policy scales asKL, with an explicit bound on the KL for best of n policies. We show that thisKL bound holds if the reference policy’s reward has sub-gaussian tails. For best of n policies, the KL bound applies to any f-divergence through a reduction to exponential order statistics using the Rényi representation. Tighter control can be achieved with Rényi divergence if additional tail information is known. Finally, we demonstrate how these bounds transfer to golden rewards, resulting in decreased golden reward improvement due to proxy reward overestimation and approximation errors.