Snelle Gecontroleerde Generatie uit Taalmodellen met Adaptieve Gewogen Afkeuringssteekproef

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

Samenvatting

De dominante benadering voor het genereren van taalmodellen onder een bepaalde beperking is lokaal beperkte decodering (LCD), waarbij tokens incrementeel worden gesamplet op elk tijdstip zodat de beperking nooit wordt geschonden. Dit wordt typisch bereikt door middel van tokenmaskering: het doorlopen van de vocabulaire en het uitsluiten van niet-conforme tokens. Er zijn twee belangrijke problemen met deze aanpak. (i) Het evalueren van de beperking op elk token kan buitensporig duur zijn -- de vocabulaire van taalmodellen overschrijdt vaak 100.000 tokens. (ii) LCD kan de globale verdeling over strings verstoren, waarbij tokens alleen worden gesamplet op basis van lokale informatie, zelfs als deze leiden tot doodlopende paden. Dit werk introduceert een nieuw algoritme dat beide problemen aanpakt. Ten eerste, om te voorkomen dat een beperking op de volledige vocabulaire bij elke stap van de generatie wordt geëvalueerd, stellen we een adaptief afwijzingssampling-algoritme voor dat typisch ordes van grootte minder beperkingsevaluaties vereist. Ten tweede laten we zien hoe dit algoritme kan worden uitgebreid om schattingen met lage variantie en zonder bias van belangrijkheidsgewichten te produceren tegen zeer geringe extra kosten -- schattingen die veilig kunnen worden gebruikt binnen eerder voorgestelde sequentiële Monte Carlo-algoritmen om het kortzichtige gedrag van lokale beperkingshandhaving te corrigeren. Door uitgebreide empirische evaluatie in domeinen zoals tekst-naar-SQL, moleculaire synthese, doelafleiding, patroonherkenning en JSON, laten we zien dat onze aanpak superieur is aan state-of-the-art baselines, waarbij een bredere klasse van beperkingen wordt ondersteund en zowel de looptijd als de prestaties worden verbeterd. Aanvullende theoretische en empirische analyses tonen aan dat de looptijdefficiëntie van onze methode wordt gedreven door het dynamische gebruik van berekeningen, waarbij deze schaalt met de divergentie tussen het onbeperkte en beperkte taalmodel, en als gevolg daarvan zijn de looptijdverbeteringen groter voor betere modellen.

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.