Technische Samenvatting: Contrastieve Distributiematching voor Geamortiseerde Sequentiële Monte Carlo in Discrete Diffusie

Probleemstelling

Discrete diffusiemodellen zijn krachtige raamwerken geworden voor het genereren van gestructureerde categorische data, zoals tekst, DNA-sequenties en eiwitten. Een centrale toepassing van deze modellen is reward-uitlijning, waarbij het doel is om te bemonsteren uit een doeldistributie die de vooraf getrainde prior naar hogere waarden van een downstream rewardfunctie (bijvoorbeeld menselijke voorkeur of eiwitontwerpbaarheid) beïnvloedt.

Hoewel Twisted Sequential Monte Carlo (SMC) een asymptotisch exacte methode biedt om te bemonsteren uit dergelijke reward-georiënteerde distributies, staat de toepassing ervan in discrete ruimtes voor een fundamentele bottleneck: de schatting van de optimale twistfunctie.

• In continue domeinen kan de twistfunctie efficiënt worden benaderd met de Tweedie-formule of gradient-based guidance.
• In discrete domeinen dwingt het ontbreken van een analoog verband tot een afhankelijkheid van kostbare Monte Carlo (MC) benaderingen om de twistfunctie bij elke inferentiestap te schatten.
• Deze afhankelijkheid creëert een zware computere last, vooral wanneer de rewardfunctie zelf duur is om te evalueren (bijvoorbeeld eiwitvouwingssimulaties of API-gebaseerde LLM-voorkeursscore). Het vergroten van het aantal MC-steekproeven ($M$) verbetert de nauwkeurigheid, maar verhoogt de inferentietijd lineair, waardoor de aanpak voor veel real-world toepassingen onpraktisch wordt.

Methodologie

De auteurs stellen Contrastieve Distributiematching (CDM) voor, een raamwerk dat is ontworpen om de kosten van SMC-inferentie te amortiseren door een geparametriseerde twistfunctie te leren via een contrastief leerdoel.

1. Doel van Contrastieve Distributiematching

In plaats van de twistfunctie te leren via directe regressie (minimaliseren van de Mean Squared Error tegen MC-schattingen), minimaliseert CDM de forward Kullback-Leibler (KL) divergentie tussen de doeldistributie $p^*_t$ en een geparametriseerde distributie $p^\phi_t$ die wordt gemoduleerd door de geleerde twistfunctie $\psi^\phi_t$.

Het doel wordt geformuleerd als:
$$\mathcal{L}_{\text{CDM}}(\phi) = \mathbb{E}_{t}\left[ \mathcal{D}_{\text{KL}}\left( p^*_t(\mathbf{x}_t) \big\| p^\phi_t(\mathbf{x}_t) \right) \right]$$

De gradient van deze loss vertoont een contrastieve structuur:
$$-\nabla_\phi \mathcal{L} = \mathbb{E}_{t}\left[ \underbrace{\mathbb{E}_{p^*_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Positieve Term}} - \underbrace{\mathbb{E}_{p^\phi_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Negatieve Term}} \right]$$

• Positieve Term: Verhoogt de log-twist op steekproeven getrokken uit de optimale doeldistributie (hoog-reward gebieden).
• Negatieve Term: Verlaagt de log-twist op steekproeven getrokken uit de huidige benadering (suboptimale gebieden).

2. Efficiënt Trainingschema

Om de inefficiëntie van het bemonsteren uit de doeldistributie $p^*_t$ voor de positieve term aan te pakken, maken de auteurs gebruik van een uniek kenmerk van diffusiemodellen: de gesloten-vorm forward kernel.

• Forward-gebaseerde Gradientsschatting: In plaats van onafhankelijk te bemonsteren uit elke tussentijdse target $p^*_t$, gebruikt de methode eerst SMC om schone positieve steekproeven ($\mathbf{x}_0$) te trekken uit de doeldistributie. Vervolgens maakt het gebruik van de gesloten-vorm forward kernel $p^{\text{base}}(\mathbf{x}_t | \mathbf{x}_0)$ om meerdere positieve steekproeven op elk willekeurig tussentijdstip $t$ te genereren tegen verwaarloosbare kosten.
• Buffer-gebaseerd Trainen: Er wordt een buffer van schone positieve steekproeven bijgehouden en hergebruikt over meerdere tijdstappen en gradient-updates, wat het aantal dure reward-evaluaties tijdens het trainen aanzienlijk reduceert.
• Negatief Bemonsteren: Negatieve steekproeven worden getrokken via Importance Sampling (IS) met behulp van de huidige geparametriseerde twist, wat computatie-efficiënt is.

