Sample-efficient evidence estimation of score based priors for model selection

Deze paper introduceert een nieuwe, sample-efficiënte methode om de modelevidentie van diffusiemodel-priors te schatten via posterior-steekproeven, waardoor de beste prior kan worden geselecteerd voor complexe inverse problemen zoals black hole-imaging.

De kern

Stel je voor dat je probeert een wazige, onherkenbare foto te herstellen. Je hebt een stukje van de foto (de meting), maar de rest is weg of vervormd. Om de volledige foto te reconstrueren, heb je een "gok" nodig over hoe de rest eruit zou moeten zien. In de wereld van wiskunde en kunstmatige intelligentie noemen we die gok een prior (een voorafgaande veronderstelling).

Het probleem is: welke gok is de beste? Als je kiest voor een prior die zegt "alle foto's zijn van katten", maar je probeert een foto van een auto te herstellen, krijg je een rare, vervormde auto die op een kat lijkt. Dat is een slechte keuze.

De auteurs van dit paper, Frederic Wang en Katherine Bouman, hebben een nieuwe manier bedacht om te bepalen welke "gok" (prior) het beste past bij de data die je hebt. Ze noemen hun methode DiME.

Hier is een uitleg in simpele taal, met een paar creatieve vergelijkingen:

1. Het Probleem: De "Gokkeuze"

Stel je voor dat je in een donkere kamer staat en je hoort een geluid. Je wilt weten wat er in de kamer is.

• Optie A: Je denkt: "Het is waarschijnlijk een hond." (Je prior is een hond).
• Optie B: Je denkt: "Het is waarschijnlijk een kat." (Je prior is een kat).
• Optie C: Je denkt: "Het is waarschijnlijk een draaimolen." (Je prior is een draaimolen).

Als het geluid een blaf is, past Optie A het beste. Maar hoe weet je dat zeker zonder de lichten aan te doen? Je moet een manier vinden om te "meten" hoe goed je gok past bij het geluid. In de wetenschap heet dit model evidence. Hoe hoger de "evidence", hoe waarschijnlijker het is dat je de juiste gok hebt gedaan.

2. De Uitdaging: De "Onmogelijke Rekening"

Vroeger was het heel moeilijk om deze "evidence" te berekenen voor de nieuwste en slimste AI-modellen (die Diffusion Models heten). Het was alsof je probeerde het totale gewicht van een wolk te berekenen door elke waterdruppel afzonderlijk te wegen. Het kostte te veel tijd en rekenkracht, of de resultaten waren onnauwkeurig.

Bestaande methoden waren als iemand die probeert een berg te beklimmen door blindelings elke steen te voelen. Als de berg (de wiskundige ruimte) erg steil en complex is, raken ze vast.

3. De Oplossing: DiME (De "Tijdsreiziger")

De auteurs hebben DiME bedacht. Dit werkt als een slimme tijdsreiziger.

Stel je voor dat je een foto hebt die volledig in ruis (witte statische) is veranderd. Een Diffusion Model kan deze ruis stap voor stap terugdraaien naar een heldere foto. Tijdens dit proces passeert het model duizenden tussenstadia: eerst heel wazig, dan iets minder wazig, dan bijna scherp.

DiME kijkt niet alleen naar het begin en het einde. Het kijkt naar alle tussenstappen die het model al maakt.

• De Analogie: Stel je voor dat je een boot vaart van een stormachtige zee (ruis) naar een rustige haven (de heldere foto). De boot maakt honderden kleine golfbewegingen onderweg.
• De oude methoden probeerden de reis te berekenen door alleen naar de start en finish te kijken, of door duizenden bootjes te laten varen.
• DiME kijkt naar de golven van één bootje dat je al hebt. Het meet hoe "nat" de boot wordt onderweg. Als de boot onderweg veel water opvangt (wat betekent dat de prior niet goed past bij de meting), dan is de "evidence" laag. Als de boot droog blijft (de prior past perfect), is de "evidence" hoog.

Het slimme aan DiME is dat het geen extra rekenkracht nodig heeft. Het gebruikt de tussenstappen die het AI-model al maakt om de foto te maken. Het is alsof je de temperatuur van de kamer meet terwijl je al naar de thermostaat kijkt, in plaats van een nieuwe thermometer te kopen.

4. Waarom is dit geweldig?

De auteurs hebben hun methode getest op verschillende dingen:

