Emparejamiento de Distribuciones Contrastivo para Monte Carlo Secuencial Amortizado en Difusión Discreta

Resumen

Los modelos de difusión discretos han emergido como marcos poderosos para generar datos categóricos estructurados. Sin embargo, muestrear eficientemente a partir de distribuciones sesgadas por recompensa sigue siendo un desafío fundamental. Si bien el Monte Carlo Secuencial Torcido (SMC, por sus siglas en inglés) ofrece exactitud asintótica para esta tarea, estimar la función de torsión óptima en espacios de estados discretos requiere costosas aproximaciones Monte Carlo, lo que genera un grave cuello de botella computacional durante la inferencia. Para superar esta limitación, introducimos la Coincidencia de Distribuciones Contrastiva (CDM), un novedoso marco que amortiza el costo de la inferencia SMC aprendiendo una función de torsión parametrizada mediante muestras positivas y negativas. Para un entrenamiento eficiente, reformulamos el estimador del gradiente aprovechando los núcleos directos de forma cerrada de los modelos de difusión discretos. En la práctica, evaluar nuestra función de torsión aprendida genera menos de un 5% de sobrecarga computacional adicional en comparación con una sola pasada directa del modelo base. Mediante amplias evaluaciones empíricas, demostramos que CDM supera consistentemente a las líneas base existentes bajo tiempo real comparable. Validamos la efectividad y versatilidad de nuestro enfoque en una variedad de aplicaciones diversas, que incluyen la generación de texto tóxico, el diseño de secuencias de ADN reguladoras, la capacidad de diseño de proteínas y la alineación de modelos de lenguaje grandes mediante difusión.

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.