Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

Résumé Technique : Flow Matching Multi-Marginal avec Interpolants Appris par Adversaire

Énoncé du Problème
Le papier traite du défi consistant à apprendre la dynamique temporelle d'un système à partir de seules observations échantillonnées (instantanés) à des points temporels discrets, sans accès aux trajectoires réelles (ground-truth). Il s'agit d'une tâche centrale dans des domaines tels que la biologie (par exemple, le séquençage de l'ARN de cellule unique, le suivi cellulaire) et la médecine. Formellement, l'objectif est d'inférer un champ de vecteurs dépendant du temps $v_t$ régissant une équation différentielle ordinaire (ODE) $dx_t = v_t(x_t)dt$, de telle sorte que l'application d'intégration pousse une distribution marginale initiale $q_0$ vers une cible $q_1$ tout en satisfaisant les contraintes marginales intermédiaires $p_{t_i} = q_{t_i}$ pour un ensemble de pas de temps $0 = t_1 < \dots < t_K = 1$.

Les approches existantes, telles que l'application du Flow Matching (FM) standard à des paires consécutives de marginales ou l'utilisation d'interpolants par splines linéaires ou cubiques, aboutissent souvent à des courbes d'interpolation non lisses. Ces problèmes de non-lissage entraînent une variance élevée du gradient dans l'objectif d'apprentissage et des difficultés pour marginaliser les champs de vecteurs résultants, particulièrement dans des contextes de haute dimension ou lorsque les données sont bruitées.

Méthodologie : ALI-CFM
Les auteurs proposent ALI-CFM (Adversarially Learnt Interpolants with Conditional Flow Matching), une méthode en deux étapes qui découple l'apprentissage des interpolants de l'apprentissage du champ de vecteurs.

