Technische Samenvatting: Geometrie-bewuste Post-hoc Onzekerheidskwantificering in Operator Learning

Probleemstelling

Neurale operatoren zijn opgekomen als krachtige surrogaten voor het oplossen van partiële differentiaalvergelijkingen (PDE's) in domeinen zoals vloeistofdynamica en vaste mechanica. De standaard deterministische aard van deze operatoren beperkt echter hun bruikbaarheid in wetenschappelijke toepassingen die een strikte onzekerheidskwantificering (UQ) vereisen, met name bij het omgaan met geometrische variabiliteit. Bestaande UQ-benaderingen vallen over het algemeen uiteen in twee categorieën:

1. Train-time methoden: Technieken zoals Deep Ensembles of Bayesiaanse benaderingen (bijv. MC Dropout) die de onzekerheid in netwerkparameters modelleren. Deze zijn vaak computationeel te kostbaar voor grootschalige operator learning op ongestructureerde meshes.
2. Post-hoc methoden: Benaderingen zoals de Laplace-benadering die probabilistische componenten toevoegen aan een bevroren model. Hoewel deze efficiënter zijn, vertrouwen veel methoden op lokale gewichtsruimte-benaderingen of generieke feature maps die er niet in slagen de specifieke geometrie-bewuste representaties te vatten die door de operator zijn geleerd.

Er bestaat een kritische kloof in het benutten van de intrinsieke, fysica-bewuste latente representaties van een vooraf getrainde neurale operator om geometrie-bewuste onzekerheden te construeren zonder het model opnieuw te trainen of misgespecificeerde priors in oneindig-dimensionale ruimtes op te leggen.

Methodologie: REEF-GP

De auteurs stellen REEF-GP voor (Residual on Embedded Features Gaussian Process), een post-hoc UQ-framework dat de discrepantie tussen een bevroren neurale operator en de ware oplossing direct in de geleerde feature-ruimte van de operator modelleert.

Kernframework

De methode behandelt de vooraf getrainde neurale operator $G_\theta$ als een deterministisch computermodel. De ware oplossing $v(x)$ op een query-coördinaat $x$ wordt gemodelleerd als:
$$v(x) = G_\theta(u, a)(x) + \delta(u, x, h(u, a, x)) + \epsilon$$
waarbij $\delta$ de model inadequaatheid (residue) vertegenwoordigt en $\epsilon$ observationele ruis is. In tegen tegenstelling tot standaard KOH (Kennedy–O'Hagan) frameworks die de operator als een black box behandelen, benut REEF-GP de interne structuur van $G_\theta$.

Belangrijke Technische Componenten

1. Geometrie-bewuste Diepe Kernel:

   • In plaats van een feature map vanaf nul te leren, gebruikt REEF-GP de verborgen lagen van de bevroren operator (specifiek Transolver) die al vervormde, fysica-bewuste representaties van de invoergeometrie coderen.
   • Het selecteert een subset van informatieve lagen (vroege, middelste en late) op basis van Centered Kernel Alignment (CKA) om een uitgebreide staat $h(u, a, x)$ te construeren.
   • Een diepe kernel wordt gedefinieerd over het product van ruimtelijke coördinaten, invoerfuncties en deze latente embeddings: $k_\phi(\bar{z}, \bar{z}') = k_{space} \cdot k_{fun} \cdot k_{latent}$. De latente kernel gebruikt een spectraal-genormaliseerde MLP om de embeddings te vervormen, wat een niet-stationaire prior creëert die zich aanpast aan lokale fysieke regimes.

2. Probabilistische Discrepantie Modellering:

   • Het residue $\delta$ wordt gemodelleerd als een zero-mean Gaussian Process (GP) met een heteroscedastische ruisterm.
   • De ruisvariantie is input-afhankelijk en wordt geparametriseerd door een spectraal-genormaliseerd dense netwerk om feature collapse te voorkomen en stabiliteit te waarborgen.

3. Schaalbare Training en Inferentie:

   • Stochastische Subset Optimalisatie: Om de $O((MN)^3)$ complexiteit van exacte GP-training te hanteren (waarbij $M$ het aantal samples is en $N$ het aantal discretisatiepunten), gebruikt REEF-GP een stochastische subset-strategie geïnspireerd op Kernel Flows (KF). Het optimaliseert hyperparameters met behulp van mini-batches van residuen, waardoor geheugenbeperkingen worden losgekoppeld van de volledige mesh-resolutie.
   • Generalized Product of Experts (gPoE): Voor inferentie conditioneert de methode $K$ onafhankelijke GP-experts op willekeurige subsets van de data. De uiteindelijke voorspellende distributie wordt verkregen door deze experts te aggregeren in precisieruimte, waarbij de geheugenvoetafdruk tijdens de testtijd wordt begrensd tot $O(N_s^2)$ per expert in plaats van de volledige datasetgrootte.

Belangrijkste Bijdragen

1. Geometrie-bewuste Diepe Kernel: Het ontwerp van een kernel die gebruikmaakt van de interne, bevroren representaties van een vooraf getrainde neurale operator, waardoor het niet nodig is om een aparte feature map te leren of een prior op te leggen die mogelijk misgespecificeerd is voor de regulariteit van de operator.
2. Schaalbare Post-Hoc UQ: De ontwikkeling van training- en procedure voor inferentie op basis van stochastische subset-optimalisatie en gPoE, wat GP-gebaseerde UQ praktisch maakt voor operator learning datasets met miljoenen punten.
3. Empirische Validatie: Demonstratie van competitieve kalibratie over vijf 2D en 3D PDE-benchmarks (inclusief Airfoil, ShapeNet Car en Ahmed Body), met robuustheid onder geometrische distributieverschuivingen.

Resultaten

Het paper evalueert REEF-GP tegen zes baselines: Deep Ensembles, MC Dropout, Probabilistic Neural Operator (PNO), Input Perturbation, Laplace Approximation (LUNO-LA) en Deep Vecchia Ensemble (DVE-spatial).

• Voorspellende Nauwkeurigheid: REEF-GP behoudt de voorspellende nauwkeurigheid van de basale deterministische operator (Transolver) over alle datasets, met relatieve $L_2$-fouten die identiek zijn aan de baseline. In tegenstelling hiertoe brengen train-time methoden (MC Dropout, PNO) en input perturbatie aanzienlijke nauwkeurigheidsverliezen met zich mee.
• Onzekerheidskalibratie: Op 3D benchmarks (ShapeNet Car) bereikt REEF-GP Negative Log-Likelihood (NLL) scores die competitief zijn met Deep Ensembles (bijv. 2.84 vs. 4.16 voor Ensembles) en presteert het significant beter dan andere post-hoc methoden zoals LUNO-LA (33.55).
• Ruimtelijke Coherentie: De door REEF-GP gegenereerde onzekerheidskaarten zijn ruimtelijk betekenisvol en concentreren zich rond fysiek significante regio's met hoge fouten, zoals schokfronten op vleugels (airfoils) en de voorbumper van auto-geometrieën.
• Geometrische OOD Robuustheid: Onder geometrische distributieverschuivingen (gemeten via Maximum Mean Discrepancy), behoudt REEF-GP een stabiele kalibratie. Terwijl andere methoden zoals LUNO-LA en Perturbation een significante NLL-degradatie vertonen naarmate de testgeometrieën afwijken van de trainingsdistributie, blijft de NLL van REEF-GP relatief stabiel.
• Computationele Kosten: REEF-GP brengt een trainingsoverhead met zich mee die een orde van grootte lager is dan die van Deep Ensembles. Hoewel de evaluatiegeheugenvoetafdruk hoger is dan bij sommige baselines vanwege de kernelmatrices, blijft het binnen de grenzen van gangbare GPU's.

Betekenis en Claims

De auteurs beweren dat REEF-GP aantoont dat nauwkeurige en schaalbare post-hoc UQ voor neurale operatoren direct in de geleerde feature-ruimte van de operator kan worden bereikt. Door de geometrie-bewuste embeddings van een bevroren operator te hergebruiken, biedt de methode een praktisch alternatief voor parameter-gecentreerde benaderingen (zoals Deep Ensembles) die te duur zijn voor grootschalige problemen.

Het paper benadrukt dat de geometrie die een neurale operator leert bij het oplossen van een PDE inherent veel van de informatie bevat die nodig is om de eigen onzekerheid te kwantificeren. Deze benadering biedt een complementaire as voor parameter-gebaseerde UQ, wat een natuurlijk startpunt biedt voor onzekerheidskwantificering op steeds grotere, geometrie-bewuste operatoren zonder de prohibitieve kosten van hertraining of ensemble-methoden. De auteurs merken beperkingen op met betrekking tot backbones die geen punt-voor-punt ruimtelijke correspondentie behouden (bijv. standaard FNO's) en de huidige focus op stationaire, single-output problemen, wat zij identificeren als natuurlijke richtingen voor toekomstig werk.