Optimal Design for Reward Modeling in RLHF

TL;DR

The paper formalizes dataset design for RLHF reward modeling as a simple-regret minimization problem under a contextual linear reward with Bradley–Terry dueling feedback. It introduces ODPO, an offline optimal-design method using Frank–Wolfe to select the most informative prompt–completion pairs, and provides a near-optimal simple-regret bound with a matching lower bound, all without online feedback. The approach relies on a linear embedding of prompt–completion pairs and a projected MLE for the reward parameter, yielding a one-shot, offline training pipeline. The results establish theoretical guarantees for data-efficient reward modeling in RLHF and offer practical guidance for constructing human-label datasets.

Abstract

Reinforcement Learning from Human Feedback (RLHF) has become a popular approach to align language models (LMs) with human preferences. This method involves collecting a large dataset of human pairwise preferences across various text generations and using it to infer (implicitly or explicitly) a reward model. Numerous methods have been proposed to learn the reward model and align a LM with it. However, the costly process of collecting human preferences has received little attention and could benefit from theoretical insights. This paper addresses this issue and aims to formalize the reward training model in RLHF. We frame the selection of an effective dataset as a simple regret minimization task, using a linear contextual dueling bandit method. Given the potentially large number of arms, this approach is more coherent than the best-arm identification setting. We then propose an offline framework for solving this problem. Under appropriate assumptions - linearity of the reward model in the embedding space, and boundedness of the reward parameter - we derive bounds on the simple regret. Finally, we provide a lower bound that matches our upper bound up to constant and logarithmic terms. To our knowledge, this is the first theoretical contribution in this area to provide an offline approach as well as worst-case guarantees.

Figures (2)

• Figure 1: Illustration of ODPO among the whole RLHF framework.
• Figure 2: Note that on this figure, $a_2$ and $a_3$ are optimal, alghough playing both of them for the duel will provide feedback of very low value since they lie in the same region of $\mathbb{R}^d$. This is why a duel between $a_1$ and another arm is of greater interest in our exploration setup: sampling good arms is not the optimal strategy, hence the link with pure exploration.

Theorems & Definitions (24)

• Lemma 1
• Lemma 2
• Theorem 1
• Corollary 1
• Theorem 2
• Lemma 3
• Theorem 3
• Lemma 4
• proof : Proof of \ref{['lemma:bounding_sigmoid']}
• Lemma 4