Resumen Técnico: GANs Escalables con Transformers (GAT)

1. Declaración del Problema

Aunque la escalabilidad ha impulsado avances recientes en el modelado generativo (particularmente en las familias autoregresivas y de difusión), las Redes Generativas Antagónicas (GANs) permanecen poco exploradas en este aspecto. Los intentos existentes de escalar GANs a menudo dependen de modelos únicos de alta capacidad con un ajuste extensivo específico de la tarea, en lugar de demostrar una escalabilidad genuina y sistemática. Además, el entrenamiento estándar de GANs enfrenta obstáculos específicos al escalar:

1. Subutilización de las Capas Tempranas: En generadores grandes basados en transformers, las capas tempranas a menudo se vuelven inactivas, contribuyendo marginalmente a la síntesis de imágenes.
2. Inestabilidad de la Optimización: Aumentar ingenuamente la profundidad y el ancho del modelo manteniendo los hiperparámetros (específicamente las tasas de aprendizaje) idénticos conduce a fallos de convergencia. Los modelos más grandes exhiben cambios de salida por paso más grandes bajo la misma tasa de aprendizaje, desestabilizando la dinámica antagónica sensible.

2. Metodología

Los autores proponen Transformadores Generativos Antagónicos (GAT), un marco que combina dos ingredientes clave para la escalabilidad: el entrenamiento en un espacio latente compacto de un Autoencoder Variacional (VAE) y el uso de arquitecturas puramente basadas en transformers tanto para el generador como para el discriminador.

2.1 Arquitectura

• Espacio Latente: El modelo opera en el espacio latente de un VAE preentrenado y congelado (específicamente SD-VAE), reduciendo la carga computacional mientras preserva la fidelidad perceptual.
• Generador: Una arquitectura pura de Vision Transformer (ViT). Toma un código latente $z$ y una condición $c$, los procesa a través de una pila de bloques de transformers (bloques GAT) y genera imágenes mediante una capa de desfragmentación (decodificador lineal). Emplea normalización adaptativa y LayerScale para la estabilidad.
• Discriminador: Una columna vertebral ViT con un token [cls] dedicado para la clasificación, que también utiliza LayerScale.

2.2 Innovaciones Clave para la Escalabilidad

Para abordar los modos de fallo específicos de escalar GANs, los autores introducen dos mecanismos principales:

A. Guía de Imagen Perturbada por Ruido Multinivel (MNG)
Para evitar que las capas tempranas se vuelvan inactivas, el generador se divide en $K$ etapas, cada una produciendo una salida intermedia $\hat{x}_k$.

• Jerarquía de Ruido: Cada salida intermedia se perturba con ruido gaussiano con una intensidad que disminuye monótonamente con la profundidad ($\alpha_1 < \alpha_2 < \dots < \alpha_K = 1$).
• Supervisión: Todas las salidas intermedias perturbadas se envían al discriminador. Esto obliga a las capas tempranas a aprender estructuras gruesas bajo ruido fuerte y a las capas posteriores a refinar detalles finos bajo ruido débil.
• Efecto: Esto fomenta una trayectoria de refinamiento de lo grueso a lo fino, asegurando que todas las capas contribuyan activamente al proceso de síntesis sin requerir una jerarquía de imágenes explícita (a diferencia de MSG-GAN).

B. Escalamiento de Tasa de Aprendizaje Consciente del Ancho
Para mantener dinámicas de entrenamiento estables a través de diferentes tamaños de modelo, los autores proponen una regla de escalamiento para la tasa de aprendizaje ($\eta$).

• Principio: A medida que aumenta el ancho del modelo (dimensión de canal $C$), la norma cuadrada esperada de las entradas crece linealmente, causando actualizaciones de salida por paso más grandes para una tasa de aprendizaje fija.
• Regla: La tasa de aprendizaje debe disminuir inversamente con la dimensión del canal para mantener constante la magnitud de las actualizaciones funcionales:
  $$\eta_{adapt} = \eta_{base} \cdot \frac{C_{base}}{C_{model}}$$
• Efecto: Esto previene la divergencia en modelos más grandes y elimina la necesidad de ajuste manual de hiperparámetros para cada escala.

C. Objetivos Adicionales

• Alineación de Representaciones (REPA): El discriminador se alinea con un Modelo de Fundación de Visión congelado (VFM, por ejemplo, DINOv2) utilizando una pérdida de similitud en datos reales. Esto fomenta que el discriminador aprenda características semánticamente ricas, proporcionando un mejor feedback antagónico.
• Pérdida Antagónica: El marco utiliza una pérdida de emparejamiento relativista con penalizaciones de gradiente de dos lados (R3GAN aproximado).

3. Resultados Clave

Los autores validan GAT a través de tamaños de modelo desde Pequeño (S) hasta Extra Grande (XL) en el conjunto de datos ImageNet-256.

• Rendimiento Estado del Arte: El modelo GAT-XL/2 logra un FID de 2.18 en ImageNet-256 en solo 60 épocas. Esto representa una mejora significativa sobre bases fuertes de 1-NFE (un paso) como MeanFlow-XL/2 (FID 3.43) y requiere 4 veces menos épocas que otras bases fuertes (por ejemplo, GigaGAN requiere 480 épocas).
• Escalabilidad: El rendimiento mejora monótonamente a medida que aumenta el tamaño del modelo (S $\to$ XL). Existe una fuerte correlación negativa (-0.95) entre el costo de inferencia (GFLOPs) y el FID, y una correlación similar para la computación total de entrenamiento, siguiendo una ley de potencia: $FID(C) \approx 3.52 \times 10^5 \cdot C^{-0.456}$.
• Eficiencia: GAT conserva la ventaja de inferencia de un solo paso de las GANs (1 NFE), ofreciendo una aceleración de aproximadamente 200 veces en el tiempo de inferencia en comparación con modelos de difusión de múltiples pasos (por ejemplo, DiT) con calidad comparable.
• Estudios de Ablación:
  • Eliminar MNG conduce a capas tempranas inactivas y rendimiento degradado.
  • Usar una tasa de aprendizaje fija a través de escalas causa divergencia en modelos más grandes o convergencia lenta en modelos más pequeños.
  • El objetivo REPA mejora significativamente el rendimiento, demostrando que las técnicas de modelos de difusión se transfieren efectivamente a las GANs.
• Generalización: El método escala a resoluciones más altas (ImageNet-512) y diferentes tokenizadores (FLUX-e2e), y soporta generación de texto a imagen (MS-COCO) y generación incondicional (FFHQ).

4. Significado y Afirmaciones

El artículo afirma que GAT establece exitosamente un marco de GAN escalable que rivaliza con el rendimiento de los modelos de difusión y autoregresivos modernos, manteniendo al mismo tiempo las ventajas únicas de las GANs:

1. Inferencia de Un Solo Paso: Generación de alta eficiencia sin denoising iterativo.
2. Manipulación Latente: La capacidad de realizar interpolación latente suave y edición semántica, que permanece coherente incluso a grandes escalas.
3. Eficiencia de Datos: Lograr resultados estado del arte con significativamente menos épocas de entrenamiento en comparación con las bases de difusión.

Los autores posicionan este trabajo como un paso hacia "escalar GANs", demostrando que con las elecciones arquitectónicas correctas (transformers puros) y estrategias de optimización (MNG y tasas de aprendizaje conscientes del ancho), el aprendizaje antagónico puede hacerse robusto y escalable, desafiando la noción de que las GANs son inherentemente limitadas en su potencial de escalado en comparación con los modelos de difusión.