Résumé technique : Flotter avec confiance (FMwC)

Énoncé du problème

Les modèles génératifs, en particulier les modèles d'appariement de flux (Flow Matching, FM), ont atteint des performances de pointe dans les domaines du langage, de la vision et des sciences. Cependant, ils produisent des sorties (par exemple, des citations hallucinées, des figures anatomiquement incorrectes ou des cristaux thermodynamiquement instables) sans indication native de leur fiabilité. Les solutions existantes pour la quantification de l'incertitude, telles que les ensembles profonds (Deep Ensembles) ou le MC-Dropout, reposent sur l'exécution de $k$ modèles distincts ou de $k$ trajectoires stochastiques par échantillon. Cette approche engendre un coût computationnel multiplié par $k$, rendant la confiance par échantillon prohibitivement coûteuse pour un déploiement à grande échelle. De plus, il a été démontré que les méthodes basées sur la vraisemblance sont peu fiables pour distinguer les entrées in-distribution des entrées out-of-distribution dans les modèles génératifs profonds. La question centrale abordée est de savoir si les modèles d'appariement de flux peuvent apprendre à quantifier leur propre incertitude pour les sorties générées sans augmenter les coûts d'inférence ni nécessiter l'accès aux vraisemblances.

Méthodologie : Appariement de flux avec confiance (FMwC)

FMwC est une approche en un seul passage qui équipe les modèles d'appariement de flux d'un score de confiance par échantillon et au niveau de la trajectoire, au même coût computationnel que l'appariement de flux déterministe standard. La méthode intègre le Dropout Adaptatif Variationnel (VAD) avec une propagation des moments sous forme close à travers la trajectoire de l'Équation Différentielle Ordinaire (ODE).

1. Formulation variationnelle

Au lieu d'apprendre un seul champ de vitesse déterministe $v_{t,\theta}(x_t)$, FMwC traite les poids du réseau $\omega$ comme des variables aléatoires tirées d'une postérieure variationnelle $q_\psi(\omega)$. Le modèle maximise la borne inférieure de l'évidence (ELBO) :
$$\text{ELBO}(\psi) = \mathbb{E}_{\omega \sim q_\psi} [\log \mathcal{L}(D | \omega)] - \text{KL}[q_\psi(\omega) \parallel \pi(\omega)]$$
où le terme de vraisemblance correspond à la perte de régression standard d'appariement de flux conditionnel. Crucialement, le gradient de la vraisemblance logarithmique espérée sous la postérieure coïncide avec le gradient de la perte de régression standard évaluée sur des échantillons de poids, permettant un entraînement sans simulation.

2. Postérieure dépendante de l'état via VAD

Pour garantir que l'incertitude soit sensible à la géométrie locale du flux, FMwC emploie le Dropout Adaptatif Variationnel (VAD). Dans ce cadre, la postérieure sur les poids de la couche $l$ est une gaussienne $q_\psi(\omega_l | x, t) = \mathcal{N}(\theta_l, \alpha_l(x, t)^2 \theta_l^2)$, où l'échelle de dropout $\alpha_l(x, t)$ est prédite par un petit réseau d'inférence $E_\gamma$. Cela permet à la dispersion de la postérieure de s'élargir dans les régions d'ambiguïté géométrique (où le champ de vitesse est difficile à ajuster) et de se resserrer ailleurs, appris uniquement via l'ELBO.

3. Propagation analytique de la variance

L'innovation centrale de FMwC est la propagation analytique de la variance du champ de vitesse à travers à la fois les couches du réseau neuronal et la trajectoire d'intégration de l'ODE en un seul passage avant :

• À travers le réseau : En utilisant l'astuce de reparamétrisation locale, la moyenne ($\mu_l$) et la variance ($\sigma_l^2$) des activations sont propagées couche par couche. Pour les couches linéaires, la récurrence de la variance est sous forme close. Pour les non-linéarités, les moments sont calculés analytiquement en utilisant la quadrature de Gauss-Hermite ou une linéarisation de Taylor, évitant ainsi le besoin d'échantillonnage Monte Carlo.
• À travers l'ODE : À chaque étape d'Euler, le réseau sort la vitesse moyenne de la postérieure $\mu_t$ et sa variance $\sigma_t^2$. L'intégrateur avance la position en utilisant $\mu_t$, tandis que $\sigma_t^2$ est enregistrée. Cela produit une trajectoire de variance $\{\sigma_t^2\}_{t=0}^T$ plutôt qu'un scalaire unique.

4. Lectures de confiance

La trajectoire de variance peut être agrégée en scores de confiance scalaires pour les tâches en aval :

