Distributionally Robust Reinforcement Learning with Human Feedback
Debmalya Mandal, Paulius Sasnauskas, Goran Radanovic
TL;DR
This work tackles the fragility of RLHF under distribution shifts by introducing distributionally robust RLHF (DRO-RLHF) that hardens both reward estimation and policy optimization against shifts in the prompt distribution. It develops minibatch SGD-based algorithms for distributionally robust reward estimation, robust policy optimization via a weighted natural policy gradient, and a robust DPO variant, accompanied by convergence guarantees under standard linear/log-linear assumptions. Empirical results on the Unified-Feedback suite show that DRO-RLHF improves out-of-distribution performance across reward models, PPO-style policy optimization, and DPO, with especially strong gains in reasoning tasks. The study demonstrates a practical, theoretically grounded pathway to more reliable LLM alignment under realistic distribution shifts, with potential extensions to other divergence measures and text-specific metrics.
Abstract
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.