Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

Resumen Técnico: Flujo de Emparejamiento Multi-Marginal con Interpolantes Aprendidos Adversarialmente

Planteamiento del Problema
El artículo aborda el desafío de aprender la dinámica dependiente del tiempo de un sistema dados únicamente observaciones muestreadas (instantáneas o snapshots) en puntos temporales discretos, sin acceso a las trayectorias reales (ground-truth). Esta es una tarea central en campos como la biología (por ejemplo, la secuenciación de ARN de célula única, el seguimiento de células) y la medicina. Formalmente, el objetivo es inferir un campo vectorial dependiente del tiempo $v_t$ que gobierna una Ecuación Diferencial Ordinaria (EDO) $dx_t = v_t(x_t)dt$, de tal manera que el mapa de integración traslada una distribución marginal inicial $q_0$ a una distribución objetivo $q_1$, satisfaciendo al mismo tiempo las restricciones marginales intermedias $p_{t_i} = q_{t_i}$ para un conjunto de pasos temporales $0 = t_1 < \dots < t_K = 1$.

Los enfoques existentes, como la aplicación del Flow Matching (FM) estándar a pares consecutivos de marginales o el uso de interpolantes lineales por tramos o de spline cúbico, suelen dar lugar a curvas de interpolación no suaves. Estos problemas de falta de suavidad provocan una alta varianza en el gradiente del objetivo de aprendizaje y dificultades para marginalizar los campos vectoriales resultantes, particularmente en entornos de alta dimensión o cuando los datos son ruidosos.

Metodología: ALI-CFM
Los autores proponen ALI-CFM (Adversarially Learnt Interpolants with Conditional Flow Matching), un método de dos etapas que desacopla el aprendizaje de los interpolantes del aprendizaje del campo vectorial.

1. Aprendizaje Adversarial de Interpolantes (ALI):
   En lugar de forzar a los interpolantes a pasar por puntos de datos específicos (emparejamiento puntual), ALI aprende un interpolante neuronal $G_\phi(x_0, x_1, t)$ que coincide con la distribución de los marginales intermedios. El interpolante se parametriza como una interpolación lineal más un término de corrección neuronal:
   $$G_\phi(x_0, x_1, t) = (1-t)x_0 + tx_1 + t(1-t)f_\phi(x_0, x_1, t)$$
   Para asegurar que la distribución de la transformación directa (pushforward) del interpolante coincida con los marginales intermedios observados $q_{t_i}$, los autores emplean un objetivo adversarial inspirado en las GAN. Un discriminador $D_\gamma$ distingue entre muestras de la marginal verdadera $q_{t_i}$ y muestras generadas por el interpolante $G_\phi$. El generador (el interpolante) se entrena para minimizar la divergencia de Jensen-Shannon entre estas distribuciones.

   Para garantizar la unicidad de los interpolantes y evitar trayectorias arbitrariamente curvas, los autores introducen regularizadores:

   • Referencia Lineal: Penaliza la desviación de una línea recta entre los puntos finales acoplados.
   • Referencia Lineal por Tramos: Realiza una regresión del interpolante hacia una trayectoria lineal por tramos definida por acoplamientos de Transporte Óptimo (OT) de cadena de Markov.
   • Norma de la Segunda Derivada: Impone suavidad penalizando la curvatura del interpolante.

2. Marginalización mediante Flow Matching Condicional (CFM):
   Una vez entrenados los interpolantes $G_\phi$, estos se fijan. Una segunda red neuronal $u_\theta$ se entrena para aproximar el campo vectorial que genera estos interpolantes. Esto se realiza utilizando el objetivo estándar de Conditional Flow Matching, donde el campo vectorial objetivo es la derivada temporal del interpolante aprendido:
   $$L_{CFM} = \mathbb{E}_{t, (x_0, x_1) \sim \pi} \| u_\theta(G_\phi(x_0, x_1, t)) - \frac{d}{dt}G_\phi(x_0, x_1, t) \|^2$$
   El acoplamiento $\pi$ entre $q_0$ y $q_1$ puede ser independiente o un plan de Transporte Óptimo.

Contribuciones Clave

1. Nueva Técnica de Interpolación: El artículo introduce un método para aprender interpolantes utilizando un objetivo adversarial que coincide con distribuciones en lugar de puntos. Esto permite aprender interpolantes suaves que son únicos bajo supuestos leves (probados en los apéndices).
   2.os Robustez ante el Ruido y la Complejidad: El método está diseñado para manejar conjuntos de datos ruidosos y geometrías complejas que varían en el tiempo, donde los métodos de interpolación puntual (como los splines o lineales por tramos) fallan o se vuelven inestables.
2. Escalabilidad y Versatilidad: Los autores demuestran que ALI-CFM escala a datos de alta dimensión y a un gran número de distribuciones marginales ($K$), superando a las líneas base existentes en escenarios complejos específicos, mientras se mantiene competitivo en otros.

Resultados Experimentales
Los autores evalúan ALI-CFM en cuatro tareas distintas:

• Distribución de "Nudo" Sintética: En un conjunto de datos sintéticos con 1,200 distribuciones marginales que forman un nudo, ALI-CFM fue el único método capaz de capturar con precisión la geometría dependiente del tiempo. Los métodos competidores (OT-MFM, splines cúbicos, lineales por tramos) produjeron interpolantes no suaves con alta varianza de gradiente.
• Seguimiento Celular (Glioblastoma): Utilizando datos de microscopía de una célula en movimiento, ALI-CFM aprendió con éxito la dinámica espacio-temporal a partir de datos ruidosos y submuestreados. En contraste, las trayectorias de OT-CFM divergieron, y OT-MFM, al ser independiente del tiempo, no logró capturar las dependencias temporales.
• Inferencia de Trayectorias de Célula Única (scRNA-seq): En los conjuntos de datos Embryoid Body (EB), Cite-seq y Multiome, ALI-CFM se desempeñó a la par de las líneas base de vanguardia (OT-MFM, OT-CFM) en términos de la Distancia del Movimiento de la Tierra (EMD). Los autores señalan que, aunque el entrenamiento adversarial dificulta superar estrictamente a los métodos que se ajustan excesivamente (overfit) a los datos puntuales en estas métricas específicas, el método sigue siendo altamente competitivo.
• Inferencia de Coordenadas de Tumor (Transcriptómica Espacial): En un conjunto de datos de cáncer de mama con secciones de tejido ordenadas espacialmente, ALI-CFM superó significativamente a todas las líneas base (incluyendo OT-CFM, OT-MMFM y OT-MFM) en la predicción de la distribución de las regiones tumorales en secciones no utilizadas para el entrenamiento (held-out). Los autores atribuyen este éxito a la capacidad del método para manejar distribuciones multimodales y los desafíos de alinear las secciones de tejido.

Significado y Reivindicaciones
El artículo afirma que ALI-CFM ofrece un enfoque conceptualmente novedoso para el flow matching multi-marginal al pasar del emparejamiento puntual estricto al emparejamiento de distribuciones. Este cambio permite el aprendizaje de interpolantes suaves y únicos que son robustos al ruido y capaces de modelar dinámicas complejas donde los métodos anteriores fallan. Los autores destacan que su método es particularmente efectivo en escenarios con cientos o miles de distribuciones marginales y datos ruidosos, un régimen donde ningún otro algoritmo de flow matching ha demostrado un rendimiento comparable. Si bien el método está a la par con los estados del arte existentes en los bancos de pruebas estándar de célula única, su rendimiento superior en la transcriptómica espacial y en las tareas de geometría sintética compleja subraya su versatilidad y escalabilidad.