Reward Weighted Classifier-Free Guidance as Policy Improvement in Autoregressive Models
For practitioners using autoregressive models (e.g., in drug discovery), this provides a method to adapt to new reward functions at test time without retraining, and accelerates RL fine-tuning.
The paper shows that reward weighted classifier-free guidance (RCFG) can approximate tilting the sampling distribution by the Q function in autoregressive models, enabling test-time optimization of novel reward functions without retraining. Applied to molecular generation, RCFG optimizes new rewards at test time and, when distilled into the base policy, significantly speeds up convergence for standard RL.
Consider an auto-regressive model that produces outputs x (e.g., answers to questions, molecules) each of which can be summarized by an attribute vector y (e.g., helpfulness vs. harmlessness, or bio-availability vs. lipophilicity). An arbitrary reward function r(y) encodes tradeoffs between these properties. Typically, tilting the model's sampling distribution to increase this reward is done at training time via reinforcement learning. However, if the reward function changes, re-alignment requires re-training. In this paper, we show that a reward weighted classifier-free guidance (RCFG) can act as a policy improvement operator in this setting, approximating tilting the sampling distribution by the Q function. We apply RCFG to molecular generation, demonstrating that it can optimize novel reward functions at test time. Finally, we show that using RCFG as a teacher and distilling into the base policy to serve as a warm start significantly speeds up convergence for standard RL.