Génération rapide et contrôlée à partir de modèles de langage avec échantillonnage de rejet pondéré adaptatif

Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

April 7, 2025

Auteurs: Benjamin Lipkin, Benjamin LeBrun, Jacob Hoover Vigly, João Loula, David R. MacIver, Li Du, Jason Eisner, Ryan Cotterell, Vikash Mansinghka, Timothy J. O'Donnell, Alexander K. Lew, Tim Vieira

cs.AI

Résumé

L'approche dominante pour générer des modèles de langage soumis à certaines contraintes est le décodage localement contraint (LCD), qui consiste à échantillonner des tokens de manière incrémentale à chaque étape temporelle de sorte que la contrainte ne soit jamais violée. Typiquement, cela est réalisé par masquage des tokens : en parcourant le vocabulaire et en excluant les tokens non conformes. Cette approche présente deux problèmes majeurs. (i) Évaluer la contrainte sur chaque token peut être prohibitivement coûteux — les vocabulaires des modèles de langage dépassent souvent 100 000 tokens. (ii) Le LCD peut fausser la distribution globale sur les chaînes de caractères, en échantillonnant des tokens uniquement sur la base d'informations locales, même s'ils mènent à des impasses. Ce travail introduit un nouvel algorithme qui résout ces deux problèmes. Premièrement, pour éviter d'évaluer une contrainte sur l'intégralité du vocabulaire à chaque étape de génération, nous proposons un algorithme d'échantillonnage par rejet adaptatif qui nécessite généralement des ordres de grandeur moins d'évaluations de contraintes. Deuxièmement, nous montrons comment cet algorithme peut être étendu pour produire des estimations à faible variance et non biaisées des poids d'importance à un coût supplémentaire très faible — des estimations qui peuvent être utilisées de manière fiable dans des algorithmes de Monte Carlo séquentiels précédemment proposés pour corriger le comportement myope de l'application locale des contraintes. Grâce à une évaluation empirique approfondie dans les domaines de la génération de texte vers SQL, de la synthèse moléculaire, de l'inférence d'objectifs, de la correspondance de motifs et du JSON, nous montrons que notre approche est supérieure aux méthodes de référence actuelles, supportant une classe plus large de contraintes et améliorant à la fois le temps d'exécution et les performances. Des analyses théoriques et empiriques supplémentaires montrent que l'efficacité en temps d'exécution de notre méthode est due à son utilisation dynamique des calculs, évoluant avec la divergence entre le modèle de langage non contraint et contraint, et par conséquent, les gains en temps d'exécution sont plus importants pour les meilleurs modèles.

English

The dominant approach to generating from language models subject to some constraint is locally constrained decoding (LCD), incrementally sampling tokens at each time step such that the constraint is never violated. Typically, this is achieved through token masking: looping over the vocabulary and excluding non-conforming tokens. There are two important problems with this approach. (i) Evaluating the constraint on every token can be prohibitively expensive -- LM vocabularies often exceed 100,000 tokens. (ii) LCD can distort the global distribution over strings, sampling tokens based only on local information, even if they lead down dead-end paths. This work introduces a new algorithm that addresses both these problems. First, to avoid evaluating a constraint on the full vocabulary at each step of generation, we propose an adaptive rejection sampling algorithm that typically requires orders of magnitude fewer constraint evaluations. Second, we show how this algorithm can be extended to produce low-variance, unbiased estimates of importance weights at a very small additional cost -- estimates that can be soundly used within previously proposed sequential Monte Carlo algorithms to correct for the myopic behavior of local constraint enforcement. Through extensive empirical evaluation in text-to-SQL, molecular synthesis, goal inference, pattern matching, and JSON domains, we show that our approach is superior to state-of-the-art baselines, supporting a broader class of constraints and improving both runtime and performance. Additional theoretical and empirical analyses show that our method's runtime efficiency is driven by its dynamic use of computation, scaling with the divergence between the unconstrained and constrained LM, and as a consequence, runtime improvements are greater for better models.