Geração Rápida e Controlada a partir de Modelos de Linguagem com Amostragem por Rejeição Ponderada Adaptativa

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

Resumo

A abordagem dominante para geração a partir de modelos de linguagem sujeitos a alguma restrição é a decodificação localmente restrita (LCD, do inglês Locally Constrained Decoding), que amostra tokens incrementalmente em cada passo de tempo de forma que a restrição nunca seja violada. Normalmente, isso é alcançado por meio de mascaramento de tokens: iterando sobre o vocabulário e excluindo tokens que não se conformam à restrição. Há dois problemas importantes nessa abordagem. (i) Avaliar a restrição em cada token pode ser proibitivamente caro — os vocabulários dos modelos de linguagem frequentemente excedem 100.000 tokens. (ii) A LCD pode distorcer a distribuição global sobre as sequências, amostrando tokens com base apenas em informações locais, mesmo que eles levem a caminhos sem saída. Este trabalho introduz um novo algoritmo que aborda ambos os problemas. Primeiro, para evitar avaliar a restrição em todo o vocabulário a cada passo da geração, propomos um algoritmo de rejeição adaptativa que tipicamente requer ordens de magnitude menos avaliações de restrição. Segundo, mostramos como esse algoritmo pode ser estendido para produzir estimativas de baixa variância e não enviesadas de pesos de importância a um custo adicional muito pequeno — estimativas que podem ser usadas de forma sólida dentro de algoritmos de Monte Carlo sequencial previamente propostos para corrigir o comportamento míope da aplicação local de restrições. Por meio de uma extensa avaliação empírica em domínios como texto-para-SQL, síntese molecular, inferência de objetivos, correspondência de padrões e JSON, mostramos que nossa abordagem é superior às baselines de última geração, suportando uma classe mais ampla de restrições e melhorando tanto o tempo de execução quanto o desempenho. Análises teóricas e empíricas adicionais mostram que a eficiência de tempo de execução do nosso método é impulsionada pelo uso dinâmico de computação, escalando com a divergência entre o modelo de linguagem não restrito e o restrito, e, como consequência, as melhorias no tempo de execução são maiores para modelos melhores.

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.