Technische Zusammenfassung: Flow Matching with Confidence (FMwC)

Problemstellung

Generative Modelle, insbesondere Flow-Matching (FM)-Modelle, haben in den Bereichen Sprache, Vision und Wissenschaft state-of-the-art-Leistung erzielt. Allerdings erzeugen sie Ausgaben (z. B. halluzinierte Zitate, anatomisch inkorrekte Abbildungen oder thermodynamisch instabile Kristalle) ohne native Angabe der Zuverlässigkeit. Bestehende Lösungen zur Unsicherheitsquantifizierung, wie Deep Ensembles oder MC-Dropout, beruhen darauf, $k$ separate Modelle oder $k$ stochastische Trajektorien pro Stichprobe auszuführen. Dieser Ansatz verursacht eine $k\times$ höhere Rechenkosten, was eine pro-Stichproben-Konfidenz für den großflächigen Einsatz prohibitiv teuer macht. Darüber hinaus haben sich likelihood-basierte Methoden als unzuverlässig erwiesen, um in-Distribution- von out-of-Distribution-Eingaben in tiefen generativen Modellen zu unterscheiden. Die zentrale Frage lautet, ob Flow-Matching-Modelle lernen können, ihre eigene Unsicherheit für generierte Ausgaben zu quantifizieren, ohne die Inferenzkosten zu erhöhen oder Zugriff auf Likelihoods zu benötigen.

Methodik: Flow Matching with Confidence (FMwC)

FMwC ist ein Single-Pass-Ansatz, der Flow-Matching-Modelle mit einer pro-Stichproben-, Trajektorien-Level-Konfidenznote ausstattet, und zwar zu denselben Rechenkosten wie das standardmäßige deterministische Flow Matching. Die Methode integriert Variational Adaptive Dropout (VAD) mit einer geschlossenen Momentenfortpflanzung durch die Ordinary-Differential-Equation (ODE)-Trajektorie.

1. Variationale Formulierung

Anstatt ein einziges deterministisches Geschwindigkeitsfeld $v_{t,\theta}(x_t)$ zu lernen, behandelt FMwC die Netzwerkgewichte $\omega$ als Zufallsvariablen, die aus einer variationalen Posterior-Verteilung $q_\psi(\omega)$ gezogen werden. Das Modell maximiert die 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)]$$
wobei der Likelihood-Term dem standardmäßigen regressionsbasierten Flow-Matching-Loss entspricht. Entscheidend ist, dass der Gradient des erwarteten Log-Likelihoods unter der Posterior-Verteilung mit dem Gradienten des standardmäßigen Regressionslosses übereinstimmt, der auf Gewichtsproben ausgewertet wird, was ein simulationsfreies Training ermöglicht.

2. Zustandsabhängige Posterior-Verteilung via VAD

Um sicherzustellen, dass die Unsicherheit empfindlich auf die lokale Geometrie des Flows reagiert, setzt FMwC Variational Adaptive Dropout (VAD) ein. In diesem Rahmen ist die Posterior-Verteilung über die Gewichte in Schicht $l$ eine Gauß-Verteilung $q_\psi(\omega_l | x, t) = \mathcal{N}(\theta_l, \alpha_l(x, t)^2 \theta_l^2)$, wobei die Dropout-Skala $\alpha_l(x, t)$ von einem kleinen Inferenznetzwerk $E_\gamma$ vorhergesagt wird. Dies ermöglicht es, die Streuung der Posterior-Verteilung in Regionen geometrischer Mehrdeutigkeit (wo das Geschwindigkeitsfeld schwer zu approximieren ist) zu verbreitern und anderswo zu verengen, wobei dies ausschließlich über die ELBO gelernt wird.

3. Analytische Varianzfortpflanzung

Die Kerninnovation von FMwC besteht darin, die Varianz des Geschwindigkeitsfeldes analytisch sowohl durch die Netzwerkschichten als auch durch die ODE-Integrations-Trajektorie in einem einzigen Vorwärtsdurchgang fortzupflanzen:

• Durch das Netzwerk: Unter Verwendung des lokalen Reparameterisierungstricks werden der Mittelwert ($\mu_l$) und die Varianz ($\sigma_l^2$) der Aktivierungen schichtweise fortgepflanzt. Für lineare Schichten ist die Varianzrekursion in geschlossener Form gegeben. Für Nichtlinearitäten werden Momente analytisch unter Verwendung von Gauss-Hermite-Quadratur oder Taylor-Linearisierung berechnet, wodurch die Notwendigkeit von Monte-Carlo-Sampling entfällt.
• Durch die ODE: Bei jedem Euler-Schritt gibt das Netzwerk die Posterior-Mittelwert-Geschwindigkeit $\mu_t$ und ihre Varianz $\sigma_t^2$ aus. Der Integrator bewegt die Position unter Verwendung von $\mu_t$, während $\sigma_t^2$ aufgezeichnet wird. Dies ergibt eine Varianz-Trajektorie $\{\sigma_t^2\}_{t=0}^T$ anstelle eines einzelnen Skalars.

4. Konfidenz-Auslesungen

Die Varianz-Trajektorie kann zu skalaren Konfidenzwerten für nachgelagerte Aufgaben aggregiert werden:

• Endpoint Confidence: Die Varianz am letzten Schritt ($t=1$).
• Integrated Confidence: Die Summe der Standardabweichungen der Geschwindigkeit entlang der Trajektorie.
• Temporal Ratio: Das Verhältnis der Varianz in der späten Fensterphase zur Varianz in der frühen Fensterphase, was hilft, Stichproben zu unterscheiden, die konvergieren (niedrige späte Varianz), von solchen, die mehrdeutig bleiben.
• Learned Heads: Ein kleiner überwachter Head kann die vollständige Varianzsequenz auf eine aufgaben spezifische Note abbilden.

