Resumen Técnico: Flujo con Confianza (FMwC)

Enunciado del Problema

Los modelos generativos, particularmente los modelos de Flow Matching (FM), han logrado un rendimiento de vanguardia en dominios de lenguaje, visión y ciencias. Sin embargo, producen salidas (por ejemplo, citas alucinadas, figuras anatómicamente incorrectas o cristales termodinámicamente inestables) sin ninguna indicación nativa de fiabilidad. Las soluciones existentes para la cuantificación de la incertidumbre, como los Ensembles Profundos o MC-Dropout, dependen de ejecutar $k$ modelos separados o $k$ trayectorias estocásticas por muestra. Este enfoque incurre en un costo computacional de $k\times$, haciendo que la confianza por muestra sea prohibitivamente costosa para el despliegue a gran escala. Además, se ha demostrado que los métodos basados en verosimilitud son poco fiables para distinguir entradas dentro de la distribución de entradas fuera de la distribución en modelos generativos profundos. La pregunta central abordada es si los modelos de flow matching pueden aprender a cuantificar su propia incertidumbre para las salidas generadas sin aumentar los costos de inferencia ni requerir acceso a verosimilitudes.

Metodología: Flow Matching con Confianza (FMwC)

FMwC es un enfoque de paso único que equipa a los modelos de flow matching con una puntuación de confianza por muestra y a nivel de trayectoria al mismo costo computacional que el flow matching determinista estándar. El método integra el Dropout Adaptativo Variacional (VAD) con la propagación de momentos en forma cerrada a través de la trayectoria de la Ecuación Diferencial Ordinaria (ODE).

1. Formulación Variacional

En lugar de aprender un único campo de velocidades determinista $v_{t,\theta}(x_t)$, FMwC trata los pesos de la red $\omega$ como variables aleatorias extraídas de una posterior variacional $q_\psi(\omega)$. El modelo maximiza la Cota Inferior de la Evidencia (ELBO):
$$\text{ELBO}(\psi) = \mathbb{E}_{\omega \sim q_\psi} [\log \mathcal{L}(D | \omega)] - \text{KL}[q_\psi(\omega) \parallel \pi(\omega)]$$
donde el término de verosimilitud corresponde a la pérdida de regresión estándar de flow matching condicional. Crucialmente, el gradiente de la log-verosimilitud esperada bajo la posterior coincide con el gradiente de la pérdida de regresión estándar evaluada en muestras de pesos, permitiendo un entrenamiento sin simulación.

2. Posterior Dependiente del Estado mediante VAD

Para asegurar que la incertidumbre sea sensible a la geometría local del flujo, FMwC emplea el Dropout Adaptativo Variacional (VAD). En este marco, la posterior sobre los pesos en la capa $l$ es una Gaussiana $q_\psi(\omega_l | x, t) = \mathcal{N}(\theta_l, \alpha_l(x, t)^2 \theta_l^2)$, donde la escala de dropout $\alpha_l(x, t)$ es predicha por una pequeña red de inferencia $E_\gamma$. Esto permite que la dispersión de la posterior se amplíe en regiones de ambigüedad geométrica (donde el campo de velocidades es difícil de ajustar) y se estreche en otros lugares, aprendido únicamente mediante la ELBO.

3. Propagación Analítica de la Varianza

La innovación central de FMwC es propagar la varianza del campo de velocidades analíticamente a través de las capas de la red neuronal y la trayectoria de integración de la ODE en una sola pasada hacia adelante:

• A través de la Red: Utilizando el truco de reparametrización local, la media ($\mu_l$) y la varianza ($\sigma_l^2$) de las activaciones se propagan capa por capa. Para las capas lineales, la recursión de la varianza es en forma cerrada. Para las no linealidades, los momentos se calculan analíticamente utilizando cuadratura de Gauss-Hermite o linealización de Taylor, evitando la necesidad de muestreo Monte Carlo.
• A través de la ODE: En cada paso de Euler, la red devuelve la velocidad media de la posterior $\mu_t$ y su varianza $\sigma_t^2$. El integrador avanza la posición utilizando $\mu_t$, mientras que $\sigma_t^2$ se registra. Esto produce una trayectoria de varianza $\{\sigma_t^2\}_{t=0}^T$ en lugar de un único escalar.

4. Lecturas de Confianza

La trayectoria de varianza puede agregarse en puntuaciones de confianza escalares para tareas posteriores:

