Envelope sampling improves model training with better reward estimates
How to post-train on a surrogate: Envelope sampling mitigates reward hacking
Machine LearningArtificial Intelligence
Summary
When training language models to respond well, it's costly to get true human feedback, so cheaper automatic judges are used instead. But these judges can be unreliable and cause the model to learn bad behaviors, a problem called reward hacking. The authors propose a new sampling method, envelope sampling, that helps recalibrate these automatic judges using a small set of expert reviews. Their method makes the model learn better and avoids reward hacking, as shown in experiments with clinical text and a specialized behavior test.
What this means in practice
- •For clinical ai developers: Improve clinical note generation models by recalibrating reward signals with envelope sampling to reduce undesirable outputs.
- •For ai safety teams: Use envelope sampling to detect and mitigate reward hacking in AI systems that rely on imperfect automated judges.
Authors
Sanjit Dandapanthula, Shuvom Sadhuka, Samir Khan, Michael Oberst, Aaditya Ramdas, Alexandra Chouldechova
Abstract
Large language models (LLMs) are commonly post-trained against LLM judges and other cheap surrogates because the true reward, such as human preference, is too expensive to query at scale. This practice often leads to reward hacking, where reinforcement learning against a miscalibrated surrogate leads to undesirable side effects. In this work, we study a setting in which a small number $n$ of model outputs are annotated with ground-truth labels (e.g., from expert review) and used to recalibrate the LLM judge before optimizing against it. Prior approaches to judge recalibration are costly or heuristic, and it is known that on-policy sampling fails when the surrogate is miscalibrated on a rare set of outputs. In this work, we propose envelope sampling, a theoretically-grounded method for judge recalibration that seeks to minimize an upper bound on the regret of the post-trained model under the assumption that the human reward and re-calibrated reward lie in an $L^2$ ball around the judge. We give practical algorithms to sample from the envelope by rejection or by fine-tuning against a modified reward, and experiments on clinical note generation and on a controlled sycophancy task show that recalibrating on envelope samples mitigates reward hacking where recalibrating on base-model samples does not.