Generazione Controllata Rapida da Modelli Linguistici con Campionamento di Reiezione Ponderato Adattivo

Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

April 7, 2025

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

Abstract

L'approccio dominante per generare da modelli linguistici soggetti a qualche vincolo è il decoding localmente vincolato (LCD), che campiona incrementamente i token a ogni passo temporale in modo che il vincolo non venga mai violato. Tipicamente, ciò si ottiene attraverso il mascheramento dei token: iterando sul vocabolario ed escludendo i token non conformi. Ci sono due problemi importanti con questo approccio. (i) Valutare il vincolo su ogni token può essere proibitivamente costoso -- i vocabolari dei modelli linguistici spesso superano i 100.000 token. (ii) Il LCD può distorcere la distribuzione globale sulle stringhe, campionando token basandosi solo su informazioni locali, anche se questi portano a percorsi senza uscita. Questo lavoro introduce un nuovo algoritmo che affronta entrambi questi problemi. Innanzitutto, per evitare di valutare un vincolo sull'intero vocabolario a ogni passo della generazione, proponiamo un algoritmo di campionamento per rifiuto adattivo che tipicamente richiede ordini di grandezza meno valutazioni del vincolo. In secondo luogo, mostriamo come questo algoritmo possa essere esteso per produrre stime a bassa varianza e non distorte dei pesi di importanza a un costo aggiuntivo molto ridotto -- stime che possono essere utilizzate in modo solido all'interno di algoritmi di Monte Carlo sequenziali precedentemente proposti per correggere il comportamento miope dell'applicazione locale dei vincoli. Attraverso una valutazione empirica estesa nei domini di text-to-SQL, sintesi molecolare, inferenza di obiettivi, pattern matching e JSON, dimostriamo che il nostro approccio è superiore rispetto alle baseline state-of-the-art, supportando una classe più ampia di vincoli e migliorando sia il tempo di esecuzione che le prestazioni. Ulteriori analisi teoriche ed empiriche mostrano che l'efficienza in termini di tempo di esecuzione del nostro metodo è guidata dal suo uso dinamico del calcolo, scalando con la divergenza tra il modello linguistico non vincolato e quello vincolato, e di conseguenza, i miglioramenti nel tempo di esecuzione sono maggiori per modelli migliori.

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.