Coevolving Representations in Joint Image-Feature Diffusion

1. El Problema

El modelado generativo de imágenes mediante modelos de difusión ha avanzado significativamente, especialmente al combinar latentes de baja nivel (de Autoencoders Variacionales, VAE) con características semánticas de alto nivel extraídas de codificadores visuales preentrenados. Sin embargo, los enfoques existentes presentan una limitación fundamental:

• Espacio de Representación Fijo: En los métodos actuales (como ReDi o REPA), el espacio de representación semántica se construye de forma independiente del objetivo generativo y permanece fijo durante todo el entrenamiento. Típicamente, se utiliza una proyección PCA predefinida o un autoencoder ligero para reducir la dimensionalidad de las características semánticas.
• Suboptimalidad: Dado que este espacio no se adapta a la tarea específica de síntesis de imágenes, puede no capturar la estructura semántica más relevante para la generación, limitando la calidad final de la muestra y la velocidad de convergencia.
• Pregunta Central: ¿Debe el espacio de representación que guía la difusión permanecer fijo, o debería adaptarse conjuntamente con el modelo generativo?

2. Metodología: CoReDi

Los autores proponen CoReDi (Coevolving Representation Diffusion), un marco donde el espacio de representación semántica evoluciona durante el entrenamiento junto con el modelo de difusión.

Componentes Clave del Enfoque:

1. Proyección Adaptable: En lugar de usar una proyección fija (como PCA), se introduce una capa lineal entrenable, $g_\phi(\cdot)$, que proyecta las características del codificador visual congelado a un espacio de menor dimensión. Esta proyección se optimiza junto con los parámetros del modelo de difusión.
2. Desafío de la Degeneración: Optimizar directamente la proyección junto con el objetivo de difusión lleva a soluciones degeneradas (colapso de características), donde el modelo minimiza la pérdida trivialmente alterando el objetivo en lugar de aprender una representación útil.

Tres Ingredientes para una Coevolución Estable:

Para evitar el colapso y lograr una evolución estable, CoReDi integra tres mecanismos críticos:

• Gradiente Detenido (Stop-Gradient): Se aplica un operador de gradiente detenido al objetivo limpio de la representación ($\tilde{\mathbf{z}}_0$) dentro de la función de pérdida de difusión. Esto evita que el modelo trivialmente modifique el objetivo para reducir la pérdida, forzando al modelo a aprender a denoising tanto la imagen como la representación de manera conjunta pero controlada.
• Normalización por Lotes (Batch Normalization): Se aplica inmediatamente después de la proyección aprendida. Esto estabiliza la escala de las características, preserva el programa de ruido previsto y actúa como un regularizador implícito para evitar que todas las muestras colapsen en un mismo punto de características.
• Regularización Explícita contra el Colapso: Incluso con los dos puntos anteriores, se observa que los canales de características pueden volverse redundantes. Se proponen tres estrategias de regularización para mantener la diversidad e información:
  • Regularización de Varianza de Características: Penaliza canales con baja varianza, asegurando que cada canal aporte información significativa.
  • Regularización de Ortogonalidad: Penaliza la desviación de la matriz de pesos de la proyección respecto a la identidad, fomentando direcciones de proyección distintas.
  • Regularización de Covarianza: Penaliza las entradas fuera de la diagonal en la matriz de covarianza de los canales, descorrelacionando las características.

Extensión al Espacio de Píxeles:

El marco se extiende más allá de los latentes de VAE hacia la difusión en espacio de píxeles (basado en la arquitectura DeCo). En este escenario, las características semánticas coevolucionan directamente con los píxeles crudos, eliminando el cuello de botella de reconstrucción del VAE y mejorando la fidelidad de la imagen.

3. Contribuciones Principales

1. Marco CoReDi: Introducción de un marco donde el espacio de representación semántica se aprende y evoluciona conjuntamente con el modelo de difusión, en lugar de ser estático.
2. Análisis de Estabilidad: Identificación y validación empírica de los tres ingredientes necesarios (Stop-Gradient, BN, Regularización) para evitar soluciones degeneradas en la optimización conjunta.
3. Mejora en Múltiples Escenarios: Demostración de que las representaciones coevolutivas mejoran la generación tanto en espacios latentes (VAE) como en espacios de píxeles, logrando una convergencia más rápida y mayor calidad de muestra.
4. Descubrimiento Estructural: Evidencia de que las representaciones aprendidas desarrollan una organización espacial más rica y estructurada a medida que avanza el entrenamiento, superando a las proyecciones PCA fijas.

4. Resultados Experimentales

Los experimentos se realizaron principalmente en el conjunto de datos ImageNet (256x256) utilizando arquitecturas SiT (Diffusion Transformer) y DeCo.

• Espacio Latente (VAE):

  • Convergencia Rápida: CoReDi alcanza un FID (Fréchet Inception Distance) de 3.3 con la arquitectura XL/2 en solo 2M de iteraciones, igualando el rendimiento de ReDi que requiere 4M de iteraciones (convergencia ~2x más rápida).
  • Calidad Superior: En la arquitectura B/2, CoReDi logra un FID de 16.4 en 400K iteraciones, superando significativamente a ReDi (21.4) y SiT base (33.0).
  • Generalización: Los beneficios se mantienen consistentes al usar diferentes codificadores visuales (DINOv2, MOCOv3, SigLIPv2, MAE), demostrando que la proyección aprendida extrae información más relevante para la generación que la PCA fija.

• Espacio de Píxeles:

  • CoReDi aplicado sobre DeCo-L/16 logra un FID de 21.5 en 200K iteraciones, superando al modelo base DeCo (31.3 en el mismo número de iteraciones) y demostrando una aceleración de convergencia de 2x.

• Análisis de Estabilidad:

  • Sin Stop-Gradient, el FID se degrada drásticamente (de 24.7 a 50.8).
  • Sin Batch Normalization, el entrenamiento colapsa completamente (FID > 220).
  • Sin Regularización, se observa colapso de características (FID 37.2), peor que la proyección PCA fija.

• Estructura Espacial:

  • Las métricas de auto-similitud espacial (LDS, CDS, RMSC) muestran que las representaciones de CoReDi desarrollan una estructura espacial coherente y organizada a medida que avanza el entrenamiento, superando a las proyecciones fijas.

5. Significado e Impacto

El trabajo CoReDi representa un cambio de paradigma en la integración de características semánticas en modelos de difusión. Al demostrar que el espacio de representación debe ser dinámico y adaptable en lugar de estático, el artículo:

• Mejora la eficiencia: Reduce el tiempo de entrenamiento necesario para alcanzar el estado del arte.
• Aumenta la fidelidad: Permite una síntesis de imágenes de mayor calidad al alinear la estructura semántica directamente con los objetivos de generación.
• Amplía el alcance: Muestra que estas técnicas son aplicables tanto a modelos latentes eficientes como a modelos de píxeles de alta fidelidad.
• Proporciona estabilidad: Ofrece un conjunto de herramientas (Stop-Gradient, BN, Regularización) que resuelven el problema fundamental de la optimización conjunta de representaciones, abriendo la puerta a futuras investigaciones en espacios de representación adaptativos.

En resumen, CoReDi establece que la "coevolución" de la representación y el generador es una estrategia superior a la simple alineación de representaciones predefinidas, logrando una síntesis de imágenes más robusta, rápida y estructuralmente coherente.