Technische Samenvatting: Metropolis-gecorrigeerde Diffusiemodellen (MADM)

Probleemstelling

Het afnemen van steekproeven uit op scores gebaseerde diffusiemodellen brengt inherent een vertekening met zich mee die voortkomt uit twee bronnen: de benaderingsfout van de geleerde scorefunctie en de tijdsdiscretisatiefout die wordt geïntroduceerd door numerieke integratoren. In het standaard Predictor-Corrector-raamwerk verplaatst een predictorstap (vaak een Probability Flow ODE) een steekproef tussen ruisniveaus, gevolgd door een correctorstap (typisch het Unadjusted Langevin Algorithm, ULA) om de steekproef op het huidige ruisniveau te verfijnen.

Hoewel ULA de scorefunctie gebruikt om verplaatsingen voor te stellen, is het zelf een vertekende steekproefnemer omdat het een continu Langevin-diffusieproces discretiseert. Bijgevolg introduceert de "corrector" in het Predictor-Corrector-schema zijn eigen discretisatiebias, waardoor hij de doelrandverdeling niet invariant laat. Standaard Metropolis-Hastings (MH) of Barker accept-afkeurmechanismen, die dergelijke vertekening in Markov Chain Monte Carlo (MCMC) corrigeren, kunnen niet direct worden toegepast op diffusiemodellen. Dit komt doordat MH de verhouding van doelkansen, $p_t(\tilde{x})/p_t(x)$, moet evalueren, wat onberekenbaar is in diffusiemodellen waar alleen de scorefunctie $s_\phi(x, t) \approx \nabla_x \log p_t(x)$ beschikbaar is.

Methodologie

De auteurs stellen Metropolis-gecorrigeerde Diffusiemodellen (MADM) voor, een raamwerk dat vertekende ULA-correctors vervangt door accept-afkeur gecorrigeerde Langevin-stappen. De kerninnovatie is een methode om de benodigde acceptatiekans te berekenen met uitsluitend de scorefunctie, waardoor de behoefte aan de onberekenbare kansverhouding wordt omzeild.

1. Identiteit van de Kansverhouding

De auteurs stellen vast dat de logaritme van de verhouding van de doelkansen kan worden uitgedrukt als een lijnintegraal van de scorefunctie langs het pad tussen de huidige toestand $x$ en de voorgestelde $\tilde{x}$:
$$\log \left( \frac{p_t(\tilde{x})}{p_t(x)} \right) = \int_0^1 \langle s(x + u(\tilde{x} - x), t), \tilde{x} - x \rangle \, du$$
Deze identiteit maakt het mogelijk om de kansverhouding te schatten met uitsluitend score-evaluaties.

2. Exacte Correctie via Bernoulli Factory

Om een exacte correctie te bereiken zonder de kans te evalueren, stellen de auteurs een algoritme voor dat gebaseerd is op een twee-munten Bernoulli factory (Gonçalves et al., 2017).

• Aanneming: Er bestaat een bovengrens $C(x, \tilde{x})$ zodat de absolute waarde van de integrand in de lijnintegraal begrensd is door $C$.
• Mechanisme: Het algoritme construeert een onbevooroordeelde schatter $W$ van de exponentiële van de integraal (geschaald met de bovengrens) met behulp van een Poisson-getruncateerde machtreeks.
• Procedure: Het algoritme gebruikt een "twee-munten" procedure om een Barker accept-afkeur beslissing te simuleren. Het verwierpt eerst met een vaste kans $\alpha'$, en als het niet wordt verworpen, draait het een munt met kans $W$. Als de munt slaagt, wordt de voorstel aanvaard; anders begint de ronde opnieuw.
• Resultaat: Deze procedure levert een exact Barker-gecorrigeerde Langevin-corrector op die de doelverdeling $p_t$ invariant laat, mits de score exact is en de bovengrensaanneming geldt.

3. Efficiënte Benadering via de Regel van Simpson

Omdat de exacte Bernoulli factory-methode een willekeurig en potentieel groot aantal score-evaluaties vereist, stellen de auteurs een schaalbare benadering voor.

