Posterior Augmented Flow Matching

A continuación se presenta un resumen técnico detallado del artículo "Posterior-Augmented Flow Matching" (PAFM).

1. Planteamiento del Problema: Supervisión Escasa y Colapso del Flujo

El artículo aborda una limitación fundamental en el Flow Matching (FM), una clase de modelos generativos de tiempo continuo que aprenden un campo vectorial dependiente del tiempo para transportar muestras desde un prior simple (por ejemplo, ruido gaussiano) hasta una distribución de datos compleja.

• El Problema Central: En el FM estándar, el entrenamiento depende de una supervisión escasa, uno a uno. Para cualquier estado intermedio dado $z_t$ a lo largo de una trayectoria que conecta una muestra de ruido $\epsilon$ y un punto de datos objetivo $z$, el modelo es supervisado para predecir la velocidad hacia ese único objetivo específico $z$.
• Realidad de Alta Dimensión: En espacios de alta dimensión (como imágenes), un único punto intermedio $z_t$ podría residir válidamente en muchas trayectorias diferentes que conducen a diversas imágenes objetivo válidas que satisfacen la misma condición (por ejemplo, "una foto de un perro").
• Consecuencias:
  • Alta Varianza: El gradiente en cualquier $z_t$ específico se estima a partir de una única señal ruidosa.
  • Colapso del Flujo: El modelo tiende a memorizar emparejamientos específicos de origen-objetivo en lugar de aprender el mapeo continuo verdadero. Esto conduce a un "colapso del flujo", donde entradas diversas se mapean a salidas excesivamente similares o confusas, fallando en generalizar.

2. Metodología: Flow Matching Aumentado con Posterior (PAFM)

Los autores proponen PAFM, una generalización teóricamente fundamentada del FM que reemplaza la supervisión de objetivo único con una esperanza sobre la distribución posterior de completaciones de objetivos válidos.

Formulación Teórica

En lugar de minimizar la pérdida contra un único objetivo $z$, PAFM minimiza la pérdida esperada sobre un conjunto de objetivos plausibles $\{z_j\}$ extraídos de la distribución posterior $p(z_j | z_t, y)$, donde $y$ es la condición.

• Objetivo: La nueva función de pérdida agrega gradientes de múltiples objetivos candidatos, ponderados por su probabilidad posterior.
• Estimador Insesgado: Los autores demuestran (Teorema 1) que PAFM es un estimador insesgado del objetivo original de FM. En esperanza, $L_{PAFM} = L_{FM}$.
• Reducción de Varianza: Crucialmente, el Teorema 2 demuestra que al agregar información de $K > 1$ trayectorias plausibles por punto intermedio, PAFM acota inferiormente la varianza del gradiente en comparación con el FM estándar. El factor de reducción de la varianza está relacionado con el Tamaño de Muestra Efectivo (ESS) del esquema de muestreo.

Implementación Práctica

Dado que la posterior verdadera es intratable, PAFM factoriza el problema en dos aproximaciones utilizando Muestreo por Importancia:

1. Probabilidad del Intermedio: $p(z_t | z_j)$: ¿Qué tan plausible es que el objetivo $z_j$ haya generado el intermedio actual $z_t$? Esto se calcula utilizando la trayectoria de probabilidad condicional conocida (Gaussiana).
2. Probabilidad de la Condición: $p(y | z_j)$: ¿Qué tan bien coincide el objetivo $z_j$ con la condición $y$?
   • Condicionado por clase: Determinista (1 si las etiquetas coinciden, 0 en caso contrario).
   • Condicionado por texto: Aproximado utilizando modelos de visión-idioma (por ejemplo, CLIP/SigLIP) o asumido uniforme si los candidatos están pre-filtrados.

Algoritmo de Entrenamiento:
Para cada paso de entrenamiento, el modelo muestrea un intermedio $z_t$. En lugar de un objetivo, recupera un conjunto de $K$ objetivos candidatos $\{z_j\}$. El modelo predice la velocidad hacia un promedio ponderado de estos objetivos, donde los pesos son proporcionales a $p(z_t | z_j) \cdot p(y | z_j)$.

