Resumen Técnico: Autoencoders de Representación Laminada para la Difusión Eficiente (FlatDINO)

Planteamiento del Problema

Los avances recientes en la generación de imágenes han demostrado que los modelos de difusión pueden operar directamente sobre características de parches de aprendizaje auto-supervisado (SSL) (por ejemplo, de DINOv2) en lugar de latentes en el espacio de píxeles. Este enfoque, ejemplificado por RAE (Zheng et al., 2025), evita el compromiso entre reconstrucción y generación inherente a los VAE entrenados en píxeles, aprovechando la estructura semánticamente significativa ya codificada en las características de SSL.

Sin embargo, persiste un cuello de botella computacional significativo. Los codificadores SSL estándar como DINOv2 producen rejillas bidimensionales densas de embeddings de parches de alta dimensión (por ejemplo, 256 parches de 768 dimensiones para una imagen de 256×256). Esta representación es comparable en tamaño a la señal original, lo que no ofrece ninguna ventaja computacional sobre la difusión en el espacio de píxeles. La rejilla densa contiene una redundancia espacial significativa porque los parches vecinos comparten un contenido semántico sustancial, lo que hace que la difusión sobre estas características sea innecesariamente costosa en términos de FLOPs y memoria.

Metodología

Los autores proponen FlatDINO, un autoencoder variacional (VAE) diseñado para comprimir los densos embeddings de parches bidimensionales de DINOv2 en una secuencia compacta unidimensional de tokens continuos.

1. Arquitectura

• Codificador (Encoder): Un Vision Transformer (ViT-B) que toma los 256 embeddings de parches de un DINOv2 ViT-B/14 congelado (con registros) como entrada. El codificador antepone $T$ tokens de registro aprendibles a la secuencia de parches. Tras el procesamiento, solo se retienen estos $T$ tokens de registro como la representación latente, descartando la estructura espacial 2D original.
• Decodificador (Decoder): Un ViT-L que utiliza tokens de registro aprendibles como consultas (queries) para atender a la secuencia latente comprimida y reconstruir los 256 embeddings de parches originales de DINOv2.
• Objetivo de Entrenamiento: El modelo se entrena utilizando un objetivo de VAE estándar con una pérdida de error cuadrático medio (MSE) de reconstrucción y una penalización de divergencia KL. El peso de la KL ($\beta$) se normaliza por la dimensionalidad latente total ($T \times d$) para asegurar una presión de regularización constante entre diferentes configuraciones.
• Generación Latente: Un modelo generativo (específicamente un DiT-XL con flow matching) se entrena sobre el espacio latente comprimido de FlatDINO. El modelo aprende un campo de velocidad para transportar muestras desde un prior Gaussiano estándar hacia la distribución de los datos a lo largo de trayectorias rectas.

2. Decisiones de Diseño Clave

• Compresión 1D: Inspirado en TiTok (Yu et al., 2024a), el método comprime la rejilla 2D en una secuencia 1D. Los autores encontraron que 32 tokens proporcionan un equilibrio favorable entre la fidelidad de la reconstrucción y la longitud de la secuencia.
• Reducción de la Dimensión de Características: Mientras que el número de tokens se reduce de 256 a 32 (una reducción de 8×), la dimensión de características por token también se comprime (por ejemplo, de 768 a 128) para reducir aún más la dimensionalidad total.
• Flow Matching: El modelo generativo utiliza flow matching en lugar de difusión estándar, operando directamente sobre la secuencia latente 1D con embeddings posicionales aprendidos.

Contribuciones Clave

1. FlatDINO: El primer método que comprime características de parches auto-supervisadas en una representación latente unidimensional, descartando efectivamente la estructura espacial 2D original mientras preserva el contenido semántico.
2. Ganancias de Eficiencia: Demuestra que los modelos de difusión entrenados en latentes de FlatDINO logran una calidad de generación competitiva reduciendo significativamente los costos computacionales. Específicamente, la reducción de 8× en la longitud de la secuencia conduce a una reducción de 8× en los FLOPs por paso de propagación hacia adelante (forward pass) y hasta una reducción de 4.5× en los FLOPs por paso de entrenamiento en comparación con la difusión sobre características de DINOv2 sin comprimir.
3. Análisis de la Tokenización: A través de estudios de ablación de tokens, los autores muestran que el modelo aprende campos receptivos localizados y de tipo "mancha" (blob-like) para la mayoría de los tokens, particionando efectivamente la imagen espacialmente. Identifican una "transición de fase" donde la reducción de tokens por debajo de un cierto umbral (por ejemplo, de 32 a 16) obliga al modelo a adoptar estrategias de codificación menos eficientes (franjas horizontales), lo que provoca una caída abrupta en la calidad.

Resultados

Los experimentos se realizaron en ImageNet 256×256 utilizando una arquitectura DiT-XL.

• Calidad de Generación: Un DiT-XL entrenado en latentes de FlatDINO (32×128) logra un gFID de 1.85 con guía de clasificación libre (CFG) y 3.34 sin CFG. Esto es competitivo con RAE (que logra 1.41 gFID con CFG sobre características sin comprimir) y otros modelos de difusión latente de última generación, a pesar de la compresión significativa.
• Eficiencia Computacional:
  • Inferencia: FlatDINO requiere aproximadamente 8× menos FLOPs por paso de propagación hacia adelante en comparación con operar sobre parches completos de DINOv2.
  • Entrenamiento: Debido a la reducción en la longitud de la secuencia y la inclusión de la sobrecarga del codificador, FlatDINO requiere 4.5× menos FLOPs por paso de entrenamiento para modelos DiT-XL en comparación con la línea base de RAE.
• Reconstrucción: La configuración 32×128 logra un rFID de 0.77 (frente a 0.62 para la línea base de RAE sin comprimir), lo que indica una reconstrucción de alta fidelidad del espacio de características.
• Ablación de Tokens: El estudio revela que 32 tokens son suficientes para mantener campos receptivos localizados. Reducir a 16 tokens provoca un cambio cualitativo hacia una codificación de franjas horizontales y una caída significativa en la calidad de la reconstrucción.

Significado y Reivindicaciones

El artículo posiciona a FlatDINO como un paso hacia la síntesis de imágenes eficiente y con base semántica. Los autores afirman que, al explotar la redundancia espacial inherente en las rejillas de parches 2D, es posible lograr una reducción de 8× en el número de tokens sin sacrificar los beneficios de convergencia de operar en un espacio latente semánticamente organizado.

El trabajo sugiere que la difusión en representaciones semánticas comprimidas es un enfoque práctico para entornos con restricciones de recursos. Los autores reconocen que sus resultados actuales son preliminares y que el modelo no ha convergido completamente (entrenado por 600 épocas frente a 800 para RAE). Atribuyen la brecha restante en la calidad de la generación a una duración de entrenamiento insuficiente y a la necesidad de recetas de difusión específicamente adaptadas para latentes semánticos comprimidos. Se propone como trabajo futuro investigar esquemas de entrenamiento extendidos y autoencoders SSL optimizados para cerrar esta brecha.