• Confiance au point final : La variance à l'étape finale ($t=1$).
• Confiance intégrée : La somme des écarts-types de vitesse le long de la trajectoire.
• Ratio temporel : Le ratio de la variance de la fenêtre tardive sur la variance de la fenêtre précoce, qui aide à distinguer les échantillons qui convergent (faible variance tardive) de ceux qui restent ambigus.
• Têtes apprises : Une petite tête supervisée peut mapper la séquence complète de variance vers un score spécifique à la tâche.

Contributions clés

1. Estimateur de confiance en un seul passage : FMwC est la première méthode à fournir un score de confiance par échantillon calibré pour l'appariement de flux au coût d'un seul passage avant déterministe, correspondant au budget d'inférence du FM standard. Elle y parvient en propageant analytiquement un bruit appris et dépendant de l'entrée à travers l'ODE.
2. Améliorations de l'inférence sans réentraînement : Le score de confiance permet des applications pratiques sans réentraînement :
   • Filtrage : Améliore de manière monotone la qualité des échantillons en éliminant les sorties à faible confiance.
   • Édition : Identifie l'instant précis ($t^*$) où le modèle s'engage dans une structure (pic de variance), permettant des éditions ciblées respectant les contraintes (par exemple, des échanges chimiques dans les cristaux) en injectant du bruit à $t^*$.
   • Pas adaptatif : Concentre les étapes de calcul de l'ODE dans les régions de forte variance (ambiguïté), récupérant des échantillons de haute qualité avec des budgets de pas faibles où un pas uniforme échoue.
3. Interprétabilité géométrique : L'article établit que le score de confiance FMwC corrèle avec la divergence locale du champ de vitesse appris. Cela offre une fenêtre sur le processus génératif, reliant l'incertitude épistémique à l'instabilité géométrique (points de bifurcation) du flux.

Résultats expérimentaux

Les auteurs ont évalué FMwC sur trois modalités : estimation de densité 2D (Checkerboard), génération d'images (MNIST) et génération de cristaux inorganiques de novo (Appariement de flux Riemannien).

• Qualité des échantillons : FMwC préserve ou améliore la qualité des échantillons par rapport au FM standard. Sur Checkerboard et MNIST, il a obtenu des taux de mauvais placement et des scores FID inférieurs, respectivement, tout en maintenant une qualité comparable sur les cristaux.
• Coût d'inférence : FMwC engendre une surcharge négligeable sur les gros modèles de base (par exemple, ~1,003$\times$ pour MNIST, ~1,08$\times$ pour les Cristaux) par rapport à la base déterministe. En revanche, les bases d'ensembles et de MC-Dropout nécessitent $5\times$ le calcul.
• Performance de filtrage : Le filtrage basé sur les scores FMwC a obtenu des scores de surface sous la courbe Précision-Rappel (AUPRC) comparables aux ensembles $k=5$ et au MC-Dropout, mais au coût d'une seule trajectoire. Par exemple, sur le jeu de données Checkerboard, le score de ratio temporel a atteint un AUPRC de 0,55, égalant la base d'ensemble (0,55) avec 1/5 du calcul.
• Succès de l'édition : Lors de l'édition de cristaux en injectant du bruit au pic de variance ($t^*$), FMwC a atteint des taux de succès significativement plus élevés (par exemple, 39 % contre 22 % pour les échanges chimiques) par rapport à la sélection aléatoire d'instant, démontrant que les régions de forte variance correspondent à des points de décision critiques.
• Corrélation avec la divergence : La trajectoire de variance $\sigma_t^2$ a montré une forte corrélation avec l'amplitude de la divergence du champ de vitesse ($|\nabla \cdot v_{t,\theta}|$). Sur le jeu de données Crystal, la corrélation intégrée a atteint un $\rho$ de Spearman de 0,96, validant l'interprétation géométrique du score de confiance.

Importance et revendications

L'article revendique que FMwC transforme l'appariement de flux d'un générateur d'échantillons en un générateur d'échantillons couplés à une mesure de fiabilité. En apprenant une postérieure dépendante de l'entrée sur les poids et en propageant la variance analytiquement, FMwC rend la génération consciente de la confiance par défaut pratique pour les applications scientifiques et critiques pour la sécurité. Les auteurs soulignent que le score de confiance est une propriété géométrique intrinsèque du flux unique entraîné, distincte du désaccord inter-modèle mesuré par les ensembles. Cela permet une guidance « chirurgicale » du processus génératif, ciblant les moments d'instabilité sans les coûts prohibitifs des méthodes existantes de quantification de l'incertitude. Ce travail suggère que la compréhension et l'exploitation de la géométrie du champ de vitesse appris peuvent conduire à des modèles génératifs plus fiables, interprétables et contrôlables.