Adapting Self-Supervised Representations as a Latent Space for Efficient Generation

1. El Problema

Los modelos generativos modernos, como los Modelos de Difusión Latente (LDM) y los modelos basados en flujo (Flow Matching), han demostrado una gran capacidad para sintetizar imágenes de alta calidad. Sin embargo, enfrentan dos desafíos principales:

• Costo Computacional: Estos modelos operan en espacios latentes que aún mantienen una estructura de cuadrícula 2D (por ejemplo, 32x32 tokens), lo que no explota completamente las redundancias espaciales inherentes a las imágenes naturales. Esto requiere arquitecturas complejas con mecanismos de atención (como Transformers) que son costosos de entrenar e inferir.
• Compromiso entre Semántica y Reconstrucción: Los enfoques que utilizan representaciones auto-supervisadas (SSL) pre-entrenadas (como DINO o CLIP) a menudo se centran en la alineación de características semánticas, pero carecen de los detalles de bajo nivel necesarios para una reconstrucción fiel de píxeles. Por otro lado, los tokenizadores discretos (como VQGAN) o secuencias largas (como TiTok) pueden ser ineficientes o perder la suavidad geométrica del espacio latente.

2. Metodología: RepTok

Los autores proponen RepTok (Representation Tokenizer), un marco que representa una imagen completa utilizando un único token latente continuo derivado de un transformador de visión auto-supervisado (SSL).

Componentes Clave del Pipeline:

1. Codificador SSL Ajustado (Fine-tuned SSL Encoder):

   • Se utiliza un encoder SSL pre-entrenado (ej. DINOv2, MAE, CLIP) que permanece mayormente congelado.
   • Adaptación Crítica: Solo se ajusta (fine-tune) el embedding del token de clase ([cls]). Este token, que originalmente agrupa información semántica de alto nivel, se enriquece con detalles de bajo nivel necesarios para la reconstrucción.
   • El resto de los pesos del encoder se mantienen fijos para preservar la geometría suave y semánticamente estructurada del espacio latente original.

2. Decodificador Generativo:

   • Un decodificador ligero (en los experimentos principales, un MLP-Mixer sin mecanismos de atención) se entrena para mapear el token latente único de vuelta al espacio de píxeles (o a un espacio latente de baja dimensión pre-entrenado).
   • Se utiliza Flow Matching como objetivo de entrenamiento estándar para la generación y reconstrucción.

3. Pérdida de Similitud Coseno (Cosine-Similarity Loss):

   • Para evitar que el token ajustado se desvíe demasiado del espacio latente original (lo cual arruinaría la capacidad de generación), se introduce una pérdida de regularización:
     $$\mathcal{L}_{cos}(x) = \lambda (1 - \cos(z, z_{frozen}))$$
   • Donde $z$ es el token ajustado y $z_{frozen}$ es la salida del encoder congelado. El hiperparámetro $\lambda$ controla el equilibrio entre la fidelidad de la reconstrucción (necesita detalles nuevos) y la preservación de la estructura geométrica del espacio SSL (necesario para la generación).

4. Generación en un Solo Token:

   • Dado que la imagen se comprime en un solo vector continuo, el modelo generativo no necesita modelar interacciones token-a-token. Esto permite eliminar la atención (self-attention) y utilizar arquitecturas puramente basadas en MLP (como MLP-Mixer), reduciendo drásticamente la complejidad computacional.

3. Contribuciones Clave

• Token Único Continuo: Demostraron que un solo token continuo, derivado de un encoder SSL finamente ajustado, es suficiente para codificar tanto la semántica de alto nivel como los detalles espaciales necesarios para una reconstrucción fiel, eliminando la redundancia espacial de las cuadrículas 2D.
• Eficiencia Computacional Extrema: Al eliminar la atención y reducir el espacio latente a un solo token, el costo de entrenamiento se reduce en más de un 90% en comparación con baselines basados en Transformers (como DiT o SiT), manteniendo un rendimiento competitivo.
• Preservación de la Geometría SSL: Mediante la pérdida de similitud coseno, logran mantener la suavidad del espacio latente pre-entrenado, lo que facilita la interpolación y la generación estable sin necesidad de pérdidas adversarias o perceptuales adicionales.
• Escalabilidad a Texto-Imagen (T2I): El marco se extiende naturalmente a la síntesis texto-imagen mediante la incorporación de atención cruzada para las condiciones de texto, logrando un rendimiento zero-shot competitivo en MS-COCO con un presupuesto de entrenamiento extremadamente limitado (<20 horas en 4 GPUs A100).

4. Resultados

• Generación Condicionada por Clase (ImageNet-1K):
  • RepTok alcanza un gFID de 1.88 (con CFG) en ImageNet 256x256.
  • Logra este rendimiento con un costo total de entrenamiento de ~34.5 K PFLOPs, en comparación con los ~212.5 K PFLOPs requeridos por SiT-XL/2 para resultados similares.
  • Reduce el costo computacional por iteración en un orden de magnitud (de ~118 GFLOPs/iter a ~23 GFLOPs/iter).
• Reconstrucción:
  • Obtiene un rFID de 1.85 y un PSNR de 14.94, superando a métodos que utilizan solo códigos semánticos puros (como RCG) y siendo competitivo con tokenizadores discretos más complejos.
• Interpolación en el Espacio Latente:
  • Las interpolaciones lineales entre tokens muestran transiciones suaves tanto en semántica como en detalles espaciales (forma, tamaño, rotación), validando la preservación de la geometría del espacio SSL.
• Síntesis Texto-Imagen:
  • En MS-COCO, el modelo alcanza un FID competitivo en un entorno zero-shot tras solo 200k iteraciones, demostrando una alta eficiencia de muestreo y entrenamiento.

5. Significado e Impacto

El trabajo RepTok representa un cambio de paradigma en el modelado generativo latente:

1. Simplificación Arquitectónica: Demuestra que la complejidad de los modelos generativos (atención, múltiples tokens) puede trasladarse casi en su totalidad a la etapa de compresión pre-entrenada, permitiendo que el generador sea extremadamente simple y eficiente.
2. Uso Potenciado de SSL: Transforma los modelos SSL de meros extractores de características o guías de entrenamiento en codificadores latentes completos y funcionales para la generación.
3. Sostenibilidad: Al reducir el costo de entrenamiento y la huella de carbono asociada a la generación de imágenes, RepTok ofrece una ruta viable para desarrollar modelos generativos de alta calidad con recursos limitados, haciendo la investigación más accesible y sostenible.

En resumen, RepTok valida que un espacio latente compacto, suave y semánticamente rico, derivado de un solo token SSL ajustado, es suficiente para lograr una generación de imágenes de alta fidelidad con una eficiencia sin precedentes.