Technische Samenvatting: Align & Invert

Probleemstelling

Inverse problemen streven ernaar een schoon signaal $x_0$ te herstellen uit ruisbevangen of gedegradeerde metingen $y$. Hoewel vooraf getrainde diffusie- en stromingsgebaseerde modellen krachtige a priori-kennis zijn geworden voor deze taken, hebben ze vaak moeite om gedetailleerde afbeeldingen te produceren in scenario's met ernstige degradatie (bijvoorbeeld extreme super-resolutie, bewegingsonscherpte) of complexe scènes. Een specifieke beperking doet zich voor bij Latent Diffusion Models (LDM's), waarbij de niet-lineariteiten van de encoder-decoder-architectuur artefacten kunnen introduceren tijdens het omgekeerde steekproefproces. Bovendien vereisen bestaande methoden vaak een groot aantal discretisatiestappen om hoogwaardige reconstructies te bereiken, en ontbreekt het hen aan sterke semantische beperkingen tijdens de inferentie om de trajectorie te sturen naar perceptueel realistische oplossingen.

Methodologie

Het artikel stelt Align & Invert voor, een raamwerk dat Representation Alignment (REPA) integreert in het inferentieproces van diffusie- en stromingsgebaseerde oplossingsmethoden voor inverse problemen.

Kernconcept

De methode dwingt alignering af tussen de interne representaties van het generatieve model (diffusie/stroming) en de kenmerken van een vooraf getrainde, zelftoezichtende visuele encoder (specifiek DINOv2). Hoewel REPA eerder werd gebruikt als trainingsregularisator om de generatieve kwaliteit te verbeteren, past dit werk het aan voor inferentie in inverse problemen waar grondwahrheidssignalen niet beschikbaar zijn.

Technische Implementatie

1. Proxy-construktie: Aangezien de grondwahrheid $x_0$ onbekend is, construeert de methode een proxy-representatie $c_{proxy}$ om de ontbrekende doelkenmerken te vervangen.

   • Initieel wordt $c_{proxy}$ afgeleid uit de DINOv2-kenmerken van de beschikbare meting $y$.
   • Naarmate het omgekeerde diffusieproces vordert, wordt de proxy geleidelijk bijgewerkt naar de DINOv2-kenmerken van de huidige ontroostte schatting van het model, $f_{DINOv2}(D(E[z_0 | z_t]))$.
   • Deze strategie vertrouwt op de waargenomen robuustheid van DINOv2-kenmerken tegenover veelvoorkomende beelddegradaties.

2. Kantelende verdeling: Het steekproefproces wordt geleid door een gekantelde verdeling die de aligneringsterm als beloning incorporeert:
   $$\tilde{p}(x_t | y) \propto p(y | x_t) p_t(x_t) \exp(\lambda \cdot \text{REPA}(\bar{x}_0, h_t))$$
   Hierbij enforceert $p(y | x_t)$ de consistentie van de meting, is $p_t(x_t)$ de onvoorwaardelijke prior, en stimuleert de REPA-term alignering tussen de tussentijdse verborgen toestanden $h_t$ van het diffusiemodel en de proxy-DINOv2-kenmerken.

3. Algoritme-integratie: De aanpak wordt geïntegreerd in bestaande state-of-the-art oplossingsmethoden (bijvoorbeeld Latent DPS, ReSample) door een gradiëntstap op basis van het REPA-doel toe te voegen tijdens de omgekeerde steekproeflus. Deze stap aligneert de latente diffusietoestand met de semantische kenmerken van het doel.

4. Aanpasbaarheid: Voor modellen die oorspronkelijk niet met REPA zijn voorgetraind, leert het raamwerk een lichtgewicht post-hoc projectiekop ($g_\phi$) om de interne representaties van het diffusiemodel af te beelden op de DINOv2-kenmerkenruimte. Deze mapping wordt geoptimaliseerd terwijl het diffusiemodel bevroren blijft, waardoor de aanpassing rekenkundig goedkoop is.

Theoretische inzichten

Het artikel biedt twee belangrijke theoretische resultaten die de REPA-regularisator karakteriseren:

