Résumé technique : Flow Matching pour les données de comptage (count-FM)

Énoncé du problème

Les données de comptage de haute dimension, omniprésentes dans des domaines tels que le séquençage de l'ARN de cellules uniques (scRNA-seq) et l'analyse des trains de potentiels d'action neuronaux, posent des défis uniques pour la modélisation générative. Ces types de données se caractérisent par leur discrétion, leur parcimonie et des structures de corrélation complexes. Les approches génératives profondes existantes peinent souvent dans ce domaine :

• Approches à états catégoriels : Les méthodes traitant les comptages comme des variables catégorielles (par exemple, des tokens) nécessitent des dimensions de sortie qui évoluent avec la plage de comptage maximale. Cela entraîne une inefficacité paramétrique et traite les comptages adjacents comme des états non liés, ignorant la géométrie inhérente de l'espace de comptage.
• Transformations continues : Les approches qui mappent les comptages vers des espaces continus via une déquantification ou des représentations latentes remplacent la masse de probabilité discrète par des densités continues, risquant d'estomper la structure discrète sous-jacente.
• Limites des modèles de comptage existants : Bien que des modèles de saut spécifiques aux comptages existent (par exemple, Poisson-JUMP), ils ne sont souvent pas conçus pour le transport entre des distributions de comptage arbitraires ou manquent des propriétés d'entraînement et de génération efficaces des méthodes basées sur les flots modernes.

L'article répond au besoin d'un cadre génératif opérant directement dans l'espace de comptage, respectant la géométrie discrète des données et permettant un transport efficace entre des distributions source et cible arbitraires, sans la surcharge paramétrique des modèles à états catégoriels.

Méthodologie : count-FM

Les auteurs proposent count-FM, un cadre de flow matching adapté aux données de comptage, basé sur un processus de naissance-mort en temps continu avec des sauts unitaires locaux.

1. Processus de naissance-mort en temps continu

Le modèle définit la transition entre une distribution source $p_0$ et une distribution cible $p_1$ comme un processus de saut de Markov en temps continu (CTMC) sur l'espace des vecteurs de comptage $\mathbb{N}_0^d$. Le processus évolue via des naissances unitaires locales ($+1$) et des morts ($-1$) pour chaque coordonnée $i$ :

• Taux de naissance : $\lambda_{t,i}(x)$
• Taux de mort : $\mu_{t,i}(x)$
• Dynamique : Sur un petit intervalle $h$, la probabilité de transition de l'état $x$ vers $x + e_i$ est $h \lambda_{t,i}(x)$, et vers $x - e_i$ est $h \mu_{t,i}(x)$. La mort est interdite lorsque le comptage est nul.

Cette paramétrisation est hautement efficace : au lieu de prédire une distribution catégorielle complète sur tous les niveaux de comptage possibles (ce qui évolue avec le comptage maximal), le modèle ne prédit que deux taux par dimension ($2d$ sorties), indépendamment de la plage de comptage.

2. Pont binomial conditionnel et entraînement

Pour entraîner le modèle, les auteurs définissent un pont conditionnel entre des extrémités fixes $(x_0, x_1)$. Pour chaque coordonnée $i$, le pont est construit en utilisant une distribution binomiale :
$$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$
Ce pont garantit que la moyenne conditionnelle se déplace linéairement de $x_0$ à $x_1$ tout en maintenant des états intermédiaires à valeurs entières.

En imposant la préservation de la masse (l'équation de Kolmogorov vers l'avant) sur ce pont conditionnel, les auteurs dérivent des expressions fermées pour les taux de naissance et de mort conditionnels :
$$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$$
où $(\cdot)^+$ désigne la rectification positive.

3. Objectif d'entraînement

Le modèle est entraîné en minimisant une divergence de Kullback-Leibler (KL) dans l'espace des chemins entre le processus de pont conditionnel et le processus appris par le modèle. Cela se décompose en un objectif d'appariement de taux local :
$$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$$
où $\ell(u, v) = v - u \log v$ est un terme de KL généralisé. Cela permet un entraînement sans simulation des taux de transition marginaux en apprenant les taux conditionnels identifiables.

4. Échantillonnage et couplage

• Échantillonnage : Les échantillons sont générés en simulant le processus de naissance-mort appris vers l'avant de $t=\epsilon_t$ à $t=1-\epsilon_t$ en utilisant une discrétisation de saut local du premier ordre.
• Couplage des extrémités : L'article explore deux stratégies de couplage pour l'appariement des échantillons source et cible :
  • Couplage indépendant : Appariement aléatoire ($p_0(x_0)p_1(x_1)$).
  • Couplage par Transport Optimal (OT) : Appariement OT par mini-lots utilisant un coût de Poisson symétrique. Cela tend à induire des transitions avec une courbure plus faible, améliorant l'interprétabilité et l'efficacité de l'échantillonnage.

5. Génération conditionnelle

