Information Theoretic Guarantees For Policy Alignment In Large Language Models

TL;DR

This paper provides information-theoretic bounds on policy alignment in large language models, showing how reward improvements relate to divergence measures under tail assumptions and extending results to various divergences and reward proxies.

Contribution

It establishes new upper bounds on reward improvements for policy alignment using $f$-divergences, including Rényi divergence, under tail assumptions, and connects proxy and true rewards.

Findings

Reward improvement scales with $\

Reward bounds hold under sub-gaussian tail assumptions.

Bounds extend to any $f$-divergence via order statistics and data processing inequality.

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 with respect to an $f$-divergence such as the $\mathsf{KL}$ divergence. The best of $n$ alignment policy selects a sample from the reference policy that has the maximum reward among $n$ independent samples. For both cases (policy alignment and best of $n$), recent works showed empirically that the reward improvement of the aligned policy on the reference one scales like $\sqrt{\mathsf{KL}}$, with an explicit bound in $n$ on the $\mathsf{KL}$ for the best of $n$ policy. We show in this paper that the $\sqrt{\mathsf{KL}}$ information theoretic upper bound holds if the reward under the reference policy has sub-gaussian tails. Moreover, we prove for the best of $n$ policy, that the $\mathsf{KL}$ upper bound can be obtained for any $f$-divergence via a reduction to exponential order statistics owing to the R\'enyi representation of order statistics, and a data processing inequality. If additional information is known on the tails of the aligned policy we show that tighter control on the reward improvement can be obtained via the R\'enyi divergence. Finally we demonstrate how these upper bounds transfer from proxy rewards to golden rewards which results in a decrease in the golden reward improvement due to overestimation and approximation errors of the proxy reward.