Appariement de distribution contrastif pour le Monte Carlo séquentiel amorti dans la diffusion discrète

Résumé

Les modèles de diffusion discrets sont devenus des cadres puissants pour générer des données catégorielles structurées. Cependant, échantillonner efficacement à partir de distributions inclinées par récompense reste un défi fondamental. Bien que le Monte Carlo séquentiel tordu (SMC) offre une exactitude asymptotique pour cette tâche, l'estimation de la fonction de torsion optimale dans des espaces d'états discrets nécessite des approximations de Monte Carlo coûteuses, créant un goulot d'étranglement computationnel sévère lors de l'inférence. Pour surmonter cette limitation, nous introduisons l'Appariement de Distributions par Contraste (CDM), un nouveau cadre qui amortit le coût de l'inférence SMC en apprenant une fonction de torsion paramétrée à l'aide d'échantillons positifs et négatifs. Pour un entraînement efficace, nous reformulons l'estimateur du gradient afin d'exploiter les noyaux directs de forme fermée des modèles de diffusion discrets. En pratique, l'évaluation de notre fonction de torsion apprise n'entraîne qu'un surcoût computationnel inférieur à 5% par rapport à une seule passe avant du modèle de base. À travers des évaluations empiriques approfondies, nous démontrons que CDM surpasse systématiquement les bases de référence existantes à temps réel égal. Nous validons l'efficacité et la polyvalence de notre approche sur un ensemble diversifié d'applications, incluant la génération de texte toxique, la conception de séquences d'ADN régulatrices, la concevabilité des protéines et l'alignement des grands modèles de langage par diffusion.

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.