UniCompress: Token Compression for Unified Vision-Language Understanding and Generation

1. El Problema

Los modelos unificados recientes (como UNITOK, VILA-U, VARGPT) buscan manejar tanto la comprensión como la generación de imágenes dentro de un único marco autoregresivo, codificando las imágenes en tokens discretos. Sin embargo, este enfoque presenta una ineficiencia crítica en el uso de tokens:

• Sobrecarga Computacional: Los tokenizadores estándar (basados en VQ-VAE o VQGAN) convierten una imagen de $512 \times 512$ en aproximadamente 1024 tokens ($32 \times 32$). Esto genera secuencias visuales muy largas que aumentan drásticamente el consumo de memoria, el costo de entrenamiento y la latencia de inferencia.
• Limitaciones de Despliegue: Esta ineficiencia impide el uso de estos modelos en sistemas con recursos limitados, como la IA encarnada (embodied AI).
• El Dilema de la Compresión: Las soluciones ingenuas, como el downsampling o la poda uniforme de tokens, funcionan bien para tareas de comprensión (reconocimiento), pero degradan severamente la generación de imágenes (más del 15% de pérdida de calidad). La generación requiere tokens espacialmente consistentes y de alta fidelidad para reconstruir detalles finos, lo cual es sensible a la eliminación de información.
• Costo de Reentrenamiento: Entrenar un nuevo tokenizador más eficiente desde cero suele requerir reentrenar todo el modelo de lenguaje (LLM) downstream, lo cual es prohibitivamente costoso.

2. Metodología: UNICOMPRESS

Los autores proponen UNICOMPRESS, un algoritmo de compresión de tokens "plug-and-play" (plug-in) diseñado para reducir la cantidad de tokens visuales (hasta 4x) sin sacrificar el rendimiento en comprensión ni en generación. La arquitectura se integra alrededor de un tokenizador discreto existente y consta de tres módulos ligeros:

1. Extracción de Tokens Globales (Global Token Extraction):

   • Utiliza un conjunto pequeño de tokens meta globales aprendibles ($N_g$, típicamente 4) que interactúan con todos los tokens de la imagen mediante una atención cruzada unidireccional.
   • Estos tokens capturan la semántica holística de la escena (relaciones entre objetos, layout general) sin alterar los tokens de imagen originales.

2. Compresión de Tokens (Image Token Compression):

   • Aplica un agrupamiento por promedio (Average Pooling) sobre la cuadrícula de tokens visuales en parches no superpuestos (ej. $2 \times 2$ o $4 \times 4$).
   • Esto reduce la secuencia visual densa a una secuencia comprimida más corta, preservando la estructura espacial gruesa.
   • La secuencia multimodal resultante se estructura como: [IMG_BOS] + Tokens Globales + [IMG_SEP] + Tokens Locales Comprimidos + [IMG_EOS].

3. Descompresión Autoregresiva Guiada (Autoregressive Decompression):

   • Este es el componente clave para la generación. Un decodificador autoregresivo reconstruye la cuadrícula de tokens densa original a partir de los tokens comprimidos y los tokens globales.
   • Los tokens globales actúan como anclas semánticas que guían al decodificador para refinar texturas y bordes locales, mitigando la pérdida de detalles que ocurre con la compresión uniforme.

Flujo de Entrenamiento (Dos Etapas):

1. Entrenamiento del Tokenizador: Se entrena el tokenizador junto con los módulos de compresión y descompresión para minimizar la pérdida de reconstrucción de la imagen. El LLM permanece congelado.
2. Ajuste Fino del LLM: Se congela el tokenizador comprimido y se realiza un ajuste fino ligero del modelo de lenguaje sobre los tokens comprimidos para ambas tareas (comprensión y generación).

3. Contribuciones Clave

• Identificación del Cuello de Botella: Demuestran que la eficiencia de tokens es un obstáculo fundamental para los modelos unificados y que la compresión ingenua perjudica desproporcionadamente la generación.
• Marco Modular Plug-in: Proponen UNICOMPRESS, un marco que se integra en modelos unificados existentes sin modificar la arquitectura del LLM ni requerir reentrenamiento completo.
• Mecanismo de Descompresión Guiada: Introducen el uso de tokens meta globales aprendibles combinados con un decodificador autoregresivo, lo que permite recuperar la fidelidad espacial necesaria para la generación.
• Resultados Empíricos Robustos: Validan el método en múltiples modelos unificados (UNITOK, VILA-U, VARGPT, UNIFORK, etc.), logrando reducciones significativas de tokens manteniendo el rendimiento.

4. Resultados Experimentales

Los experimentos se realizaron en una amplia gama de benchmarks de comprensión (GQA, MME, POPE, TextVQA, etc.) y generación (MJHQ-30K, medido por FID y CLIPScore).

• Reducción de Tokens: Logran una reducción de 4x en los tokens visuales (ej. de 256 a 64 tokens).
• Rendimiento en Comprensión: La caída de precisión es mínima (≤ 3 puntos porcentuales en la mayoría de los benchmarks). En algunos casos, el rendimiento se mantiene casi idéntico al modelo original.
• Rendimiento en Generación:
  • La degradación en calidad de imagen (FID) es moderada (ej. aumento de ~1.5 a 2 puntos en FID).
  • La alineación texto-imagen (CLIPScore) se mantiene competitiva.
  • Comparado con la poda ingenua, UNICOMPRESS evita la caída catastrófica (>15%) en generación.
• Eficiencia:
  • Latencia de Inferencia: Reducción de hasta un 41.8% en la latencia de generación.
  • Tiempo de Entrenamiento: Reducción de hasta un 15.4% en el tiempo de entrenamiento.
  • Ejemplo: En UNITOK, el tiempo de inferencia de generación bajó de 32.25 min a 18.96 min.

Estudios de Ablación:

• La compresión basada en pooling es superior a la selección de tokens (Top-k) o convoluciones estriadas para mantener la estructura espacial.
• Los tokens meta globales aprendibles son superiores a un token CLS o al promedio simple de la imagen, especialmente para la generación, ya que proporcionan una guía semántica más fuerte.
• La descompresión autoregresiva es esencial; sin ella, las imágenes reconstruidas presentan texturas suavizadas y artefactos.

5. Significado e Impacto

UNICOMPRESS representa un paso crucial hacia la viabilidad práctica de los modelos multimodales unificados en entornos con recursos limitados.

• Despliegue Real: Al reducir drásticamente la latencia y el costo computacional sin sacrificar significativamente la calidad, habilita el uso de estos modelos en dispositivos con restricciones de memoria y potencia (como robots o dispositivos móviles).
• Eficiencia Integral: A diferencia de métodos anteriores que solo optimizan el entrenamiento o la comprensión, UNICOMPRESS ofrece ganancias tanto en entrenamiento como en inferencia, y tanto para comprensión como para generación.
• Escalabilidad: Proporciona una ruta para escalar modelos unificados manteniendo la calidad, pero operando con presupuestos computacionales mucho menores, demostrando que las representaciones visuales compactas son suficientes para tareas complejas si se gestionan correctamente mediante mecanismos de reconstrucción guiada.