Resumen Técnico: MergeTok

Declaración del Problema

Los tokenizadores visuales actuales para la generación de imágenes están bifurcados en dos familias distintas, cada una de las cuales sufre limitaciones estructurales complementarias:

1. Tokenizadores basados en VAE continuos: Aunque ofrecen una reconstrucción de alta fidelidad, dependen de espacios latentes que carecen de una organización intrínseca. Bajo una fuerte compresión, los factores semánticos se entrelazan dentro de representaciones latentes densas, lo que resulta en códigos débilmente desentrelazados que dificultan el control semántico.
2. Tokenizadores basados en VQ discretos: Aunque permiten la generación autorregresiva, enfrentan cuellos de botella de optimización. La naturaleza no diferenciable de la cuantización conduce a actualizaciones de gradiente dispersas e irregulares, lo que provoca que muchas entradas del libro de códigos (codebook) reciban una supervisión insuficiente. Esto aumenta el riesgo de colapso del libro de códigos y reduce la eficiencia de la representación.

Estas limitaciones crean una tensión persistente entre la fidelidad de reconstrucción, la estructura semántica, la eficiencia y la estabilidad del entrenamiento.

Metodología: Marco de Trabajo MergeTok

MergeTok propone un tokenizador de doble rama unificado que optimiza conjuntamente la tokenización continua (VAE) y la discreta (VQ) dentro de una única arquitectura de codificador-decodificador. La innovación central es el uso de la Fusión de Tokens (ToMe - Token Merging) como un puente semántico entre las dos ramas.

Descripción de la Arquitectura

Dado un input de imagen $X$, un codificador CNN compartido $E_c$ produce un mapa de características $Z$. Este mapa es procesado por dos ramas paralelas:

1. Rama VAE (Continua):

   • Fusión de Tokens (Token Merging): Un codificador basado en atención equipado con módulos ToMe comprime la secuencia de tokens de $L$ tokens a $K$ tokens ($K < L$) mediante la fusión de tokens similares. Este proceso genera una representación condensada $Z^{(vae)}_K$ y un mapa de origen $S$ que registra la asignación de los tokens originales a los grupos fusionados.
   • Alineación Semántica: Los tokens fusionados se alinean con un modelo profesor de estilo DINO (también equipado con ToMe) utilizando una pérdida de distancia de coseno ($L_{align}$). Esto regulariza el espacio latente continuo hacia representaciones semánticamente conscientes y desentrelazadas.
   • Reconstrucción: Un decodificador híbrido utiliza tanto los tokens condensados como el mapa de origen $S$ para "desfusionar" los grupos y reconstruir la imagen en el espacio de píxeles.

2. Rama VQ (Discreta):

   • Cuantización Consciente de Grupos: Esta rama omite la operación ToMe durante el pase hacia adelante (forward pass) para preservar la longitud completa de la secuencia ($L$ tokens). Sin embargo, utiliza el mapa de origen $S$ generado por la rama VAE como un prior estructural.
   • Regularización: El mapa de origen particiona los $L$ tokens en $K$ grupos semánticos. Se aplican dos pérdidas conscientes de grupo a las asignaciones del libro de códigos:
     • Diversidad Intra-grupo ($L_{div}$): Maximiza la entropía dentro de cada grupo para prevenir el colapso hacia un solo código.
     • Exclusividad Inter-grupo ($L_{cons}$): Minimiza el solapamiento entre grupos para asegurar que diferentes grupos semánticos utilicen subconjuntos de códigos distintos.
   • Reconstrucción: Un decodificador VQ estándar reconstruye la imagen a partir de los tokens cuantizados.

Estrategia de Entrenamiento

• Objetivo Unificado: El modelo se entrena de extremo a extremo con una pérdida total que combina regularizadores de reconstrucción, adversariales y de fusión consciente:
  $$L_{total} = \lambda_{vae}L_{VAE} + \lambda_{vq}L_{VQ} + \lambda_{ali}L_{align} + \lambda_{div}L_{div} + \lambda_{cons}L_{cons}$$