• Confianza del Punto Final: La varianza en el paso final ($t=1$).
• Confianza Integrada: La suma de las desviaciones estándar de la velocidad a lo largo de la trayectoria.
• Ratio Temporal: La relación entre la varianza de la ventana tardía y la de la ventana temprana, lo que ayuda a distinguir muestras que convergen (baja varianza tardía) de aquellas que permanecen ambiguas.
• Cabezas Aprendidas: Una pequeña cabeza supervisada puede mapear la secuencia completa de varianzas a una puntuación específica de la tarea.

Contribuciones Clave

1. Estimador de Confianza de Paso Único: FMwC es el primer método que proporciona una puntuación de confianza por muestra calibrada para flow matching al costo de una sola pasada determinista hacia adelante, igualando el presupuesto de inferencia del FM estándar. Logra esto propagando analíticamente el ruido aprendido dependiente de la entrada a través de la ODE.
2. Mejoras en la Inferencia sin Entrenamiento: La puntuación de confianza habilita aplicaciones prácticas sin reentrenamiento:
   • Filtrado: Mejora monótonamente la calidad de las muestras descartando salidas de baja confianza.
   • Edición: Identifica el instante exacto ($t^*$) donde el modelo se compromete con una estructura (pico de varianza), permitiendo ediciones dirigidas y respetuosas de las restricciones (por ejemplo, intercambios químicos en cristales) inyectando ruido en $t^*$.
   • Paso Adaptativo: Concentra los pasos de cálculo de la ODE en regiones de alta varianza (ambigüedad), recuperando muestras de alta calidad con presupuestos de pasos bajos donde el paso uniforme falla.
3. Interpretabilidad Geométrica: El artículo establece que la puntuación de confianza de FMwC se correlaciona con la divergencia local del campo de velocidades aprendido. Esto proporciona una ventana al proceso generativo, vinculando la incertidumbre epistémica con la inestabilidad geométrica (puntos de bifurcación) del flujo.

Resultados Experimentales

Los autores evaluaron FMwC en tres modalidades: estimación de densidad 2D (Checkerboard), generación de imágenes (MNIST) y generación de cristales inorgánicos de novo (Flow Matching Riemanniano).

• Calidad de la Muestra: FMwC preserva o mejora la calidad de la muestra en comparación con el FM estándar. En Checkerboard y MNIST, logró tasas de colocación incorrecta y puntuaciones FID más bajas, respectivamente, manteniendo una calidad comparable en cristales.
• Costo de Inferencia: FMwC incurre en una sobrecarga insignificante en backbones más grandes (por ejemplo, ~1.003$\times$ para MNIST, ~1.08$\times$ para Cristales) en comparación con la línea base determinista. Por el contrario, las líneas base de ensembles y MC-Dropout requieren $5\times$ el cómputo.
• Rendimiento de Filtrado: El filtrado basado en puntuaciones de FMwC logró puntuaciones del Área bajo la Curva de Precisión-Recall (AUPRC) comparables a ensembles de $k=5$ y MC-Dropout, pero al costo de una sola trayectoria. Por ejemplo, en el conjunto de datos Checkerboard, la puntuación de ratio temporal logró un AUPRC de 0.55, igualando la línea base de ensemble (0.55) con 1/5 del cómputo.
• Éxito en la Edición: Al editar cristales inyectando ruido en el pico de varianza ($t^*$), FMwC logró tasas de éxito significativamente más altas (por ejemplo, 39% frente a 22% para intercambios químicos) en comparación con la selección aleatoria de instantes, demostrando que las regiones de alta varianza corresponden a puntos de decisión críticos.
• Correlación con la Divergencia: La trayectoria de varianza $\sigma_t^2$ mostró una fuerte correlación con la magnitud de la divergencia del campo de velocidades ($|\nabla \cdot v_{t,\theta}|$). En el conjunto de datos Crystal, la correlación integrada alcanzó un Spearman $\rho$ de 0.96, validando la interpretación geométrica de la puntuación de confianza.

Significado y Afirmaciones

El artículo afirma que FMwC transforma el flow matching de un generador de muestras en un generador de muestras emparejadas con una medida de fiabilidad. Al aprender una posterior dependiente de la entrada sobre los pesos y propagar la varianza analíticamente, FMwC hace que la generación consciente de la confianza sea un comportamiento predeterminado práctico para aplicaciones científicas y críticas para la seguridad. Los autores enfatizan que la puntuación de confianza es una propiedad geométrica intrínseca del flujo único entrenado, distinta del desacuerdo inter-modelo medido por los ensembles. Esto permite una guía "quirúrgica" del proceso generativo, apuntando a momentos de inestabilidad sin los costos prohibitivos de los métodos existentes de cuantificación de la incertidumbre. El trabajo sugiere que comprender y aprovechar la geometría del campo de velocidades aprendido puede conducir a modelos generativos más fiables, interpretables y controlables.