Technische Samenvatting: Geamortiseerde Gids voor Beeldinpainting met Voorgeïmplementeerde Diffusiemodellen (AID)

Probleemstelling

Beeldinpainting is een conditionele generatietask waarbij het doel is ontbrekende gebieden van een afbeelding te reconstrueren, terwijl consistentie met zichtbare pixels wordt behouden. Hoewel voorgeïmplementeerde diffusiemodellen deze capaciteit aanzienlijk hebben verbeterd, staan bestaande benaderingen voor een dichotomie in efficiëntie en aanpasbaarheid:

1. Opdrachtspecifieke training: Methoden zoals LaMa, MAT en BrushNet trainen dedicated modellen of grote aanpassingsmodules. Hoewel dit de inferentiekosten amortiseert, vereist dit aanzienlijke trainingtijd en voegt het aanzienlijke modelcomplexiteit toe.
2. Per-instantie optimalisatie: Methoden zoals RePaint, DPS en MCG houden de voorgeïmplementeerde diffusie-ruggengraat vast, maar passen deze afzonderlijk aan voor elke gemaskerde afbeelding tijdens het testen. Dit behoudt het voorgeïmplementeerde model, maar brengt herhaalde optimalisatiekosten met zich mee tijdens de implementatie, waarbij elke afbeelding wordt behandeld als een nieuw inferentieprobleem.

Er is een gat voor een "tussenoplossing" die inferentie amortiseert zonder een groot opdrachtspecifiek model te trainen of per-instantie optimalisatie tijdens het testen uit te voeren.

Methodologie: Geamortiseerde Inpainting met Diffusie (AID)

De auteurs stellen AID voor, een raamwerk dat de voorgeïmplementeerde diffusie-ruggengraat vasthoudt en een kleine, herbruikbare gidsmodule offline traint. Deze module wordt vervolgens bij het testen hergebruikt voor alle gemaskerde afbeeldingen zonder verdere optimalisatie.

1. Controlerend Formulering

Het probleem wordt geformuleerd als een deterministisch controleprobleem.

• Dynamiek: Het omgekeerde diffusieproces wordt beheerst door een stroom van waarschijnlijkheid-ODE aangedreven door een vaste voorgeïmplementeerde scorefunctie $\hat{S}$.
• Gids: Een opdrachtsafhankelijke controlemoot $u(t, x; \xi)$ wordt toegevoegd aan de drift, waarbij $\xi = (M, y)$ de waarneembare masker en zichtbare pixels voorstelt. De controle is niet afhankelijk van het verborgen schone beeld $x^\dagger$.
• Doel: Het doel is om een terminale verlies $\Psi$ te minimaliseren (balancerend tussen trouw aan zichtbare pixels en het ontbrekende gebied) plus een lopende kosten die de grootte van het gidsveld bestraffen. Dit leidt tot een Hamilton-Jacobi-Bellman (HJB) vergelijking waarbij de optimale controle $u^*$ deterministisch is.

2. Theoretische Brug: Deterministisch naar Stochastisch

Het direct oplossen van de hoog-dimensionale deterministische HJB-vergelijking is onuitvoerbaar. De auteurs introduceren een auxiliair gerandomiseerd controleprobleem met behulp van Gaussische beleidslijnen om het probleem leersaam te maken via Continuous-Time Reinforcement Learning (CT-RL).

• Auxiliair Probleem: In plaats van een deterministische controle, gebruikt het systeem een Gaussisch beleid $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$, waarbij het gemiddelde $\mu$ de kandidaat-gidswet is.
• Optimizer-bewarende Brug: Het artikel bewijst twee kritieke stellingen:
  1. Voldoende: De optimale deterministische gids $u^*$ kan worden opgeheven naar een optimaal Gaussisch beleid door het beleidsgemiddelde in te stellen op $u^*$.
  2. Noodzakelijk: Elk optimaal Gaussisch beleid voor het auxiliaire probleem moet een gemiddelde hebben dat de optimale deterministische gids $u^*$ exact herstelt.
• Betekenis: Deze brug zorgt ervoor dat het leren van een stochastisch beleid (wat gradiëntgebaseerde updates mogelijk maakt) de doel-deterministische gids die bij implementatie wordt gebruikt, niet verandert. De variantie van het Gaussische beleid is vast en dient alleen als leermiddel, niet als entropie-regularisator.