1. Apprentissage Adversaire des Interpolants (ALI) :
   Au lieu de forcer les interpolants à passer par des points de données spécifiques (appariement ponctuel), ALI apprend un interpolant neural $G_\phi(x_0, x_1, t)$ qui correspond à la distribution des marginales intermédiaires. L'interpolant est paramétré comme une interpolation linéaire plus un terme de correction neural :
   $$G_\phi(x_0, x_1, t) = (1-t)x_0 + tx_1 + t(1-t)f_\phi(x_0, x_1, t)$$
   Pour garantir que la distribution de poussée (pushforward) de l'interpolant corresponde aux marginales intermédiaires observées $q_{t_i}$, les auteurs emploient un objectif adversaire inspiré des GAN. Un discriminateur $D_\gamma$ distingue les échantillons de la vraie marginale $q_{t_i}$ des échantillons générés par l'interpolant $G_\phi$. Le générateur (l'interpolant) est entraîné pour minimiser la divergence de Jensen-Shannon entre ces distributions.

   Pour assurer l'unicité des interpolants et empêcher des trajectoires arbitrairement courbes, les auteurs introduisent des régulariseurs :

   • Référence Linéaire : Pénalise l'écart par rapport à une ligne droite entre les points terminaux couplés.
   • Référence Linéaire par Morceaux : Effectue une régression de l'interpolant vers un chemin linéaire par morceaux défini par des couplages de transport optimal (OT) de type chaîne de Markov.
   • Norme de la Seconde Dérivée : Impose la lissé par la pénalisation de la courbure de l'interpolant.

2. Marginalisation via le Flow Matching Conditionnel (CFM) :
   Une fois les interpolants $G_\phi$ entraînés, ils sont fixés. Un second réseau neural $u_\theta$ est entraîné pour approximer le champ de vecteurs qui génère ces interpolants. Cela est fait en utilisant l'objectif standard de Conditional Flow Matching (CFM), où le champ de vecteurs cible est la dérivée temporelle de l'interpolant appris :
   $$L_{CFM} = \mathbb{E}_{t, (x_0, x_1) \sim \pi} \| u_\theta(G_\phi(x_0, x_1, t)) - \frac{d}{dt}G_\phi(x_0, x_1, t) \|^2$$
   Le couplage $\pi$ entre $q_0$ et $q_1$ peut être indépendant ou un plan de Transport Optimal (OT).

Contributions Clés

1. Technique d'Interpolation Novatrice : Le papier introduit une méthode pour apprendre des interpolants en utilisant un objectif adversaire qui correspond aux distributions plutôt qu'aux points. Cela permet d'obtenir des trajectoires lisses et uniques sous des hypothèses légères (prouvées dans les appendices).
2. Robustesse au Bruit et à la Complexité : La méthode est conçue pour gérer des jeux de données bruités et des géométries temporelles complexes où les méthodes d'interpolation par points (comme les splines ou le linéaire par morceaux) échouent ou deviennent instables.
3. Scalabilité et Polyvalence : Les auteurs démontrent que ALI-CFM passe à l'échelle pour des données de haute dimension et un grand nombre de distributions marginales ($K$), surpassant les bases de référence existantes dans certains scénarios complexes tout en restant compétitif dans d'autres.

Résultats Expérimentaux
Les auteurs évaluent ALI-CFM sur quatre tâches distinctes :

• Distribution Synthétique en "Nœud" (Knot) : Sur un jeu de données synthétiques avec 1 200 marginales formant un nœud, ALI-CFM était la seule méthode capable de capturer précisément la géométrie dépendante du temps. Les méthodes concurrentes (OT-MFM, splines cubiques, linéaire par morceaux) ont produit des interpolants non lisses avec une variance de gradient élevée.
• Suivi Cellulaire (Glioblastome) : En utilisant des données de microscopie d'une cellule en mouvement, ALI-CFM a réussi à apprendre la dynamique spatio-temporelle à partir de données bruitées et sous-échantillonnées. En revanche, les trajectoires d'OT-CFM ont divergé, et l'OT-MFM indépendant du temps n'a pas réussi à capturer les dépendances temporelles.
• Inférence de Trajectoire de Cellule Unique (scRNA-seq) : Sur les jeux de données Embryoid Body (EB), Cite-seq et Multiome, ALI-CFM a performé au même niveau que les bases de référence de pointe (OT-MFM, OT-CFM) en termes de distance de l'acheteur de terre (Earth Mover's Distance - EMD). Les auteurs notent que, bien que l'entraînement adversaire rende difficile le fait de surpasser strictement les méthodes qui font du surapprentissage sur les données ponctuelles dans ces métriques spécifiques, la méthode reste hautement compétitive.
• Inférence de Coordonnées Tumorales (Transcriptomique Spatiale) : Sur un jeu de données de cancer du sein avec des sections de tissus ordonnées spatialement, ALI-CFM a significativement surpassé toutes les bases de référence (y compris OT-CFM, OT-MMFM et OT-MFM) dans la prédiction de la distribution des régions tumorales dans les sections non testées (held-out). Les auteurs attribuent ce succès à la capacité de la méthode à gérer les distributions multimodales et les défis de l'alignement des sections de tissus.

Signification et Revendications
Le papier affirme que ALI-CFM offre une approche conceptuellement nouvelle pour le flow matching multi-marginal en passant d'un assignement ponctuel strict à un appariement de distribution. Ce changement permet l'apprentissage d'interpolants lisses et uniques, robustes au bruit et capables de modéliser des dynamiques complexes là où les méthodes précédentes échouent. Les auteurs soulignent que leur méthode est particulièrement efficace dans des scénarios impliquant des centaines ou des milliers de distributions marginales et des données bruitées, un régime dans lequel aucun autre algorithme de flow matching n'a démontré de performance comparable. Bien que la méthode soit au même niveau que l'état de l'art sur les benchmarks standards de cellule unique, sa performance supérieure sur la transcriptomique spatiale et les tâches de géométrie synthétique complexe souligne sa polyvalence et sa scalabilité.