MacTok: Robust Continuous Tokenization for Image Generation

Resumen Técnico: MacTok

1. El Problema: Colapso Posterior en Tokenizadores Continuos

Los modelos generativos modernos (difusión, autoregresivos) dependen de tokenizadores de imágenes para mapear píxeles a espacios latentes comprimidos. Existen dos enfoques principales:

• Discretos (VQ-VAE): Utilizan un código finito, pero sufren de errores de cuantización y flujo de gradiente limitado.
• Continuos (KL-VAE): Utilizan un espacio latente continuo regularizado por una divergencia KL hacia una prior gaussiana. Aunque ofrecen representaciones suaves, sufren de un problema crítico bajo compresión fuerte: el colapso posterior.

En el colapso posterior, la regularización KL fuerza excesivamente la distribución latente a acercarse a la prior isotrópica, haciendo que el codificador ignore la información de la imagen de entrada. Como resultado, el espacio latente pierde información semántica discriminativa, lo que lleva a una reconstrucción pobre y a una generación de baja fidelidad, especialmente cuando se utilizan pocos tokens (ej. 64 o 128).

2. Metodología: MacTok

Los autores proponen MacTok (Masked Augmenting 1D Continuous Tokenizer), un marco diseñado para prevenir el colapso posterior y aprender representaciones latentes robustas y semánticamente estructuradas mediante dos estrategias principales:

A. Enmascaramiento de Imagen (Image Masking)
Para forzar al modelo a preservar información crítica, MacTok entrena el tokenizador reconstruyendo imágenes completas a partir de entradas parciales (enmascaradas). Se utilizan dos esquemas complementarios con igual probabilidad:

1. Enmascaramiento Aleatorio: Se ocultan parches aleatorios de la imagen antes de la codificación. Esto obliga a los tokens latentes a inferir información global y estructural a partir del contexto visible.
2. Enmascaramiento Semántico Guiado por DINO: Utilizando características de un modelo preentrenado (DINOv2), se identifican los parches más informativos (basados en la similitud coseno con el token de clasificación) y se enmascaran selectivamente. Esto obliga al modelo a reconstruir regiones semánticamente críticas, transfiriendo priores semánticos del espacio de imagen al espacio latente.

Nota: El enmascaramiento en los tokens latentes (drop-out) solo retrasa el colapso temporalmente, mientras que el enmascaramiento en la imagen (input) estabiliza el aprendizaje de forma consistente.

B. Alineación de Representaciones (Global y Local)
Para asegurar que el espacio latente mantenga coherencia semántica, MacTok introduce una pérdida de alineación que conecta los tokens latentes con las características de DINOv2:

• Alineación Local: Los tokens latentes expandidos se alinean con los parches locales de DINOv2.
• Alineación Global: La representación latente global (promedio) se alinea con el token de clasificación de DINOv2.
  Esto crea un espacio latente bien estructurado donde conceptos semánticos similares se agrupan, mejorando tanto la reconstrucción como la generación.

Arquitectura:
MacTok utiliza una arquitectura basada en Vision Transformer (ViT) para el codificador y decodificador, procesando secuencias 1D de tokens latentes (flexibles en longitud) junto con tokens de imagen y máscaras.

3. Contribuciones Clave

1. Identificación y Solución del Colapso: Demostraron que el colapso posterior en tokenizadores continuos es una forma de sobre-regularización y que el flujo de información a través de la reconstrucción de entradas incompletas (enmascaramiento) estabiliza el aprendizaje latente.
2. MacTok: Introducción de un tokenizador continuo 1D que combina enmascaramiento aleatorio/semántico y alineación de representaciones para prevenir el colapso y preservar características discriminativas.
3. Eficiencia y Calidad: Logran un equilibrio superior entre fidelidad de reconstrucción y calidad generativa utilizando una compresión extrema (solo 64 o 128 tokens), superando a métodos discretos y continuos existentes.

4. Resultados Experimentales

Los experimentos se realizaron en el conjunto de datos ImageNet a resoluciones de 256×256 y 512×512, utilizando generadores como SiT-XL y LightningDiT.

• Rendimiento en Generación (gFID):

  • 256×256: Logró un gFID de 1.44 con 128 tokens (usando SiT-XL y CFG), superando a modelos de vanguardia.
  • 512×512: Logró un gFID de 1.52 con 128 tokens, estableciendo un nuevo estado del arte (SOTA) en esta resolución con tan pocos tokens.
  • Comparado con SoftVQ-VAE (que usa 64 tokens), MacTok mejora significativamente el gFID (ej. 1.52 vs 2.21 en 512×512).

• Eficiencia de Tokens:

  • MacTok reduce el uso de tokens en hasta 64 veces en comparación con métodos que requieren 4096 tokens (como MDTv2) o 1024 tokens, manteniendo una fidelidad competitiva.
  • Con solo 64 tokens, MacTok supera a modelos que requieren 256 o más tokens en métricas de generación.

• Reconstrucción:

  • Alcanza un rFID de 0.43 (128 tokens) y 0.75 (64 tokens) en 256×256, superando a tokenizadores basados en VQ que requieren muchos más tokens para resultados similares.

• Análisis del Espacio Latente:

  • Las visualizaciones (UMAP) muestran que MacTok evita el colapso, creando un espacio latente estructurado y separable, a diferencia de las distribuciones isotrópicas y sin información de los KL-VAE colapsados.
  • La precisión de "linear probing" confirma que las características latentes retienen mejor la información semántica.

5. Significado e Impacto

MacTok representa un avance significativo en la tokenización de imágenes para la generación:

• Viabilidad de la Compresión Extrema: Demuestra que es posible comprimir imágenes en muy pocos tokens (64-128) sin sacrificar la calidad, lo cual es crucial para reducir los costos computacionales en modelos autoregresivos y de difusión.
• Superación del Colapso: Ofrece una solución robusta al problema del colapso posterior en modelos continuos, eliminando la necesidad de un ajuste fino delicado de hiperparámetros (como el peso de la pérdida KL) que caracterizaba a los métodos anteriores.
• Integración Semántica: Al integrar priores semánticos (vía DINO) directamente en el proceso de tokenización mediante enmascaramiento y alineación, MacTok genera un espacio latente que es intrínsecamente más rico y útil para tareas posteriores, estableciendo un nuevo estándar para la tokenización continua de alta fidelidad.