Décodage Augmenté par Récompense : Génération de Texte Contrôlée Efficace avec un Modèle de Récompense Unidirectionnel

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

October 14, 2023

Auteurs: Haikang Deng, Colin Raffel

cs.AI

Résumé

Bien que les grands modèles de langage se soient avérés efficaces dans un large éventail d'applications en aval, ils génèrent souvent des textes problématiques ou dépourvus d'un attribut souhaité. Dans cet article, nous présentons le décodage augmenté par récompense (Reward-Augmented Decoding, RAD), une procédure de génération de texte qui utilise un petit modèle de récompense unidirectionnel pour encourager un modèle de langage à produire des textes possédant certaines propriétés. Plus précisément, RAD utilise le modèle de récompense pour évaluer les générations au fur et à mesure de leur production et réajuste les probabilités d'échantillonnage pour favoriser les tokens à forte récompense. En utilisant un modèle de récompense unidirectionnel, RAD peut mettre en cache les activations des étapes de génération précédentes afin de réduire la surcharge computationnelle. À travers des expériences sur la génération de textes non toxiques et contrôlés en termes de sentiment, nous démontrons que RAD obtient les meilleurs résultats parmi les méthodes qui modifient uniquement la procédure de génération et rivalise avec les performances des méthodes de pointe impliquant un réentraînement du modèle de langage. Nous validons en outre que RAD est efficace sur des modèles de langage très volumineux tout en induisant une surcharge computationnelle minimale.

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.