FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

Resumen Técnico: FreeBridge – Puentes de Schrödinger Variacionales para la Dinámica de Transición Celular

Planteamiento del Problema
Los ensayos de imagen de alto contenido son esenciales para cuantificar las respuestas celulares a perturbaciones químicas y genéticas. Sin embargo, existe una limitación fundamental: las células están fijadas químicamente en el momento de la adquisición de la imagen, lo que hace que las trayectorias temporales individuales sean inobservables. En consecuencia, el modelado de perturbaciones se reduce a la inferencia del transporte estocástico entre las poblaciones de control y tratadas, donde solo se conocen las distribuciones marginales finales.

Aunque los modelos generativos recientes (por ejemplo, modelos de difusión, flow matching) logran una fuerte alineación entre los puntos finales de control y perturbación, la consistencia de los puntos finales por sí sola es insuficiente para determinar la evolución intermedia. Infinitos procesos estocásticos pueden conectar marginales idénticos; sin restricciones adicionales, estos modelos pueden atravesar regiones del espacio latente no sustentadas por morfologías de células individuales observadas, lo que conduce a estados intermedios biológicamente implausibles. Los métodos existentes a menudo optimizan objetivos a nivel de distribución que dejan los estados celulares intermedios subdeterminados, pudiendo satisfacer métricas como FID o KID mientras generan artefactos.

Metodología: El Marco de Trabajo FreeBridge
Los autores proponen FreeBridge, un marco de transporte estocástico con restricciones formulado como un "Motor Celular" (Cell Engine) que separa explícitamente la definición de estados celulares válidos del mecanismo de transición. El método opera bajo supervisión de solo puntos finales y consta de tres componentes principales:

1. Unidad de Estado (Geometría de Célula Única Explícita):

   • En lugar de tratar las imágenes como mezclas a nivel de población, FreeBridge define estados atómicos como representaciones de célula única segmentadas por instancia (usando Cellpose).
   • Un codificador fijo y preentrenado $\phi$ incrusta estos recortes en un espacio latente $z = \phi(x)$. Este codificador permanece congelado durante el aprendizaje del transporte, estableciendo una geometría latente fija definida antes de aprender la regla de transición.
   • Se construye un banco de soporte empírico $Z_{emp}$ a partir de las incrustaciones latentes de las células observadas. Este banco define la "variedad" (manifold) de las morfologías biológicamente válidas.

2. Regla de Transición (Transporte de Puente de Schrödinger):

   • La transición entre la distribución de control $\mu_z$ y la distribución perturbada $\nu_{c,z}$ se modela como una ecuación diferencial estocástica (SDE) controlada: $dZ_t = u_t(Z_t) dt + \sigma dW_t$.
   • El objetivo es encontrar un campo de deriva $u_t$ que transporte la distribución inicial a la distribución objetivo en $t=1$.
   • A diferencia del transporte óptimo estándar, esta formulación incorpora un costo de estado específico de la tarea para regularizar la trayectoria.

3. Regularización Geométrica (Restricción de Soporte Empírico):

   • Para evitar que el proceso estocástico atraviese "vacíos" en el espacio latente, FreeBridge introduce un costo de soporte $V(z) = \lambda \cdot d_{emp}(z)^2$, donde $d_{emp}(z)$ es la distancia al vecino más cercano en el banco de soporte empírico $Z_{emp}$.
   • El objetivo total minimiza la suma de la energía cinética (magnitud de la deriva) y el costo de soporte:
     $$\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]$$
   • Esta regularización ancla los estados intermedios a la morfología empíricamente observada sin requerir trayectorias pareadas a nivel de célula.

Contribuciones Clave

• Formulación: Introduce una formulación de Puente de Schrödinger para el modelado de transiciones de célula única que desacopla explícitamente la especificación del estado del transporte estocástico.
• Fundamentación Geométrica: Propone un término de regularización novedoso basado en la estimación de densidad no paramétrica (distancia al vecino más cercano al banco de soporte empírico) para restringir las trayectorias intermedias dentro de regiones biológicamente plausibles.
• Protocolo Unificado: Establece un protocolo de evaluación unificado a través de múltiples conjuntos de datos (BBBC021, RxRx1, JUMP) utilizando recortes de célula única segmentados por instancia, asegurando una comparación justa entre métodos.
• Desempeño: Demuestra que la restricción de la geometría mejora no solo la fidelidad del punto final, sino también la retención semántica de los mecanismos de perturbación (Mecanismo de Acción o MoA).

Resultados
Evaluado en los conjuntos de datos BBBC021, RxRx1 y JUMP, FreeBridge demuestra lo siguiente:

• Fidelidad del Punto Final: Logra puntuaciones de Distancia de Incepción de Fréchet (FID) y Distancia de Incepción de Kernel (KID) competitivas o mejoradas en comparación con las líneas base de vanguardia (PhenDiff, IMPA, CellFlux). En BBBC021, alcanza el FID más bajo (15.35) y el KID más bajo (1.10).
• Consistencia Intermedia: Reduce significativamente la Tasa de Violación de Soporte ($R_{viol}$), una métrica que mide la proporción de estados intermedios que caen fuera del soporte morfológico empírico. En BBBC021, FreeBridge reduce $R_{viol}$ a 0.11, comparado con 0.38 para PhenDiff y 0.31 para CellFlux.
• Retención Semántica: Los estudios de ablación muestran que eliminar la regularización geométrica disminuye la precisión de la clasificación del Mecanismo de Acción (MoA) (de 72.08% a 65.08% en BBBC021), indicando que restringir la evolución intermedia ayuda a preservar las señales mecánicas.
• Análisis Cualitativo: Las visualizaciones confirman que FreeBridge preserva mejor las características morfológicas específicas de cada compuesto (por ejemplo, la compactación nuclear, la contracción citoplasmática) que las líneas base, que a menudo exhiben suavizado de textura o patrones incompletos.

Significado y Reivindicaciones
El artículo sostiene que el modelado de transiciones biológicas requiere más que la simple alineación de fronteras; requiere la fundamentación geométrica explícita de la evolución intermedia. Al anclar el transporte estocástico a una variedad celular empíricamente soportada, FreeBridge proporciona un enfoque con principios para generar dinámicas de perturbación biológicamente interpretables. Los autores afirman que esta formulación ofrece una base robusta para la simulación de "células virtuales", abordando la brecha crítica donde los modelos alineados con los puntos finales fallan en capturar estados celulares intermedios realistas. El trabajo destaca que, sin restricciones geométricas, los modelos generativos pueden producir puntos finales estadísticamente válidos que atraviesan regiones del espacio latente biológicamente no sustentadas.

Limitaciones
Los autores reconocen que FreeBridge depende de un codificador congelado; los factores morfológicos ausentes en el preentrenamiento del codificador no pueden recuperarse mediante el aprendizaje de transporte por sí solo. Además, el cálculo de la distancia de soporte escala cuadráticamente ($O(B^2)$) con el tamaño del lote (minibatch), lo que puede requerir una búsqueda de vecino más cercano aproximada para bancos de soporte más grandes.