FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

Résumé Technique : FreeBridge – Ponts de Schrödinger Variationnels pour la Dynamique de Transition Cellulaire

Énoncé du Problème
Les essais d'imagerie à haut contenu sont essentiels pour quantifier les réponses cellulaires aux perturbations chimiques et génétiques. Cependant, une limitation fondamentale existe : les cellules sont chimiquement fixées au moment de l'acquisition de l'image, ce qui rend les trajectoires temporelles individuelles inobservables. Par conséquent, la modélisation des perturbations se réduit à l'inférence d'un transport stochastique entre les populations de contrôle et de traitement, où seules les distributions marginales finales sont connues.

Bien que les modèles génératifs récents (ex. : modèles de diffusion, flow matching) parviennent à un fort alignement entre les points terminaux de contrôle et de perturbation, la cohérence des points terminaux seule est insuffisante pour déterminer l'évolution intermédiaire. Une infinité de processus stochastiques peut relier des marginales identiques ; sans contraintes supplémentaires, ces modèles peuvent traverser des régions de l'espace latent non soutenues par les morphologies cellulaires observées, menant à des états intermédiaires biologiquement implausibles. Les méthodes existantes optimisent souvent des objectifs au niveau de la distribution, laissant les structures cellulaires intermédiaires sous-déterminées, ce qui peut satisfaire des métriques telles que FID ou KID tout en générant des artefacts.

Méthodologie : Le Cadre FreeBridge
Les auteurs proposent FreeBridge, un cadre de transport stochastique contraint formulé comme un « Moteur Cellulaire » (Cell Engine) qui sépare explicitement la définition des états cellulaires valides du mécanisme de transition. La méthode opère sous une supervision de type « uniquement par les points terminaux » et se compose de trois composantes clés :

1. Unité d'État (Géométrie Monocellulaire Explicite) :

   • Au lieu de traiter les images comme des mélanges au niveau de la population, FreeBridge définit les états atomiques comme des représentations monocellulaires segmentées par instance (utilisant Cellpose).
   • Un encodeur pré-entraîné et fixe $\phi$ projette ces extractions dans un espace latent $z = \phi(x)$. Cet encodeur est gelé pendant l'apprentissage du transport, établissant une géométrie latente fixe définie avant l'apprentissage de la règle de transition.
   • Une banque de support empirique $Z_{emp}$ est construite à partir des plongements latents des cellules observées. Cette banque définit le « manifold » (variété) des morphologies biologiquement valides.

2. Règle de Transition (Transport par Pont de Schrödinger) :

   • La transition entre la distribution de contrôle $\mu_z$ et la distribution perturbée $\nu_{c,z}$ est modélisée comme une équation différentielle stochastique (SDE) contrôlée : $dZ_t = u_t(Z_t) dt + \sigma dW_t$.
   • L'objectif est de trouver un champ de dérive $u_t$ qui transporte la distribution initiale vers la distribution cible à $t=1$.
   • Contrairement au transport optimal standard, cette formulation incorpore un coût d'état spécifique à la tâche pour régulariser le chemin.

3. Régularisation Géométrique (Contrainte de Support Empirique) :

   • Pour empêcher le processus stochastique de traverser des « vides » dans l'espace latent, FreeBridge introduit un coût de support $V(z) = \lambda \cdot d_{emp}(z)^2$, où $d_{emp}(z)$ est la distance au plus proche voisin dans la banque de support empirique $Z_{emp}$.
   • L'objectif total minimise l'énergie cinétique (magnitude de la dérive) et le coût de support :
     $$\min_u \mathbb{E} \left[ \int_0^1 \frac{1}{2\sigma^2}\|u_t(Z_t)\|^2 dt + \lambda \int_0^1 d_{emp}(Z_t)^2 dt \right]$$
   • Cette régularisation ancre les états intermédiaires à la morphologie empiriquement observée sans nécessier de trajectoires appariées au niveau de la cellule.

Contributions Clés

• Formulation : Introduit une formulation de pont de Schrödinger pour la modélisation des transitions monocellulaires qui sépare explicitement la spécification de l'état du transport stochastique.
• Ancrage Géométrique : Propose un terme de régularisation novateur basé sur l'estimation de densité non paramétrique (distance du plus proche voisin à une banque de support empirique) pour contraindre les trajectoires intermédiaires dans des régions biologiquement plausibles.
• Protocole Unifié : Établit un protocole d'évaluation unifié à travers plusieurs jeux de données (BBBC021, RxRx1, JUMP) en utilisant des extractions monocellulaires segmentées par instance, garantissant une comparaison équitable entre les méthodes.
• Performance : Démontre que la contrainte de la géométrie améliore non seulement la fidélité des points terminaux, mais aussi la rétention sémantique des mécanismes de perturbation (Mechanism-of-Action).

Résultats
Évalué sur les jeux de données BBBC021, RxRx1 et JUMP, FreeBridge démontre les points suivants :

• Fidélité des Points Terminaux : Atteint des scores FID (Fréchet Inception Distance) et KID (Kernel Inception Distance) compétitifs ou améliorés par rapport aux bases de référence de l'état de l'art (PhenDiff, IMPA, CellFlux). Sur BBBC021, il obtient le FID le plus bas (15,35) et le KID le plus bas (1,10).
• Cohérence Intermédiaire : Réduit significativement le Taux de Violation de Support ($R_{viol}$), une métrique mesurant la proportion d'états intermédiaires tombant en dehors du support morphologique empirique. Sur BBBC021, FreeBridge réduit $R_{viol}$ à 0,11, contre 0,38 pour PhenDiff et 0,31 pour CellFlux.
• Rétention Sémantique : Les études d'ablation montrent que la suppression de la régularisation géométrique diminue la précision de la classification du Mécanisme d'Action (MoA) (passant de 72,08 % à 65,08 % sur BBBC021), indiquant que la contrainte de l'évolution intermédiaire aide à préserver les signaux mécanistes.
• Analyse Qualitative : Les visualisations confirment que FreeBridge préserve mieux les caractements morphologiques spécifiques aux composés (ex. : compaction nucléaire, rétrécissement cytoplasmique) que les autres méthodes, qui présentent souvent un lissage de texture ou des motifs incomplets.

Signification et Revendications
L'article soutient que la modélisation des transitions biologiques nécessite plus qu'un simple alignement des frontières ; elle nécessite l'ancrage géométrique explicite de l'évolution intermédiaire. En ancrant le transport stochastique à un manifold cellulaire empiriquement supporté, FreeBridge offre une approche rigoureuse pour générer des dynamiques de perturbation biologiquement interprétables. Les auteurs affirment que cette formulation offre une base robuste pour la simulation de « cellules virtuelles », comblant la lacune critique où les modèles alignés sur les points terminaux échouent à capturer des états cellulaires intermédiaires réalistes. Le travail souligne que, sans contraintes géométriques, les modèles génératifs peuvent produire des points terminaux statistiquement valides qui traversent des régions de l'espace latent biologiquement non soutenues.

Limites
Les auteurs reconnaissent que FreeBridge dépend d'un encodeur gelé ; les facteurs morphologiques absents du pré-entraînement de l'encodeur ne peuvent être récupérés par le seul apprentissage du transport. De plus, le calcul de la distance de support est de complexité quadratique ($O(B^2)$) par rapport à la taille du mini-batch, ce qui pourrait nécessiter une recherche de plus proches voisins approximative pour des banques de support plus larges.