Resumen Técnico: Una Pasa No Basta: Refinamiento Latente Recursivo para Modelos Generativos

Declaración del Problema

A pesar de los avances significativos en la generación de imágenes, el campo enfrenta un cuello de botella crítico en la evaluación. La métrica dominante, la Distancia Fréchet Inception (FID), confunde la fidelidad de las muestras con la cobertura de modos, a menudo saturándose en valores bajos de un solo dígito. Esta saturación enmascara un problema persistente: los modelos pueden lograr puntuaciones FID bajas generando un pequeño conjunto de imágenes nítidas y casi idénticas (colapso de modos) mientras fallan en cubrir la distribución completa de los datos.

Los paradigmas generativos actuales exhiben una compensación:

• Modelos de Difusión y de Ajuste de Flujo: Logran alta fidelidad pero a menudo requieren muestreo iterativo (múltiples evaluaciones de funciones neuronales). Las variantes destiladas de un solo paso aceleran la inferencia pero sacrifican consistentemente la diversidad y la cobertura. Además, los solucionadores de difusión cubren sistemáticamente de manera insuficiente las regiones de baja densidad de la distribución de datos.
• Redes Generativas Antagónicas (GANs): Ofrecen inferencia de un solo paso pero son propensas al colapso de modos, concentrándose en modos de alta fidelidad mientras descartan otros.
• Estimación de Máxima Verosimilitud Implícita (IMLE): Garantiza la cobertura de modos por construcción (asegurando que cada imagen de entrenamiento tenga una muestra generada cercana) y soporta inferencia de un solo paso. Sin embargo, los métodos basados en IMLE han quedado históricamente rezagados en la calidad de las muestras. Esta limitación se debe parcialmente a la arquitectura del generador: la red de mapeo en los modelos IMLE basados en StyleGAN es un Perceptrón Multicapa (MLP) estándar procesado en una sola pasada hacia adelante. Esto obliga a la red a determinar todos los aspectos del código de estilo (identidad, estructura, textura, detalle fino) simultáneamente, haciendo que la salida sea altamente sensible a las imprecisiones en el mapeo latente inicial.

Metodología: Mapeador de Tokens Recursivo (RTM)

Los autores proponen reemplazar la red de mapeo MLP de un solo paso en los generadores basados en estilo con un Mapeador de Tokens Recursivo (RTM). RTM adapta la arquitectura del "Modelo Recursivo Pequeño" (TRM) al entorno generativo, permitiendo el refinamiento iterativo del espacio latente sin aumentar la cantidad de parámetros.

Arquitectura Central

• Estructura Recursiva: En lugar de una pila profunda de capas independientes, RTM utiliza un bloque compartido $f$ aplicado iterativamente. El proceso involucra dos niveles de recursión:
  • Ciclos Internos ($L$): Un estado interno $Z_L$ se actualiza $L$ veces por paso, condicionado por un estado externo y el vector de ruido original.
  • Pasos de Refinamiento ($H$): El estado externo $Z_H$ se actualiza $H$ veces.
• Mecanismo: El mapeador produce primero una estimación gruesa del vector de estilo $w$ y la corrige progresivamente a lo largo de múltiples ciclos. Los ciclos tempranos establecen estructuras gruesas (identidad, pose), mientras que los ciclos posteriores refinan la textura y el color.
• Eficiencia de Parámetros: Debido a que el mismo bloque $f$ se reutiliza, aumentar $H$ o $L$ incrementa la profundidad computacional efectiva sin añadir parámetros.
• Optimización de Gradiente Corto: Para gestionar la memoria durante el entrenamiento, solo el paso final de refinamiento se diferencia a través de; los pasos intermedios se desconectan. Esto preserva los beneficios representacionales de la recursión profunda mientras mantiene los costos de memoria de entrenamiento comparables a un único bloque de alimentación directa.
• Integración: RTM está diseñado como un reemplazo "plug-and-play" para la red de mapeo tanto en generadores basados en IMLE (específicamente RS-IMLE) como en StyleGANs entrenados adversarialmente (StyleGAN2 y StyleGAN2-ADA). La arquitectura del decodificador permanece sin cambios.

Contribuciones Clave

