Towards a Theoretical Understanding to the Generalization of RLHF

2026-01-23

AI Summary

This work investigates the generalization ability of Reinforcement Learning from Human Feedback (RLHF) in high-dimensional large language models. Departing from conventional analyses that rely on the consistency of maximum likelihood estimation, the study establishes, for the first time, a generalization bound under an end-to-end RLHF framework by leveraging the algorithmic stability perspective. By introducing a linear reward model and a feature coverage condition, and analyzing both Gradient Ascent (GA) and Stochastic Gradient Ascent (SGA) algorithms, the authors prove that when the feature coverage condition holds, the generalization error of the empirical optimal policy converges at a rate of $O(n^{-1/2})$. This result further extends to policies obtained via gradient-based optimization, thereby providing theoretical justification for the generalization performance of practical RLHF implementations.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Reinforcement Learning from Human Feedback (RLHF) and its variants have emerged as the dominant approaches for aligning Large Language Models with human intent. While empirically effective, the theoretical generalization properties of these methods in high-dimensional settings remain to be explored. To this end, we build the generalization theory on RLHF of LLMs under the linear reward model, through the framework of algorithmic stability. In contrast to the existing works built upon the consistency of maximum likelihood estimations on reward model, our analysis is presented under an end-to-end learning framework, which is consistent with practice. Concretely, we prove that under a key \textbf{feature coverage} condition, the empirical optima of policy model have a generalization bound of order $\mathcal{O}(n^{-\frac{1}{2}})$. Moreover, the results can be extrapolated to parameters obtained by gradient-based learning algorithms, i.e., Gradient Ascent (GA) and Stochastic Gradient Ascent (SGA). Thus, we argue that our results provide new theoretical evidence for the empirically observed generalization of LLMs after RLHF.

Problem

• RLHF
• generalization
• Large Language Models
• algorithmic stability
• high-dimensional settings

Innovation

• algorithmic stability
• generalization bound
• feature coverage
• RLHF
• end-to-end learning