Resumen Técnico: FLUX – Emparejamiento de Flujo Longitudinal Consciente de la Geometría con Mezcla de Expertos

1. Enunciado del Problema

Muchos sistemas biológicos, como los circuitos neuronales y las poblaciones celulares en diferenciación, evolucionan a través de dinámicas locales continuas mientras cambian entre regímenes latentes definidos por el aprendizaje, el contexto del estímulo o las etapas del desarrollo. Sin embargo, estos sistemas se observan con frecuencia solo como instantáneas longitudinales no emparejadas. En escenarios como la imagen de calcio de campo amplio o la secuenciación de ARN de células individuales, no es posible rastrear las mismas células o neuronas individuales a lo largo de los puntos temporales; en su lugar, los investigadores observan una secuencia de distribuciones marginales ($\mu_0, \mu_1, \dots, \mu_{T-1}$).

Esta configuración presenta dos desafíos acoplados:

1. Restricciones Geométricas: Las mediciones biológicas de alta dimensión a menudo se concentran cerca de variedades no lineales curvas de baja dimensión. Los interpolantes euclidianos estándar entre extremos no emparejados pueden pasar por regiones de baja densidad, violando las restricciones biofísicas y fallando en respetar la geometría intrínseca del espacio de estados.
2. Descubrimiento de Regímenes: El mecanismo de transporte subyacente (el "flujo") a menudo cambia con el tiempo. Por ejemplo, la actividad neuronal puede pasar de una expectativa general de recompensa por estímulo a una discriminación específica, o las células pueden transicionar de programas pluripotentes a comprometidos con una linaje. Estos son cambios en las dinámicas locales (regímenes), no meramente en agrupaciones espaciales. Los métodos existentes a menudo se centran en la precisión del transporte pero fallan en descomponer el campo de velocidad para identificar estos regímenes dinámicos latentes.

2. Metodología: FLUX

Los autores introducen FLUX (Emparejamiento de Flujo para datos longitudinales no emparejados con miZcla de Expertos), un marco que modela conjuntamente el transporte continuo y el descubrimiento de regímenes no supervisado. El método opera en tres etapas secuenciales:

A. Trayectorias Condicionales Conscientes de la Geometría

FLUX extiende el Emparejamiento de Flujo Métrico a configuraciones longitudinales. En lugar de utilizar trayectorias euclidianas rectas entre marginales adyacentes, FLUX aprende una métrica dependiente de los datos a partir de observaciones agrupadas (incluyendo muestras no etiquetadas).

• Aprendizaje de Métrica: Un backend de geometría (utilizando métricas de núcleo profundo RBF o RBF-MLP) estima la puntuación de distancia a la variedad, definiendo una métrica $G(x)$ que penaliza las trayectorias que abandonan la variedad de datos.
• Red de Doblado: Una red $B$ parametriza una trayectoria aproximada consciente de la geometría $\gamma_k(\alpha)$ entre los extremos $x_k$ y $x_{k+1}$. Esta trayectoria se dobla para permanecer en la variedad, proporcionando puntos de entrenamiento $z_{k,\alpha}$ y tangentes objetivo $\dot{z}_{k,\alpha}$ que respetan la geometría local.

B. Campos de Velocidad con Cambio de Régimen

FLUX descompone el campo de velocidad global $v_\theta$ en una Mezcla de Expertos (MoE).

• Campos Vectoriales de Expertos: El modelo aprende $M$ campos vectoriales de expertos dispersos $f_m$.
• Enrutador: Una red enrutadora $g(t, x)$ mapea el tiempo y el estado a logits de expertos.
• Mecanismo de Enrutamiento: Durante el entrenamiento, un mecanismo Gumbel-Softmax Directo selecciona un experto activo para cada muestra, permitiendo la optimización diferenciable de decisiones de enrutamiento discretas. En la inferencia, el régimen se determina mediante la asignación dura determinista (argmax).
• Objetivo: El campo de velocidad se entrena para coincidir con las tangentes de las trayectorias conscientes de la geometría. El enrutador se regulariza para fomentar asignaciones dispersas y temporalmente coherentes, evitando el colapso de expertos.

C. Pipeline de Entrenamiento

El entrenamiento procede en tres etapas para garantizar la estabilidad:

1. Aprendizaje de Métrica: Entrenar el backend de geometría en muestras agrupadas.
2. Entrenamiento de la Red de Doblado: Entrenar la red de doblado con la métrica congelada para generar interpolantes conscientes de la geometría.
3. Entrenamiento del Campo de Velocidad: Congelar las redes de geometría y doblado; entrenar los expertos de velocidad y el enrutador en las trayectorias y tangentes generadas.

