Technische Samenvatting: Flow Annealing Posterior Sampling (FAPS)

Probleemstelling

Het artikel behandelt de uitdaging van principiële regressie voor stochastische processen, een taak die centraal staat in wetenschappelijke inverse problemen, data-assimilatie en onzekerheidskwantificering. In deze settings is het doel om een posterieure verdeling af te leiden over een onbekende functie of fysiek veld op basis van schaarse en ruisige observaties. Hoewel klassieke Gaussian Process (GP) regressie een Bayesiaanse oplossing biedt voor Gaussische priors en eenvoudige observatiemodellen, zijn veel wetenschappelijke processen niet-Gaussisch, hoogdimensionaal en worden ze indirect waargenomen via fysieke metingen (bijv. oplossingen van Partiële Differentiaalvergelijkingen, of PDE's).

Bestaande benaderingen kampen met beperkingen:

• Neural Processes leren geamortiseerde conditionele voorspellers, maar missen vaak een rigoureuze posterieure consistentie bij variërende query-discretisaties en scheiden het herbruikbare leren van de prior niet expliciet van de inferentie tijdens de testtijd.
• Neural Operators bieden efficiënte PDE-oplossingskaarten, maar zijn doorgaans deterministisch en leveren geen posterieure verdelingen uit schaarse data.
• Diffusie-gebaseerde Posterior Samplers (bijv. DAPS, FunDPS) combineren geleerde priors met observatie-begeleiding, maar zijn vaak diffusie-gebaseerd, computationeel duur en primair gericht op specifieke PDE-settings in plaats van het verenigen van regressie en inverse problemen.

Er is behoefte aan een algemeen posterior sampling-framework dat gepre-trainde functie-ruimte priors benut om likelihood-geleide inferentie uit te voeren voor zowel directe stochastische-procesregressie als indirecte PDE-inverse problemen zonder expliciete evaluatie van de prior-densiteit.

Methodologie: Flow Annealing Posterior Sampling (FAPS)

FAPS is een functionele posterior sampling-framework gebouwd op gepre-trainde Operator Flow Matching (OFM) priors. Het verenigt stochastische-procesregressie en PDE-inverse problemen via een ontkoppelde annealing-procedure.

1. Functionele Flow-Matching Prior

De methode maakt gebruik van een gepre-trainde OFM-prior, die een probabiliteitsflow leert die een referentie Gaussische maat $\gamma = \mathcal{N}(0, C_0)$ transporteert naar een doelverdeling $\mu_1$. Dit wordt bereikt via een neurale operator $v_\theta$ die getraind is om het conditionele snelheidveld $v_t(u_t | u_0, u_1) = u_1 - u_0$ te benaderen langs een ruisig recht interpolatiepad. De resulterende flow-map $\Phi_{s \to t}$ maakt het mogelijk om samples te transporteren van de referentie-maat naar het doelproces.

2. Verenigde Posterior Formulering

FAPS behandelt zowel functionele regressie als PDE-inverse problemen onder één enkel observatiemodel:
$$y = A(u_1) + \epsilon$$
waarbij $u_1$ de onbekende functie is, $A$ de observatie-operator is, en $\epsilon$ de ruis is.

• Voor functionele regressie is $A$ een directe punt-evaluatie of een masking-operator.
• Voor PDE-inverse problemen is $A$ een compositie van een schaarse observatie-operator en een gepre-trainde neurale-operator surrogate $G_\phi$ die de forward PDE-oplossingskaart representeert.

3. De Annealing en Re-bridging Procedure

FAPS sampleert uit de posterior $\pi_1^y$ door te alterneren tussen drie stappen over een annealing-schema $0 = t_0 < t_1 < \dots < t_K = 1$:

1. Flow Transport (Naar Eindpunt): Gegeven een "ruisige" staat $u_{t_k}$ op annealing-niveau $t_k$, integreert de methode de gepre-trainde flow naar het eindpunt om een "schone" anker $\hat{u}_k = \Phi_{t_k \to 1}(u_{t_k})$ te verkrijgen. Dit anker dient als een prior-consistente schatting van de schone functie.
2. Likelihood-Guided Langevin Correctie: Rondom het anker $\hat{u}_k$ voert FAPS $L$ stappen van Langevin-dynamica uit om de sample te corrigeren op basis van de observatie-likelihood $p(y|u_1)$. De update-regel is:
   $$u_1^{(\ell+1)} = u_1^{(\ell)} + \eta \left[ -\frac{u_1^{(\ell)} - \hat{u}_k}{\lambda_k^2} + C_X \nabla_{u_1} \log p(y|u_1^{(\ell)}) \right] + \sqrt{2\eta} \zeta^{(\ell)}$$
   Cruciaal is dat de gradiënt wordt gepreconditioneerd door een low-rank covariantie-matrix $C_X$, die offline is geschat uit het doelproces. Deze preconditioner verspreidt schaarse observatie-informatie langs dominante functionele correlaties, wat de updates in ongeobserveerde regio's aanzienlijk verbetert.
3. Re-bridging: De gecorrigeerde schone sample wordt opnieuw gebridged naar het volgende annealing-niveau $t_{k+1}$ met behulp van de bridge-kernel geïnduceerd door het flow-pad, wat effectief ruis toevoegt aan de staat om de volgende iteratie voor te bereiden.

Dit proces herhaalt zich tot $t_K = 1$, wat posterieure samples oplevert zonder ooit expliciet de prior-densiteit te hoeven evalueren.

Belangrijkste Bijdragen

1. Eerste Verenigde Functionele Flow-Matching Posterior Sampler: FAPS wordt gepresenteerd als het eerste framework dat stochastische-procesregressie en PDE-inverse problemen verenigt met behulp van gepre-trainde functionele flow-matching priors. Het biedt een gemeenschappelijke likelihood-geleide sampler voor zowel directe als indirecte observaties.
2. Low-Rank Covariance-Preconditioned Langevin Correctie: De methode introduceert een low-rank covariantie-preconditioner voor de Langevin-correctiestap. Dit maakt gebruik van dominante correlaties in de functie-ruimte, wat de posterieure updates onder schaarse observaties aanzienlijk verbetert vergeleken met ongepreconditioneerde methoden.
3. Flexibele en Computationeel Efficiënte Inferentie: FAPS is "plug-and-play", waarbij gebruik wordt gemaakt van gepre-trainde priors en forward surrogates. Het ondersteunt variabele query-discretisaties en ruisige observaties zonder hertraining. Het vermijdt expliciete evaluatie van de prior-densiteit, waardoor het aanzienlijk geheugen- en tijd-efficiënter is dan likelihood-gebaseerde sampling met flow-priors of diffusie-gebaseerde samplers.
4. Sterke Empirische Prestaties: Uit uitgebreide experimenten blijkt dat FAPS state-of-the-art prestaties bereikt in functionele regressie (het overtreft varianten van Neural Processes) en competitieve of superieure prestaties levert in PDE-inverse problemen vergeleken met diffusie-gebaseerde baselines, met lagere kosten voor sampling tijdens de testtijd.

Experimentele Resultaten

Het artikel evalueert FAPS over drie settings:

• Gaussian Process Regressie: Op 1D Matérn en Gibbs-kernel benchmarks, komt FFS nauw overeen met referentie posterieure gemiddelden en onzekerheden over variërende query-resoluties (128, 512, 1024). Het presteert beter dan Neural Processes (CNP, ANP, NDP, FlowNP) en directe OFM-sampling in Sliced Wasserstein Distance (SWD) en Maximum Mean Discrepancy (MMD), vooral bij hogere resoluties waar baselines degraderen.
• Niet-Gaussische Functionele Regressie: Op Navier-Stokes, Black Hole imaging, en Global Climate datasets, reconstrueert FAPS dominante ruimtelijke structuren en produceert het coherente posterieure samples. Het bereikt de beste of bijna-beste kalibratie (SSR) en reconstructie-metrieken (CRPS, PSNR, SSIM) vergeleken met baselines, die vaak overmatig gladstrijken of onzekerheid onderschatten.
• PDE-Inverse Problemen: Op Darcy flow, Poisson, Helmholtz, en Navier-Stokes inverse benchmarks (128x128 resolutie, 128 ruisige observaties), behaalt FAPS (zowel FNO als UNet varianten) competitieve of beste resultaten in CRPS, PSNR, en SSIM vergeleken met DiffusionPDE, FunDPS, en DDIS.
  • Efficiëntie: FAPS is aanzienlijk sneller dan diffusie-gebaseerde baselines (bijv. 1.73x sneller dan DiffusionPDE, 2.37x sneller dan DDIS) en vereist aanzienlijk minder GPU-geheugen (door het vermijden van de hoge geheugenoverhead van expliciete likelihood-evaluatie in flow-priors of grote diffusie-stappen).
  • Zero-Shot Super-Resolutie: FAPS demonstreert het vermogen om inverse inferentie uit te voeren op query-meshes (bijv. 160x160) die niet gezien zijn tijdens de training van de prior, een capaciteit die mogelijk wordt gemaakt door het functionele karakter van de OFM-prior.

Betekenis en Claims

De auteurs stellen dat FAPS een nieuw paradigma vestigt voor onzekerheid-bewuste inferentie in wetenschappelijk machine learning. Door gepre-trainde functionele flow-matching priors om te zetten in likelihood-geleide posterior samplers, biedt FAPS een herbruikbaar, flexibel en computationeel efficiënt alternatief voor bestaande methoden.

Het artikel erkent bescheiden dat de praktische sampler steunt op een benaderde lokale Gaussische correctie en dat een volledige niet-asymptotische convergentietheorie voor de praktische sampler nog toekomstig werk is. De auteurs bieden echter empirisch bewijs en een analyse van foutpropagatie die suggereert dat lokale transitiefouten additief accumuleren en dat de methode convergeert naar de doel-posterior naarmate het annealing-mesh en de sampling-parameters worden verfijnd.

Uiteindelijk suggereert het werk dat functionele flow-matching priors een robuust fundament kunnen vormen voor schaalbare Bayesiaanse inferentie in complexe wetenschappelijke domeinen, waarbij regressie en inverse modellering worden verenigd onder één enkel, efficiënt framework.