Контрастивное согласование распределений для амортизированного последовательного метода Монте-Карло в дискретной диффузии

Аннотация

Дискретные диффузионные модели стали мощными инструментами для генерации структурированных категориальных данных. Однако эффективная выборка из распределений, смещённых по вознаграждению, остаётся фундаментальной задачей. Хотя скрученный последовательный метод Монте-Карло (SMC) обеспечивает асимптотическую точность для этой задачи, оценка оптимальной функции скручивания в дискретных пространствах состояний требует дорогостоящих аппроксимаций методом Монте-Карло, что приводит к серьёзному вычислительному узкому месту на этапе инференса. Чтобы преодолеть это ограничение, мы предлагаем контрастивное согласование распределений (CDM) — новый подход, который амортизирует вычислительные затраты SMC-инференса путём обучения параметризованной функции скручивания на положительных и отрицательных примерах. Для эффективного обучения мы переформулируем оценку градиента, используя замкнутые прямые ядра дискретных диффузионных моделей. На практике вычисление обученной функции скручивания приводит к дополнительным вычислительным затратам менее 5% по сравнению с однократным прямым проходом базовой модели. На основе обширных эмпирических экспериментов мы демонстрируем, что CDM стабильно превосходит существующие базовые методы при одинаковом астрономическом времени. Мы подтверждаем эффективность и универсальность нашего подхода в различных приложениях, включая генерацию токсичных текстов, дизайн регуляторных последовательностей ДНК, проектируемость белков и согласование больших языковых моделей на основе диффузии.

English

Discrete diffusion models have emerged as powerful frameworks for generating structured categorical data. However, efficiently sampling from reward-tilted distributions remains a fundamental challenge. While Twisted Sequential Monte Carlo (SMC) offers asymptotic exactness for this task, estimating the optimal twist function in discrete state spaces necessitates costly Monte Carlo approximations, resulting a severe computational bottleneck at inference. To overcome this limitation, we introduce Contrastive Distribution Matching (CDM), a novel framework that amortizes the cost of SMC inference by learning a parameterized twist function via positive and negative samples. For efficient training, we reformulate the gradient estimator to leverage the closed-form forward kernels of discrete diffusion models. In practice, evaluating our learned twist function incurs less than 5% additional computational overhead compared to a single forward pass of the base model. Through extensive empirical evaluations, we demonstrate that CDM consistently outperforms existing baselines under matched wall-clock time. We validate the effectiveness and versatility of our approach across a diverse range of applications, including toxic text generation, regulatory DNA sequence design, protein designability, and diffusion large language model alignment.