Gecontroleerd decoderen van taalmodelen

Controlled Decoding from Language Models

October 25, 2023

Auteurs: Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Cheng, Michael Collins, Trevor Strohman, Jilin Chen, Alex Beutel, Ahmad Beirami

cs.AI

Samenvatting

Wij stellen controlled decoding (CD) voor, een nieuwe off-policy reinforcement learning-methode om de autoregressieve generatie van taalmodellen te sturen naar uitkomsten met een hoge beloning. CD lost een off-policy reinforcement learning-probleem op via een waardefunctie voor de beloning, die wij een prefix scorer noemen. De prefix scorer wordt tijdens inferentie gebruikt om de generatie te sturen naar uitkomsten met een hogere beloning. Wij tonen aan dat de prefix scorer getraind kan worden op (mogelijk) off-policy data om de verwachte beloning te voorspellen wanneer het decoderen wordt voortgezet vanuit een gedeeltelijk gedecodeerd antwoord. Empirisch demonstreren wij dat CD effectief is als controlemechanisme op een corpus van Reddit-gesprekken. Wij laten ook zien dat de modulariteit van het ontwerp van CD het mogelijk maakt om te sturen voor meerdere beloningen, waardoor een multi-objectief reinforcement learning-probleem effectief wordt opgelost zonder extra complexiteit. Tot slot tonen wij aan dat CD op een nieuwe blockwise-manier kan worden toegepast tijdens inferentie, opnieuw zonder dat er wijzigingen tijdens de training nodig zijn, waardoor de kloof tussen de populaire best-of-K-strategie en token-level reinforcement learning wordt overbrugd. Dit maakt CD een veelbelovende benadering voor het afstemmen van taalmodellen.

English

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.