Résumé Technique : FLUX – Appariement de Flux Longitudinal Conscient de la Géométrie avec Mélange d'Experts

1. Énoncé du Problème

De nombreux systèmes biologiques, tels que les circuits neuronaux et les populations cellulaires en différenciation, évoluent par le biais de dynamiques locales continues tout en basculant entre des régimes latents définis par l'apprentissage, le contexte des stimuli ou les stades développementaux. Cependant, ces systèmes sont fréquemment observés uniquement sous forme d'instantanés longitudinaux non appariés. Dans des scénarios tels que l'imagerie calcique grand champ ou le séquençage de l'ARN de cellules uniques, les mêmes cellules individuelles ou neurones ne peuvent pas être suivis d'un instant à l'autre ; au lieu de cela, les chercheurs observent une séquence de distributions marginales ($\mu_0, \mu_1, \dots, \mu_{T-1}$).

Ce cadre présente deux défis couplés :

1. Contraintes Géométriques : Les mesures biologiques de haute dimension se concentrent souvent près de variétés non linéaires courbes de faible dimension. Les interpolants euclidiens standards entre des extrémités non appariées peuvent traverser des régions de faible densité, violant ainsi les contraintes biophysiques et ne respectant pas la géométrie intrinsèque de l'espace d'état.
2. Découverte de Régimes : Le mécanisme de transport sous-jacent (le « flux ») change souvent au fil du temps. Par exemple, l'activité neuronale peut passer d'une attente large de récompense liée au stimulus à une discrimination spécifique, ou les cellules peuvent passer de programmes pluripotents à des programmes engagés dans une lignée. Il s'agit de changements dans les dynamiques locales (régimes), et non de simples clusters spatiaux. Les méthodes existantes se concentrent souvent sur la précision du transport mais échouent à décomposer le champ de vitesse pour identifier ces régimes dynamiques latents.

2. Méthodologie : FLUX

Les auteurs introduisent FLUX (Flow matching for Unpaired longitudinal data with miXture-of-experts), un cadre qui modélise conjointement le transport continu et la découverte de régimes non supervisée. La méthode opère en trois étapes séquentielles :

A. Trajectoires Conditionnelles Conscientes de la Géométrie

FLUX étend l'Appariement de Flux Métrique aux cadres longitudinaux. Au lieu d'utiliser des trajectoires euclidiennes rectilignes entre des marginales adjacentes, FLUX apprend une métrique dépendante des données à partir d'observations regroupées (incluant des échantillons non étiquetés).

• Apprentissage de la Métrique : Un backend géométrique (utilisant des métriques de noyau profond RBF ou RBF-MLP) estime le score de distance à la variété, définissant une métrique $G(x)$ qui pénalise les trajectoires quittant la variété des données.
• Réseau de Courbure : Un réseau $B$ paramétrise une trajectoire approximative consciente de la géométrie $\gamma_k(\alpha)$ entre les extrémités $x_k$ et $x_{k+1}$. Cette trajectoire se courbe pour rester sur la variété, fournissant des points d'entraînement $z_{k,\alpha}$ et des tangentes cibles $\dot{z}_{k,\alpha}$ qui respectent la géométrie locale.

B. Champs de Vitesse à Commutation de Régimes

FLUX décompose le champ de vitesse global $v_\theta$ en un Mélange d'Experts (MoE).

• Champs Vectoriels d'Experts : Le modèle apprend $M$ champs vectoriels d'experts parcimonieux $f_m$.
• Routeur : Un réseau routeur $g(t, x)$ mappe le temps et l'état vers des logits d'experts.
• Mécanisme de Routage : Pendant l'entraînement, un mécanisme Straight-Through Gumbel-Softmax sélectionne un expert actif pour chaque échantillon, permettant une optimisation différentiable des décisions de routage discrètes. Lors de l'inférence, le régime est déterminé par l'affectation dure déterministe (argmax).
• Objectif : Le champ de vitesse est entraîné pour correspondre aux tangentes des trajectoires conscientes de la géométrie. Le routeur est régularisé pour encourager des affectations parcimonieuses et temporellement cohérentes et prévenir l'effondrement des experts.

C. Pipeline d'Entraînement

L'entraînement se déroule en trois étapes pour assurer la stabilité :

