Constrained Particle Seeking: Solving Diffusion Inverse Problems with Just Forward Passes

Deze paper introduceert Constrained Particle Seeking (CPS), een nieuwe gradient-vrije methode die diffusie-modellen effectief toepast op inverse problemen zonder kennis van het waarnemingsproces, door het zoeken naar optimale deeltjes te herformuleren als een geoptimaliseerde taak met beperkingen.

De kern

Stel je voor dat je een oude, beschadigde foto probeert te herstellen. Je hebt alleen een wazig, gebroken stukje van de foto (de "waarneming") en je moet het hele originele plaatje weer reconstrueren. Dit is een klassiek probleem in de wetenschap: hoe haal je het originele signaal terug uit een vervormde meting?

Vroeger deden computers dit door een soort "gids" te gebruiken die precies wist hoe de foto beschadigd was. Maar in de echte wereld is die gids vaak niet beschikbaar, of is het te ingewikkeld om te berekenen (zoals bij het beeldvormen van een zwart gat of het voorspellen van windstromen).

Dit paper introduceert een nieuwe methode genaamd CPS (Constrained Particle Seeking). Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het oude probleem: "Het Verkeerde Spel"

Stel je voor dat je in het donker probeert een schat te vinden.

• De oude methode (zoals SCG): Je roept 100 vrienden om je heen. Ze gooien allemaal een steen in het donker. Jij kijkt welke steen het dichtst bij de schat landt, pakt die ene steen en gooit de andere 99 weg. Je herhaalt dit keer op keer.
  • Het nadeel: Je gooit enorm veel waardevolle informatie weg. Misschien landde steen #99 net iets verder weg, maar hij gaf je wel een hint over de richting! Door ze weg te gooien, ben je inefficiënt en moet je heel veel "proefballonnen" laten gaan.

2. De nieuwe methode: CPS (De Slimme Zoeker)

CPS doet iets heel anders. Het is alsof je niet alleen kijkt naar de beste steen, maar naar alle stenen die je vrienden hebben gegooid.

• De Analogie van de "Landelijke Kaart":
  In plaats van één steen te kiezen, kijkt CPS naar de patronen van alle 100 stenen die in het donker zijn geland.

  • Het zegt: "Oké, steen #1 landde hier, steen #2 daar, en steen #99 daar. Als we al deze punten met elkaar verbinden, kunnen we een lokale kaart maken van waar de schat waarschijnlijk zit."
  • Dit noemen ze een "surrogaatmodel". Het is alsof je uit de chaos van de gegooiden stenen een logische route trekt.

• De "Magische Rand" (De Beperking):
  Er is een risico: als je alleen kijkt naar de stenen die het dichtst bij de schat lijken, kun je in een valkuil stappen (bijvoorbeeld een plek die eruitziet als de schat, maar eigenlijk een leeg veld is).
  CPS lost dit op door een magische rand om het gebied te trekken. Deze rand zorgt ervoor dat je alleen zoekt op plekken waar het "waarschijnlijk" is dat de schat zit, gebaseerd op wat we al weten over hoe foto's eruitzien (de "prior").

  • Kortom: Je zoekt actief naar de beste plek, maar je mag niet buiten de veilige zone treden die door de AI is getekend.

3. Waarom is dit zo slim?

• Geen "Gids" nodig: Je hoeft niet te weten hoe de foto beschadigd is (geen wiskundige afgeleiden nodig). Je hoeft alleen maar te kijken wat er gebeurt als je probeert.
• Efficiëntie: Omdat CPS alle informatie gebruikt (zelfs van de "slechte" stenen), heeft het veel minder pogingen nodig dan de oude methoden om de juiste oplossing te vinden.
• De "Restart" Strategie: Soms loop je vast in een verkeerde richting. CPS heeft een trucje: als het even niet lukt, gooit het de hele situatie een beetje op zijn kop (een "restart") en probeert het opnieuw, maar dan met de kennis die het al heeft opgedaan. Dit voorkomt dat je vastloopt.

Wat hebben ze bewezen?

De auteurs hebben deze methode getest op drie heel verschillende dingen:

1. Foto's herstellen: (Invullen van ontbrekende stukjes, scherper maken, JPEG-fouten wegwerken).
2. Zwartgaten beelden: Het reconstrueren van een zwart gat uit heel weinig en ruisige data.
3. Stroming van vloeistoffen: Het voorspellen van hoe wind of water beweegt, gebaseerd op een paar metingen.

