Surrogate rewards improve steerability of generative models without oracle queries
Steered Not Stirred: Effective Reward Tilting using Surrogate Reward Mixing
Machine Learning
Summary
Some computer programs create things like text or images and can be guided by rewards to make better results. But sometimes the best rewards, like human opinions, are hard to get. The authors found a way to combine several easier-to-get rewards to guide the program toward what a hard-to-get reward would want. This method, called MARTINI, works with pre-trained models and uses only a little feedback from the hard reward. They showed it helps make better language, images, and protein designs without needing the hard feedback at every step.
What this means in practice
- •For language model engineers: Improve language generation outputs by steering pretrained models using surrogate reward signals without needing expensive human feedback.
- •For protein design teams: Enhance accuracy and validity in protein-ligand cofolding predictions by combining surrogate rewards to guide pretrained generative models at inference time.
Authors
Alec Sargood, Olivia Gallup, Georg Maierhofer, Noam Ghenassia, Paul Duckworth, Krisztina Sinkovics, Avishek Joey Bose
Abstract
Continuous-time generative models, from diffusion models, flow matching, and now flow maps, owe their practical value to high-fidelity generation and steerability. In many applications, a ground-truth latent ``oracle'' reward $R^*$ is unavailable or expensive to evaluate, e.g., human judgments or experimental measurements. Although the oracle reward is not directly observable, it can often be approximated through multiple cheaper surrogates that, individually, provide incomplete characterizations but collectively reveal substantial information about the downstream task. In this paper, we present a method for composing such surrogates to steer a pretrained model toward the distribution tilted by an unavailable $R^*$. Given a small feedback set of oracle rankings or numerical rewards, we introduce Mixing And Reward Tilting for INference-time Improvement (MARTINI), which learns a surrogate composition for plug-and-play use with standard inference-time steering methods. Our theory characterizes surrogate information limits and relates population recovery error to ideal-target KL at small tilt strength. For preferences, the bound additionally accounts for an irreducible mismatch between the Plackett-Luce population target and the conditional mean. Experiments demonstrate improved oracle-tilt fidelity in language generation and ImageNet alignment, and improved joint pose accuracy and validity in protein-ligand cofolding, without inference-time oracle queries.