Unbiased Diffusion Variational Inversion via Principled Posterior Matching

Probleemstelling

Berekeningsmatige beeldvormingstaken behelzen het herstellen van een onbekend signaal $x$ uit ruisbeïnvloede, indirecte metingen $y = A(x) + \eta$, waarbij $A$ een slecht gestelde voorwaartse operator is. Een kritieke uitdaging bij deze inverse problemen is niet alleen het verkrijgen van een hoogwaardige puntsschatting, maar ook het uitvoeren van betrouwbare onzekerheidskwantificatie door de volledige Bayesiaanse posteriorverdeling $p(x|y)$ te herstellen.

Bestaande op diffussie gebaseerde methoden voor het oplossen van deze inverse problemen vallen over het algemeen in drie categorieën, die allemaal te lijden hebben onder aanzienlijke theoretische beperkingen met betrekking tot de optimalisatie van de Kullback-Leibler (KL)-divergentie:

1. Gradiëntgeleide Monte Carlo (bijv. DPS): Deze methoden benaderen de onberekenbare tijd-afhankelijke waarschijnlijkheids-term, wat vaak leidt tot instabiele steekproeftrajecten en vertekende posterior-schattingen.
2. Optimalisatiegebaseerde Variatiele Inference (bijv. RED-Diff, SDS): Deze methoden gaan doorgaans uit van een degeneratieve variatiele verdeling (een Dirac-delta), waardoor ze effectief de entropieterm in het variatiele doelwit verwerpen. Dit reduceert het probleem tot Maximum A Posteriori (MAP)-schatting, wat leidt tot ernstige mode-collaps en het onvermogen om de multi-modale structuur van de posterior te vangen.
3. Geamortiseerde Inference (bijv. DAVI): Deze methoden trainen feed-forward netwerken om posteriors direct te voorspellen. Ze vertrouwen echter op de Integral KL (IKL)-divergentie, wat een heuristische benadering is in plaats van een strikte variatiele ondergrens. Deze afwijking, gecombineerd met een afhankelijkheid van gekoppelde ground-truth-data voor supervisie, beperkt hun vermogen om zero-shot te generaliseren en de exacte Bayesiaanse posterior te herstellen.

Methodologie: Principled Posterior Matching (PPM)

De auteurs stellen Principled Posterior Matching (PPM) voor, een raamwerk dat terugkeert naar de fundamenten van variatiele inference om onvertrokken optimalisatie van de exacte KL-divergentie te bereiken.

1. Theoretische Grondslag: Integratie van Fisher-divergentie
In plaats van te vertrouwen op heuristische benaderingen of aannames over degeneratieve verdelingen, formuleert PPM de exacte optimalisatie van de KL-divergentie strikt door de Fisher-divergentie te integreren langs het diffussieproces. De auteurs maken gebruik van een klassiek theoretisch resultaat dat stelt dat de KL-divergentie tussen twee verdelingen kan worden uitgebreid als de integraal van de Fisher-divergentie langs het diffussietraject.

• Kerninzicht: Het minimaliseren van de integraal van de Fisher-divergentie moedigt inherent massa-dekkend gedrag aan, wat theoretisch de mode-collaps die bij eerdere methoden werd waargenomen, mitigeert.

2. Gradiënt-equivalentiestelling
Om de optimalisatie van de Fisher-divergentie hanteerbaar te maken, leiden de auteurs een nieuwe Gradiënt-equivalentiestelling af. Deze stelling biedt een exacte, onvertrokken gradiëntschatter voor de variatiele parameters zonder backpropagation door het bevroren, voorgetrainde score-netwerk te vereisen.

• De gradiënt wordt afgeleid met behulp van een "stop-gradient" operator en een hulpscore-netwerk, wat efficiënte optimalisatie via stochastische gradiëntafdaal mogelijk maakt.

3. Unificerend Raamwerk
PPM verenigt twee complementaire paradigma's onder één enkel optimalisatiedoelwit:

• Variatiele Inference (VI): De variatiele posterior wordt gemodelleerd als een ensemble van beelddeeltjes $\{\mu_k\}$. Dit staat hoge-fideliteit, op deeltjes gebaseerde posterior-herstel toe voor specifieke waarnemingen.
• Geamortiseerde Inference (AI): De variatiele posterior wordt gemodelleerd als een reconstructienetwerk $x = g_\phi(z)$. Dit maakt snelle, single-step inferentie mogelijk door een globale mapping te leren van waarnemingen naar posterior-steekproeven.
• Architectuur: Het raamwerk maakt gebruik van een voorgetraind diffussie-score-netwerk $s_\theta$ als prior en een hulpscore-netwerk $s_\phi$ (aangepast via Low-Rank Adaptation) om de score van de variatiele verdeling $q_\phi$ te benaderen. De optimalisatie wisselt af tussen het updaten van het hulpscore-netwerk en het updaten van de variatiele parameters (deeltjes of netwerkgewichten).

