Towards P areto-Efficient RLHF : Paying Attention to a Few High-Reward Samples with Reward Dropout (2024)

Abstract: Recently, leveraging reinforcement learning (RL) to fine-tune language models (LMs), known as reinforcement learning from human feedback (RLHF), has become an important research topic. However, there is still a lack of theoretical understanding of how RLHF works, the conditions under which it succeeds or fails, and whether it guarantees optimization of both likelihood \beta(\cdot) and reward R(\cdot) objectives. To address these issues, we consider RLHF as a bi-objective problem that has the nature of a Pareto optimization, present a Pareto improvement condition that is necessary to obtain Pareto-efficient policies, and propose a simple yet powerful method named reward dropout that guarantees a Pareto improvement. To demonstrate the performance of reward dropout, two benchmark datasets commonly used in text style transfer tasks were utilized in our study: sentiment and topic datasets sourced from Yelp and AG_News, respectively. Our experiments highlight that paying attention to a few samples with higher rewards leads to greater Pareto improvements regardless of model size. We also demonstrate that the effect of reward dropout is generalizable and most effective with non-pretrained target models, saving the effort of pretraining.

## introduction the emergence of chatgpt has sparked public interest in language models (lms), leading to a surge in lm research in both academia and industry. recently, leveraging reinforcement learning (rl) to fine-tune lms, known as reinforcement learning from human feedback (rlhf), has become an important research topic. this approach aims to generate reliable sequences with desired attributes by simultaneously maximizing the reward objective r(•) and the likelihood β(•) of reference or behavior lms @xcite figure 1 : pareto improvement in rlhf. π 0 and π * are the local and global pareto optimal solutions, respectively. moving from π 0 to either 1 ⃝ or 2 ⃝ always sacrifices one objective, whereas moving to 3 ⃝, that is π * , does not. the latter case is a pareto improvement. the purple hatched area represents an expanded feasible region and the orange shaded area is a region where pareto improvement is available. despite its success and popularity, there is still a lack of theoretical understanding of how rlhf works, the conditions under which it succeeds or fails, and whether it guarantees optimization of both likelihood and reward objectives. to address these issues, we study the theoretical aspects of rlhf through the lens of a pareto optimization. specifically, we consider rlhf as a bi-objective problem that has the nature of a pareto optimization (see §2). then, we analyze the objective function of rlhf from pareto optimization perspectives and present a pareto improvement condition that is necessary to obtain pareto-efficient policies (see §4). based on the analysis, we propose a simple yet powerful method called reward dropout (see §5) and evaluate it on two benchmark datasets with six control attributes (see §6 and §7). the contributions of our study can be summarized as follows: ## analysis analyzing the region of feasible solutions can help us to identify pareto-dominant policies. in this section, we derive the gradient of eq (1), examine the feasible region of rlhf, and identify the condition necessary to achieve pareto-efficient policies. ## reward dropout in this section, we introduce a simple yet powerful method called the reward dropout. this method aims to achieve a pareto-efficient rlhf by enforcing pareto improvements. rationale since the entropy rewardln π θ (τ ) is always positive, eq ( 5) ensures that r(τ ) remains positive for all τ , leading to pareto improvement. consequently, we can obtain pareto-efficient policies by considering only samples where the bi-objective reward r(τ ) + ln β(τ ) exceeds any positive real-valued threshold δ: r(τ ) + ln β(τ ) > δ where δ ∈ r + 0 . (6) implementation the principle behind eq ( 6) is straightforward. first, set a threshold δ greater than zero. second, exclude samples with bi-objective rewards below δ from each training batch. in practice, this means retaining only a few bi-objective rewards above δ and setting the rest to zero. to elaborate, bi-objective rewards are sorted in ascending order within each batch, divided into equal intervals, and those below δ are set to zero. we refer to this technique as reward dropout because samples are dropped out based on their rewards being set to zero. see algorithm 1 for details. ## experiments in this study, we demonstrate the performance of reward dropout on benchmark datasets and test its validity across different configurations. ## limitations & future works our work raises important questions for future rlhf studies. for example, since rlhf can be framed as a pareto optimization problem, it would be fascinating to explore whether traditional algorithms for pareto optimization could be applied to rlhf. specifically, while we assumed λ = 1 for analytical convenience, future research should focus on optimizing λ as a variable or finding ways to naturally cancel it out during optimization. furthermore, theorem 4.2 and our results prove that reward dropout guarantees policy updates in the direction of pareto improvement. however, this does not ensure convergence to a pareto-efficient policy. therefore, proving the existence of paretoefficient policies and theoretically analyzing the minimum number of training iterations, i.e., computational complexity, required to achieve them are compelling topics for future research. we hope to address these topics in subsequent studies. our evaluation leaves some areas unexplored. it is necessary to validate the performance of reward dropout across more attributes and language models. although we evaluated models up to 7b parameters, assessing the effect of reward dropout on much larger, state-of-the-art models would be an intriguing future direction. additionally, it would be valuable to experiment with reward dropout in conjunction with auxiliary methods such as reject sampling or self-training mechanism @xcite @xcite . ## conclusion in this study, we established a theoretical foundation for rlhf from a bi-objective perspective, proposed a simple yet powerful method called reward dropout and empirically demonstrated its effectiveness in various aspects. as the first study on pareto-efficient rlhf, we hope our work will help address major rlhf challenges, such as misalignment, the curse of recursion, and catastrophic forgetting, based on pareto optimization. as described in §5, we applied reward dropout in an off-policy manner (see @xcite and algorithm 1). this could raise concerns about the generalizability of reward dropout. specifically, pareto improvement, indicated by an increase in π, might only occur in drop-in samples. since the sample distribution is governed by the behavior model, this suggests that the pareto improvement in the target models could be heavily dependent on the sampling distribution of the behavior model, leading to concerns that the drop-out samples will never benefit from reward dropout and the effect of reward dropout cannot be generalized.