• Muestreo Consciente de la Granularidad: Durante el entrenamiento, la ratio de fusión (y, por ende, el número de tokens mantenidos $K$) se muestrea dinámicamente de una distribución discreta centrada alrededor de un hiperparámetro $K^*$. Esto expone al codificador a variaciones en el ancho del cuello de botella, fomentando la robustez a través de distintas granularidades semánticas.
• Inferencia: En la inferencia, el tokenizador siempre produce una secuencia fija de 256 tokens. Los generadores descendentes (tanto autorregresivos como de difusión) operan sobre esta secuencia fija y son agnósticos a la estructura interna de fusión.

Contribuciones Clave

1. Tokenizador de Doble Rama Unificado: MergeTok es la primera arquitectura que optimiza conjuntamente los latentes continuos y discretos bajo un único objetivo, utilizando la fusión de tokens como una interfaz semántica para reforzar la estabilidad del entrenamiento de VAE y las ventajas estructurales de VQ.
2. Restricciones de Entrenamiento Conscientes de la Fusión: Los autores introducen dos objetivos novedosos derivados del proceso de fusión:
   • Alineación de tokens fusionados: Mejora la estructura semántica en el espacio latente continuo.
   • Cuantización consciente de grupos: Estabiliza el entrenamiento de VQ y mejora la utilización del libro de códigos mediante la diversidad intra-grupo y la exclusividad inter-grupo.
3. Muestreo de Ratio de Fusión Consciente de la Granularidad: Una estrategia de tiempo de entrenamiento que muestrea diferentes granularidades de tokens, permitiendo al modelo aprender una tokenización unificada que es robusta ante diversas tasas de compresión.

Resultados Experimentales

Los experimentos se realizaron en ImageNet-256 (resolución 256×256) con un presupuesto fijo de 256 tokens.

• Fidelidad de Reconstrucción:
  • Discreto (VQ): MergeTok-BL logró un rFID de 0.78, superando a bases sólidas como GigaTok-BL (0.81) y MergeVQ-GR (1.12). Con un discriminador DINO congelado, MergeTok-BL* alcanzó un rFID de 0.48.
  • Continuo (VAE): MergeTok-BL logró un rFID de 0.47, superando ligeramente a RAE (0.49) y venciendo significativamente a SoftVQ-L (0.61) y DC-AE-f32 (0.69).
• Calidad de la Representación:
  • MergeTok-BL alcanzó un 78.3% de precisión en el sondeo lineal (linear probing), el más alto entre los tokenizadores discretos listados, superando a UniTok (70.8%) y MergeVQ-GR (77.9%). Esto indica una organización semántica superior.
• Generación Descendente:
  • Autorregresiva (Rama VQ): Emparejado con LlamaGen-XXL, MergeTok-BL logró un gFID de 1.93 (sin CFG) y 2.13 (con CFG), superando a GigaTok-BL (2.03).
  • Difusión/Flujo (Rama VAE): Emparejado con SiT-XL/1, MergeTok-BL logró un gFID de 1.79 (sin CFG) y 1.29 (con CFG), igualando a las mejores entradas de la tabla.

Significado y Reivindicaciones

El artículo afirma que MergeTok demuestra que una única arquitectura puede dotar a los tokenizadores visuales tanto de una organización semántica robusta como de una discreción amigable para el generador.

• Unificación: Resuelve el compromiso entre la alta fidelidad de reconstrucción de los VAE y la estructura discreta requerida para la generación autorregresiva.
• Puente Semántico: Al aprovechar la fusión de tokens como un prior estructural, el modelo transfiere con éxito la información de agrupación semántica de la rama continua para estabilizar la rama discreta, y viceversa.
• Compatibilidad: El tokenizador produce secuencias de 256 tokens que son compatibles tanto con generadores autorregresivos (ej. LlamaGen) como de difusión/flujo (ej. DiT, SiT) sin requerir modificaciones en los generadores mismos.

Los autores posicionan a MergeTok no meramente como una herramienta de reconstrucción, sino como una interfaz unificada que soporta simultáneamente la generación autorregresiva discreta, la generación de difusión continua y el aprendizaje de representaciones semánticamente estructuradas.