Casamento de Distribuições Contrastivo para Monte Carlo Sequencial Amortizado em Difusão Discreta

Resumo

Modelos de difusão discretos emergiram como estruturas poderosas para a geração de dados categóricos estruturados. No entanto, a amostragem eficiente a partir de distribuições enviesadas por recompensa continua sendo um desafio fundamental. Embora o Monte Carlo Sequencial com Twist (SMC) ofereça exatidão assintótica para essa tarefa, a estimativa da função de twist ótima em espaços de estados discretos exige aproximações de Monte Carlo custosas, resultando em um grave gargalo computacional durante a inferência. Para superar essa limitação, introduzimos a Correspondência de Distribuição Contrastiva (CDM), uma nova estrutura que amortiza o custo da inferência SMC ao aprender uma função de twist parametrizada por meio de amostras positivas e negativas. Para um treinamento eficiente, reformulamos o estimador de gradiente de modo a aproveitar os kernels diretos de forma fechada dos modelos de difusão discretos. Na prática, a avaliação da nossa função de twist aprendida acarreta menos de 5% de custo computacional adicional em comparação com uma única passagem direta do modelo base. Por meio de extensas avaliações empíricas, demonstramos que a CDM supera consistentemente as linhas de base existentes sob tempo de execução equivalente. Validamos a eficácia e versatilidade da nossa abordagem em uma ampla gama de aplicações, incluindo geração de texto tóxico, design de sequências de DNA regulatório, designabilidade de proteínas e alinhamento de modelos de linguagem grandes baseados em difusão.

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.