ShuffleFlow: Scalable Posterior Inference for Bayesian Inverse Imaging

Technische Samenvatting: ShuffleFlow

Probleemstelling
Bayesiaanse inverse beeldvorming streeft naar het herstellen van latente beelden $x$ uit ruisende, incomplete metingen $y = A(x) + n$ door de posterieure verdeling $p(x|y)$ te schatten. Hoewel Variational Inference (VI) een fundamentele aanpak biedt voor posterieure inferentie die efficiënt samples kan genereren na optimalisatie, wordt de toepassing ervan op grootschalige beeldreconstructie gehinderd door een gebrekkige schaalbaarheid. Traditionele VI-methoden die vertrouwen op flow-gebaseerde netwerken, modelleren doorgaans de volledige beeldverdeling direct, wat resulteert in computationele kosten die slecht schalen met de beelddimensie (bijv. $O(N^4)$ voor $N \times N$ pixels). Dit leidt tot onhoudbare geheugen- en tijdbudgetten tijdens de optimalisatie. Bestaande pogingen om de schaalbaarheid te verbeteren, zoals naïeve patching- of pixelatie-strategieën, offeren vaak de beeldkwaliteit en de expressiviteit van de posterieure verdeling op. Daarentegen vereisen diffusiemodellen, hoewel krachtig, grote datasets voor training en brengen ze hoge computationele kosten met zich mee voor het genereren van hoog-resolutie, hoog-sample-aantal posterieure verdelingen vanwege hun iteratieve aard van sampling.

Methodologie: ShuffleFlow
De auteurs stellen ShuffleFlow voor, een schaalbaar VI-framework dat ontworpen is om de taak van posterieure modellering te deconstrueren in drie modulaire componenten: coördinaat-sampling, feature-encoding en posterieure schatting.

1. Coördinaat-sampling (Pixel-Unshuffling): In plaats van het volledige $N \times N$ beeld of onafhankelijke patches te modelleren, verdeelt ShuffleFlow het beeld in een stapel van $S^2$ gecorreleerde sub-beelden met behulp van een pixel-unshuffling operatie met een downsampling-factor $S$. Elk sub-beeld heeft dimensies $(N/S) \times (N/S)$. Dit reduceert de dimensionaliteit van het probleem van $N^2$ naar $N^2/S^2$.
2. Feature Encoding (Neural Field): Om ruimtelijke structuren te vangen zonder de flow-netwerken te belasten met volledige beeldreconstructietaken, gebruikt het framework een Neural Field (NF) encoder (specifiek een coördinaat-gebaseerde MLP, SIREN). De 2D-coördinaat-offsets van de sampling-locaties worden in de NF gevoed om feature-embeddings ($h_i$) voor elk sub-beeld te genereren.
3. Posterieure Schatting (Conditional Normalizing Flow): Een gedeelde Conditional Normalizing Flow (CNF) modelleert de gezamenlijke verdeling van de stapel sub-beelden. De CNF is geconditioneerd op de feature-embeddings ($h_i$) die overeenkomen met de locatie van elk sub-beeld. Cruciaal is dat het framework een door De Finetti geïnspireerde aanpak hanteert door een enkele gedeelde base-latente variabele $z_0$ te delen over alle sub-beelden. Dit maakt het mogelijk om de gezamenlijke verdeling te factoriseren als:
   $$q(x) = \left( \prod_{i=1}^{S^2} q(x_i | z_0; c_i) \right) q(z_0)$$
   waarbij $z_0$ wordt getrokken uit een standaard normale verdeling. Dit deelmechanisme correleert de sub-beelden, wat ruis-artefacten die geassocieerd worden met pixel-onafhankelijkheid voorkomt, terwijl het de omvang van het flow-netwerk aanzienlijk verkleint van $O(N^4)$ naar $O(N^4/S^4)$.

Optimalisatie en Inferentie
De CNF en NF worden gezamenlijk getraind door de Kullback-Leibler (KL) divergentie tussen de variationele verdeling en de Bayesiaanse posterieure verdeling te minimaliseren. De loss-functie bestaat uit drie termen: de data-likelihood, een prior (ofwel een klassieke Total Variation prior of een geleerde score-gebaseerde prior met bevroren gewichten), en een negatieve entropie-term om diversiteit te stimuleren. Tijdens de inferentie genereert de NF embeddings voor de sampling-coördinaten, en de CNF genereert samples geconditioneerd op deze embeddings en een gedeelde base-variabele, die vervolgens via pixel-shuffling worden samengevoegd tot een volledig beeld.

