Technische Samenvatting: ART voor Diffusion Sampling

Probleemstelling

Diffusiemodellen genereren data door ruis te transformeren naar samples via een geleerde reverse-time dynamiek. Hoewel de training van het score- (of denoising-) model goed gevestigd is, vertrouwt de samplingfase op het numeriek discretiseren van het reverse proces op een eindig rooster van tijdstappen. De huidige standaardpraktijken maken gebruik van uniforme roosters of handmatig ontworste schema's (bijv. EDM, DPM-Solver). Deze benaderingen zijn doorgaans ad hoc, vaste voorschriften die niet aanpassen aan de specifieke numerieke kenmerken van het traject. Gevolgd kan dit leiden tot een suboptimale toewijzing van computationele middelen, wat resulteert in ofwel significante discretisatiefouten (als de stappen te grof zijn) of verspilde evaluaties (als de stappen te fijn zijn in regio's waar de dynamiek vloeiend is). Het kernprobleem is het formuleren van een systematische, principiële methode voor tijdstaptoewijzing die de kwaliteit van de samples optimaliseert onder een vastgesteld computationeel budget (aantal functie-evaluaties).

Methodologie

1. Adaptive Reparameterized Time (ART)

De auteurs stellen ART voor, een continu-tijd optimaal controleframework dat de selectie van tijdstappen behandelt als een controleprobleem.

• Tijd-reparameterisatie: In plaats van het reverse proces te evolueren op een vast fysiek tijdsrooster $\tau \in [0, T]$, introduceert ART een hulp- "sampling-klok" $t \in [0, T]$. De fysieke tijd is een functie van deze klok, $\tau = \psi(t)$.
• Controle-variabele: De snelheid van de voortgang in de fysieke tijd relatief aan de sampling-klok, $\theta(t) = \dot{\psi}(t)$, wordt behandeld als de controle-variabele.
• Doelstelling: Het doel is om een doelfunctie te maximaliseren die afgeleid is van de leidende-orde Euler-discretisatiefout. De lokale fout van het Euler-schema is getoond proportioneel aan $|Q(x, \psi)| \theta^2$, waarbij $Q$ de lokale stijfheid van het probability-flow vectorveld vangt. Het controleprobleem streeft naar het minimaliseren van de geaccumuleerde fout onder een totale tijdbudgetbeperking $\int_0^T \theta(t) dt = T$.
• Formulering: Dit resulteert in een deterministisch optimaal controleprobleem waarbij de staat zowel het sample-traject als de geaccumuleerde fysieke tijd omvat, en de controle $\theta$ bepaalt hoe snel of langzaam de sampler door de fysieke tijd beweegt.

2. ART-RL: Een Gerandomiseerde Controleformulering

Het direct oplossen van het deterministische ART-probleem via de Hamilton–Jacobi–Bellman (HJB) vergelijking is onuitvoerbaar vanwege de hoge dimensionaliteit van de toestandsruimte (bijv. beelddimensies). Om dit te overwinnen, introduceren de auteurs ART-RL, een hulp- gerandomiseerd controleprobleem.

• Gaussiaanse Policies: De deterministische controle $\theta$ wordt vervangen door een stochastische policy $\pi(\lambda)$, specifiek een Gaussische distributie $N(\mu, \lambda/|Q|)$. De variantie is omgekeerd evenredig met de lokale stijfheid $|Q|$, wat zorgt voor minder willekeur in stijve regio's.
• Theoretische Equivalentie: Het artikel bewijst dat het gemiddelde van de optimale Gaussische policy in de ART-RL formulering identiek is aan de optimale controleoplossing van het oorspronkelijke deterministische ART-probleem. Hiermee wordt bewezen dat het oplossen van het gerandomiseerde probleem de oplossing van het deterministische probleem herstelt.
• Actor–Critic Leren: Door gebruik te maken van recente vooruitgang in continuous-time reinforcement learning (CTRL), leiden de auteurs implementeerbare actor–critic update-regels af.
  • Critic: Schat de waardefunctie (policy evaluatie) in met behulp van martingale-gebaseerde momentcondities.
  • Actor: Update de policy mean $\mu$ om de waarde te verbeteren (policy verbetering) op basis van de feedback van de critic.
  • Lagrange-multiplier: Een multiplier $\gamma$ wordt bijgewerkt om de totale tijdbudgetbeperking af te dwingen.

3. Implementatie en Distillatie

De geleerde policy is een feedbackfunctie van tijd, staat en geaccumuleerde tijd. Echter, empirische resultaten tonen aan dat de optimale controle $\theta$ in de praktijk primair afhangt van de tijd. De auteurs stellen een distillatiestap voor waarbij de geleerde neurale netwerkpolicy wordt vervangen door een vast, deterministisch tijd-alleen schema (het empirische gemiddelde van de controle). Dit maakt het mogelijk om het geleerde schema te implementeren als een eenvoudig vooraf berekend rooster, zonder extra overhead tijdens de inferentie vergeleken met handmatig ontwerpte schema's.

