ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

Technische Samenvatting: ART voor Diffusie-Sampling

Probleemstelling
Score-gebaseerde diffusiemodellen genereren steekproeven door een geleerde dynamiek in omgekeerde tijd te integreren. Een kritieke numerieke keuze in dit proces is de discretisatie van het tijdsinterval in een rooster van tijdstappen. Standaardbenaderingen maken doorgaans gebruik van uniforme roosters of handgemaakte schema's (bijv. EDM, DPM-Solver). De auteurs betogen dat deze vaste schema's vaak suboptimaal zijn gezien een vast budget aan functiewaarderingen (NFE). Uniforme roosters verdelen de rekenkracht gelijkmatig, zelfs wanneer de omgekeerde dynamiek aanzienlijk varieert langs het traject, terwijl handgemaakte roosters mogelijk niet adaptief zijn voor specifieke data-distributies of complexiteiten van de scorefunctie. Het doel is om de rekenkracht langs het sampling-traject systematisch te herverdelen om de geaggregeerde Euler-discretisatiefout te minimaliseren.

Methodologie
Het artikel introduceert Adaptive Reparameterized Time (ART), een regeltheoretisch kader voor optimale planning van tijdstappen.

1. Regeltheoretische Formulering (ART):

   • In plaats van roosterpunten direct te selecteren, behandelt ART de "kloksnelheid" van het diffusieproces als een regelvariabele.
   • Het introduceert een gereparameteriseerde tijdsvariabele $t$ die via een continue afbeelding $\psi(t)$ wordt gemapt naar de oorspronkelijke diffusietijd $\tau$.
   • De regel $\theta(t) = \dot{\psi}(t)$ vertegenwoordigt de instantane veranderingssnelheid van de diffusietijd.
   • Het doel is om $\theta(t)$ te kiezen om de geaggregeerde lokale Euler-discretisatiefout te minimaliseren, onder de beperking dat het totale tijdsbudget behouden blijft ($\int_0^T \theta(t) dt = T$).
   • De lokale foutproxy wordt afgeleid als evenredig met $|Q(x, \psi)| \theta(t)^2$, waarbij $Q$ afhankelijk is van de scorefunctie en haar afgeleiden.

2. Oplossing via Versterkend Leren (ART-RL):

   • Het deterministische ART-probleem is onoplosbaar om direct op te lossen vanwege de hoge-dimensionale toestandsruimte en het ontbreken van een gesloten vormoplossing voor de Hamilton-Jacobi-Bellman (HJB)-vergelijking.
   • De auteurs stellen ART-RL voor, een aanpak voor continu-tijd versterkend leren (CTRL).
   • Zij introduceren een gerandomiseerd hulp-probleem waarbij de regel $\theta$ wordt bemonsterd uit een Gaussisch beleid $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$.
   • Deze formulering maakt het gebruik van actor-critic-algoritmen mogelijk (specifiek het aanpassen van het kader van Jia en Zhou [14]) om het beleidsgemiddelde $\mu$ en de waarderingsfunctie te leren via stochastische benadering.

3. Theoretische Brug:

   • Een centrale theoretische bijdrage is de vestiging van een tweerichtingsbrug tussen de deterministische ART-optimum en het optimale Gaussische beleid in het ART-RL hulp-probleem.
   • Voldoende: De deterministische optimale regel $\mu^*$ lift naar een optimaal Gaussisch beleid waarvan het gemiddelde exact $\mu^*$ is.
   • Noodzakelijk: Elk optimaal Gaussisch beleid voor het hulp-probleem moet een gemiddelde hebben dat de deterministische ART-regel $\mu^*$ herstelt.
   • Dit resultaat rechtvaardigt het parametriseren van de actor als een Gaussisch gemiddelde en het inzetten ervan deterministisch na training, en biedt rigoureuze garanties dat het geleerde schema de ware optimum is voor het oorspronkelijke regelprobleem.