In al deze gevallen deed CPS het net zo goed als de super-complexe methoden die wel een "gids" nodig hebben, maar dan zonder die gids. En het deed het veel beter dan andere methoden die ook geen gids gebruikten.

Samenvatting in één zin

CPS is als een slimme schatzoeker die niet alleen kijkt naar de beste steen die hij vindt, maar de patronen van alle gegooiden stenen gebruikt om een kaart te tekenen, terwijl hij zich tegelijkertijd houdt aan de regels van de natuur, zodat hij de schat sneller en zekerder vindt zonder een handboek nodig te hebben.

Technische samenvatting

Probleemstelling

Inverse problemen zijn fundamenteel in vele wetenschappelijke en technische domeinen (zoals medische beeldvorming, geofysica en computervisie). Het doel is om een origineel signaal $x$ te reconstrueren uit indirecte en ruisbeïnvloede waarnemingen $y$, vaak gemodelleerd als $y = H(x) + \eta$. Dit probleem is inherent slecht gesteld (ill-posed), wat betekent dat er meerdere mogelijke oplossingen kunnen zijn.

Recente vooruitgang in generatieve modellen, met name diffusiemodellen, heeft bewezen effectief te zijn als "plug-and-play" priors voor het oplossen van deze problemen. Echter, de bestaande methoden (zoals DPS) zijn afhankelijk van de gradiënten van het waarnemingsproces $H(\cdot)$. In veel real-world scenario's is deze gradiënt niet beschikbaar, te duur om te berekenen (bijv. bij complexe numerieke simulaties), of is het model een "black-box".

Bestaande gradiëntvrije methoden (zoals SCG en EnKG) hebben hun eigen beperkingen:

• SCG (Symbolic Constrained Guidance): Gebruikt een afwijzingsstrategie (rejection sampling) waarbij bij elke stap meerdere deeltjes worden gesampleerd, maar er slechts één wordt behouden en de rest wordt verworpen. Dit is inefficiënt en gooit waardevolle lokale informatie weg.
• EnKG (Ensemble Kalman Diffusion): Vereist het onderhoud van duizenden deeltjes in hoge dimensies, wat leidt tot een enorme rekenlast.
• DPG: Voegt ruis toe aan schone samples, wat kan leiden tot afwijkingen van het data-mannequin.

Methodologie: Constrained Particle Seeking (CPS)

De auteurs introduceren CPS, een nieuwe gradiëntvrije aanpak die het inverse probleem herschrijft als een beperkt optimalisatieprobleem. In plaats van passief het beste deeltje uit een set te kiezen, zoekt CPS actief naar de optimale deeltjespositie door informatie van alle kandidaat-deeltjes te synthetiseren.

De methode bestaat uit drie hoofdstappen per tijdstap $t$ tijdens het denoising-proces:

1. Sampling van Kandidaten:
   Er worden $n$ kandidaat-deeltjes ($x_t^1, ..., x_t^n$) gesampleerd vanuit de voorwaartse overgangskern van het diffusiemodel (de onvoorwaardelijke prior).

