LoCoFlow: Local-Context Flow Matching for Single-Cell Dynamics and Fate Inference

Resumen Técnico: LoCoFlow

Declaración del Problema

La secuenciación de célula única proporciona instantáneas de alta resolución de los estados celulares en puntos temporales discretos, pero la naturaleza destructiva de la medición impide la observación de correspondencias longitudinales reales entre células individuales. La reconstrucción de trayectorias de desarrollo continuas a partir de estas instantáneas es un problema mal planteado con múltiples soluciones posibles. Además, los métodos existentes suelen tratar las células como vectores aislados, fallando al representar adecuadamente la geometría del colector (manifold) local del espacio de estados celulares o la incertidumbre inherente en la diversificación del linaje. Aunque las células con estados transcripcionales similares pueden exhibir tendencias de destino divergentes, los enfoques actuales de transporte óptimo y de ajuste de flujo (flow-matching) suelen inferir acoplamientos temporales a nivel de célula individual, careciendo de un marco robusto para modelar la evolución de los vecindarios celulares locales como unidades cohesivas.

Metodología: LoCoFlow

Los autores presentan LoCoFlow (Local-Context Flow Matching), un marco que integra la geometría de nubes de puntos locales con la incertidumbre del destino condicional para modelar la dinámica de célula única. El método opera en dos etapas principales:

1. Construcción de Parches Locales y Correspondencia Jerárquica

En lugar de emparejar células aisladas, LoCoFlow construye "parches" (vecindarios locales) superpuestos alrededor de células ancla.

• Definición de Parche: Para datos de expresión, los parches son vecindarios geodésicos en un grafo de k-vecinos más cercanos; para datos espaciales, son vecindarios de tamaño fijo en coordenadas medidas.
• Transporte Jerárquico: El método emplea un esquema de transporte óptimo (OT) de dos niveles. Primero, resuelve un problema de transporte a nivel de parche utilizando la distancia de Bures para alinear los centros y las covarianzas de los parches (capturando tanto la ubicación como la dispersión). Segundo, dentro de los pares de parches retenidos, calcula un plan de transporte a nivel de célula regularizado por entropía.
• Asignación Suave (Soft Assignment): La agregación de estos planes produce una distribución condicional de células objetivo para cada fuente, en lugar de una asignación rígida, preservando múltiples correspondencias de sucesores potenciales.

2. Flujo de Entrenamiento de Dos Etapas

• Etapa 1: Transporte de Puente Regularizado Geométricamente:
  • Utilizando los conjuntos de puntos finales emparejados derivados del transporte jerárquico, el modelo aprende "puentes" (trayectorias intermedias) entre los parches de origen y de destino.
  • Una red neuronal predice un campo de velocidad residual $f_\theta$ que, combinado con una interpolación lineal, asegura que la trayectoria pase exactamente por los puntos finales.
  • Prioris Geométricos: La función de pérdida incluye términos de energía cinética y una penalización radial que restringe la deformación relativa (expansión/contracción) del parche sin restringir la traslación común. También penaliza las velocidades ortogonales a la geometría del colector local estimado.
• Etapa 2: Destilación de Velocidad Condicional:
  • Dado que el estado futuro es desconocido durante la inferencia, el puente (que requiere ambos puntos finales) no puede usarse directamente. LoCoFlow destila las velocidades del puente en un campo de velocidad condicional $V_\omega$.
  • Latentes Jerárquicos: El modelo aprende a generar códigos de destino $c_i$ a partir de un contexto de origen $Z_0$ utilizando una estructura latente jerárquica: un latente compartido por el parche $\eta_P$ (que captura resultados a nivel de población) y latentes específicos de la célula $\eta_i$ (que capturan la variación individual).
  • Entrenamiento: Un codificador de destino aprende primero a mapear los puntos finales observados hacia los códigos $c^\star$. Luego, un generador aprende a producir estos códigos a partir del contexto de origen y los latentes. La red de velocidad se entrena para coincidir con las velocidades del puente condicionadas a estos códigos generados.
  • Inferencia: Durante la inferencia, los códigos latentes se extraen de un prior (sin observaciones futuras) y se mantienen fijos durante la integración. Esto permite el muestreo de trayectorias alternativas que representan la incertidumbre del destino.