4. Generaliseerbaarheid
De formulatie breidt zich op natuurlijke wijze uit naar een bredere familie van divergentiemaatstaven door de integraal van de Fisher-divergentie te generaliseren, waardoor PPM wordt gepositioneerd als een veelzijdige basis voor score-gebaseerde variatiele inference.

Belangrijkste Bijdragen

• Onvertrokken Optimalisatie: PPM is het eerste raamwerk dat de exacte KL-divergentie voor op diffussie gebaseerde inverse problemen strikt optimaliseert zonder terug te grijpen naar de vertrokken benaderingen (bijv. Dirac-delta-aannames of IKL-heuristieken) die bestaande methoden plagen.
• Theoretische Unificatie: Het paper bewijst dat het minimaliseren van de exacte KL via Fisher-integratie bestaande score-gebaseerde methoden verenigt als speciale gevallen en een theoretische verklaring biedt voor waarom eerdere methoden lijden aan mode-collaps (de "benaderingskloof").
• Ondersteuning voor Dubbele Paradigma's: Het biedt een unificerend trainingslogica die zowel deeltjesgebaseerde VI voor hoogwaardig posterior-herstel als geamortiseerde inference voor efficiënte, ongesuperviseerde generatie ondersteunt.
• Robuuste Onzekerheidskwantificatie: Door mode-collaps te vermijden, produceert PPM goed gekalibreerde onzekerheidsschattingen die de multi-modale aard van de posterior trouw weerspiegelen.

Experimentele Resultaten

De auteurs valideren PPM over een uitgebreide reeks benchmarks voor berekeningsmatige beeldvorming, inclusief natuurlijke afbeeldingen en wetenschappelijke beeldvormingstaken.

1. 2D Posterior-schatting (Speelgoedvoorbeelden)

• Op 2D Gaussische mengsel-priors herstelt PPM nauwkeurig alle posterior-modes. Daarentegen vertoont RED-Diff ernstige mode-collaps (convergentie naar één enkel punt), en RLSD toont meetinconsistentie of slechte prior-gevolgzaamheid.

2. Computergestuurde Fotografie (FFHQ & ImageNet)

• Taken: Inpainting van afbeeldingen, bewegingsontwarring, super-resolutie en Gaussische ruisreductie.
• Prestaties: PPM presteert beter dan state-of-the-art baselines (DPS, RED-Diff, RLSD, DAVI) in reconstructiefideliteit (PSNR, SSIM) en steekproefdiversiteit.
• Kwalitatief: Bij inpainting-taken genereert PPM diverse en plausibele invullingen (bijv. verschillende kapsels), terwijl baselines bijna identieke outputs produceren.
• Onzekerheid: De standaarddeviatie-kaarten van PPM weerspiegelen nauwkeurig positionele onzekerheid (bijv. bij objectranden), terwijl de kaarten van RED-Diff aangeven dat modes zijn ingestort en die van RLSD artefacten bevatten.

3. Toepassingen in Wetenschappelijke Beeldvorming

• Super-resolutie Fluorescentiemicroscopie: Toegepast op subcellulaire structuren (microtubuli, ER). PPM bereikt hogere fideliteit dan RED-Diff en produceert onzekerheidskaarten die overeenkomen met de beeldvormingsfysica (lagere onzekerheid bij structuurcentra, hogere bij vage randen).
• Radio Interferometrische Zwartegatbeelden: Toegepast op synthetische Event Horizon Telescope (EHT)-data. PPM herstelt trouw de ringstructuur en de heldere maansikkel van de schaduw van het zwarte gat. In tegenstelling tot RED-Diff (dat instort) en DPS (dat wazig is), biedt PPM een fysisch betekenisvolle onzekerheidsverdeling die structurele variatie vangt.

Betekenis en Aanspraken

Het paper beweert dat PPM een robust raamwerk voor wetenschappelijke beeldvorming vestigt door de fundamentele theoretische beperkingen van huidige op diffussie gebaseerde inverse oplossers aan te pakken.

• Betrouwbaarheid: Door de benaderingsbias te elimineren die leidt tot mode-collaps, maakt PPM trouwe multi-modale posterior-herstel en goed gekalibreerde onzekerheidsschattingen mogelijk, wat cruciaal is voor wetenschappelijke en medische toepassingen waar het begrijpen van de volledige oplossingsruimte even belangrijk is als de reconstructie zelf.
• Zero-Shot Generalisatie: In tegenstelling tot geamortiseerde methoden die gekoppelde trainingsdata vereisen, werkt PPM op een ongesuperviseerde manier, waardoor het toepasbaar is op wetenschappelijke domeinen waar ground-truth-data niet beschikbaar is.
• Theoretische Grondslag: Het werk biedt een strikte theoretische brug tussen Fisher-divergentie en KL-minimalisatie, en biedt een principiële alternatief voor heuristische benaderingen die het veld hebben gedomineerd.

Kortom, PPM demonstreert dat terugkeren naar de exacte optimalisatie van het variatiele doelwit, gefaciliteerd door Fisher-divergentie-integratie, superieure reconstructiekwaliteit, diversiteit en onzekerheidskwantificatie oplevert in vergelijking met bestaande vertrokken benaderingen.