Pour les tâches conditionnelles, le modèle intègre des covariables $y$ et utilise le Classifier-Free Guidance (CFG). Pendant l'entraînement, les variables de conditionnement sont aléatoirement supprimées. Lors de l'échantillonnage, les taux conditionnels et non conditionnels sont combinés via une échelle de guidage $w$ pour contrôler la force de la condition.

Contributions clés

1. Cadre count-FM : Un cadre de flow matching pour les données de comptage basé sur une dynamique de naissance-mort locale, permettant un transport directement dans l'espace de comptage avec des états intermédiaires restant à valeurs de comptage.
2. Schéma d'entraînement traitable : Une méthode pour apprendre les taux de transition marginaux via un pont binomial conditionnel, conservant l'efficacité du flow matching tout en évitant la complexité de la modélisation à états catégoriels.
3. Efficacité paramétrique : Une formulation qui prédit uniquement les taux locaux de naissance et de mort ($2d$ sorties) plutôt que des distributions catégorielles complètes, la rendant adaptée aux données de haute dimension avec de grandes plages de comptage.

Résultats expérimentaux

1. Simulation (Mélange Gamma-Poisson 2D)

• Performance : count-FM et count-FM-OT ont obtenu la meilleure qualité d'échantillon (distance de Wasserstein-2 et MMD$^2$ les plus faibles) par rapport aux bases de référence incluant D3PM, discrete-FM, tauLDR, SEDD et Poisson-JUMP.
• Efficacité : count-FM a obtenu ces résultats avec substantiellement moins de paramètres entraînables (2 372) par rapport aux bases de référence à états catégoriels (14 000–67 000+).
• Géométrie : Les états intermédiaires ont évolué de manière fluide dans l'espace de comptage. La version couplée par OT a suivi un chemin de transition visiblement plus droit et a atteint la qualité cible avec moins d'évaluations de fonctions (NFE).

2. Séquençage de l'ARN de cellules uniques (Gyrus denté)

• Génération non conditionnelle : Sur un jeu de données de 2 930 cellules et 13 913 gènes, count-FM a surpassé les bases de référence à variables latentes (scVI) et les modèles de diffusion latente (scDiffusion, scLDM) en qualité d'échantillon (W2 et MMD$^2$) tout en utilisant moins de paramètres (environ 9,8 M contre 29 M–39 M pour les concurrents).
• Transport (P12 à P35) : Le modèle a réussi à modéliser le transport développemental du jour postnatal 12 au jour 35. En utilisant un couplage OT restreint à la lignée, les chemins de transport ont respecté les structures développementales connues (par exemple, la progression de neuroblaste à cellule granulaire), produisant des distributions intermédiaires interprétables restant dans l'espace de comptage.

3. Trains de potentiels d'action neuronaux (Hippocampe/Entorhinal)

• Génération conditionnelle : Le modèle a été appliqué pour générer des comptages de potentiels d'action conditionnés par la position linéarisée et la direction de course.
• Comparaison : count-FM avec une échelle de guidage $w=1$ a surpassé un régresseur MLP déterministe et une base de référence Poisson MLP.
• Structure de corrélation : Contrairement au Poisson MLP, qui sous-estimait les corrélations de population, count-FM a mieux préservé la structure de dépendance inter-neurones.
• Compromis du guidage : $w=1$ a offert le meilleur équilibre entre l'ajustement spatial et l'étalonnage de la distribution. $w=2$ a affiné les champs de lieu mais a déformé l'étalonnage de l'amplitude.

Signification et affirmations

L'article affirme que count-FM fournit un cadre efficace pour la modélisation de données de comptage de haute dimension en permettant la génération et le transport directement dans l'espace de comptage. Sa signification principale réside dans :

• Fidélité géométrique : En modélisant le transport via des sauts unitaires locaux, la méthode respecte la géométrie discrète des données de comptage, évitant les effets d'estompage des approximations continues.
• Évolutivité : L'efficacité paramétrique la rend viable pour des applications de haute dimension (par exemple, des milliers de gènes) où les méthodes à états catégoriels deviennent prohibitives en termes de calcul.
• Interprétabilité : Les états intermédiaires dans le chemin de transport sont des échantillons de comptage valides, permettant l'inspection de distributions de transition significatives dans des applications telles que la biologie du développement.
• Polyvalence : Le cadre prend en charge la génération non conditionnelle, le transport entre des distributions arbitraires et la génération conditionnelle avec un guidage contrôlable.

Les auteurs reconnaissent certaines limites, notant que le pont binomial conditionnel est relativement simple et pourrait ne pas explorer tous les chemins intermédiaires plausibles, et que la discrétisation actuelle du premier ordre pourrait être améliorée par des schémas de simulation plus précis. Cependant, les résultats empiriques sur des simulations et des jeux de données biomédicaux démontrent que count-FM atteint une haute qualité d'échantillon avec une efficacité de modélisation supérieure aux approches concurrentes.