3. Contribuciones Clave

1. Extensión a Marginales Longitudinales: Los autores extienden el emparejamiento de flujo consciente de la geometría a cadenas ordenadas de distribuciones marginales, enfatizando el papel del aprendizaje de variedades no supervisado cuando las etiquetas de trayectoria son escasas pero están disponibles observaciones no etiquetadas.
2. Descubrimiento de Regímenes mediante Descomposición de Velocidad: Introducen un campo de velocidad con cambio de régimen donde el enrutamiento de expertos dispersos descubre regímenes dinámicos latentes directamente desde la estructura de transporte local, en lugar de agrupar observaciones estáticas a posteriori.
3. Interacción Geometría-Regímenes: El trabajo demuestra que el aprendizaje geométrico es crítico para el descubrimiento de regímenes, particularmente cuando los regímenes se definen por cambios en las dinámicas locales en lugar de una separación espacial estática.
4. Validación Biológica: El marco se valida en dos configuraciones biológicas distintas:
   • Datos Neuronales: Separando el entrenamiento temprano/intermedio del entrenamiento tardío en el aprendizaje asociativo, coincidiendo con la divergencia conductual.
   • Biología del Desarrollo: Recuperando regímenes de evolución de expresión asociados con poblaciones celulares pluripotentes versus diferenciadas en cuerpos embrionarios.

4. Resultados

El artículo evalúa FLUX en cuatro conjuntos de datos: Stanford Bunny (variedad controlada), Atractor de Lorenz (dinámicas con cambio de régimen), Imagen de Calcio de Campo Amplio (aprendizaje neuronal) y Diferenciación de Cuerpos Embrionarios.

• Stanford Bunny: FLUX preserva el transporte superficial en incrustaciones de alta dimensión donde los métodos euclidianos fallan en permanecer en la variedad.
• Atractor de Lorenz: FLUX logra una recuperación perfecta de regímenes (ARI=1.0, NMI=1.0) al identificar el cambio entre regímenes caóticos y subcríticos. Crucialmente, FLUX sin Aprendizaje de Variedad (utilizando trayectorias euclidianas) colapsa a un solo experto, fallando en descubrir el régimen a pesar de tener la misma arquitectura MoE. Esto demuestra que las trayectorias conscientes de la geometría exponen las firmas de curvatura y velocidad necesarias para la especialización de expertos.
• Datos Neuronales: FLUX identifica una transición en las dinámicas corticales que se alinea con la emergencia de discriminación conductual (CS+ vs. CS−), superando al enrutamiento euclidiano y a las líneas base de cambio dinámico.
• Cuerpo Embrionario: Mientras que el enrutamiento euclidiano funciona bien aquí (porque las etapas son espacialmente separables), FLUX con aprendizaje geométrico permanece competitivo, reforzando que la geometría es más crítica cuando los regímenes se codifican en dinámicas locales.

Hallazgos de Ablación: El artículo afirma explícitamente que el enrutamiento de mezcla de expertos por sí solo es insuficiente. Sin aprendizaje geométrico, el modelo puede ajustar el transporte local pero falla en recuperar regímenes cuando esos regímenes se codifican en las dinámicas locales en lugar de en agrupaciones estáticas.

5. Significado y Afirmaciones

Los autores afirman que la descomposición de velocidad consciente de la geometría proporciona una estrategia general para descubrir transiciones de estados biológicos latentes a partir de instantáneas longitudinales no emparejadas.

• Impacto Científico: FLUX desplaza el enfoque de optimizar puramente la precisión del transporte (minimizando la distancia de Wasserstein) hacia la recuperación de estructuras dinámicas interpretables. Sugiere que los sistemas biológicos a menudo se reorganizan en regímenes dinámicos distintos durante procesos como el aprendizaje y la diferenciación.
• Limitaciones: Los autores reconocen que FLUX es computacionalmente más costoso que el emparejamiento de flujo euclidiano debido al entrenamiento de tres redes distintas. También señalan que el número de expertos se fija de antemano, y los regímenes inferidos deben interpretarse con cautela cuando las dinámicas locales son ruidosas o están débilmente muestreadas.
• Conclusión: El artículo concluye que FLUX ofrece un marco para el modelado dinámico interpretable de instantáneas biológicas no emparejadas, cerrando la brecha entre el aprendizaje de variedades geométricas y los sistemas dinámicos con cambio de régimen.