Decodifica Controllata da Modelli Linguistici

Controlled Decoding from Language Models

October 25, 2023

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

Abstract

Proponiamo il decoding controllato (CD), un nuovo metodo di apprendimento per rinforzo off-policy per controllare la generazione autoregressiva da modelli linguistici verso esiti ad alta ricompensa. CD risolve un problema di apprendimento per rinforzo off-policy attraverso una funzione di valore per la ricompensa, che chiamiamo prefix scorer. Il prefix scorer viene utilizzato al momento dell'inferenza per orientare la generazione verso esiti con ricompensa più elevata. Dimostriamo che il prefix scorer può essere addestrato su dati (possibilmente) off-policy per prevedere la ricompensa attesa quando il decoding continua da una risposta parzialmente decodificata. Empiricamente, dimostriamo che CD è efficace come meccanismo di controllo sul corpus di conversazioni di Reddit. Mostriamo inoltre che la modularità del design di CD rende possibile controllare più ricompense, risolvendo efficacemente un problema di apprendimento per rinforzo multi-obiettivo senza alcuna complessità aggiuntiva. Infine, dimostriamo che CD può essere applicato in modo innovativo a blocchi al momento dell'inferenza, ancora una volta senza la necessità di modifiche durante l'addestramento, colmando essenzialmente il divario tra la popolare strategia best-of-K e l'apprendimento per rinforzo a livello di token. Questo rende CD un approccio promettente per l'allineamento dei modelli linguistici.

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.