Scalable Ensembling For Mitigating Reward Overoptimisation
Ahmed M. Ahmed, Rafael Rafailov, Stepan Sharkov, Xuechen Li, Sanmi Koyejo
TL;DR
The paper tackles reward overoptimisation in RLHF for large language models by proposing a scalable ensemble method that uses a shared encoder with multiple linear reward heads and a minimum operator to create a pessimistic yet efficient reward signal. Through PPO-based RLHF experiments on the AlpacaFarm dataset, the proposed multi-head reward model achieves comparable performance to a full ensemble while substantially reducing memory and training time. The findings include favorable calibration properties and guidance that a small number of heads (around three) may be optimal. This approach enables more scalable and robust alignment of instruction-following models with human preferences, without the prohibitive costs of full ensemble methods.
Abstract
Reinforcement Learning from Human Feedback (RLHF) has enabled significant advancements within language modeling for powerful, instruction-following models. However, the alignment of these models remains a pressing challenge as the policy tends to overfit the learned ``proxy" reward model past an inflection point of utility as measured by a ``gold" reward model that is more performant -- a phenomenon known as overoptimisation. Prior work has mitigated this issue by computing a pessimistic statistic over an ensemble of reward models, which is common in Offline Reinforcement Learning but incredibly costly for language models with high memory requirements, making such approaches infeasible for sufficiently large models. To this end, we propose using a shared encoder but separate linear heads. We find this leads to similar performance as the full ensemble while allowing tremendous savings in memory and time required for training for models of similar size.