Быстрое управляемое генерация из языковых моделей с адаптивным взвешенным методом отклоняющего сэмплирования

Fast Controlled Generation from Language Models with Adaptive Weighted Rejection Sampling

April 7, 2025

Авторы: 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

Аннотация

Основной подход к генерации с использованием языковых моделей с учетом определенных ограничений — это локально ограниченное декодирование (LCD), при котором на каждом шаге пошагово выбираются токены, не нарушающие заданное ограничение. Обычно это достигается с помощью маскирования токенов: перебора всего словаря и исключения токенов, не соответствующих ограничению. Однако у этого подхода есть две важные проблемы. (i) Проверка ограничения для каждого токена может быть чрезмерно затратной — словари языковых моделей часто превышают 100 000 токенов. (ii) LCD может искажать глобальное распределение строк, выбирая токены только на основе локальной информации, даже если это ведет к тупиковым путям. В данной работе представлен новый алгоритм, который решает обе эти проблемы. Во-первых, чтобы избежать проверки ограничения на всем словаре на каждом шаге генерации, мы предлагаем адаптивный алгоритм отбраковки, который обычно требует на порядки меньше проверок ограничений. Во-вторых, мы показываем, как этот алгоритм можно расширить для получения низкодисперсионных и несмещенных оценок весов важности с минимальными дополнительными затратами — оценки, которые можно надежно использовать в ранее предложенных алгоритмах последовательного Монте-Карло для коррекции близорукого поведения локального применения ограничений. Благодаря обширной эмпирической оценке в областях текста в SQL, синтеза молекул, вывода целей, сопоставления шаблонов и JSON, мы показываем, что наш подход превосходит современные базовые методы, поддерживая более широкий класс ограничений и улучшая как время выполнения, так и производительность. Дополнительные теоретические и эмпирические анализы показывают, что эффективность времени выполнения нашего метода обусловлена динамическим использованием вычислений, масштабируемым в зависимости от расхождения между неограниченной и ограниченной языковой моделью, и, как следствие, улучшения времени выполнения более значительны для более качественных моделей.

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.