1. Mapeador de Tokens Recursivo (RTM): Una arquitectura de red de mapeo novedosa que reemplaza el MLP estándar de un solo paso en los generadores basados en estilo con un proceso de refinamiento iterativo.
2. Integración con RS-IMLE: Cuando se integra con IMLE de Muestreo por Rechazo (RS-IMLE), RTM mejora la FID, la Precisión y el Recall en nueve benchmarks de pocos ejemplos, CIFAR-10 incondicional y CelebA-HQ (256×256), manteniendo la inferencia de un solo paso y el mapeo directo de latente a imagen de IMLE.
3. Generalización al Entrenamiento Adversarial: RTM mejora StyleGAN2 y StyleGAN2-ADA en CIFAR-10 y AFHQ-v1 (512×512), demostrando que los beneficios del refinamiento recursivo no son específicos del objetivo de entrenamiento IMLE.
4. Flexibilidad de Inferencia: El número de pasos de refinamiento ($H$) actúa como un hiperparámetro computacional. Los modelos pueden evaluarse con diferentes valores de $H$ en el momento de la inferencia sin reentrenamiento, permitiendo una compensación entre el costo de inferencia y la fidelidad.

Resultados Experimentales

El artículo evalúa RTM utilizando Precisión, Recall, FID y Puntuación Inception (IS) en múltiples conjuntos de datos.

• CIFAR-10 (32×32): RTM integrado con RS-IMLE logra la Precisión (0.896) y el Recall (0.773) más altos entre todos los métodos listados, incluidos los baselines de difusión y ajuste de flujo. Reduce la FID en un 30% en comparación con el baseline RS-IMLE estándar (3.97 vs. 5.69) mientras mantiene la inferencia de un solo paso.
• CelebA-HQ (256×256): RTM reduce la FID en un 27% sobre el baseline RS-IMLE (10.67 vs. 15.43) y mejora simultáneamente la Precisión (0.952) y el Recall (0.592).
• Benchmarks de Pocos Ejemplos: En nueve conjuntos de datos estándar de pocos ejemplos (por ejemplo, Obama, Perro, Anime), RTM reduce aproximadamente a la mitad la FID promedio de la reproducción RS-IMLE correspondiente (13.94 vs. 25.2).
• StyleGAN2/ADA: En CIFAR-10, RTM reduce la FID de 3.88 a 3.55 para StyleGAN2 y mantiene la FID en 2.31 para StyleGAN2-ADA mientras mejora la Precisión y la Cobertura. En AFHQ-v1, RTM mejora la FID (4.79 vs. 4.99) y el Recall (0.565 vs. 0.507).
• Profundidad vs. Recursión: Un estudio de ablación en CIFAR-10 muestra que simplemente aumentar la profundidad de un MLP no recursivo (a 32 capas) mejora la Precisión pero degrada severamente el Recall (0.550 vs. 0.664 del baseline). RTM logra una mejor FID con 13 veces menos parámetros que el MLP emparejado en profundidad, lo que sugiere que la recursión proporciona un sesgo inductivo superior para la fidelidad distribucional sin sacrificar la cobertura.

Significado y Afirmaciones

El artículo argumenta que el objetivo de la generación de imágenes debería cambiar de simplemente reducir la FID (que está cerca de la saturación) a mejorar explícitamente la diversidad y la cobertura (Recall) junto con la fidelidad (Precisión).

• Más allá de la FID: Los autores afirman que la FID por sí sola es insuficiente porque recompensa muestras nítidas y concentradas que pueden ignorar grandes porciones de la distribución de datos. La Precisión y el Recall son complementos esenciales.
• Mejora Simultánea: A diferencia de los baselines de ajuste de flujo que a menudo intercambian cobertura por FID, o de los modelos de difusión destilados que pierden diversidad en la inferencia de un solo paso, RTM mejora simultáneamente la calidad y la diversidad.
• Perspectiva Arquitectónica: El artículo postula que el refinamiento es un cálculo cualitativamente diferente a una sola pasada hacia adelante. Al permitir que el mapeador revise decisiones tempranas (estructura gruesa) a la luz de cálculos posteriores (detalles finos), RTM supera la sensibilidad del decodificador a las imprecisiones del código latente.
• Modestia: Los autores reconocen limitaciones, señalando que los experimentos se restringen a conjuntos de datos donde la búsqueda de vecinos más cercanos de IMLE es computacionalmente viable (excluyendo ImageNet). No afirman haber resuelto la generación de imágenes, sino haber proporcionado un método que equilibra mejor la compensación fidelidad-diversidad dentro de los marcos de evaluación actuales.