Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards a Theoretical Understanding to the Generalization of RLHF

Zhaochun Li, Mingyang Yi, Yue Wang, Shisheng Cui, Yong Liu

TL;DR

This work provides an end-to-end theoretical framework for RLHF by analyzing the generalization of policies learned through KL-regularized RLHF under a linear reward model. Using algorithmic stability and a feature coverage assumption, it establishes a dimension-free suboptimality bound of order $\tilde{O}(n^{-1/2})$ for empirical optima and extends the results to Gradient Ascent and Stochastic Gradient Ascent with explicit rates in terms of $n$ and $T$. The analysis reveals how sufficient data coverage leads to exact recovery, while realistic insufficient coverage yields robust generalization via a decomposed error bound. The results offer theoretical support for the empirically observed generalization capabilities of LLMs after RLHF and provide guidance for data collection and optimization strategy in practice. Overall, the framework advances understanding of how RLHF generalizes in high dimensions and under practical training dynamics, highlighting the role of feature coverage and algorithmic stability in enabling data-efficient alignment.

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.

Towards a Theoretical Understanding to the Generalization of RLHF

TL;DR

This work provides an end-to-end theoretical framework for RLHF by analyzing the generalization of policies learned through KL-regularized RLHF under a linear reward model. Using algorithmic stability and a feature coverage assumption, it establishes a dimension-free suboptimality bound of order for empirical optima and extends the results to Gradient Ascent and Stochastic Gradient Ascent with explicit rates in terms of and . The analysis reveals how sufficient data coverage leads to exact recovery, while realistic insufficient coverage yields robust generalization via a decomposed error bound. The results offer theoretical support for the empirically observed generalization capabilities of LLMs after RLHF and provide guidance for data collection and optimization strategy in practice. Overall, the framework advances understanding of how RLHF generalizes in high dimensions and under practical training dynamics, highlighting the role of feature coverage and algorithmic stability in enabling data-efficient alignment.

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 . 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.
Paper Structure (48 sections, 44 theorems, 193 equations, 2 figures)

This paper contains 48 sections, 44 theorems, 193 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. Theoretical Foundations of RLHF.
  4. Algorithmic Stability and Generalization.
  5. Formal Setting
  6. Notation.
  7. Preliminaries
  8. Suboptimality Decomposition
  9. Optimization error.
  10. Concentration error.
  11. Generalization error.
  12. Suboptimality Gap of Empirical Stationary Points
  13. RL with Sufficient Prompts
  14. RL without Sufficient Prompts
  15. Suboptimality Gap under Gradient-Based Algorithms
  16. ...and 33 more sections

Key Result

Proposition 1

For any reward function $r(x, a)$, the optimal policy $\pi^*$ that maximizes eq:expectation loss admits the following closed-form expression:

Figures (2)

  • Figure 1: The suboptimality gap under the feature coverage assumption. The parameters are empirical stationary points or obtained by gradient-based methods. Left panels (\ref{['fig:img1']}-\ref{['fig:img2']}): The suboptimality gap plotted against $n^{-\frac{1}{2}}$. Right panels (\ref{['fig:img3']}-\ref{['fig:img4']}): The suboptimality gap w.r.t. optimization steps $T$ in log-log scale.
  • Figure 2: The required samples to make empirical stationary points becomes ground truth parameters.

Theorems & Definitions (91)

  • Proposition 1: rafailov2023direct
  • Definition 1: linear reward function class
  • Definition 2: Posterior Boltzmann policy function class
  • Definition 3: Suboptimality Gap Decomposition
  • Corollary 1
  • Remark 1
  • Lemma 1
  • Definition 4: Generalization Error
  • Definition 5: Uniform Stability bousquet2002stability
  • Proposition 2: bousquet2020sharper
  • ...and 81 more