DAWN-FM: Data-Aware and Noise-Informed Flow Matching for Solving Inverse Problems

Titel: DAWN-FM: Data-Aware en Noise-Informed Flow Matching voor het Oplossen van Inverse Problemen

1. Het Probleem: Ill-geplaatste Inverse Problemen

Inverse problemen komen veel voor in velden zoals medische beeldvorming, geofysica en signaalverwerking. Het doel is om parameters van een systeem te schatten op basis van onvolledige of ruisbeïnvloede waarnemingen. Deze problemen zijn vaak ill-geplaatst (ill-posed), wat betekent dat:

• Er geen unieke oplossing bestaat.
• De oplossing extreem gevoelig is voor kleine verstoringen in de data (ruis).

Traditionele oplossingen vereisen regularisatie. Bestaande machinelearning-benaderingen, zoals voorgeprogrammeerde diffusiemodellen, hebben echter tekortkomingen:

• Ze presteren vaak slecht bij sterk ill-geplaatste problemen.
• Ze zijn niet robuust bij hoge ruisniveaus omdat ze de data-distributie als een starre prior behandelen, ongeacht de specifieke metingen.
• Ze genereren vaak één oplossing en bieden beperkte mogelijkheden voor onzekerheidskwantificatie.

2. Methodologie: DAWN-FM

De auteurs introduceren DAWN-FM (Data-AWare and Noise-Informed Flow Matching), een generatief raamwerk dat Flow Matching (FM) aanpast voor inverse problemen. In plaats van een generiek voorgeprogrammeerd model te gebruiken, wordt het model specifiek getraind voor het desbetreffende inverse probleem.

Kernconcepten:

• Flow Matching (FM): Een generatief proces dat een eenvoudige referentieverdeling (bijv. een Gaussische verdeling $\pi_0$) transformeert naar de doelverdeling (de posterior $\pi_1$) via een deterministisch proces (een differentiaalvergelijking of ODE). Het leert een tijdsafhankelijk snelheidsveld $v$ dat deze transformatie stuurt.
• Data- en Ruismelding (Embedding): Het cruciale verschil met standaard FM is dat DAWN-FM de gemeten data ($b$) en het ruisniveau ($\sigma$) expliciet in het trainingsproces integreert.
  • De data $b$ wordt getransformeerd naar de beeldruimte via de geadjungeerde operator $A^\top b$ (of een encoder $E_\eta(A^\top b)$).
  • Het ruisniveau $\sigma$ wordt als een scalar input aan het netwerk gegeven.
  • Het snelheidsnetwerk $s_\theta$ neemt dus als input: het geïnterpoleerde punt $x_t$, de tijd $t$, de data-embedding $f(b)$ en het ruisniveau $\sigma$.

Trainingsdoel en Loss-functie:
De training minimaliseert een gecombineerde loss-functie (Vergelijking 14):

1. Snelheidsfout ($L_1$): Het verschil tussen de voorspelde snelheid en de werkelijke snelheid van de trajecten.
2. Data-consistentie ($L_2$): Een misfit-term die de afwijking straft tussen de waargenomen data $b$ en de data die zou worden gegenereerd door de gereconstrueerde afbeelding ($A\hat{x}_1$). Dit zorgt ervoor dat de oplossing niet alleen plausibel is volgens de prior, maar ook consistent is met de fysieke metingen.

Inferentie en Onzekerheid:
Omdat het proces stochastisch is (afhankelijk van de initiële random noise $x_0$), kan het model meerdere mogelijke oplossingen genereren voor dezelfde invoer. Door de ODE $M$ keer te integreren met verschillende startpunten, kan men:

• De posteriormiddellijn berekenen (gemiddelde van alle oplossingen) als de beste schatting.
• De standaarddeviatie berekenen om een onzekerheidskaart te genereren, die aangeeft welke delen van de afbeelding minder betrouwbaar zijn gereconstrueerd.

3. Belangrijkste Bijdragen

1. Probleem-specifiek Training: Het model wordt getraind om direct de posterior-distributie $\pi(x|b)$ te leren voor een specifiek inverse probleem, in plaats van te vertrouwen op een generieke prior. Dit voorkomt artefacten die ontstaan wanneer een generieke prior botst met de data.
2. Robuustheid tegen Ruis: Door het ruisniveau $\sigma$ expliciet als input te geven, past het model zich aan aan verschillende ruisniveaus. Dit maakt het superieur aan voorgeprogrammeerde modellen die vaak falen bij hoge ruis.
3. Onzekerheidskwantificatie: Het framework maakt het mogelijk om een ensemble van plausibele oplossingen te genereren, waardoor de onzekerheid in de oplossing kwantificeerbaar wordt (essentieel voor medische diagnoses).
4. Integratie van Forward Model: De methode gebruikt de forward-operator (en zijn geadjungeerde) direct in de netwerkarchitectuur en de loss-functie, wat zorgt voor fysiek consistente reconstructies.

4. Resultaten

De auteurs hebben hun methode getest op twee taken: beeldontwarring (image deblurring) en tomografie.

• Beeldontwarring (MNIST, STL10, CIFAR10):

  • DAWN-FM presteerde significant beter dan bestaande methoden zoals Diffusion Posterior Sampling en InverseUNetODE.
  • Op de STL10-dataset verbeterde DAWN-FM de PSNR met ongeveer 4 dB ten opzichte van de beste concurrent.
  • De methode bleek zeer robuust bij hoge ruisniveaus (tot 20%), terwijl andere methoden snel degradeerden.
  • De onzekerheidskaarten toonden correct aan dat de randen van objecten (waar de reconstructie lastig is) de hoogste onzekerheid hebben.

• Tomografie (OrganAMNIST, OrganCMNIST):

  • DAWN-FM overtrof de klassieke Filtered Backprojection (FBP) met een ramp-filter op alle metrieken (MSE, SSIM, PSNR).
  • De methode kon anatomische structuren scherper en accurater reconstrueren.
  • De onzekerheidsanalyse toonde variabiliteit in de grootte en vorm van lobben, wat cruciale informatie biedt voor medische besluitvorming.

5. Betekenis en Conclusie

DAWN-FM vertegenwoordigt een verschuiving in de aanpak van inverse problemen binnen het machine learning-domein. In plaats van te vertrouwen op "one-size-fits-all" voorgeprogrammeerde modellen, biedt het een flexibele, data-bewuste aanpak die specifiek wordt getraind voor het probleem en de ruiscondities.

De belangrijkste implicaties zijn:

• Betere reconstructies: Hogere kwaliteit en minder artefacten, zelfs bij zeer ruisige of onvolledige data.
• Betrouwbaarheid: De mogelijkheid om onzekerheid te kwantificeren maakt de methode veiliger voor kritieke toepassingen zoals medische beeldvorming.
• Fysieke consistentie: Door de forward-operator in het trainingsproces te integreren, worden oplossingen gegarandeerd consistent met de meetdata, wat de "inverse crime" (het gebruik van identieke modellen voor synthese en reconstructie) minimaliseert en de generalisatie verbetert.

De code is beschikbaar gesteld, wat de reproduceerbaarheid en verdere ontwikkeling van deze techniek voor andere ill-geplaatste problemen faciliteert.