Optimizing Language Models with Fair and Stable Reward Composition in Reinforcement Learning (2024)

Abstract: Reinforcement learning from human feedback (RLHF) and AI-generated feedback (RLAIF) have become prominent techniques that significantly enhance the functionality of pre-trained language models (LMs). These methods harness feedback, sourced either from humans or AI, as direct rewards or to shape reward models that steer LM optimization. Nonetheless, the effective integration of rewards from diverse sources presents a significant challenge due to their disparate characteristics. To address this, recent research has developed algorithms incorporating strategies such as weighting, ranking, and constraining to handle this complexity. Despite these innovations, a bias toward disproportionately high rewards can still skew the reinforcement learning process and negatively impact LM performance. This paper explores a methodology for reward composition that enables simultaneous improvements in LMs across multiple dimensions. Inspired by fairness theory, we introduce a training algorithm that aims to reduce disparity and enhance stability among various rewards. Our method treats the aggregate reward as a dynamic weighted sum of individual rewards, with alternating updates to the weights and model parameters. For efficient and straightforward implementation, we employ an estimation technique rooted in the mirror descent method for weight updates, eliminating the need for gradient computations. The empirical results under various types of rewards across a wide range of scenarios demonstrate the effectiveness of our method.

## introduction in recent years, pretrained language models (lms) have made significant strides in the field of natural language processing, leading to their widespread use in downstream applications such as conversational agents @xcite @xcite figure 1 : an example of question answering with two types of rewards optimizing in different directions. @xcite , code generation @xcite @xcite , and machine translation @xcite . reinforcement learning from human feedback (rlhf) @xcite @xcite and reinforcement learning from ai feedback (rlaif) @xcite @xcite plays a critical role in this evolution, enhancing the models' ability to generate outputs that better align with human preferences and greatly increasing their versatility. the methods of rlhf and rlaif typically incorporate three principal stages. initially, there is supervised fine-tuning, which entails honing a foundational language model by utilizing a specialized dataset crafted for this purpose. following this, the second stage is the development of reward functions, which are designed to serve as surrogate indicators of human or ai judgments and preferences. subsequently, the language model serves as policy model and undergoes optimization via a reinforce-ment learning algorithm @xcite . it is evident that the reward model is crucial in ensuring the language model's outputs continually improves and adapts to evaluative standards, which in turn, directly impacts the efficacy of the reinforcement learning phase. in some scenarios, there may be multiple reward functions @xcite @xcite @xcite , as people may wish to assess and enhance the lm from various perspectives. however, two factors make this challenging. firstly, different reward functions evaluate text quality from various angles, but their insights are not entirely independent. secondly, it is difficult for lms to determine the specific optimization direction for each reward function since they only receive an aggregated reward score. considering the scenario depicted in figure 1 , where two reward functions are employed, let us examine the example question, "how many times have the lakers won the finals?" the language model may generate various responses, which the two reward functions assess based on factuality and completeness, providing guidance for optimizing the language model. the first response is given a factuality score of 1.0 and a completeness score of 0.3, indicating that it is entirely factual but lacks some detail. the second response receives a factuality score of 0.3 and a completeness score of 1.0, signaling that it is complete but contains inaccuracies. it becomes challenging to evaluate the overall success of the responses by merely aggregating the individual rewards because they each excel in different aspects. therefore, the development of an effective algorithm that can integrate various reward functions is of utmost significance for both research and practical applications. recent studies have investigated multiple methods for composing rewards, such as ranking @xcite , applying weightings @xcite , using welfare functions @xcite , and practicing safe reinforcement learning @xcite . while these approaches may involve complex designs and the fine-tuning of hyperparameters, there is a still high risk that the policy may overfocus on one reward function and neglect others, thereby negatively impacting the lm. to tackle the aforementioned challenges, we propose a method named fast rl (fair and stable reward reinforcement learning), which is designed for simple but effective integration of diverse re-wards. inspired by fairness theory @xcite , we have formulated a training objective that aims to minimize disparity and maximize stability across different reward functions simultaneously. drawing on the principles of distributionally robust optimization (dro) @xcite @xcite , we compute composite rewards as a weighted sum of individual rewards and transform the training objective into a max-min optimization problem. we iteratively optimize the language model and the weights, with the latter being updated via an estimation of the mirror descent method without the need for gradient computation. this strategy not only guides language models towards a more balanced, stable, and comprehensive improvement, but also offers simplicity in implementation. our contributions are summarized as follows: ## method in this section, we propose a method and its implementation that are designed to simultaneously improve the performance of lms across various reward functions. the entire training framework is depicted in figure 2 . initially, the input state is fed into the lm, and each reward function provides a score for evaluating the model outputs. these scores are then integrated using a weighted sum to obtain a composite reward, which is subsequently used to optimize the lm. concurrently, the weights are updated through an estimation based on mirror descent. ## experiment we evaluate the effectiveness of our method across various scenarios using different language models. our experiments encompass dialogue generation, question answering, and tasks aimed at mitigating harmfulness and enhancing helpfulness. our proposed method consistently outperforms the baselines across all experimental scenarios. details regarding hyperparameters are provided in the appendix. ## related work reinforcement learning for optimizing the language model. @xcite @xcite has emerged as a crucial methodology for fine-tuning language models to better reflect human intentions, as documented in various studies. its efficacy is demonstrated in downstream tasks such as summarization @xcite , story-telling (ziegler et al., 2019), following instructions, and reducing harm @xcite @xcite . however, rlhf involves gathering pairwise human-labeled data and an additional training procedure for the reward model, which can be resource-intensive. to optimize lms in a faster and more lightweight manner, recent studies have shifted toward applying reward learning from ai feedback (rlaif) @xcite @xcite . this approach leverages ai-generated feedback or provides direct reward signals, thus bypassing the need for extensive human-labeled datasets. furthermore, research by @xcite has found that reward-model-based approaches continue to hold their benefits, particularly when dealing with samples that are not well represented within the initial training preferences. these insights underscore the sustained importance of reinforcement learning in the enhancement of language models. ## conclusion in this study, we focus on the scenarios that involve complex, multi-faceted reward models for optimizing lms. given the diverse perspectives from which various reward models assess text, our aim is to develop a method that can appropriately compose different rewards, so as to ensure that lms do not excessively prioritize one perspective over others. leveraging fairness theory, we propose a method wherein the training objective is to reduce disparity and increase robustness among rewards. drawing on the principles of dro, we calculate composite rewards as a weighted sum of individual rewards, and transforms the training objective into a max-min optimization problem. the updating mechanism for the weights assigned to different rewards utilizes an estimation approach based on the mirror descent method, which is not only straightforward but also highly effective, simplifying the implementation process. the empirical results across various scenarios demonstrate the efficacy of our approach. limitations and future work. while our study yields promising results, it is not without its limitations. firstly, the absence of human evaluations is noteworthy; we have relied on gpt-4 as a standin, but this may not reflect human judgment with complete accuracy. furthermore, the potential for conflict and inaccuracy arises from using outputs of various reward models, as our current approach does not have a mechanism to distinguish between the efficacy of these models. instead, we calculate a composite reward in an effort to concurrently boost performance across all reward signals, which unfortunately may lead to less-than-ideal outcomes. in our future research, we plan to enhance our methodology by developing and implementing theoretical frameworks designed to detect and eliminate superfluous rewards.