Reward Weighted Classifier-Free Guidance as Policy Improvement in Autoregressive Models

TL;DR

This paper introduces Reward Weighted Classifier-Free Guidance (RCFG), a method for policy improvement in autoregressive models that allows optimizing new reward functions at test time without retraining.

Contribution

The paper proposes RCFG as a policy improvement operator that approximates distribution tilting via the Q function, enabling flexible test-time reward optimization.

Findings

RCFG effectively optimizes novel reward functions in molecular generation.

Using RCFG as a teacher accelerates convergence in reinforcement learning.

RCFG can adapt to changing reward functions without retraining the model.

Abstract

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.