Distributionally Robust Reinforcement Learning from Human Feedback

Reinforcement learning from human feedback (RLHF) has evolved to be one of the main methods for fine-tuning large language models (LLMs). However, existing RLHF methods are non-robust, and their performance deteriorates if the downstream task differs significantly from the preference dataset used in fine-tuning. In order to mitigate this problem, we introduce a distributionally robust RLHF for fine-tuning LLMs. In particular, our goal is to ensure that a fine-tuned model retains its performance even when the distribution of prompts significantly differs from the distribution encountered during fine-tuning. We formulate distributionally robust optimization (DRO) version of two popular fine-tuning methods -- (1) reward-based RLHF and (2) reward-free DPO (direct preference optimization). We propose a minibatch gradient descent based algorithms for both of them, and theoretically prove convergence guarantees for the algorithms. Subsequently, we evaluate our algorithms on an out-of-distribution (OOD) task by first training the model on the Unified-Feedback dataset and evaluating its performance on two different datasets. The experimental results show that our robust training improves the accuracy of the learned reward models on average, and markedly on some tasks, such as reasoning. Furthermore, we show that the robust versions of policy optimization methods, similarly improve performance on OOD tasks.

Large language models are typically fine-tuned using examples of human preferences. However, they are brittle when the questions they encounter after deployment look different from the ones they were trained on. This paper develops a new training approach that deliberately prepares the model for such shifts by, during training, focusing extra attention on the kinds of examples the model finds hardest, rather than treating every training example equally. This work tests the proposed method on two open-source models (a 2-billion-parameter Gemma model and a 7-billion-parameter Mistral model), and finds that when the trained models are evaluated on tasks drawn from datasets they were not trained on, the new method matches or beats standard training across the board. The improvements are especially striking on reasoning tasks, where accuracy jumps substantially. The work suggests a simple, drop-in change to existing fine-tuning pipelines can make language models meaningfully more reliable when the real world doesn't match the training data.