Asymptotics of Language Model Alignment

2024/04/02 by Joy Qiping Yang, Salman Salamatian, Yang, Joy Qiping +7 · 16 citations
Computer Science · #FOS: Computer and information sciences #Information Theory (cs.IT) #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Natural Language Processing Techniques

paper · pdf · doi:10.48550/arxiv.2404.01730

openalex publication_date 2024/04/02 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Let p denote a generative language model. Let r denote a reward model that returns a scalar that captures the degree at which a draw from p is preferred. The goal of language model alignment is to alter p to a new distribution ϕ that results in a higher expected reward while keeping ϕ close to p. A popular alignment method is the KL-constrained reinforcement learning (RL), which chooses a distribution ϕ_Δ that maximizes E_ϕΔ r(y) subject to a relative entropy constraint KL(ϕ_Δ|| p) ≤ Δ. Another simple alignment method is best-of-N, where N samples are drawn from p and one with highest reward is selected. In this paper, we offer a closed-form characterization of the optimal KL-constrained RL solution. We demonstrate that any alignment method that achieves a comparable trade-off between KL divergence and reward must approximate the optimal KL-constrained RL solution in terms of relative entropy. To further analyze the properties of alignment methods, we introduce two simplifying assumptions: we let the language model be memoryless, and the reward model be linear. Although these assumptions may not reflect complex real-world scenarios, they enable a precise characterization of the asymptotic behavior of both the best-of-N alignment, and the KL-constrained RL method, in terms of information-theoretic quantities. We prove that the reward of the optimal KL-constrained RL solution satisfies a large deviation principle, and we fully characterize its rate function. We also show that the rate of growth of the scaled cumulants of the reward is characterized by a proper Renyi cross entropy. Finally, we show that best-of-N is asymptotically equivalent to KL-constrained RL solution by proving that their expected rewards are asymptotically equal, and concluding that the two distributions must be close in KL divergence.

Cited by

• Test-Time Alignment of LLMs via Sampling-Based Optimal Control in pre-logit space
• The Best of N Worlds: Aligning Reinforcement Learning with Best-of-N Sampling via max@k Optimisation
• Certified Self-Consistency: Statistical Guarantees and Test-Time Training for Reliable Reasoning in LLMs
• Optimal Stopping vs Best-of-N for Inference Time Optimization
• Best-of-Majority: Minimax-Optimal Strategy for Pass@k Inference Scaling
• Rethinking Large Language Model Distillation: A Constrained Markov Decision Process Perspective
• Inference-Time Reward Hacking in Large Language Models
• SPECS: Faster Test-Time Scaling through Speculative Drafts
• Theoretical Tensions in RLHF: Reconciling Empirical Success with Inconsistencies in Social Choice Theory
• Tokenized Bandit for LLM Decoding and Alignment
• Psi-Sampler: Initial Particle Sampling for SMC-Based Inference-Time Reward Alignment in Score Models
• Soft Best-of-n Sampling for Model Alignment
• Learning to Choose or Choosing to Learn: Best-of-N vs. Supervised Fine-Tuning for Bit String Generation
• AdaBoN: Adaptive Best-of-N Alignment
• Tail-Aware Information-Theoretic Bounds for LLM Alignment under Heavy-Tailed Rewards
• Restoring Calibration for Aligned Large Language Models: A Calibration-Aware Fine-Tuning Approach

Related

• Beyond RLHF: A Unified Theoretical Framework of Alignment
• Averaging log-likelihoods in direct alignment
• A Principled Loss Function for Direct Language Model Alignment
• Generalist Reward Models: Found Inside Large Language Models
• Towards Reliable Alignment: Uncertainty-aware RLHF