4. Distillatie:

   • Na offline training wordt het stochastische beleid gedistilleerd tot een deterministisch, enkel op tijd gebaseerd rooster. Dit verwijdert de noodzaak voor per-stap actor-netwerk-evaluaties of variantieberekeningen tijdens inferentie, en garandeert geen extra rekenkosten in vergelijking met standaard-schema's.

Belangrijkste Bijdragen

• Formulering: Stelt ART voor, het eerste regeltheoretische kader voor diffusie-tijdstappenplanning via adaptieve tijdsreparameterisatie.
• Theorie: Vestigt de tweerichtings-equivalentie tussen een deterministisch optimalisatieprobleem en een actor-critic-leringsprocedure in continue tijd met Gaussische beleidsregels. Dit transformeert heuristische RL-toepassingen in een principiële route voor het vinden van deterministische optima.
• Algoritme: Ontwikkelt ART-RL, een actor-critic-algoritme dat adaptieve tijdschema's datagedreven leert.
• Generalisatie: Demonstreert dat het geleerde schema kan worden gedistilleerd en overgedragen tussen verschillende datasets en aantallen tijdstappen zonder hertraining.

Experimentele Resultaten
De auteurs evalueren ART-RL binnen de officiële EDM-pijplijn over CIFAR-10, AFHQv2, FFHQ en ImageNet.

• CIFAR-10: ART-RL presteert consistent beter dan Uniform, DPM-Solver log-SNR en standaard EDM-schema's over een breed scala aan NFE-budgetten.
  • Bij lage budgetten (bijv. NFE=3) produceert ART-RL herkenbare objecten waar baselines dicht bij ruis blijven.
  • Bij NFE=35 bereikt ART-RL een FID van 1,82, een verbetering ten opzichte van EDM (1,85) en DPM-logSNR (1,89).
  • Statistische significantie wordt bevestigd via meerdere seeds, waarbij de slechtste run van ART-RL de beste run van EDM verslaat.
• Generalisatie:
  • Interpolatie/Extrapolatie: Een schema geleerd bij $K=18$ stappen op CIFAR-10 behoudt zijn prestatievoordeel wanneer geïnterpoleerd of geëxtrapoleerd naar andere stap-aantallen ($K \in \{4, 9, 12, 15, 20\}$).
  • Cross-Dataset Overdracht: Het op CIFAR-10 gedistilleerde schema wordt toegepast op AFHQv2, FFHQ en ImageNet zonder hertraining. Het verbetert consistent de FID ten opzichte van EDM en DPM-logSNR, met name bij lage en gemiddelde NFE's.
• Eendimensionale Analyse: In een setting met een analytische scorefunctie reduceert ART-RL de Wasserstein-afstand ($W_2$) aanzienlijk in vergelijking met handgemaakte schema's, wat bevestigt dat de verbetering voortkomt uit betere tijdsallocatie in plaats van fouten in de schatting van de score.

Betekenis en Claims
Het artikel claimt de eerste te zijn die diffusie-tijdstappenplanning formuleert als een regelprobleem via adaptieve tijdsreparameterisatie, en de eerste die een rigoureuze tweerichtingsbrug biedt tussen een deterministische optimum en een actor-critic-leringsprocedure.

De auteurs benadrukken dat hun methode datagedreven en geamortiseerd is:

• Datagedreven: Het schema wordt geleerd uit de objectfunctie en data-dynamiek, niet handgemaakt.
• Geamortiseerde Overdracht: De kosten voor offline training (ongeveer 1–2 uur op een Colab T4 voor CIFAR-10) zijn een eenmalige uitgave. Het resulterende gedistilleerde deterministische rooster is herbruikbaar over verschillende datasets en stap-aantallen met geen extra inferentiekosten, en komt overeen met de runtime van standaard-schema's zoals Uniform of EDM.

Het werk positioneert adaptieve tijdsreparameterisatie als een beginnende maar veelbelovende richting voor systematisch schema-ontwerp in generatieve diffusiemodellen, en biedt een regeltheoretisch startpunt voor toekomstig onderzoek.