Utility-inspired Reward Transformations Improve Reinforcement Learning Training of Language Models (2501.06248v2)

Abstract: Current methods that train LLMs with reinforcement learning feedback, often resort to averaging outputs of multiple rewards functions during training. This overlooks crucial aspects of individual reward dimensions and inter-reward dependencies that can lead to sub-optimal outcomes in generations. In this work, we show how linear aggregation of rewards exhibits some vulnerabilities that can lead to undesired properties of generated text. We then propose a transformation of reward functions inspired by economic theory of utility functions (specifically Inada conditions), that enhances sensitivity to low reward values while diminishing sensitivity to already high values. We compare our approach to the existing baseline methods that linearly aggregate rewards and show how the Inada-inspired reward feedback is superior to traditional weighted averaging. We quantitatively and qualitatively analyse the difference in the methods, and see that models trained with Inada-transformations score as more helpful while being less harmful.

• The paper introduces a novel reward transformation inspired by economic utility theory to better balance multiple reward signals during RL training.
• It leverages Inada conditions and CRRA functions to adjust reward sensitivity, prioritizing underperforming dimensions for improved language model performance.
• Experimental results using the Gemma 2B model on Anthropic's datasets demonstrate enhanced outcomes in both helpfulness and harmlessness with the proposed method.

Overview of "Utility-inspired Reward Transformations Improve Reinforcement Learning Training of LLMs"

The paper addresses a significant aspect of fine-tuning LLMs using reinforcement learning from human feedback (RLHF) by focusing on the transformation of reward functions. The authors critique the prevalent strategy of linearly aggregating multiple reward signals for LLM training, which often fails to capture the intricacies of individual reward components. This can culminate in suboptimal model performance, particularly when these models generate text without adequately balancing factors like helpfulness and harmlessness.

To enhance the training process, the authors introduce a novel approach inspired by the economic theory of utility functions, specifically the Inada conditions, which are designed to handle diminishing returns in utility. This reward transformation prioritizes improvements in dimensions where the model's performance is critically below a desired threshold while diminishing attention on dimensions already performing satisfactorily. By leveraging constant relative risk aversion (CRRA) utility functions, the proposed method adjusts the sensitivity of rewards to retrain LLMs more effectively.

Methodology and Results

The Inada Reward Transformation (IRT) is presented as a superior method over linear aggregation methods for balancing conflicting dimensions of rewards. IRT operates by transforming individual reward values before aggregation, aligning them closer to economic principles of human utility optimization. Empirically, models fine-tuned with IRT exhibit improved performances in generating more helpful and less harmful text.

Experimental evaluations are conducted using the Gemma 2B model on Anthropic's Helpfulness and Harmlessness datasets, incorporating feedback from learned reward models. The use of LLMs as evaluators further underpins the practicality of the proposed method, showing preferred outcomes from models enhanced with IRT over those relying on traditional linear combinations of reward functions.

Implications and Future Work

The implications of this research are both practical and theoretical. Practically, the introduction of utility-inspired transformations in the reward aggregation process can reduce safety risks and enhance alignment of AI systems with complex human values, making them more robust against harmful behavior. Theoretically, it paves the way for further exploration into economic theories for improving AI design. It suggests a viable path for integrating advanced economic concepts with AI training methods, potentially leading to AI systems capable of handling nuanced human feedback more effectively.

Future research may focus on extending this approach to larger, state-of-the-art models and exploring other economic principles that might enhance AI alignment. Additional studies could also refine the parameter choices for IRT, as well as investigate learnable thresholds and penalty factors for dynamic adaptability in diverse contexts. Integrating this method with other alignment techniques may improve both safety and utility, contributing to the broader field of AI ethics and governance.