VIPaint: Image Inpainting with Pre-Trained Diffusion Models via Variational Inference

1. Probleemstelling

Het paper adresseert de uitdaging van image inpainting (het invullen van ontbrekende of gemaskerde gebieden in een afbeelding) met behulp van vooraf getrainde Diffusiemodellen (DMs), met name Latent Diffusion Models (LDMs) zoals Stable Diffusion.

• De Kernmoeilijkheid: Hoewel DMs uitblinken in het genereren van nieuwe afbeeldingen uit ruis, is het conditioneren ervan op beschadigde of gedeeltelijk waargenomen afbeeldingen (inpainting) niet triviaal. Het doel is om te stalen uit de ware conditionele posteriorverdeling $p(x|y)$, waarbij $y$ de waargenomen (gemaskerde) afbeelding is en $x$ de complete afbeelding.
• Beperkingen van Bestaande Methoden:
  • Heuristische/Iteratieve Methoden: Benaderingen zoals RePaint, CoPaint en DPS (Diffusion Posterior Sampling) vertrouwen vaak op iteratieve verfijning of "time-travel" heuristieken. Deze falen frequent om een globale structuur vast te leggen wanneer grote gebieden zijn gemaskerd, wat leidt tot oncoherente of onrealistische resultaten.
  • Variationale Methoden (REDdiff): Eerdere pogingen tot variational inference (bijv. REDdiff) optimaliseren vaak een enkele afbeelding op het ruisvrije niveau ($t=0$) met behulp van een eenvoudige Gaussische benadering. Dit faalt in het vastleggen van onzekerheid, is vatbaar voor optimalisatie-instabiliteit (vooral in regimes met lage ruis) en heeft moeite om te generaliseren naar LDMs.
  • Trainingsvereisten: Veel state-of-the-art methoden vereisen het fine-tunen van het generatieve model op specifieke datasets of maskers, wat rekenkundig duur is en "zero-shot" toepassing op nieuwe taken verhindert.

2. Methodologie: VIPaint

De auteurs stellen VIPaint voor, een hieraarchische variational inference (VI) algoritme dat het vooraf getrainde diffusiemodel behandelt als een prior. In plaats van een enkele afbeelding te optimaliseren, optimaliseert VIPaint een verdeling over latente variabelen op meerdere ruisniveaus.

A. Hieraarchische Variationale Posterior

VIPaint definieert een niet-Gaussische, Markoviaanse benadering van de ware diffusieposterior. Het introduceert een set van $K$ "keypoints" (intermediate timesteps) $\{h(1), \dots, h(K)\}$ binnen de reverse diffusietrajectorie (waarbij $h(K)$ hoge ruis is en $h(1)$ lage ruis).

De variationale verdeling wordt gedefinieerd als:
$$q_\lambda(z_{h(1):h(K)}) = q_\lambda(z_{h(K)}) \prod_{i=1}^{K-1} q_\lambda(z_{h(i)} | z_{h(i+1)})$$

• Structuur:
  • Hoge Ruis ($h(K)$): Gemodelleerd als een gefactoriseerde Gaussische verdeling.
  • Intermediaire/Lage Ruis: Gemodelleerd als een Gaussische verdeling waarbij het gemiddelde een convexe combinatie is van de prior diffusievoorspelling (van het vooraf getrainde model) en een geleerde variational parameter specifiek voor de afbeelding.
  • Variantie: De variantie wordt geleerd en is dimensiespecifiek, waardoor het model dynamisch de onzekerheid kan verhogen in gemaskerde gebieden terwijl het zekerheid behoudt in waargenomen gebieden.

B. De Doelfunctie

VIPaint optimaliseert een variational lower bound op de marginale likelihood van de observatie $y$. De verliesfunctie bestaat uit drie termen:

1. Reconstructieverlies: Zorgt ervoor dat de gegenereerde steekproeven op het laagste ruisniveau ($h(1)$) consistent zijn met de waargenomen pixels $y$. Dit wordt benaderd met behulp van een één-staps denoising-voorspelling.
2. Priorverlies: Regulariseert het hoogste ruisniveau ($h(K)$) om uit te lijnen met de globale semantiek die is gecodeerd door de diffusieprior in het regime met hoge ruis.
3. Hieraarchisch Verlies: Zorgt voor consistentie tussen aangrenzende keypoints in de hiërarchie, waardoor mid-level afbeeldingsdetails en structuur worden vastgelegd.

C. Optimalisatie en Steekproefneming