1. Apprentissage de la Métrique : Entraînement du backend géométrique sur des échantillons regroupés.
2. Entraînement du Réseau de Courbure : Entraînement du réseau de courbure avec la métrique figée pour générer des interpolants conscients de la géométrie.
3. Entraînement du Champ de Vitesse : Gel de la géométrie et des réseaux de courbure ; entraînement des experts de vitesse et du routeur sur les trajectoires et les tangentes générées.

3. Contributions Clés

1. Extension aux Marginales Longitudinales : Les auteurs étendent l'appariement de flux conscient de la géométrie aux chaînes ordonnées de distributions marginales, en soulignant le rôle de l'apprentissage de variété non supervisé lorsque les étiquettes de trajectoire sont rares mais que des observations non étiquetées sont disponibles.
2. Découverte de Régimes par Décomposition de Vitesse : Ils introduisent un champ de vitesse à commutation de régimes où le routage parcimonieux d'experts découvre des régimes dynamiques latents directement à partir de la structure de transport locale, plutôt que de regrouper des observations statiques a posteriori.
3. Interaction Géométrie-Régimes : Le travail démontre que l'apprentissage géométrique est crucial pour la découverte de régimes, en particulier lorsque les régimes sont définis par des changements dans les dynamiques locales plutôt que par une séparation spatiale statique.
4. Validation Biologique : Le cadre est validé sur deux contextes biologiques distincts :
   • Données Neurales : Séparation de l'entraînement précoce/intermédiaire de l'entraînement tardif dans l'apprentissage associatif, coïncidant avec une divergence comportementale.
   • Biologie du Développement : Récupération de régimes d'évolution de l'expression associés aux populations cellulaires pluripotentes versus différenciées dans des embryoides.

4. Résultats

L'article évalue FLUX sur quatre jeux de données : Stanford Bunny (variété contrôlée), Attracteur de Lorenz (dynamiques à commutation de régimes), Imagerie Calcique Grand Champ (apprentissage neuronal) et Différenciation d'Embryoides.

• Stanford Bunny : FLUX préserve le transport de surface dans des plongements de haute dimension où les méthodes euclidiennes échouent à rester sur la variété.
• Attracteur de Lorenz : FLUX atteint une récupération parfaite des régimes (ARI=1.0, NMI=1.0) en identifiant la commutation entre les régimes chaotiques et sous-critiques. Crucialement, FLUX sans Apprentissage de Variété (utilisant des trajectoires euclidiennes) s'effondre vers un seul expert, échouant à découvrir le régime malgré la même architecture MoE. Cela prouve que les trajectoires conscientes de la géométrie exposent les signatures de courbure et de vitesse nécessaires à la spécialisation des experts.
• Données Neurales : FLUX identifie une transition dans les dynamiques corticales qui correspond à l'émergence de la discrimination comportementale (CS+ vs CS−), surpassant le routage euclidien et les dynamiques de commutation de base.
• Embryoides : Bien que le routage euclidien fonctionne bien ici (car les stades sont spatialement séparables), FLUX avec apprentissage géométrique reste compétitif, renforçant l'idée que la géométrie est la plus critique lorsque les régimes sont encodés dans les dynamiques locales.

Résultats d'Ablation : L'article indique explicitement que le routage par mélange d'experts seul est insuffisant. Sans apprentissage géométrique, le modèle peut ajuster le transport local mais échoue à récupérer les régimes lorsque ces régimes sont encodés dans les dynamiques locales plutôt que dans des clusters statiques.

5. Importance et Revendications

Les auteurs affirment que la décomposition de vitesse consciente de la géométrie fournit une stratégie générale pour découvrir les transitions d'états biologiques latentes à partir d'instantanés longitudinaux non appariés.

• Impact Scientifique : FLUX déplace l'accent de l'optimisation pure de la précision du transport (minimisation de la distance de Wasserstein) vers la récupération de structures dynamiques interprétables. Il suggère que les systèmes biologiques se réorganisent souvent en régimes dynamiques distincts lors de processus tels que l'apprentissage et la différenciation.
• Limites : Les auteurs reconnaissent que FLUX est plus coûteux en calcul que l'appariement de flux euclidien en raison de l'entraînement de trois réseaux distincts. Ils notent également que le nombre d'experts est fixé à l'avance, et que les régimes inférés doivent être interprétés avec prudence lorsque les dynamiques locales sont bruyantes ou faiblement échantillonnées.
• Conclusion : L'article conclut que FLUX offre un cadre pour la modélisation dynamique interprétable d'instantanés biologiques non appariés, comblant le fossé entre l'apprentissage de variété géométrique et les systèmes dynamiques à commutation de régimes.