Theoretical guarantees on the best-of-n alignment policy

2024-01-03

arXiv.org

AI Summary

The widely adopted best-of-$n$ strategy in generative model alignment has long relied on a flawed theoretical foundation: the classical KL divergence analytical formula is mistakenly treated as an exact expression, whereas it only constitutes an upper bound; similarly, the win rate lacks a rigorous, tight theoretical bound. Method: We first disprove the exactness of the classical KL formula, formally establishing it as a loose upper bound, and propose a novel KL estimator with provably tight estimation guarantees. We further derive, for the first time, the strict upper bound $n/(n+1)$ on win rate and demonstrate its tightness via construction. Our analysis integrates theoretical derivation, probabilistic inequality techniques, and large-scale numerical experiments. Results: We establish an exact theoretical trade-off between KL divergence and win rate, showing that near-optimal performance is achievable even for $n < 1000$. This work provides the first tight, reliable, and empirically verifiable theoretical framework for alignment strategy design.

Technology Category

• Game Theory and Economic Paradigms: Adversarial Learning
• Machine Learning: Online Learning & Bandits
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Uses of LLMs and GenAI for marketplace design, bidding, and strategic interactions
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

A simple and effective method for the inference-time alignment of generative models is the best-of-$n$ policy, where $n$ samples are drawn from a reference policy, ranked based on a reward function, and the highest ranking one is selected. A commonly used analytical expression in the literature claims that the KL divergence between the best-of-$n$ policy and the reference policy is equal to $log (n) - (n-1)/n.$ We disprove the validity of this claim, and show that it is an upper bound on the actual KL divergence. We also explore the tightness of this upper bound in different regimes, and propose a new estimator for the KL divergence and empirically show that it provides a tight approximation. We also show that the win rate of the best-of-$n$ policy against the reference policy is upper bounded by $n/(n+1)$ and derive bounds on the tightness of this characterization. We conclude with analyzing the tradeoffs between win rate and KL divergence of the best-of-$n$ alignment policy, which demonstrate that very good tradeoffs are achievable with $n<1000$.

Problem

• Strategy Evaluation
• Generative Models
• Probability Estimation

Innovation

• Optimal n Strategies
• Improved Estimation Method
• Winning Probability Analysis