Technische Samenvatting: Flowing with Confidence (FMwC)

Probleemstelling

Generatieve modellen, met name Flow Matching (FM)-modellen, hebben state-of-the-art prestaties bereikt op het gebied van taal, visie en wetenschappelijke domeinen. Ze produceren echter outputs (bijvoorbeeld gehallucineerde citaten, anatomisch incorrecte figuren of thermodynamisch instabiele kristallen) zonder native indicatie van betrouwbaarheid. Bestaande oplossingen voor onzekerheidskwantificatie, zoals Deep Ensembles of MC-Dropout, vertrouwen op het uitvoeren van $k$ aparte modellen of $k$ stochastische trajecten per steekproef. Deze aanpak brengt een $k\times$ hogere rekenkosten met zich mee, waardoor per-steekproef vertrouwen voor grootschalige implementatie onbetaalbaar duur wordt. Bovendien is aangetoond dat likelihood-gebaseerde methoden onbetrouwbaar zijn voor het onderscheiden van in-distribution van out-of-distribution inputs in diepe generatieve modellen. De centrale vraag die wordt behandeld, is of flow-matching modellen kunnen leren hun eigen onzekerheid te kwantificeren voor gegenereerde outputs zonder de inferentiekosten te verhogen of toegang tot likelihoods te vereisen.

Methodologie: Flow Matching with Confidence (FMwC)

FMwC is een single-pass aanpak die flow-matching modellen uitrust met een per-steekproef, traject-niveau vertrouwensscore tegen dezelfde rekenkosten als standaard deterministische flow matching. De methode integreert Variational Adaptive Dropout (VAD) met gesloten-vorm momentpropagatie door de Ordinary Differential Equation (ODE)-traject.

1. Variationale Formulering

In plaats van een enkel deterministisch snelheidsveld $v_{t,\theta}(x_t)$ te leren, behandelt FMwC de netwerkgewichten $\omega$ als stochastische variabelen getrokken uit een variational posterior $q_\psi(\omega)$. Het model maximaliseert de Evidence Lower Bound (ELBO):
$$\text{ELBO}(\psi) = \mathbb{E}_{\omega \sim q_\psi} [\log \mathcal{L}(D | \omega)] - \text{KL}[q_\psi(\omega) \parallel \pi(\omega)]$$
waarbij de likelihood-term overeenkomt met de standaard conditionele flow-matching regressieverliesfunctie. Cruciaal is dat de gradient van de verwachte log-likelihood onder de posterior overeenkomt met de gradient van de standaard regressieverliesfunctie, geëvalueerd op gewichtsteekproeven, wat simulatie-vrije training mogelijk maakt.

2. Standaard-afhankelijke Posterior via VAD

Om ervoor te zorgen dat de onzekerheid gevoelig is voor de lokale geometrie van de flow, maakt FMwC gebruik van Variational Adaptive Dropout (VAD). In dit kader is de posterior over gewichten in laag $l$ een Gaussische verdeling $q_\psi(\omega_l | x, t) = \mathcal{N}(\theta_l, \alpha_l(x, t)^2 \theta_l^2)$, waarbij de dropout-schaal $\alpha_l(x, t)$ wordt voorspeld door een klein inferentienetwerk $E_\gamma$. Hierdoor kan de posterior-spreiding verbreden in gebieden met geometrische ambiguïteit (waar het snelheidsveld moeilijk te fitten is) en zich elders vertragen, uitsluitend geleerd via de ELBO.

3. Analytische Variantiepropagatie

De kerninnovatie van FMwC is het analytisch doorgeven van de variantie van het snelheidsveld door zowel de neurale netwerklagen als de ODE-integratietraject in één enkele forward pass:

• Door het Netwerk: Met behulp van de lokale reparameterization-trick worden het gemiddelde ($\mu_l$) en de variantie ($\sigma_l^2$) van activaties laag-voor-laag doorgegeven. Voor lineaire lagen is de variantierecursie in gesloten vorm. Voor niet-lineariteiten worden momenten analytisch berekend met behulp van Gauss-Hermite-kwadratuur of Taylor-linearisatie, waardoor Monte Carlo-steekproeven overbodig zijn.
• Door de ODE: Bij elke Euler-stap geeft het netwerk de posterior-gemiddelde snelheid $\mu_t$ en zijn variantie $\sigma_t^2$ uit. De integrator werkt de positie bij met $\mu_t$, terwijl $\sigma_t^2$ wordt geregistreerd. Dit levert een traject van variantie $\{\sigma_t^2\}_{t=0}^T$ op in plaats van een enkele scalair.

4. Vertrouwensuitkomsten

De variantietraject kan worden geaggregeerd tot scalar vertrouwensscores voor downstream taken:

• Eindpuntvertrouwen: De variantie bij de laatste stap ($t=1$).
• Geïntegreerd Vertrouwen: De som van de standaardafwijkingen van de snelheid langs het traject.
• Temporele Ratio: De verhouding van de variantie in het late venster tot de variantie in het vroege venster, wat helpt bij het onderscheiden van steekproeven die convergeren (lage late variantie) van die welke ambigu blijven.
• Gestudeerde Koppen: Een kleine supervised kop kan de volledige variantieserie afbeelden op een taak-specifieke score.

