Generación Controlada Rápida a partir de Modelos de Lenguaje con Muestreo de Rechazo Ponderado Adaptativo

Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

April 7, 2025

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

Resumen

El enfoque dominante para generar a partir de modelos de lenguaje sujetos a alguna restricción es el decodificado localmente restringido (LCD, por sus siglas en inglés), que consiste en muestrear tokens de manera incremental en cada paso de tiempo de forma que la restricción nunca se viole. Normalmente, esto se logra mediante el enmascaramiento de tokens: recorriendo el vocabulario y excluyendo los tokens que no cumplen con la restricción. Este enfoque presenta dos problemas importantes. (i) Evaluar la restricción en cada token puede ser prohibitivamente costoso, ya que los vocabularios de los modelos de lenguaje suelen superar los 100,000 tokens. (ii) El LCD puede distorsionar la distribución global sobre las cadenas, muestreando tokens basados únicamente en información local, incluso si estos conducen a caminos sin salida. Este trabajo introduce un nuevo algoritmo que aborda ambos problemas. En primer lugar, para evitar evaluar una restricción sobre el vocabulario completo en cada paso de generación, proponemos un algoritmo de muestreo por rechazo adaptativo que generalmente requiere órdenes de magnitud menos evaluaciones de restricción. En segundo lugar, mostramos cómo este algoritmo puede extenderse para producir estimaciones de baja varianza y sin sesgo de los pesos de importancia a un costo adicional muy pequeño, estimaciones que pueden usarse de manera sólida dentro de algoritmos previamente propuestos de Monte Carlo secuencial para corregir el comportamiento miope de la aplicación local de restricciones. A través de una extensa evaluación empírica en dominios como texto-a-SQL, síntesis molecular, inferencia de objetivos, coincidencia de patrones y JSON, demostramos que nuestro enfoque es superior a los métodos de referencia más avanzados, admitiendo una clase más amplia de restricciones y mejorando tanto el tiempo de ejecución como el rendimiento. Análisis teóricos y empíricos adicionales muestran que la eficiencia en tiempo de ejecución de nuestro método se debe a su uso dinámico de la computación, escalando con la divergencia entre el modelo de lenguaje no restringido y el restringido, y como consecuencia, las mejoras en tiempo de ejecución son mayores para modelos mejores.

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.