Schnelle kontrollierte Generierung aus Sprachmodellen mit adaptiv gewichtetem Rejection Sampling

Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

April 7, 2025

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

Zusammenfassung

Der dominante Ansatz zur Generierung aus Sprachmodellen unter Berücksichtigung einer bestimmten Einschränkung ist das lokal eingeschränkte Decodieren (LCD), bei dem Token in jedem Zeitschritt inkrementell so abgetastet werden, dass die Einschränkung niemals verletzt wird. Typischerweise wird dies durch Token-Masking erreicht: Es wird über das Vokabular iteriert und nicht konforme Token ausgeschlossen. Es gibt zwei wichtige Probleme bei diesem Ansatz. (i) Die Bewertung der Einschränkung für jedes Token kann unverhältnismäßig teuer sein – die Vokabulare von Sprachmodellen überschreiten oft 100.000 Token. (ii) LCD kann die globale Verteilung über Zeichenketten verzerren, indem Token nur basierend auf lokalen Informationen abgetastet werden, selbst wenn sie in Sackgassen führen. Diese Arbeit stellt einen neuen Algorithmus vor, der beide Probleme adressiert. Erstens schlagen wir, um die Bewertung einer Einschränkung auf dem gesamten Vokabular in jedem Generierungsschritt zu vermeiden, einen adaptiven Rejection-Sampling-Algorithmus vor, der typischerweise um Größenordnungen weniger Einschränkungsbewertungen erfordert. Zweitens zeigen wir, wie dieser Algorithmus erweitert werden kann, um mit sehr geringem zusätzlichem Aufwand Schätzungen mit geringer Varianz und unverzerrten Gewichten zu erzeugen – Schätzungen, die sicher innerhalb zuvor vorgeschlagener sequentieller Monte-Carlo-Algorithmen verwendet werden können, um das kurzsichtige Verhalten der lokalen Einschränkungsdurchsetzung zu korrigieren. Durch umfangreiche empirische Bewertungen in den Bereichen Text-zu-SQL, Molekülsynthese, Zielinferenz, Mustererkennung und JSON-Domänen zeigen wir, dass unser Ansatz den state-of-the-art Baselines überlegen ist, eine breitere Klasse von Einschränkungen unterstützt und sowohl die Laufzeit als auch die Leistung verbessert. Zusätzliche theoretische und empirische Analysen zeigen, dass die Laufzeiteffizienz unserer Methode durch ihre dynamische Nutzung von Berechnungen angetrieben wird, die mit der Divergenz zwischen dem uneingeschränkten und dem eingeschränkten Sprachmodell skaliert, und als Konsequenz sind die Laufzeitverbesserungen für bessere Modelle größer.

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.