Decodificação Controlada em Modelos de Linguagem

Controlled Decoding from Language Models

October 25, 2023

Autores: 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

Resumo

Propomos o decodificação controlada (CD, do inglês "controlled decoding"), um novo método de aprendizado por reforço off-policy para controlar a geração autoregressiva de modelos de linguagem em direção a resultados de alta recompensa. O CD resolve um problema de aprendizado por reforço off-policy por meio de uma função de valor para a recompensa, que chamamos de avaliador de prefixo. O avaliador de prefixo é utilizado no momento da inferência para direcionar a geração em direção a resultados de maior recompensa. Mostramos que o avaliador de prefixo pode ser treinado com dados (possivelmente) off-policy para prever a recompensa esperada quando a decodificação é continuada a partir de uma resposta parcialmente decodificada. Demonstramos empiricamente que o CD é eficaz como mecanismo de controle no corpus de conversas do Reddit. Também mostramos que a modularidade do design do CD torna possível controlar múltiplas recompensas, resolvendo efetivamente um problema de aprendizado por reforço multiobjetivo sem complexidade adicional. Por fim, mostramos que o CD pode ser aplicado de forma inovadora em blocos durante a inferência, novamente sem a necessidade de alterações no treinamento, essencialmente preenchendo a lacuna entre a popular estratégia best-of-K e o aprendizado por reforço em nível de token. Isso torna o CD uma abordagem promissora para o alinhamento de modelos de linguagem.

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.