Semantic-Aware Prefix Learning for Token-Efficient Image Generation

1. El Problema

Aunque los modelos generativos de imágenes han avanzado significativamente, la mayoría de los tokenizadores visuales existentes (que mapean imágenes de alta dimensión a representaciones latentes compactas) se entrenan principalmente con objetivos dominados por la reconstrucción. Esto genera dos limitaciones clave:

• Alineación Semántica Débil: Las representaciones latentes aprendidas a menudo están mal ancladas en conceptos de alto nivel (semántica), ya que los objetivos de reconstrucción tienden a priorizar detalles visuales de bajo nivel.
• Semántica como Regularización Auxiliar: Los enfoques recientes que intentan mejorar la alineación semántica suelen tratar las señales semánticas como una regularización auxiliar o un objetivo de alineación "suave", en lugar de hacerlas funcionalmente necesarias para el aprendizaje de la representación. Como resultado, la estructura global y los conceptos de alto nivel no se codifican de manera robusta en los tokens, lo que limita la eficiencia y el control en la generación posterior.

2. Metodología Propuesta

Los autores proponen dos componentes principales: SMAP (un tokenizador) y CARD (un generador híbrido).

A. SMAP: Tokenizador Consciente de Prefijos Semánticos

SMAP es un tokenizador basado en consultas (query-based) en 1D que integra la semántica directamente en la formación de tokens y la reconstrucción. Su diseño se basa en dos mecanismos innovadores:

1. Inyección de Semántica como Elementos de Secuencia:

   • A diferencia de métodos anteriores que usan encoders semánticos externos, SMAP inyecta incrustaciones condicionales (embedding) derivadas de etiquetas de clase directamente en la secuencia de tokens del codificador y decodificador.
   • Estas incrustaciones se colocan entre los parches visuales y las consultas latentes aprendibles, permitiendo que la información semántica interactúe mediante autoatención para moldear la formación de los tokens latentes.

2. Estrategia de Eliminación de Tokens Finales (Tail Token Dropping):

   • Para forzar a que la semántica sea indispensable, SMAP introduce un entrenamiento perturbado donde, en cada iteración, se elimina aleatoriamente una parte de la cola de la secuencia de tokens latentes (los tokens finales).
   • Se entrena con diferentes "presupuestos de tokens" (longitudes de prefijo $k$). En el caso extremo ($k=0$), el decodificador debe reconstruir la imagen utilizando únicamente las incrustaciones condicionales (semánticas) y los tokens de máscara.
   • Esto obliga a que la información semántica y los prefijos latentes tempranos absorban la responsabilidad de codificar la estructura global y la identidad de la categoría, mientras que los tokens restantes refinan los detalles específicos de la instancia.

B. CARD: Generador Híbrido Causal Autoregresivo-Difusión

Basado en el espacio latente estructurado por SMAP, se propone CARD para la generación de imágenes:

• Arquitectura Híbrida: Combina un módulo autoregresivo causal (Transformador) que modela las dependencias estructurales de alto nivel entre los tokens latentes, seguido de un módulo de refinamiento continuo basado en Difusión (Flow Matching).
• Compartición de Condiciones Semánticas: Una característica clave es que CARD reutiliza directamente las incrustaciones semánticas aprendidas durante el pre-entrenamiento de SMAP como entrada de condición para la generación. Esto elimina la necesidad de un encoder de clase adicional y asegura la consistencia entre el espacio de tokenización y el espacio de generación.

3. Contribuciones Clave

1. Identificación de una Limitación Central: Se señala que la alineación semántica actual es insuficiente porque no hace que la semántica sea un componente funcional necesario para la reconstrucción, sino solo una señal de alineación.
2. SMAP (Tokenizador): Un tokenizador 1D que trata la información semántica como invariantes preservadas en el prefijo. Mediante la inyección condicional y la eliminación de tokens, logra secuencias de tokens ordenadas por información y ancladas semánticamente, donde el prefijo contiene la identidad de la categoría y la estructura global.
3. CARD (Generador): Un modelo generativo híbrido que explota la estructura de ordenamiento de información de SMAP, demostrando que un tokenizado semánticamente fundamentado mejora tanto la reconstrucción como la generación downstream bajo presupuestos de tokens compactos.

4. Resultados Experimentales

Los experimentos se realizaron principalmente en el conjunto de datos ImageNet-1K (256x256):

• Calidad de Reconstrucción: SMAP supera consistentemente a los tokenizadores baselines (como TiTok y TA-TiTok) en métricas de FID de reconstrucción (rFID) e Inception Score (IS), tanto en configuraciones discretas (VQ) como continuas (KL, SoftVQ).
  • Ejemplo: Con 128 tokens y variante KL, SMAP logra un rFID de 0.75 frente a 1.02 de la versión sin truncamiento.
• Efecto de la Eliminación de Tokens: La estrategia de "tail token dropping" mejora significativamente la calidad de reconstrucción, especialmente con presupuestos de tokens bajos, demostrando que fuerza al modelo a aprender una representación semántica robusta.
• Generación de Imágenes:
  • El sistema SMAP + CARD logra un rendimiento competitivo con modelos mucho más grandes y que usan muchos más tokens.
  • Con un presupuesto de solo 128 tokens, el modelo SMAP(KL) + CARD (1.1B parámetros) alcanza un gFID de 1.85 (con CFG) e IS de 325.1.
  • Esto es superior o comparable a métodos que utilizan 256, 576 o incluso 4096 tokens, demostrando una eficiencia extrema.
• Análisis de Roles: Las visualizaciones confirman que las incrustaciones condicionales (C) por sí solas pueden reconstruir la estructura global y la identidad de la categoría, mientras que los tokens latentes (Z) añaden los detalles específicos de la instancia. La combinación cruzada de C y Z de diferentes imágenes funciona correctamente, validando la separación de roles.

5. Significado e Impacto

Este trabajo es significativo porque cambia el paradigma de cómo se aprenden las representaciones latentes para la generación de imágenes:

• De "Alineación" a "Necesidad Funcional": Demuestra que la semántica no debe ser solo una señal de regularización, sino un componente estructural obligatorio que soporta la reconstrucción incluso cuando los tokens visuales son escasos.
• Eficiencia de Tokens: Permite generar imágenes de alta calidad con un número muy reducido de tokens (128), lo que reduce drásticamente la carga computacional en la fase de generación autoregresiva.
• Coherencia End-to-End: Al compartir el espacio de incrustaciones semánticas entre el tokenizador y el generador, se logra una mejor alineación entre la compresión de la imagen y su síntesis, superando la brecha común entre el entrenamiento del tokenizador y el del generador.

En resumen, SMAP y CARD establecen un nuevo estándar para la tokenización eficiente, donde la comprensión semántica es el cimiento sobre el cual se construyen los detalles visuales, permitiendo una generación de imágenes más rápida, controlable y de mayor calidad.