On the Error-Correcting Effects of Stochasticity in Discrete Diffusion

Probleemstelling
Discrete diffusiemodellen hebben concurrerende prestaties behaald in tekst- en beeldgeneratie, maar staan voor een fundamentele afweging tussen bemonsteringsefficiëntie en steekproefkwaliteit. Hoewel sterk deterministische overgangen (zoals die in Discrete Probability Flow, of DPF) snelle convergentie mogelijk maken, zijn ze vatbaar voor foutaccumulatie, wat leidt tot een lagere uiteindelijke steekproefkwaliteit. Omgekeerd kunnen meer stochastische overgangen fouten corrigeren, maar convergeren ze langzamer. Bestaande intuïties over deze afweging zijn grotendeels heuristisch; bijvoorbeeld worden op masking gebaseerde overgangen beschouwd als onomkeerbaar en foutgevoelig, terwijl uniforme overgangen gedacht worden een zichzelf corrigerend gedrag te vertonen. Echter, eerder werk heeft niet systematisch gekarakteriseerd hoe de mate van stochastischheid deze afweging beheerst, noch heeft het een principiële verklaring gegeven voor de onderliggende mechanismen van foutpropagatie en -correctie in discrete settingen.

Methodologie en Theoretische Analyse
De auteurs presenteren een systematische studie van hoe stochastischheid in Markov-overgangen de bemonsteringsafweging in discrete diffusie beheerst. Hun analyse maakt gebruik van informatie-theoretische hulpmiddelen om foutpropagatie te karakteriseren:

1. Foutcorrigerend Mechanisme: Het paper identificeert dat redundante overgangen tussen staatparen, die symmetrisch waarschijnlijkheidsmassa uitwisselen zonder de tijdmarginaalverdelingen te veranderen, een foutcorrigerend effect induceren. Met behulp van KL-divergentiecontractie-analyse bewijzen de auteurs dat deze overgangen bemonsteringsfouten op bewijswaardige wijze kunnen contracteren.
2. Theoretische Afweging: De analyse onthult een dubbele rol voor stochastischheid:
   • In het regime met een laag aantal functiewaarderingen (NFE) verhoogt het verhogen van stochastischheid de discretisatiefouten door frequentere staatveranderingen.
   • In het regime met een hoog NFE, waar discretisatiefouten verwaarloosbaar zijn, mengen stochastische overgangen waarschijnlijkheidsmassa over staten, contracteren verdelingsfouten en mitigeren modelbenaderingsfouten.
3. Contractiecoëfficiënten: De auteurs leiden contractiecoëfficiënten af voor verschillende corruptieprocessen. Zij tonen aan dat voor kleine statruimten, uniforme diffusie een sterkere foutcontractie vertoont dan masking-diffusie. Omgekeerd bereikt masking-diffusie de slechtst mogelijke contractiesnelheid, waarbij het in het slechtste geval faalt om fouten te contracteren (specifiek wanneer de posterior niet volledig gedragen wordt over alle waarden).
4. DCRS-algoritme: Gemotiveerd door deze bevindingen, stellen de auteurs Discrete Churn and Restart Sampling (DCRS) voor, een trainingsvrij inferentie-algoritme. DCRS wisselt af tussen:
   • Bijna-deterministische terugwaartse updates: Om snelle convergentie te garanderen.
   • Gecontroleerde injectie van stochastischheid: Om geaccumuleerde fouten te corrigeren.
     Dit wordt bereikt via twee mechanismen:
   • Opnieuw starten (Restarting): Herstelt periodiek de bemonsteringstrajectorie met behulp van overgangen van het voorwaartse proces (die gesloten-vorm uitdrukkingen hebben) om globale stochastischheid in te brengen zonder extra neurale netwerkbewerkingen.
   • Churnen: Injecteert lokale stochastischheid binnen een tijdsvenster door een kleine voorwaartse stap gevolgd door een terugwaartse stap uit te voeren, wat lokale foutcorrectie biedt.

Belangrijkste Resultaten
Het paper evalueert DCRS op synthetische datasets en schaalbare benchmarks (CIFAR10, CelebA, LM1B, OpenWebText):

• Beeldgeneratie: Op beelddatasets verbetert DCRS de snelheid-kwaliteit-afweging in het regime met lage NFE aanzienlijk. Het bereikt vergelijkbare steekproefkwaliteit als standaard bemonsteraars (zoals $\tau$-leaping) met aanzienlijk minder bemonsteringsstappen (bijvoorbeeld tot een 10$\times$ reductie in stappen op CelebA terwijl concurrerende kwaliteit behouden blijft).
• Taalgeneratie: Op taalbenchmarks zijn de resultaten genuanceerder. Hoewel DCRS voordelen toont, observeert het paper dat het verminderen van stochastischheid een beperkte impact heeft op de uiteindelijke steekproefkwaliteit voor bepaalde modellen (bijvoorbeeld UDLM). De auteurs merken op dat de effectiviteit van foutcorrigerende stochastischheid sterk afhankelijk is van het corruptieproces (uniform versus masking) en de bemonsteringsprocedure (bijvoorbeeld het gebruik van nucleus sampling).
• Gecontroleerde Experimenten: Experimenten op 1D-datasets en datasets met mengsels van Gaussians valideren de theoretische voorspellingen. Zij tonen aan dat bijna-deterministische bemonstering sneller convergeert maar gevoelig is voor modelfouten, terwijl stochastische bemonstering deze fouten in de loop van de tijd corrigeert. DCRS slaagt erin deze concurrerende effecten in evenwicht te brengen.

Betekenis en Claims
Het paper claimt de eerste systematische karakterisering te bieden van hoe stochastischheid de afweging tussen bemonsteringssnelheid en -kwaliteit in discrete diffusiemodellen beheerst. De primaire bijdragen zijn:

1. Theoretisch Inzicht: Het stelt vast dat stochastischheid de generatie verbetert door fouten "weg te mengen", fungerend als een mechanisme om modelbenaderingsfouten te corrigeren ten koste van langzamere bemonstering. Dit wordt geformaliseerd door informatie-theoretische analyse die aantoont dat redundante overgangen KL-divergentie contracteren.
2. Algoritmische Innovatie: Het introduceert DCRS, een methode die expliciet gebruikmaakt van de foutcorrigerende effecten van stochastischheid door voorwaartse en terugwaartse processen af te wisselen. Dit maakt snelle generatie mogelijk in het regime met lage NFE zonder dat modelhertraining vereist is.
3. Unificatie van Principes: Het werk verenigt het begrip van foutcorrectie in discrete diffusie, en toont aan dat de sterkte van correctie kritiek afhankelijk is van het type en de mate van stochastischheid die haalbaar is onder het gekozen corruptieproces.

De auteurs concluderen dat het beheersen van stochastischheid een fundamenteel ontwerpprincipe is voor discrete diffusie-inferentie, hoewel zij erkennen dat het afstemmen van het niveau en het schema van stochastischheid een noodzakelijke stap blijft. Zij suggereren dat toekomstig werk zich kan richten op adaptieve methoden voor het optimaliseren van deze afweging.