• Optimalisatie: De variational parameters $\lambda$ (gemiddelden, varianties en menggewichten) worden per afbeelding geoptimaliseerd met behulp van gradient descent (Adam). Cruciaal is dat dit gebeurt zonder annealing (random time sampling wordt gebruikt voor het verlies), wat een onbevooroordeelde bound biedt en de instabiliteit vermijdt die wordt gezien in methoden die annealen van $T$ naar $0$.
• Steekproefneming: Na optimalisatie worden diverse reconstructies gegenereerd door:
  1. Ancestral Sampling: Steekproeven nemen uit de geoptimaliseerde hierarchische posterior van $h(K)$ naar $h(1)$.
  2. Verfijning: Verdere verfijning van de steekproef van $h(1)$ naar $t=0$ met behulp van het vooraf getrainde denoising-netwerk, geleid door de likelihood-gradiënt (vergelijkbaar met DPS maar startend vanuit een beter geïnitieerde posterior).

3. Belangrijkste Bijdragen

1. Hieraarchische Niet-Gaussische Benadering: In tegenstelling tot eerdere VI-methoden die de posterior benaderen op een enkel punt ($t=0$), modelleert VIPaint de posterior over een hiërarchie van ruisniveaus. Dit stelt het in staat om zowel globale semantiek (bij hoge ruis) als lokale details (bij lage ruis) gelijktijdig vast te leggen.
2. Zero-Shot Adaptatie: VIPaint vereist geen fine-tuning van het vooraf getrainde diffusiemodel. Het past het model aan op nieuwe maskers en afbeeldingen door variational parameters voor elke query te optimaliseren, waardoor het toepasbaar is op elke vooraf getrainde DM (pixel-based of latent).
3. Stabiliteit en Onzekerheidsmodellering: Door het regime met lage ruis te vermijden tijdens de hoofdoptimalisatiefase (waar denoising-netwerken niet-glad en instabiel zijn) en expliciet variantie te modelleren, produceert VIPaint diverse, plausibele imputaties die de onzekerheid van de ontbrekende data weerspiegelen.
4. Generalisatie naar LDMs: De methode is expliciet ontworpen om te werken met Latent Diffusion Models, waardoor high-resolution inpainting (bijv. 1024x1024) mogelijk is met lagere rekenkosten dan pixel-based benaderingen.

4. Experimentele Resultaten

De auteurs hebben VIPaint geëvalueerd op ImageNet64, ImageNet256, LSUN-Churches en Stable Diffusion 3.5 (tekst-geconditioneerd).

• Kwantitatieve Prestaties:

  • Grote Maskers: VIPaint presteert significant beter dan baselines (DPS, RePaint, CoPaint, REDdiff) op grote gemaskerde gebieden (40-80% van de afbeelding) in termen van PSNR, SSIM en LPIPS.
  • Tekst-geconditioneerde Modellen: Op Stable Diffusion 3.5 bereikt VIPaint superieure PSNR (26,69 vs 25,22 voor DPS) en LPIPS (0,050 vs 0,061), waarbij het globaal coherente afbeeldingen produceert die fijne details behouden.
  • Andere Inverse Problemen: De methode toont ook sterke prestaties bij deblurring en super-resolution, en presteert beter dan gespecialiseerde baselines zoals PSLD en ReSample.

• Kwalitatieve Prestaties:

  • Visuele resultaten tonen aan dat VIPaint globaal coherente structuren genereert (bijv. consistente objectvormen en belichting) waar andere methoden wazige of oncoherente artefacten produceren.
  • Het genereert succesvol meerdere diverse plausibele aanvullingen voor dezelfde gemaskerde afbeelding, wat zijn vermogen demonstreert om de multi-modale aard van de posterior te modelleren.

• Efficiëntie:

  • VIPaint is rekenkundig efficiënt. Voor $K=2$ zijn slechts ongeveer 150 denoising-netwerkaanroepen nodig voor optimalisatie (vergeleken met 1000+ voor op steekproeven gebaseerde methoden) en is het sneller dan veel baselines terwijl het hogere kwaliteit levert.

5. Betekenis

VIPaint vertegenwoordigt een aanzienlijke vooruitgang in inverse problemen met diffusiemodellen. Door te verschuiven van heuristische steekproefneming of optimalisatie op een enkel punt naar een principiële hierarchische variational inference framework, lost het het kritieke probleem van globale incoherentie bij inpainting met grote maskers op.

Het vermogen om zero-shot te werken met state-of-the-art foundation modellen (zoals Stable Diffusion 3.5) zonder hertraining, maakt het een zeer praktisch hulpmiddel voor real-world taken zoals afbeeldingredactie, restauratie en voltooiing. Het overbrugt de kloof tussen de theoretische strengheid van variational inference en de praktische mogelijkheden van moderne generatieve AI.