1. Wiskundige puzzels: Waar ze het antwoord al wisten, bleek DiME het exact juiste antwoord te geven.
2. Digitale foto's (MNIST): Ze lieten de AI kiezen tussen een prior van het cijfer '6' of '9'. Als ze een wazige '6' lieten zien, koos DiME altijd voor '6', zelfs als andere methoden dachten dat het een '9' was.
3. Zwarte gaten (M87):* Dit is het coolste deel. Ze gebruikten het om te kijken welke theorie over zwarte gaten het beste paste bij de echte foto's van het Event Horizon Telescope.
   • Ze hadden een prior gebaseerd op simpele wiskundige modellen.
   • Ze hadden een prior gebaseerd op complexe fysica (GRMHD).
   • Ze hadden een prior getraind op gezichten of sterrenbeelden.
   • Het resultaat: DiME zei: "De complexe fysica-theorie (GRMHD) past het beste bij de echte foto van het zwarte gat." Dit bevestigt dat onze huidige wetenschappelijke modellen over zwarte gaten goed zijn, maar laat ook zien dat er nog ruimte is voor verbetering.

Samenvatting

DiME is een slimme, snelle manier om te zeggen: "Welke van deze AI-voorkennis (priors) is het meest waarschijnlijk correct voor deze specifieke meting?"

In plaats van duizenden extra berekeningen te doen, gebruikt het de "reistijd" van het AI-model zelf om een antwoord te geven. Het helpt wetenschappers om de beste modellen te kiezen voor het reconstrueren van beelden, van medische scans tot foto's van zwarte gaten, zonder dat ze blindelings hoeven te gokken.

Kortom: Het is de kompasnaald die je vertelt welke route je moet nemen in een labyrint van onzekerheid, zonder dat je het hele labyrint hoeft te doorlopen.

Technische samenvatting

Probleemstelling

In Bayesiaanse inverse beeldproblemen (zoals beeldherstel of reconstructie) speelt de keuze van de prior-verdeling $p(x)$ een cruciale rol. Een slecht gekozen prior kan leiden tot ernstige bias in de reconstructie, vooral als de werkelijke afbeelding buiten het draagvlak van de prior valt. De ideale manier om de beste prior te selecteren is via model evidence (de marginale waarschijnlijkheid $p(y | M)$), die aangeeft hoe goed een model $M$ de waarnemingen $y$ verklaart.

Het probleem is echter dat het berekenen van deze model evidence voor diffusiemodellen (de huidige state-of-the-art voor data-gedreven priors) onberekenbaar (intractable) is. Bestaande schatters voor model evidence hebben meestal twee grote nadelen:

1. Ze vereisen duizenden steekproeven of nauwkeurige evaluaties van de ongebruikelijke prior-dichtheid of de "clean prior score" ($\nabla_x \log p(x)$).
2. Diffusiemodellen zijn vaak onnauwkeurig of slecht geconditioneerd bij het schatten van de score van de schone prior, wat leidt tot vertekende schattingen of trage convergentie bij bestaande methoden.

Methodologie: DiME

De auteurs stellen DiME (Diffusion Model Evidence) voor, een schatter voor model evidence die specifiek is ontworpen voor diffusie-priors. De kern van de methode is het integreren over de tijdmarginalen van posterior-steekproefmethoden.

Kernprincipes:

• Integratie over tijdmarginalen: In plaats van te vertrouwen op de schone prior-score, integreert DiME langs het pad van de posterior-verdeling $p(x_t | y)$ tijdens het reverse diffusion-proces.
• Gebruik van tussentijdse steekproeven: De methode maakt gebruik van de grote hoeveelheid intermediaire samples die natuurlijk gegenereerd worden tijdens het posterior-sampling proces (bijvoorbeeld met de DAPS-methode). Hierdoor zijn slechts een handvol posterior-steekproeven nodig (bijv. 20) voor een nauwkeurige schatting.
• Formulering: De log-evidence wordt geschat via de relatie:
  $$\log p(y) = \mathbb{E}_{x_0 \sim p(x_0|y)}[\log p(y | x_0)] - D_{KL}(p(x_0 | y) || p(x_0))$$
  De KL-divergentie tussen de posterior en de prior wordt geschat door een som te nemen over de tijdstappen, waarbij de kwadratische norm van de likelihood-score $\|\nabla_{x_t} \log p(y | x_t)\|^2$ wordt gebruikt.

Technische Innovaties:

1. Verbeterde Covariantie Benadering: Voor de Gaussian Approximation DAPS-methode stellen de auteurs een verbeterde covariantie-benadering voor die kennis van de prior $p(x_0)$ integreert. Dit voorkomt dat de variantie bij hoge ruisniveaus wordt overschat, wat een veelvoorkomend probleem is bij eerdere heuristieken.
2. Onbevooroordeelde Schatters voor de Score: Omdat het direct berekenen van $\nabla_{x_t} \log p(y | x_t)$ onmogelijk is, gebruiken ze twee onbevooroordeelde schatters ($\Theta_{high}$ en $\Theta_{low}$) die afhankelijk zijn van het ruisniveau.
   • Bij hoge ruis wordt gebruikgemaakt van de afstand tussen een gesamplede $\tilde{x}_0$ en de verwachting $\mathbb{E}[x_0|x_t]$.
   • Bij lage ruis wordt de likelihood-score direct benaderd.
   • Om een onbevooroordeelde schatting van de gekwalificeerde score te krijgen, worden er twee onafhankelijke steekproeven $\tilde{x}_0^{(1)}$ en $\tilde{x}_0^{(2)}$ genomen per tijdstap, zodat het product $\Theta^{(1)T}\Theta^{(2)}$ de verwachte waarde correct benadert zonder de bias van het kwadrateren van een enkele schatter.

Belangrijkste Bijdragen

1. DiME Schatter: Een nieuwe, sample-efficiënte methode om model evidence te schatten voor diffusie-priors zonder de prior-score of -dichtheid te hoeven kennen.
2. Generalisatie: De methode is afgeleid voor standaard marginalen (zoals gebruikt in DAPS) en uitgebreid naar willekeurige marginalenpaden (zoals in PnP-DM).
3. Praktische Implementatie: Een praktische integratie met de Decoupled Annealing Posterior Sampling (DAPS) methode, inclusief verbeterde covariantie-benaderingen en strategieën voor het schatten van de likelihood-score.
4. Validatie op Real-World Data: Toepassing op een complex, niet-convex probleem: de reconstructie van het zwart gat M87* met data van de Event Horizon Telescope.

Resultaten

De auteurs testen DiME op drie niveaus:

1. Gaussische Mengsels (Analytische Ground Truth):

   • DiME levert bijna onbevooroordeelde schattingen van de evidence, vergelijkbaar met geavanceerde baselines zoals Sequential Monte Carlo (SMC) en Thermodynamic Integration (TI).
   • Belangrijk: DiME presteert goed zonder de ware prior-score te gebruiken, terwijl baselines die dit wel doen, vaak vertekende resultaten geven bij uit-distributie (OOD) data.
   • De verbeterde covariantie-benadering elimineert de bias die ontstaat bij de originele DAPS-heuristiek.

2. Model Selectie op MNIST (Niet-convexe Problemen):

   • Bij fase-retrieval problemen (Gaussisch en Fourier) selecteert DiME consistent de juiste prior uit een set van 10 diffusiemodellen (één per MNIST-cijfer) op basis van één ruisbevatte meting.
   • Baselines (zoals SMC met een geleerde score) falen vaak in het selecteren van het juiste model, wat aantoont dat methoden die afhankelijk zijn van een schone prior-score niet geschikt zijn voor diffusie-gebaseerde modelselectie.

3. M87 Zwart Gat Imaging:*

   • DiME wordt toegepast op echte observaties van het Event Horizon Telescope.
   • Model Selectie: De methode concludeert dat een prior getraind op synthetische GRMHD-simulaties (General-Relativistic Magnetohydrodynamics) de hoogste waarschijnlijkheid heeft, vergeleken met priors getraind op RIAF-simulaties, algemene ruimtebeelden, gezichten of MNIST-cijfers.
   • Model Validatie: Door de evidence van de M87*-observaties te vergelijken met de verdeling van evidence van gesimuleerde GRMHD-data, wordt vastgesteld dat M87* statistisch gezien "in-distribution" is ten opzichte van de GRMHD-prior (p-waarde 0.209). Dit bevestigt de geldigheid van het huidige fysische model, terwijl er nog ruimte is voor verfijning.

Significantie

Dit paper is significant omdat het een fundamentele beperking van diffusiemodellen in de wetenschappelijke beeldvorming oplost: het gebrek aan een betrouwbare manier om model evidence te berekenen.

• Wetenschappelijke Impact: Het stelt onderzoekers in staat om niet alleen reconstructies te maken, maar ook om principieel modelselectie en validatie uit te voeren. Dit is cruciaal in gebieden zoals astrofysica, waar het bepalen van de juiste fysische prior (bijv. GRMHD vs. RIAF) direct invloed heeft op de interpretatie van waarnemingen.
• Efficiëntie: De methode is extreem sample-efficiënt (slechts ~20 steekproeven nodig), wat het computatiekosten drastisch verlaagt ten opzichte van bestaande methoden die duizenden steekproeven vereisen.
• Robuustheid: DiME is robuust tegen uit-distributie data en vereist geen toegang tot de onderliggende prior-score, wat het toepasbaar maakt voor complexe, geleerde priors waar de score onnauwkeurig kan zijn.

Kortom, DiME maakt diffusie-priors niet alleen bruikbaar voor reconstructie, maar ook voor rigoureuze statistische inferentie en modelvalidatie in kritieke wetenschappelijke toepassingen.