3. Estrategias de Selección de Objetivos

El artículo explora varias estrategias prácticas para construir el conjunto de candidatos $\{z_j\}$:

1. Vecinos Más Cercanos en el Espacio Latente (Principal): Utilizando FAISS para recuperar los $K$ vecinos más cercanos de la imagen de entrenamiento en el espacio latente. Esto asegura proximidad geométrica y similitud semántica.
2. Aumentaciones Aleatorias: Crear objetivos aplicando recortes aleatorios, redimensionamientos y rotaciones a la imagen fuente. Esto aumenta la diversidad espacial.
3. Muestreo de Momentos VAE: Muestrear $K$ veces desde la distribución gaussiana (momento) del codificado VAE de la imagen fuente. Esto proporciona pequeñas perturbaciones alrededor del codificado medio.

4. Contribuciones Clave

• Generalización Teórica: Se demuestra que PAFM es una generalización insesgada del FM que reduce la varianza del gradiente al agregar supervisión sobre toda la posterior de objetivos compatibles.
• Mitigación del Colapso del Flujo: Al proporcionar una señal de supervisión más densa, PAFM evita que el modelo memorice emparejamientos específicos, lo que conduce a campos vectoriales más robustos.
• Flexibilidad: El marco es agnóstico a cómo se seleccionan los candidatos, permitiendo la integración con diversas modalidades de datos y arquitecturas.
• Eficiencia: El método introduce una sobrecarga computacional negligible porque los conjuntos de candidatos pueden pre-calcularse (por ejemplo, mediante índices FAISS) y cargarse durante el entrenamiento.

5. Resultados Experimentales

Los autores evaluaron PAFM en diferentes escalas, arquitecturas y conjuntos de datos:

• Conjuntos de Datos:
  • ImageNet-1K (Condicionado por clase): Probado en modelos SiT-B/2 y SiT-XL/2.
  • CC12M (Texto a Imagen): Probado en la arquitectura MMDiT.
• Métricas: FID50K (Distancia de Inception Fréchet), sFID, Precisión, Recuperación.
• Mejoras de Rendimiento:
  • ImageNet: PAFM mejoró el FID en hasta 3.4 puntos (por ejemplo, SiT-XL/2 mejoró de 11.14 a 9.85 FID).
  • CC12M: PAFM redujo el FID de 10.37 (FM) a 9.45.
  • Consistencia: Se observaron mejoras en todas las escalas de modelos (SiT-B/2, SiT-XL/2) y arquitecturas (SiT, MMDiT).
• Varianza del Gradiente: El análisis empírico mostró que PAFM redujo la varianza del gradiente por mini-lote en aproximadamente 4 veces en comparación con FM, lo que llevó a dinámicas de entrenamiento más estables.
• Costo Computacional:
  • El rendimiento (throughput) disminuyó solo un 6.6%.
  • El uso máximo de memoria aumentó un 0.4%.
  • Los FLOPs por paso permanecieron sin cambios.

6. Significado

PAFM representa un paso significativo hacia adelante en el modelado generativo al abordar el cuello de botella de la "supervisión escasa" inherente a los modelos basados en flujo y difusión.

• Robustez: Mejora fundamentalmente la estabilidad del entrenamiento de modelos generativos de alta dimensión, reduciendo el riesgo de colapso de modos y artefactos.
• Eficiencia: Logra mejoras de rendimiento de vanguardia con computación adicional mínima, lo que lo convierte en un reemplazo práctico y directo para los objetivos estándar de Flow Matching.
• Generalizabilidad: El enfoque es aplicable tanto a la generación condicionada por clase como a la generación de texto a imagen, lo que sugiere que podría beneficiar a una amplia gama de modelos generativos continuos más allá de las arquitecturas específicas probadas.