Decodificación Aumentada por Recompensas: Generación Controlada de Texto Eficiente con un Modelo de Recompensa Unidireccional

Reward-Augmented Decoding: Efficient Controlled Text Generation With a Unidirectional Reward Model

October 14, 2023

Autores: Haikang Deng, Colin Raffel

cs.AI

Resumen

Si bien los modelos de lenguaje de gran escala han demostrado ser efectivos en una amplia gama de aplicaciones posteriores, a menudo generan texto que es problemático o carece de un atributo deseado. En este artículo, presentamos la Decodificación Aumentada por Recompensas (Reward-Augmented Decoding, RAD), un procedimiento de generación de texto que utiliza un pequeño modelo de recompensa unidireccional para incentivar a un modelo de lenguaje a generar texto con ciertas propiedades. Específicamente, RAD utiliza el modelo de recompensa para puntuar las generaciones a medida que se producen y reescala las probabilidades de muestreo para favorecer tokens con alta recompensa. Al emplear un modelo de recompensa unidireccional, RAD puede almacenar en caché las activaciones de pasos previos de generación para reducir la sobrecarga computacional. A través de experimentos sobre la generación de texto no tóxico y controlado por sentimiento, demostramos que RAD supera a los métodos que solo modifican el procedimiento de generación y se equipara al rendimiento de los métodos más avanzados que implican reentrenar el modelo de lenguaje. Además, validamos que RAD es efectivo en modelos de lenguaje muy grandes mientras incurre en una sobrecarga computacional mínima.

English

While large language models have proven effective in a huge range of downstream applications, they often generate text that is problematic or lacks a desired attribute. In this paper, we introduce Reward-Augmented Decoding (RAD), a text generation procedure that uses a small unidirectional reward model to encourage a language model to generate text that has certain properties. Specifically, RAD uses the reward model to score generations as they are produced and rescales sampling probabilities to favor high-reward tokens. By using a unidirectional reward model, RAD can cache activations from prior generation steps to decrease computational overhead. Through experiments on generating non-toxic and sentiment-controlled text, we demonstrate that RAD performs best among methods that change only the generation procedure and matches the performance of state-of-the-art methods that involve re-training the language model. We further validate that RAD is effective on very large language models while incurring a minimal computational overhead.