Kontrastives Verteilungsmatching für amortisiertes sequentielles Monte Carlo in diskreter Diffusion

Zusammenfassung

Diskrete Diffusionsmodelle haben sich als leistungsfähige Rahmenwerke zur Erzeugung strukturierter kategorialer Daten erwiesen. Dennoch stellt die effiziente Stichprobenentnahme aus belohnungsgewichteten Verteilungen eine grundlegende Herausforderung dar. Während Twisted Sequential Monte Carlo (SMC) asymptotische Exaktheit für diese Aufgabe bietet, erfordert die Schätzung der optimalen Twist-Funktion in diskreten Zustandsräumen kostspielige Monte-Carlo-Näherungen, was zu einem erheblichen rechnerischen Engpass während der Inferenz führt. Um diese Einschränkung zu überwinden, führen wir Contrastive Distribution Matching (CDM) ein, ein neuartiges Rahmenwerk, das die Kosten der SMC-Inferenz amortisiert, indem es eine parametrisierte Twist-Funktion mithilfe positiver und negativer Stichproben lernt. Für ein effizientes Training formulieren wir den Gradientenschätzer neu, um die geschlossenen Vorwärtskerne diskreter Diffusionsmodelle zu nutzen. In der Praxis verursacht die Auswertung unserer erlernten Twist-Funktion weniger als 5 % zusätzlichen rechnerischen Mehraufwand im Vergleich zu einem einzigen Vorwärtsdurchlauf des Basismodells. Durch umfangreiche empirische Auswertungen zeigen wir, dass CDM unter vergleichbarer Wanduhrzeit durchgängig bessere Ergebnisse als bestehende Baselines erzielt. Wir validieren die Effektivität und Vielseitigkeit unseres Ansatzes über eine Vielzahl von Anwendungen hinweg, darunter die Erzeugung toxischer Texte, das Design regulatorischer DNA-Sequenzen, die Designfähigkeit von Proteinen und das Alignment von Diffusions-Großsprachmodellen.

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.