Belangrijkste Bijdragen

1. Single-Pass Vertrouwensschatting: FMwC is de eerste methode die een gekalibreerde per-steekproef vertrouwensscore biedt voor flow matching tegen de kosten van een enkele deterministische forward pass, wat overeenkomt met het inferentiebudget van standaard FM. Dit wordt bereikt door analytisch geleerde, input-afhankelijke ruis door de ODE te propageren.
2. Inferentieverbeteringen zonder Training: De vertrouwensscore maakt praktische toepassingen mogelijk zonder hertraining:
   • Filteren: Verbetert monotoon de steekproefkwaliteit door outputs met laag vertrouwen te verwijderen.
   • Bewerken: Identificeert het exacte tijdstip ($t^*$) waarop het model zich vastlegt op een structuur (variantiepiek), waardoor gerichte, aan de beperkingen voldane bewerkingen mogelijk zijn (bijvoorbeeld chemische swaps in kristallen) door ruis in te brengen bij $t^*$.
   • Adaptief Stappen: Concentreert ODE-berekeningsstappen in gebieden met hoge variantie (ambiguïteit), waardoor hoogwaardige steekproeven worden hersteld bij lage stapbudgetten waar uniforme stappen falen.
3. Geometrische Interpretatieerbaarheid: Het artikel stelt vast dat de FMwC-vertrouwensscore correleert met de lokale divergentie van het geleerde snelheidsveld. Dit biedt een venster in het generatieve proces, waarbij epistemische onzekerheid wordt gekoppeld aan de geometrische instabiliteit (bifurcatiepunten) van de flow.

Experimentele Resultaten

De auteurs hebben FMwC geëvalueerd over drie modaliteiten: 2D-dichtheidsschatting (Checkerboard), beeldgeneratie (MNIST) en de novo generatie van anorganische kristallen (Riemannian Flow Matching).

• Steekproefkwaliteit: FMwC behoudt of verbetert de steekproefkwaliteit ten opzichte van standaard FM. Op Checkerboard en MNIST werden respectievelijk lagere misplaatsingspercentages en FID-scores bereikt, terwijl vergelijkbare kwaliteit werd behouden voor kristallen.
• Inferentiekosten: FMwC brengt verwaarloosbare overhead met zich mee op grotere backbones (bijvoorbeeld ~1,003$\times$ voor MNIST, ~1,08$\times$ voor Kristallen) ten opzichte van de deterministische baseline. Daarentegen vereisen ensemble- en MC-Dropout-baselines $5\times$ de rekencapaciteit.
• Filterprestaties: Filteren op basis van FMwC-scores behaalde Area Under the Precision-Recall Curve (AUPRC)-scores die vergelijkbaar waren met $k=5$ ensembles en MC-Dropout, maar tegen de kosten van een enkel traject. Bijvoorbeeld, op de Checkerboard-dataset behaalde de temporele ratioscore een AUPRC van 0,55, wat overeenkomt met de ensemble-baseline (0,55) met 1/5e van de rekencapaciteit.
• Succes van Bewerking: Bij het bewerken van kristallen door ruis in te brengen bij de variantiepiek ($t^*$), behaalde FMwC aanzienlijk hogere succespercentages (bijvoorbeeld 39% versus 22% voor chemische swaps) in vergelijking met willekeurige tijdstipselectie, wat aantoont dat gebieden met hoge variantie overeenkomen met kritieke beslispunten.
• Correlatie met Divergentie: De variantietraject $\sigma_t^2$ toonde een sterke correlatie met de grootte van de divergentie van het snelheidsveld ($|\nabla \cdot v_{t,\theta}|$). Op het Crystal-dataset bereikte de geïntegreerde correlatie een Spearman $\rho$ van 0,96, wat de geometrische interpretatie van de vertrouwensscore valideert.

Betekenis en Claims

Het artikel beweert dat FMwC flow matching transformeert van een generator van steekproeven naar een generator van steekproeven gekoppeld aan een maatstaf van betrouwbaarheid. Door een input-afhankelijke posterior over gewichten te leren en variantie analytisch te propageren, maakt FMwC vertrouwensbewuste generatie een praktisch standaard voor wetenschappelijke en veiligheidskritische toepassingen. De auteurs benadrukken dat de vertrouwensscore een intrinsieke geometrische eigenschap is van de enkel getrainde flow, onderscheiden van de inter-model-onenigheid die door ensembles wordt gemeten. Dit maakt "chirurgische" begeleiding van het generatieve proces mogelijk, gericht op momenten van instabiliteit zonder de onbetaalbare kosten van bestaande onzekerheidskwantificatiemethoden. Het werk suggereert dat het begrijpen en benutten van de geometrie van het geleerde snelheidsveld kan leiden tot betrouwbaardere, interpreteerbaardere en controleerbaardere generatieve modellen.