1. Divergentie-minimalisatie: Het maximaliseren van het REPA-doel blijkt equivalent te zijn aan het minimaliseren van de Maximum Mean Discrepancy (MMD) tussen de DINOv2-kenmerkenverdelingen van de gereconstrueerde afbeelding en de grondwahrheid, tot op termen die rekening houden met kenmerk-misalignering en proxy-benaderingsfouten.
2. Contractie van de latente toestand: Onder regulariteitsaannamen met betrekking tot de Jacobianen van de diffusie-encoder en de projectiekop induceert een enkele REPA-update een contractie in de latente diffusietoestandsruimte. Specifiek stuurt het de huidige latente toestand $z_t$ dichter naar de latente representatie van het schone beeld $z^*$, mits de huidige toestand zich binnen een lokaal buurtpunt van de oplossing bevindt.

Experimentele resultaten

De auteurs hebben de methode geëvalueerd op vier inverse problemen: Super-resolutie (4×, 8×, 16×), Box Inpainting, Gaussische ontscherping en Bewegingsontscherping, met behulp van de ImageNet- en FFHQ-datasets.

• Perceptuele kwaliteit: Het integreren van REPA verbeterde consistent de perceptuele metrieken (LPIPS en FID) voor alle taken en oplossingsmethoden (Latent DPS, ReSample, Latent DAPS). Bijvoorbeeld, bij 4× super-resolutie op ImageNet verlaagde Latent DPS + REPA de LPIPS van 0,238 naar 0,217 en de FID van 123,82 naar 86,88.
• Pixelmetrieken: Standaard pixel-voor-pixel metrieken (PSNR, SSIM) bleven vergelijkbaar met de baselines, wat de bekende perceptie-distortie-trade-off weerspiegelt waarbij REPA semantische realisme prioriteert boven pixel-perfecte nauwkeurigheid.
• Efficiëntie: Een significante bevinding is dat REPA oplossingsmethoden in staat stelt om met aanzienlijk minder discretisatiestappen hetzelfde niveau van perceptuele kwaliteit te bereiken. Bij super-resolutie en Gaussische ontscherping behaalde de methode vergelijkbare LPIPS-scores met ongeveer 4× minder stappen dan de baseline, wat leidde tot een verminderde totale runtime ondanks de extra berekening per stap.
• Robuustheid: De methode bleef effectief onder ernstige degradaties (bijvoorbeeld 16× upscaling, hoge intensiteit van onscherpte) en bij toepassing op onafhankelijk voorgetrainde modellen via de geleerde projectiekop.

Betekenis en claims

Het artikel claimt dat representatie-alignering een effectief hulpmiddel vormt voor het verbeteren van generatieve oplossingsmethoden voor inverse problemen tijdens de inferentie.

• Semantische geleiding: Door interne modelrepresentaties te aligneren met robuuste visuele kenmerken (DINOv2), biedt de methode een sterke semantische prior die de reconstructie stuurt naar perceptueel betekenisvolle oplossingen, zelfs zonder grondwahrheid-toezicht.
• Theoretische rechtvaardiging: Het werk verduidelijkt waarom alignering werkt, en toont aan dat het fungeert als een mechanisme om divergentie in de kenmerkenruimte te minimaliseren en de latente toestand te contracteren naar de oplossing van het schone beeld.
• Algemene toepasbaarheid: Het raamwerk is algemeen, compatibel met diverse diffusie-/stromingsgebaseerde oplossingsmethoden, en aanpasbaar aan modellen die oorspronkelijk niet met alignering zijn getraind via een lichtgewicht post-hoc mapping.
• Efficiëntie: Het adresseert de rekenkundige bottleneck van diffusiesteekproeven door convergentie te versnellen, waardoor hoogwaardige reconstructie mogelijk is met minder steekproefstappen.

De auteurs concluderen dat deze aanpak de kloof overbrugt tussen de generatieve capaciteiten van diffusiemodellen en het semantische begrip van zelftoezichtende encoders, en een principiële manier biedt om de reconstructiekwaliteit te verbeteren in uitdagende instellingen voor inverse problemen.