Belangrijkste Bijdragen

1. Controle-theoretisch Framework: Dit artikel is de eerste die de toewijzing van tijdstappen in diffusion sampling formuleert als een continuous-time optimaal controleprobleem (ART), waarmee het verder gaat dan heuristische roosterontwerpen.
2. Theoretische Brug (ART-RL): Het legt een rigoureuze link tussen het deterministische controleprobleem en een gerandomiseerde reinforcement learning formulering. Het bewijst dat het gemiddelde van de optimale gerandomiseerde policy het oorspronkelijke deterministische ART-probleem oplost, wat het gebruik van actor–critic algoritmen mogelijk maakt.
3. Algoritmische Ontwikkeling: De auteurs leiden specifieke policy evaluatie- en verbeteringskarakteristieken af voor tijd-reparameterisatie, wat leidt tot een praktische, data-gedreven actor–critic algoritme (Algoritme 1) voor het leren van optimale schema's.
4. Generaliseerbaarheid en Transfer: Het werk demonstreert dat een schema dat geleerd is op een specifieke dataset (CIFAR-10) zonder hertraining kan worden overgedragen naar verschillende datasets (AFHQv2, FFHQ, ImageNet), verschillende resoluties (64x64 naar 512x512), verschillende pipelines (EDM naar EDM2) en verschillende representatieruimtes (pixels naar latent).

Experimentele Resultaten

De auteurs evalueren ART-RL over diverse instellingen, waarbij de score model, solver en pipeline constant blijven en alleen het tijdstaprooster verandert:

• Laag-dimensionale Analytische Case: Gebruikmakend van een 1D diffusie met een bekende scorefunctie, presteert ART-RL consistent beter dan Uniform, DPM en EDM schema's in termen van Wasserstein-2 error, wat aantoont dat handmatig ontwerpte schema's voor afbeeldingen kunnen falen in andere domeinen.
• CIFAR-10 (EDM Pipeline):
  • Heun Sampling: ART-RL bereikt de beste Fréchet Inception Distance (FID) over alle budgetten, inclusclusief de sterkste configuratie (NFE=35), waarbij het de state-of-the-art EDM schedule overtreft (1.82 vs 1.85).
  • Euler Sampling: Vergelijkbare verbeteringen worden waargenomen, wat het voordeel bevestigt van het leren van het rooster zelfs met first-order solvers.
• MNIST (Klein Model): Zelfs met een lichtgewicht, niet-pre-trained score model en een RK4 solver, presteert ART-RL significant beter dan de baselines, wat de robuustheid van de methode ten opzichte van modelcapaciteit aantoont.
• Transfer Learning:
  • Over Stap-aantallen: Een schema geleerd bij $K=18$ wordt effectief overgedragen naar andere stap-aantallen via log-lineaire interpolatie/extrapolatie.
  • Over Datasets: Het geleerde CIFAR-10 schema verbetert de FID op AFHQv2, FFHQ en ImageNet-64 zonder hertraining.
  • Over Pipelines en Resoluties: Het schema wordt succesvol overgedragen naar de EDM2 pipeline op ImageNet-512 (latent space, hoge resolutie), waarbij het zowel de EDM als DPM baselines overtreft.

Betekenis en Claims

Het artikel claimt dat ART-RL een principieel, data-gedreven alternatief biedt voor de bestaande ad hoc, vaste heuristische roosters die gebruikt worden bij diffusion sampling.

• Herbruikbaarheid: De primaire praktische betekenis is de mogelijkheid om een herbruikbaar tijdstap-schema te leren. Eenmaal getraind en gedestilleerd op een enkele dataset (bijv. CIFAR-10), kan het schema zonder hertraining worden toegepast op een breed scala aan downstream taken (verschillende datasets, resoluties en pipelines), waardoor de offline trainingskosten worden geamortiseerd.
• Drop-in Vervanging: De methode vereist geen wijzigingen in de score model architectuur, de solver implementatie of het trainingsprotocol. Het fungeert puur als een vervanging voor het tijdstap-rooster, wat het een modulaire verbetering maakt voor bestaande diffusion pipelines.
• Theoretische Onderbouwing: In tegenstelling tot eerdere heuristische benaderingen is ART geworteld in optimal control theory en continuous-time reinforcement learning, wat een systematische manier biedt om de discretisatiefout in diffusion sampling aan te pakken.

De auteurs merken bescheiden op dat hun huidige analyse beperkt is tot probability flow ODE sampling en dat de distillatie naar tijd-alleen schema's de staat-afhankelijkheid verwijdert, wat de vraag openlaat of rijkere, op de staat geconditioneerde schema's nog verdere voordelen kunnen bieden. De huidige resultaten leggen echter een sterke fundering voor systematisch tijdstap-ontwerp in generatieve modellering.