2. Fitten van een Lokale Surrogaat:
   Omdat de gradiënt van $H(x)$ niet beschikbaar is, wordt een statistische linearisatie toegepast. De auteurs fit een lokaal lineair surrogaatmodel $H(x) \approx Ax + b$ rondom de gesampleerde deeltjes. De parameters $A$ en $b$ worden geschat door de correlatie tussen de gesampleerde deeltjes en hun geschatte schone versies (via Tweedie's formule) te analyseren.

   • Dit model integreert informatie van alle deeltjes, inclusief die welke slecht presteren, wat leidt tot een robuustere schatting dan het kiezen van één enkel deeltje.

3. Beperkte Optimalisatie (Particle Seeking):
   Het doel is om een deeltje $x_t$ te vinden dat de waarneming $y$ zo goed mogelijk benadert, terwijl het binnen het hoogstwaarschijnlijkheidsgebied van de prior blijft. Dit wordt geformuleerd als:
   $$\min_{x_t} \|y - (Ax_t + b)\|^2 \quad \text{onder de voorwaarde} \quad x_t \in S^{d-1}(\mu_t, \sigma_t\sqrt{d})$$
   De beperking zorgt ervoor dat het deeltje op een hypersfeer blijft die overeenkomt met de verdeling van de diffusie-kern (in plaats van naar het gemiddelde $\mu_t$ te trekken, wat vaak fout is in hoge dimensies).

   Door de Lagrange-multiplicator-methode toe te passen en gebruik te maken van de asymptotische benadering (waarbij de variantie $\sigma_t$ klein is), wordt een gesloten vorm oplossing afgeleid:
   $$x_t^* \approx \mu_t + \sigma_t\sqrt{d} \frac{A^\top(y - \bar{H})}{\|A^\top(y - \bar{H})\|}$$
   Hierbij is $x_t^*$ het geoptimaliseerde deeltje dat de informatie van alle kandidaten combineert.

4. Restart Strategie:
   Om cumulatieve fouten te verminderen die ontstaan door de sequentiële sampling, wordt een "Restart"-strategie geïntegreerd. Als een stap niet optimaal verloopt, wordt het proces teruggesprongen naar een eerdere stap, opnieuw met ruis verstoord, en opnieuw geprobeerd. Dit verhoogt de robuustheid aanzienlijk.

Belangrijkste Bijdragen

1. Paradigmaverschuiving: Van passieve selectie (rejection sampling) naar actieve zoektocht naar optimale deeltjes binnen een beperkte ruimte.
2. Gebruik van Alle Informatie: In tegenstelling tot eerdere methoden die deeltjes wegwerpen, gebruikt CPS de volledige ensemble om een lokaal surrogaatmodel te bouwen, waardoor zelfs "slechte" deeltjes bijdragen aan de zoekrichting.
3. Gradiëntvrij met Prestaties op Gradiëntniveau: CPS vereist geen gradiënten van $H(\cdot)$, maar bereikt prestaties die vergelijkbaar zijn met geavanceerde gradiëntgebaseerde methoden (zoals DPS en RED-diff).
4. Efficiëntie: CPS is aanzienlijk efficiënter dan EnKG, omdat het geen duizenden deeltjes hoeft te onderhouden, maar slechts een klein ensemble gebruikt om een optimale richting te berekenen.

Resultaten

De auteurs hebben CPS getest op een breed scala aan taken:

• Beeld-Inverse Problemen: Inpainting, super-resolutie, deblurring en JPEG-restauratie (waarbij de forward-modellen niet-differentieerbaar zijn).
  • Resultaat: CPS overtreft alle gradiëntvrije basismethoden (SCG, DPG, EnKG) aanzienlijk en presteert vergelijkbaar met of beter dan gradiëntgebaseerde methoden. Bijvoorbeeld, bij 95% inpainting op FFHQ behaalde CPS een PSNR van 24.90, terwijl SCG slechts 6.15 haalde.
• Wetenschappelijke Inverse Problemen:
  • Zwarte Gaten Beeldvorming: Een sterk niet-lineair probleem met schaarse observaties. CPS leverde de beste visuele en kwantitatieve resultaten (PSNR en BPSNR).
  • Vloeistof Data Assimilatie: Het reconstrueren van initiële vorticiteitsvelden uit ruwe metingen van de Navier-Stokes vergelijkingen. CPS slaagde erin patronen te herstellen waar andere methoden faalden.

Ablatie Studies:

• CPS is zeer efficiënt in het gebruik van deeltjes; zelfs met slechts 8 deeltjes presteert het goed, terwijl EnKG veel meer deeltjes nodig heeft.
• De Restart-strategie verbetert de stabiliteit en prestaties, vooral bij suboptimale initiële ruis.

Betekenis en Impact

Dit werk is significant omdat het de toepasbaarheid van diffusiemodellen voor inverse problemen uitbreidt naar scenario's waar geen gradiënten beschikbaar zijn. Veel real-world systemen (zoals atmosferische simulaties of complexe sensoren) zijn black-boxen of te duur om te differentiëren. CPS biedt een robuust, efficiënt en principieel kader om deze problemen op te lossen zonder de kwaliteit van gradiëntgebaseerde methoden in te leveren. Het demonstreert dat door slimme synthese van ensemble-informatie en beperkte optimalisatie, de inefficiëntie van eerdere gradiëntvrije benaderingen kan worden overwonnen.