Belangrijkste Bijdragen

1. Probleemdecompositie: De auteurs introduceren een nieuw modulair VI-framework dat de modellering van grote beeld-posterieure verdelingen deconstrueert in coördinaat-sampling, partiële beeldrepresentatie en lichtgewicht onzekerheidskwantificatie. Deze structuur is agnostisch voor specifieke netwerkarchitecturen, waardoor integratie van vooruitgang in neural fields en generatieve priors mogelijk is.
2. Schaalbaarheid en Efficiëntie: Door gebruik te maken van pixel-unshuffling en gedeelde latente variabelen, bereikt ShuffleFlow een massale reductie in modelgrootte en computationele middelen. Theoretische en empirische analyses tonen aan dat de methode bijna even efficiënt schaalt als pixel-wijze inferentie, waarbij een 4096$\times$ reductie in de omvang van het flow-netwerk wordt bereikt voor $256 \times 256$ beelden (met $S=8$) vergeleken met volledige beeld-flowmodellen.
3. Hoog-resolutie Posterieure Generatie: De methode demonstreert het vermogen om snel hoog-sample-aantal posterieure verdelingen te genereren. In tegen tegenstelling tot diffusie-samplers, waarbij de sampling-tijd lineair groeit met het aantal samples, kan ShuffleFlow na een korte trainingsfase efficiënt duizenden samples genereren.

Experimentele Resultaten
De auteurs evalueren ShuffleFlow op zowel lineaire (motion deblurring, compressed sensing MRI) als niet-lineaire (Fourier fase-retrieval) inverse problemen.

• Lineaire Problemen: ShuffleFlow bereikt State-of-the-Art (SOTA) prestaties onder de VI-methoden wat betreft beeldkwaliteit (PSNR, SSIM, LPIPS) en onzekerheid-calibratie (NLL, ECE), waarbij het eerdere flow-gebaseerde VI-methoden (DPI) en naïve patching-strategieën overtreft. Het presteert vergelijkbaar met leidende diffusie-samplers (bijv. DAPS) qua beeldkwaliteit, maar met aanzienlijk lagere computationele overhead.
• Niet-lineaire Problemen: In de uitdagende Fourier fase-retrieval taak, die een bimodale posterieure verdeling vertoont (rechtopstaande versus geroteerde beelden), slaagt ShuffleFlow erin de bimodaliteit te vangen. In contrast hiermee falen methoden zoals CF-NeRF en naïeve patching in het modelleren van deze complexe globale correlaties.
• Efficiëntie: In een scenario dat 10.000 posterieure samples vereist, genereert ShuffleFlow een hoog-resolutie bimodale posterieure verdeling binnen 16 minuten. In hetzelfde tijdsbestek genereren diffusie-samplers (DAPS, DPS) veel minder samples en falen ze om de tweede modus duidelijk te karakteriseren, waarbij ze ongeveer 20$\times$ meer tijd nodig hebben om een vergelijkbare resolutie te bereiken.

Betekenis en Claims
Het artikel claimt dat ShuffleFlow de kritieke schaalbaarheid-bottleneck in Variational Inference voor wetenschappelijke beeldvorming aanpakt zonder de posterieure nauwkeurigheid op te offeren. Het positioneert VI als een essentieel alternatief voor diffusiemodellen, met name in scenario's met weinig data waar diffusiemodellen niet getraind kunnen worden, en in toepassingen die een snelle exploratie van complexe, hoog-resolutie posterieure verdelingen vereisen (bijv. het identificeren van bimodale verdelingen in niet-lineaire inverse problemen). De auteurs merken op dat hoewel de methode zeer efficiënt is, deze momenteel enige overmoed (overconfidence) vertoont in niet-lineaire settings door trainingsinstabiliteit, wat zij identificeren als een doelwit voor toekomstig werk. Het framework wordt gepresenteerd als een exploratief instrument dat kan bepalen wanneer de complexiteit van de posterieure verdeling uitgebreide sample-generatie rechtvaardigt.