Resumen Técnico: Flow Matching para Datos de Conteo (count-FM)

Enunciado del Problema

Los datos de conteo de alta dimensión, prevalentes en campos como la secuenciación de ARN de células individuales (scRNA-seq) y el análisis de trenes de picos neuronales, presentan desafíos únicos para el modelado generativo. Estos tipos de datos se caracterizan por su discreción, dispersión y estructuras de correlación complejas. Los enfoques generativos profundos existentes a menudo luchan en este dominio:

• Enfoques de Estado Categórico: Los métodos que tratan los conteos como variables categóricas (por ejemplo, tokens) requieren dimensiones de salida que escalan con el rango máximo de conteo. Esto conduce a una ineficiencia de parámetros y trata los conteos adyacentes como estados no relacionados, ignorando la geometría inherente del espacio de conteo.
• Transformaciones Continuas: Los enfoques que mapean los conteos a espacios continuos mediante descuantificación o representaciones latentes reemplazan la masa de probabilidad discreta con densidades continuas, potencialmente difuminando la estructura discreta subyacente.
• Limitaciones de los Modelos de Conteo Existentes: Aunque existen algunos modelos de salto específicos para conteos (por ejemplo, Poisson-JUMP), a menudo no están diseñados para el transporte entre distribuciones de conteo arbitrarias o carecen de las propiedades eficientes de entrenamiento y generación de los métodos basados en flujo modernos.

El artículo aborda la necesidad de un marco generativo que opere directamente en el espacio de conteo, respete la geometría discreta de los datos y permita un transporte eficiente entre distribuciones fuente y objetivo arbitrarias sin la sobrecarga de parámetros de los modelos de estado categórico.

Metodología: count-FM

Los autores proponen count-FM, un marco de flow matching adaptado para datos de conteo basado en un proceso de nacimiento-muerte de tiempo continuo con saltos unitarios locales.

1. Proceso de Nacimiento-Muerte de Tiempo Continuo

El modelo define la transición entre una distribución fuente $p_0$ y una distribución objetivo $p_1$ como un proceso de salto de Markov de tiempo continuo (CTMC) en el espacio de vectores de conteo $\mathbb{N}_0^d$. El proceso evoluciona mediante nacimientos unitarios locales ($+1$) y muertes ($-1$) para cada coordenada $i$:

• Tasa de Nacimiento: $\lambda_{t,i}(x)$
• Tasa de Muerte: $\mu_{t,i}(x)$
• Dinámica: En un intervalo pequeño $h$, la probabilidad de transitar del estado $x$ a $x + e_i$ es $h \lambda_{t,i}(x)$, y a $x - e_i$ es $h \mu_{t,i}(x)$. La muerte no se permite cuando un conteo es cero.

Esta parametrización es altamente eficiente: en lugar de predecir una distribución categórica completa sobre todos los niveles de conteo posibles (lo cual escala con el conteo máximo), el modelo solo predice dos tasas por dimensión ($2d$ salidas), independientemente del rango de conteo.

2. Puente Binomial Condicional y Entrenamiento

Para entrenar el modelo, los autores definen un puente condicional entre puntos fijos $(x_0, x_1)$. Para cada coordenada $i$, el puente se construye utilizando una distribución Binomial:
$$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$
Este puente asegura que la media condicional se mueva linealmente de $x_0$ a $x_1$ mientras mantiene estados intermedios de valor entero.

Al imponer la preservación de masa (la ecuación de Kolmogorov hacia adelante) en este puente condicional, los autores derivan expresiones en forma cerrada para las tasas de nacimiento y muerte condicionales:
$$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$$
donde $(\cdot)^+$ denota rectificación positiva.

3. Objetivo de Entrenamiento

El modelo se entrena minimizando una divergencia Kullback-Leibler (KL) en el espacio de trayectorias entre el proceso de puente condicional y el proceso aprendido del modelo. Esto se descompone en un objetivo de ajuste de tasas local:
$$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$$
donde $\ell(u, v) = v - u \log v$ es un término KL generalizado. Esto permite un entrenamiento libre de simulación de las tasas de transición marginales aprendiendo las tasas condicionales identificables.

4. Muestreo y Acoplamiento

• Muestreo: Las muestras se generan simulando el proceso de nacimiento-muerte aprendido hacia adelante desde $t=\epsilon_t$ hasta $t=1-\epsilon_t$ utilizando una discretización de salto local de primer orden.
• Acoplamiento de Puntos Finales: El artículo explora dos estrategias de acoplamiento para emparejar muestras fuente y objetivo:
  • Acoplamiento Independiente: Emparejamiento aleatorio ($p_0(x_0)p_1(x_1)$).
  • Acoplamiento de Transporte Óptimo (OT): Emparejamiento por mini-lotes usando un costo de Poisson simétrico. Esto tiende a inducir transiciones con menor curvatura, mejorando la interpretabilidad y la eficiencia de muestreo.