Hauptbeiträge

1. Single-Pass-Konfidenzschätzer: FMwC ist die erste Methode, die eine kalibrierte pro-Stichproben-Konfidenznote für Flow Matching zu den Kosten eines einzigen deterministischen Vorwärtsdurchgangs bereitstellt und dabei das Inferenzbudget des standardmäßigen FM einhält. Dies wird erreicht, indem erlernte, eingabeabhängige Rauschsignale analytisch durch die ODE fortgepflanzt werden.
2. Training-freie Verbesserungen der Inferenz: Die Konfidenznote ermöglicht praktische Anwendungen ohne Nachtraining:
   • Filterung: Verbessert die Probenqualität monoton durch Verwerfen von Ausgaben mit niedriger Konfidenz.
   • Bearbeitung: Identifiziert den exakten Zeitschritt ($t^*$), an dem das Modell eine Struktur festlegt (Varianzspitze), und ermöglicht gezielte, einschränkungsrespektierende Bearbeitungen (z. B. chemische Tauschvorgänge in Kristallen) durch Einbringen von Rauschen bei $t^*$.
   • Adaptive Schrittweite: Konzentriert ODE-Rechenschritte in Regionen hoher Varianz (Mehrdeutigkeit) und gewinnt hochwertige Proben bei niedrigen Schrittbudgets zurück, bei denen eine einheitliche Schrittweite versagt.
3. Geometrische Interpretierbarkeit: Die Arbeit zeigt, dass die FMwC-Konfidenznote mit der lokalen Divergenz des erlernten Geschwindigkeitsfeldes korreliert. Dies bietet einen Einblick in den generativen Prozess und verbindet epistemische Unsicherheit mit der geometrischen Instabilität (Bifurkationspunkte) des Flows.

Experimentelle Ergebnisse

Die Autoren bewerteten FMwC über drei Modalitäten hinweg: 2D-Dichteschätzung (Checkerboard), Bildgenerierung (MNIST) und de-novo-Generierung anorganischer Kristalle (Riemannian Flow Matching).

• Probenqualität: FMwC erhält oder verbessert die Probenqualität im Vergleich zum standardmäßigen FM. Auf Checkerboard und MNIST erreichte es jeweils niedrigere Fehlplatzierungsraten und FID-Scores, während es auf Kristallen eine vergleichbare Qualität beibehielt.
• Inferenzkosten: FMwC verursacht bei größeren Backbones einen vernachlässigbaren Overhead (z. B. ~1,003$\times$ für MNIST, ~1,08$\times$ für Kristalle) im Vergleich zur deterministischen Baseline. Im Gegensatz dazu erfordern Ensemble- und MC-Dropout-Baselines eine $5\times$ höhere Rechenleistung.
• Filterungsleistung: Die Filterung basierend auf FMwC-Scores erzielte Area Under the Precision-Recall Curve (AUPRC)-Scores, die mit $k=5$ Ensembles und MC-Dropout vergleichbar waren, jedoch zu den Kosten einer einzigen Trajektorie. Zum Beispiel erreichte der Temporal-Ratio-Score auf dem Checkerboard-Datensatz einen AUPRC von 0,55 und entsprach damit der Ensemble-Baseline (0,55) bei einem Fünftel der Rechenleistung.
• Erfolg bei der Bearbeitung: Beim Bearbeiten von Kristallen durch Einbringen von Rauschen am Varianzpeak ($t^*$) erreichte FMwC signifikant höhere Erfolgsquoten (z. B. 39 % gegenüber 22 % für chemische Tauschvorgänge) im Vergleich zur zufälligen Auswahl von Zeitschritten, was zeigt, dass Regionen hoher Varianz kritischen Entscheidungspunkten entsprechen.
• Korrelation mit Divergenz: Die Varianz-Trajektorie $\sigma_t^2$ zeigte eine starke Korrelation mit der Größe der Divergenz des Geschwindigkeitsfeldes ($|\nabla \cdot v_{t,\theta}|$). Auf dem Crystal-Datensatz erreichte die integrierte Korrelation einen Spearman-$\rho$ von 0,96 und validierte damit die geometrische Interpretation der Konfidenznote.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass FMwC Flow Matching von einem Generator von Proben in einen Generator von Proben transformiert, die mit einem Maß an Vertrauenswürdigkeit gepaart sind. Durch das Erlernen einer eingabeabhängigen Posterior-Verteilung über Gewichte und die analytische Fortpflanzung der Varianz macht FMwC konfidenzbewusste Generierung zu einem praktischen Standard für wissenschaftliche und sicherheitskritische Anwendungen. Die Autoren betonen, dass die Konfidenznote eine intrinsische geometrische Eigenschaft des einzelnen trainierten Flows ist, die sich von der zwischen Modellen bestehenden Diskrepanz unterscheidet, die durch Ensembles gemessen wird. Dies ermöglicht eine „chirurgische" Steuerung des generativen Prozesses, die Instabilitätsmomente gezielt anspricht, ohne die prohibitiven Kosten bestehender Unsicherheitsquantifizierungsmethoden. Die Arbeit legt nahe, dass das Verständnis und die Nutzung der Geometrie des erlernten Geschwindigkeitsfeldes zu zuverlässigeren, interpretierbareren und kontrollierbaren generativen Modellen führen können.