Contrastieve Distributie-Matching voor Geamortiseerd Sequentieel Monte Carlo in Discrete Diffusie

Samenvatting

Discrete diffusiemodellen zijn uitgegroeid tot krachtige raamwerken voor het genereren van gestructureerde categorische gegevens. Het efficiënt bemonsteren van beloning-afwijkende verdelingen blijft echter een fundamentele uitdaging. Hoewel Twisted Sequential Monte Carlo (SMC) asymptotische exactheid biedt voor deze taak, vereist het schatten van de optimale twistfunctie in discrete toestandsruimten kostbare Monte Carlo-benaderingen, wat leidt tot een ernstige computationele bottleneck tijdens de inferentie. Om deze beperking te overwinnen, introduceren we Contrastive Distribution Matching (CDM), een nieuw raamwerk dat de kosten van SMC-inferentie amortiseert door een geparametriseerde twistfunctie te leren via positieve en negatieve monsters. Voor efficiënte training herformuleren we de gradientschatter om de gesloten-vorm voorwaartse kernels van discrete diffusiemodellen te benutten. In de praktijk kost het evalueren van onze geleerde twistfunctie minder dan 5% extra computationele overhead vergeleken met een enkele voorwaartse doorgang van het basismodel. Door middel van uitgebreide empirische evaluaties tonen we aan dat CDM consequent beter presteert dan bestaande baselines onder gelijke wandkloktijd. We valideren de effectiviteit en veelzijdigheid van onze aanpak in een breed scala aan toepassingen, waaronder toxische tekstgeneratie, regulerend DNA-sequentieontwerp, eiwitontwerpbaarheid en diffusie grote-taalmodel-afstemming.

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.