• Methode: De lijnintegraal wordt deterministisch benaderd met Simpsons 1/3-regel (een Newton-Cotes kwadratuurregel). Dit vereist het evalueren van de score op de huidige toestand, de voorstel en het middelpunt.
• Acceptatie: De resulterende schatting wordt ingevoerd in de standaard Metropolis-Hastings acceptatiekans.
• Eigenschappen: Hoewel dit een kleine vertekening introduceert door kwadratuurtruncatie, bereikt de benadering een nauwkeurigheid van orde $O(h^{5/2})$ in de stapgrootte $h$. Het is computationeel goedkoop en vereist slechts één extra score-evaluatie in het middelpunt per voorstel.

Belangrijkste Bijdragen

1. MADM-raamwerk: Introductie van een Metropolis-gecorrigeerd correctorraamwerk voor diffusiemodellen dat systematisch vertekende ULA-stappen vervangt.
2. Exacte Methode: De eerste exacte methode voor het aanpassen van Langevin-correcties in diffusiemodellen met behulp van een twee-munten Bernoulli factory, die uitsluitend de scorefunctie en een bovengrens op de lijnintegraal vereist.
3. Schaalbare Benadering: Een praktische benadering volgens de regel van Simpson die hoge-orde nauwkeurigheid ($O(h^{5/2})$) bereikt met bijna nul marginale kosten (één extra score-evaluatie).
4. Empirische Validatie: Demonstratie van verbeterde steekproefkwaliteit op zowel synthetische als beelddatasets.

Resultaten

De auteurs evalueren MADM op synthetische 2D-datasets en standaard beeldbenchmarks (CIFAR-10, FFHQ, AFHQv2, ImageNet-64) met behulp van vooraf getrainde modellen van Karras et al. (2022).

• Synthetische Datasets: Op 2D-datasets (spiraal, trechter, Sierpinski-driehoek, windmolen) produceerde de standaard ancestral sampler (equivalent aan een predictor met een ULA-corrector) steekproeven met significante uitschieters buiten de ondersteuning van de ware dataverdeling. MADM-correctorstappen verplaatsten deze steekproeven succesvol terug naar de data-ondersteuning, waardoor uitschieters werden geëlimineerd.
• Beelddatasets:
  • Voor onvoorwaardelijke modellen (CIFAR-10, FFHQ, AFHQv2) leverde het toepassen van één MADM-correctorstap per predictorstap consistente, zij het bescheiden, verbeteringen op in de Fréchet Inception Distance (FID) in vergelijking met deterministische baselines (Heun- en Euler ODE-oplossers).
  • Voor het klasse-voorwaardelijke ImageNet-64-model waren de winsten duidelijker. Het gebruik van MADM met een Euler-predictor leverde grotere FID-verbeteringen op dan het gebruik van een hogere-orde Heun-predictor zonder correctie. Dit suggereert dat MH-aanpassingen effectiever zijn in het mitigeren van discretisatiefouten dan het simpelweg verhogen van de orde van de ODE-oplosser wanneer het aantal stappen groot is.
• Computationele Kosten: De exacte methode (Algoritme 2) bleek computationeel duur vanwege het willekeurige aantal vereiste score-evaluaties. Bijgevolg maakten de experimenten voornamelijk gebruik van de benadering volgens de regel van Simpson (Algoritme 3), die consistente winsten bood met beheersbare overhead.

Betekenis en Claims

Het artikel claimt de eerste exacte methode te bieden voor het aanpassen van Langevin-correcties in diffusiemodellen, waarmee de inherente vertekening van ULA-correctors wordt aangepakt. Door de scorefunctie te benutten om kansverhoudingen te schatten via lijnintegralen, maakt MADM het gebruik van rigoureuze MCMC-acceptatiecriteria mogelijk in generatieve modellering zonder toegang te vereisen tot de onberekenbare datakans.

De auteurs positioneren MADM als een natuurlijke uitbreiding van het Predictor-Corrector-raamwerk die de steekproefkwaliteit verbetert. Zij merken op dat hoewel de exacte methode theoretisch significant is, de benadering volgens de regel van Simpson een praktische, schaalbare oplossing biedt die consistent de FID op beeldbenchmarks verbetert. Het werk gaat uit van een perfecte schatting van de score, waarbij dit wordt erkend als een inherente beperking van diffusiemodellen en niet als een specifiek gebrek aan hun methode. Toekomstig werk wordt voorgesteld om zich te richten op het ontwikkelen van efficiënte varianten van de exacte methode of diens subroutines.