5. Generación Condicional

Para tareas condicionales, el modelo incorpora covariables $y$ y utiliza Guía sin Clasificador (CFG). Durante el entrenamiento, las variables de condicionamiento se eliminan aleatoriamente. En el muestreo, las tasas condicionales y no condicionales se combinan mediante una escala de guía $w$ para controlar la fuerza de la condición.

Contribuciones Clave

1. Marco count-FM: Un marco de flow matching para datos de conteo basado en dinámicas locales de nacimiento-muerte, permitiendo transporte directamente en el espacio de conteo con estados intermedios que permanecen con valores de conteo.
2. Esquema de Entrenamiento Tractable: Un método para aprender tasas de transición marginales mediante un puente binomial condicional, reteniendo la eficiencia del flow matching mientras evita la complejidad del modelado de estado categórico.
3. Eficiencia de Parámetros: Una formulación que predice solo tasas locales de nacimiento y muerte ($2d$ salidas) en lugar de distribuciones categóricas completas, haciéndola adecuada para datos de alta dimensión con grandes rangos de conteo.

Resultados Experimentales

1. Simulación (Mezcla Gamma-Poisson 2D)

• Rendimiento: count-FM y count-FM-OT lograron la mejor calidad de muestra (menor distancia de Wasserstein-2 y MMD$^2$) en comparación con líneas base que incluyen D3PM, discrete-FM, tauLDR, SEDD y Poisson-JUMP.
• Eficiencia: count-FM logró estos resultados con sustancialmente menos parámetros entrenables (2,372) en comparación con las líneas base de estado categórico (14,000–67,000+).
• Geometría: Los estados intermedios evolucionaron suavemente en el espacio de conteo. La versión acoplada con OT siguió una trayectoria de transición visiblemente más recta y alcanzó la calidad objetivo con menos evaluaciones de función (NFEs).

2. ARN de Células Individuales (Giro Dentado)

• Generación No Condicional: En un conjunto de datos de 2,930 células y 13,913 genes, count-FM superó a las líneas base de variables latentes (scVI) y modelos de difusión latente (scDiffusion, scLDM) en calidad de muestra (W2 y MMD$^2$) mientras utilizaba menos parámetros (aprox. 9.8M vs. 29M–39M para competidores).
• Transporte (P12 a P35): El modelo modeló con éxito el transporte de desarrollo desde el día postnatal 12 hasta el día 35. Utilizando un acoplamiento OT restringido por linaje, las rutas de transporte respetaron las estructuras de desarrollo conocidas (por ejemplo, la progresión de neuroblasto a célula granular), produciendo distribuciones intermedias interpretables que permanecieron en el espacio de conteo.

3. Trenes de Picos Neuronales (Hipocampo/Entorrinal)

• Generación Condicional: El modelo se aplicó para generar conteos de picos condicionados a la posición linealizada y la dirección de carrera.
• Comparación: count-FM con escala de guía $w=1$ superó a un regresor MLP determinista y a una línea base Poisson MLP.
• Estructura de Correlación: A diferencia del MLP Poisson, que subestimó las correlaciones poblacionales, count-FM preservó mejor la estructura de dependencia entre neuronas.
• Compensación de Guía: $w=1$ proporcionó el mejor equilibrio entre ajuste espacial y calibración distribucional. $w=2$ afiló los campos de lugar pero distorsionó la calibración de amplitud.

Significado y Afirmaciones

El artículo afirma que count-FM proporciona un marco efectivo para modelar datos de conteo de alta dimensión al permitir la generación y el transporte directamente en el espacio de conteo. Su principal significado radica en:

• Fidelidad Geométrica: Al modelar el transporte mediante saltos unitarios locales, el método respeta la geometría discreta de los datos de conteo, evitando los efectos de difuminado de las aproximaciones continuas.
• Escalabilidad: La eficiencia de parámetros lo hace viable para aplicaciones de alta dimensión (por ejemplo, miles de genes) donde los métodos de estado categórico se vuelven computacionalmente prohibitivos.
• Interpretabilidad: Los estados intermedios en la ruta de transporte son muestras de conteo válidas, permitiendo la inspección de distribuciones transicionales significativas en aplicaciones como la biología del desarrollo.
• Versatilidad: El marco soporta generación no condicional, transporte entre distribuciones arbitrarias y generación condicional con guía controlable.

Los autores reconocen limitaciones, señalando que el puente binomial condicional es relativamente simple y podría no explorar todas las rutas intermedias plausibles, y que la discretización de primer orden actual podría mejorarse con esquemas de simulación más precisos. Sin embargo, los resultados empíricos en conjuntos de datos de simulación y biomédicos demuestran que count-FM logra una alta calidad de muestra con una mayor eficiencia de modelado que los enfoques competidores.