Decodificación Controlada en Modelos de Lenguaje

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

Resumen

Proponemos decodificación controlada (CD, por sus siglas en inglés), un novedoso método de aprendizaje por refuerzo fuera de política para controlar la generación autoregresiva de modelos de lenguaje hacia resultados de alta recompensa. CD resuelve un problema de aprendizaje por refuerzo fuera de política mediante una función de valor para la recompensa, a la que denominamos evaluador de prefijos. Este evaluador de prefijos se utiliza en tiempo de inferencia para guiar la generación hacia resultados de mayor recompensa. Demostramos que el evaluador de prefijos puede entrenarse con datos (posiblemente) fuera de política para predecir la recompensa esperada cuando la decodificación continúa a partir de una respuesta parcialmente decodificada. Empíricamente, mostramos que CD es efectivo como mecanismo de control en el corpus de conversaciones de Reddit. También demostramos que la modularidad del diseño de CD permite controlar múltiples recompensas, resolviendo eficazmente un problema de aprendizaje por refuerzo multiobjetivo sin complejidad adicional. Finalmente, mostramos que CD puede aplicarse de manera novedosa en bloques durante la inferencia, nuevamente sin necesidad de cambios en el entrenamiento, cerrando esencialmente la brecha entre la popular estrategia de mejor-de-K y el aprendizaje por refuerzo a nivel de tokens. Esto convierte a CD en un enfoque prometedor para la alineación de modelos de lenguaje.

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.