3. Parametrizatie

Om ervoor te zorgen dat de geleerde twistfunctie geen inferentie-overhead introduceert, wordt de twistfunctie $\psi^\phi_t$ geparametriseerd als een lichtgewicht scalair hoofd dat is gekoppeld aan de laatste feature-laag van de vooraf getrainde diffusie-ruggengraat. Hierdoor kan het model zowel logits als de twist-schatting in één forward pass produceren, wat minder dan 5% computere last toevoegt ten opzichte van het basismodel.

Belangrijkste Bijdragen

1. CDM-raamwerk: Introductie van een nieuw SMC-amortisatie-raamwerk voor discrete diffusie dat de kosten van het toepassen van de twistfunctie op inferentietijd reduceert tot een operatie met constante tijd.
2. Diffusie-natief Trainen: Ontwerp van een trainschema dat gebruikmaakt van het gesloten-vorm forward-proces van diffusiemodellen om efficiënt trainen met dure rewardfuncties mogelijk te maken, waarbij de steekproefinefficiëntie van standaard regressie-gebaseerde benaderingen wordt vermeden.
3. Veelzijdigheid en Synergie: Aantonen dat CDM agnostisch is ten opzichte van de voorstel-distributie, waardoor het kan worden gecombineerd met fijngefineerde voorstellen (bijvoorbeeld d1, DRAKES) voor synergetische prestatiewinst.
4. Superieure Doelfunctie: Validatie dat het contrastief leerdoel superieure prestaties en snellere convergentie oplevert in vergelijking met standaard regressie-gebaseerde twist-doelen (Soft Value).

Experimentele Resultaten

De auteurs hebben CDM geëvalueerd over vier uiteenlopende toepassingen: giftige tekstgeneratie, regulatief DNA-sequentieontwerp, eiwitontwerpbaarheid en uitlijning van diffusie-grote taalmodellen (dLLM).

• Schalingsprestaties: In alle taken overtrof CDM consistent de baselines (Best-of-N, standaard SMC met variërende $M$, en Soft Value) wanneer ze werden vergeleken op basis van wandkloktijd.
• Dure Rewards: In de taak voor eiwitontwerpbaarheid, waar reward-evaluatie de dominante bottleneck is, behaalde CDM superieure schaling in vergelijking met zowel Best-of-N als Soft Value, terwijl standaard SMC door inferentie-overhead niet effectief kon schalen.
• Compatibiliteit met Fijnfinen: Bij integratie met fijngefineerde voorstellen (d1 en DRAKES) verbeterde CDM de prestaties verder en verminderde het mode-collapse-problemen die vaak worden waargenomen bij fijngefineerde modellen, waarbij diversiteit behouden bleef terwijl hoge rewards werden bereikt.
• Trainings-efficiëntie: CDM toonde snellere convergentie in wandkloktijd in vergelijking met Soft Value, wat wordt toegeschreven aan de effectiviteit van het contrastief leerdoel met negatieve steekproeven.

Betekenis en Claims

De paper claimt dat CDM de primaire bottleneck in twisted SMC-inferentie voor discrete diffusie aanpakt: de afhankelijkheid van kostbare Monte Carlo-schatting. Door de twistfunctie te leren via een contrastief doel dat gebruikmaakt van het diffusie-forward-proces, maakt de methode efficiënte, schaalbare reward-uitlijning mogelijk.

De auteurs benadrukken dat hun aanpak orthogonaal is aan voorstel-fijnfinen; het kan worden gecombineerd met bestaande fijnfinemethoden om verdere prestatiewinst te behalen. Bovendien blijft de methode van toepassing zelfs wanneer de rewardfunctie niet-differentieerbaar of computationeel duur is, een scenario waarin gradient-based guidance en standaard SMC moeite hebben. De resultaten suggereren dat CDM een nieuwe Pareto-grens vestigt voor de afweging tussen reward-kwaliteit en inferentietijd in discrete generatieve modellering.