Towards a Theoretical Understanding to the Generalization of RLHF

TL;DR

This paper develops a theoretical framework to understand the generalization capabilities of Reinforcement Learning from Human Feedback (RLHF) in training large language models, focusing on linear reward models and end-to-end learning.

Contribution

It introduces a new generalization theory for RLHF based on algorithmic stability, applicable to gradient-based algorithms, and provides bounds under a feature coverage condition.

Findings

Empirical optima of policy models have a generalization bound of order O(n^{-1/2})

Results extend to gradient-based learning algorithms like GA and SGA

Provides theoretical support for observed generalization in LLMs after RLHF

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.