Controlled Decoding from Language Models
Authors: name: Sidharth Mudgal; name: Jong Lee; name: Harish Ganapathy; name: YaGuang Li; name: Tao Wang; name: Yanping Huang; name: Zhifeng Chen; name: Heng-Tze Cheng; name: Michael Collins; name: Trevor Strohman; name: Jilin Chen; name: Alex Beutel; name: Ahmad Beirami
Published: 2023-10-25 22:00:05

Abstract: We propose controlled decoding (CD), a novel off-policy reinforcement
learning method to control the autoregressive generation from language models
towards high reward outcomes. CD solves an off-policy reinforcement learning
problem through a value function for the reward, which we call a prefix scorer.
The prefix scorer is used at inference time to steer the generation towards
higher reward outcomes. We show that the prefix scorer may be trained on
(possibly) off-policy data to predict the expected reward when decoding is
continued from a partially decoded response. We empirically demonstrate that CD
is effective as a control mechanism on Reddit conversations corpus. We also
show that the modularity of the design of CD makes it possible to control for
multiple rewards, effectively solving a multi-objective reinforcement learning
problem with no additional complexity. Finally, we show that CD can be applied
in a novel blockwise fashion at inference-time, again without the need for any
training-time changes, essentially bridging the gap between the popular
best-of-K strategy and token-level reinforcement learning. This makes CD a
promising approach for alignment of language models.