3. Geamortiseerd Actor-Critic Algorithm

Gebaseerd op de theoretische brug ontwikkelen de auteurs een volledig datagedreven actor-critic algoritme:

• Parametrisatie: De actor (gidsgemiddelde $\hat{\mu}_\phi$) en criticus (waardefunctie $\hat{V}_\theta$) worden geparametriseerd door neurale netwerken.
• Leren: Het algoritme gebruikt martingale orthogonaliteitsvoorwaarden afgeleid van de auxiliaire HJB-vergelijking. Het update de actor en criticus met behulp van stochastische benadering op basis van gesamplede trajecten uit de opdrachtverdeling.
• Implementatie: Eenmaal getraind, wordt de gidsmodule deterministisch ingezet. Voor een nieuwe gemaskerde afbeelding gebruikt het systeem het geleerde gemiddelde-gids $\hat{\mu}_\phi$ binnen de omgekeerde-tijdoplosser, waarbij het beleidsruis wordt verworpen.

Belangrijkste Bijdragen

1. Methodologie: AID introduceert een gecontroleerd-gidsraamwerk dat inferentie amortiseert door een lichtgewicht gidsmodule offline te trainen, die vervolgens over opdrachten wordt hergebruikt zonder per-instantie optimalisatie.
2. Theorie: Het artikel vestigt een rigoureuze "beleids-equivalentiebrug" tussen het deterministische gidsprobleem en een auxiliaire Gaussische formulering. Het bewijst dat het auxiliaire probleem optimizer-bewarend is, wat het gebruik van stochastische actor-critic-methoden rechtvaardigt om een deterministisch controleveld te leren.
3. Efficiëntie: De aanpak voegt minder dan 1% trainbare parameters toe ten opzichte van de voorgeïmplementeerde ruggengraat, terwijl de kwaliteit-snelheid trade-off aanzienlijk wordt verbeterd.

Empirische Resultaten

De auteurs hebben AID geëvalueerd op AFHQv2, FFHQ (met behulp van de pixel-ruimte EDM-pijplijn) en ImageNet (met behulp van de latent-ruimte EDM2-pijplijn).

• Prestaties vs. Vaste-Ruggengraat Baselines: AID presteert consistent beter dan sterke vaste-ruggengraat baselines (Unguided, Replacement, MCG, DPS) over meerdere metrieken (PSNR, SSIM, LPIPS) en maskertypes (vrij-vorm, centrum, strip).
• Prestaties vs. RePaint: AID overtreft RePaint (een sterke iteratieve optimalisatie baseline) in kwaliteit, terwijl slechts ~10% van het bemonsteringsbudget wordt gebruikt (bijv. 35 stappen versus 350 stappen voor RePaint).
• Prestaties vs. LatentPaint: Op ImageNet overtreft AID LatentPaint (een herbruikbare latent-ruimte aanpassingsmethode) op alle metrieken, terwijl het een gidsmodule gebruikt met minder dan 45% van de trainbare parameters.
• Latentie: AID ondersteunt een laag-latentie werkpunt (12 stappen) dat de inferentietijd verder verlaagt terwijl het prestaties behoudt die superieur zijn aan niet-RePaint baselines en concurrerend blijven met RePaint.
• Generalisatie: Dezelfde getrainde gidsmodule transferreert effectief naar ongezette maskertypes en verschillende bemonsteringsbudgetten zonder hertraining.

Betekenis en Claims

Het artikel beweert dat AID het eerste geamortiseerde gecontroleerd-gidsraamwerk is voor beeldinpainting dat:

1. De voorgeïmplementeerde diffusie-ruggengraat bevroren houdt.
2. Alleen een kleine, herbruikbare gidsmodule traint.
3. Een exacte optimizer-bewarende brug biedt van stochastisch actor-critic leren naar het deterministische gidsveld dat bij implementatie wordt gebruikt.

De auteurs positioneren dit werk als een stap in de richting van het gebruik van controle- en RL-gebaseerde raamwerken om theoretisch onderbouwde methoden te bieden voor het bestuderen en verbeteren van diffusiemodellen. Ze benadrukken dat continuous-time RL de basis vormt voor zowel de rigoureuze analyse als het praktische algoritmeontwerp, en een principieel alternatief biedt voor heuristische per-instantie optimalisatie of zware opdrachtspecifieke hertraining.