Contribuciones Clave

1. Representación de Contexto Local: Cambia la unidad fundamental del modelado dinámico de vectores de expresión aislados a estados de población local (nubes de puntos), modelando explícitamente la evolución de la geometría del vecindario.
2. Transporte Jerárquico: Introduce un mecanismo de transporte óptimo acoplado a nivel de parche y de célula que respeta tanto el movimiento global de las poblaciones celulares como sus cambios estructurales internos (ej., heterogeneidad, expansión).
3. Incertidumbre del Destino Condicional: Desarrolla un proceso de destilación de dos etapas que separa las restricciones geométricas de la trayectoria de la incertidumbre del resultado, permitiendo el muestreo de múltiples trayectorias de destino plausibles sin observaciones futuras.
4. Marco de Análisis de Perturbaciones: Proporciona un mecanismo para distinguir entre los efectos de la célula focal y los efectos del vecindario (fondo) en las simulaciones de perturbación génica, revelando cómo el contexto local moldea las respuestas predichas.

Resultados Experimentales

El artículo evalúa LoCoFlow en tres conjuntos de datos: diferenciación de cuerpos embrioides (EB), reprogramación de iPSC, desarrollo de C. elegans y regeneración cerebral de Axolote.

• Reconstrucción de Datos No Vistos (Held-Out): En los conjuntos de datos de EB e iPSC, LoCo la achieve distancias de Wasserstein ($W_1, W_2$) más bajas en comparación con líneas base como OT-CFM, MFM y PACE en la mayoría de los entornos de alta dimensión (ej., PCA50). Sin embargo, el artículo señala que LoCoFlow a veces presenta una Discrepancia de la Media Máxima (MMD) más alta, lo que sugiere que las predicciones pueden estar más concentradas en regiones específicas del espacio de estados celulares.
• Predicción de Destino (Datos LARRY): Utilizando datos de linaje rastreado LARRY, LoCoFlow reduce el error bruto medio de Brier de 0.7990 (scDiffEq) a 0.6620, lo que indica una mejora en la precisión de la predicción de los vectores de probabilidad de destino completos.
• Respuestas a Perturbaciones:
  • Parches Geométricos: Al perturbar genes en parches definidos geométricamente, las intervenciones de solo fondo y de todos los miembros producen respuestas alineadas (ej., desplazando el destino de monocitos a neutrófilos).
  • Minibatches Aleatorios: En contraste, los modelos que utilizan parches ensamblados aleatoriamente muestran respuestas opuestas entre la célula focal y el fondo, resaltando la importancia del contexto geométrico.
• Interpolación Espacial (Axolote): LoCoFlow mejora la predicción espacial ($W_2$) para la regeneración cerebral de Axolote en comparación con OT-CFM, aunque los autores señalan que las distribuciones espaciales predichas siguen siendo visualmente similares al punto temporal de origen, lo que sugiere la necesidad de un alineamiento espacial adicional.

Significación y Reivindicaciones

El artículo afirma que LoCoFlow proporciona un marco robusto para estudiar cómo la estructura de la población celular local moldea las trayectorias inferidas, representando explícitamente la incertidumbre en el destino celular. Los autores enfatizan que:

• El Contexto Local es Crítico: Los distintos patrones de respuesta observados entre los parches geométricos y los minibatches aleatorios respaldan la hipótesis de que el contexto local es una representación útil para la inferencia probabilística del destino.
• Supuestos de Modelado: Los autores mantienen una postura modesta, reconociendo que los datos de instantáneas no identifican de forma única las historias celulares. Los componentes geométricos y latentes de LoCoFlow imponen supuestos de modelado en lugar de revelar relaciones de linaje reales.
• Hipótesis Testables: El marco genera hipótesis testables sobre cómo el contexto local moldea las respuestas predichas a las perturbaciones, distinguiendo entre los cambios intrínsecos de la célula y las influencias ambientales.

El trabajo no pretende resolver la ambigüedad fundamental de la inferencia de trayectorias, sino que ofrece un método que se alinea mejor con la geometría del colector de los estados celulares y proporciona una forma estructurada de cuantificar y muestrear la incertidumbre del destino.