Resumen Técnico: Compresión de Caché de Clave-Valor Consciente de Cabeza para la Generación Eficiente de Imágenes Autoregresivas

Enunciado del Problema

Los modelos de generación visual autoregresiva (AR), que representan las imágenes como secuencias de tokens visuales, han logrado resultados de alta resolución y fotorrealistas al escalar el tamaño del modelo y la cantidad de tokens. Sin embargo, este escalado introduce una sobrecarga sustancial de memoria y un rendimiento de inferencia bajo debido a la caché de Clave-Valor (KV) requerida durante la decodificación autoregresiva. Aunque trabajos recientes en procesamiento del lenguaje natural (PLN) han explorado la gestión de la caché KV (por ejemplo, selección de tokens, expulsión y cuantización), aplicar directamente estas técnicas al dominio visual es un desafío debido a la estructura espacial inherente de las imágenes.

Los métodos existentes específicos para visión, como LineAR, intentan reducir el uso de la caché reteniendo solo un subconjunto de tokens basándose en la observación de que los tokens atienden principalmente a vecinos espacialmente adyacentes. Sin embargo, estos enfoques asignan un presupuesto fijo y uniforme a cada cabeza de atención. Esta estrategia pasa por alto la heterogeneidad entre las cabezas de atención, lo que conduce a una asignación de memoria subóptima. Específicamente, algunas cabezas se centran en vecindades locales (requiriendo presupuestos pequeños), mientras que otras capturan dependencias de largo alcance, como estructuras de objetos y regiones de fondo (requiriendo presupuestos más grandes). La asignación uniforme no logra adaptarse a estos patrones distintos, lo que resulta en un desperdicio de memoria o en una calidad de generación degradada.

Metodología: HeadKV

Los autores proponen HeadKV, un marco novedoso y sin entrenamiento que realiza una compresión de caché KV consciente de la cabeza. La idea central es que las cabezas de atención en los modelos de generación de imágenes AR exhiben dos patrones distintos: Local y Global, y que el comportamiento de una cabeza es consistente a través de las posiciones de los tokens dentro de la misma capa para una entrada dada.

La tubería de HeadKV consta de tres pasos clave:

1. Llenado de Caché (Prefilling): El modelo codifica las entradas condicionales y genera una secuencia inicial de tokens visuales sin compresión. Esta etapa acumula suficiente contexto para estimar las distribuciones de atención por cabeza.
2. Agrupación de Cabezas (Clasificación): Una vez que la caché alcanza una longitud predefinida, el marco clasifica cada cabeza de atención como "Local" o "Global" utilizando un criterio de acumulación de atención basado en umbral.
   • Para cada cabeza, las puntuaciones de atención se acumulan hacia atrás desde la posición del token actual hasta alcanzar un umbral ($\tau = 0.9$).
   • Si el número de tokens necesarios para alcanzar el umbral es menor que un tamaño de ventana local predefinido, la cabeza se clasifica como Local. De lo contrario, se clasifica como Global.
   • Esta clasificación se realiza una vez durante la etapa temprana de generación y se reutiliza para todo el proceso de decodificación, sin requerir entrenamiento adicional ni estadísticas a nivel de conjunto de datos.
3. Compresión de Caché Consciente de la Cabeza: Se aplican diferentes estrategias de compresión según el tipo de cabeza:
   • Cabezas Locales: Estas cabezas atienden principalmente al contexto cercano. El marco asigna un presupuesto compacto, reteniendo solo los tokens más recientes mediante un mecanismo de ventana deslizante. Los tokens históricos se omiten por completo.
   • Cabezas Globales: Estas cabezas capturan dependencias de largo alcance. Para preservar tokens distantes informativos sin incurrir en un crecimiento lineal de la memoria, los autores diseñan una estrategia de Expulsión de Tokens Estratificada (STE).
     • Expulsión de Tokens Estratificada: Los tokens históricos candidatos se dividen en dos subconjuntos disjuntos: tokens de Corto Alcance y tokens de Largo Alcance, basándose en su distancia desde la posición actual.
     • Las puntuaciones de importancia se calculan independientemente para cada subconjunto utilizando características de consultas pasadas.
     • Los tokens con las puntuaciones más bajas se expulsan de cada subconjunto (por ejemplo, $M/2$ de corto alcance y $M/2$ de largo alcance). Esto evita que los tokens localmente dominantes eclipsen a los tokens distantes semánticamente relevantes, asegurando una preservación equilibrada tanto de los detalles locales como de la estructura global.

Contribuciones Clave

1. Análisis de la Heterogeneidad de la Atención: El artículo proporciona un análisis en profundidad de los modelos de generación de imágenes AR, identificando dos patrones de atención distintos (Local y Global) y revelando que los comportamientos de las cabezas son invariantes a través de las posiciones de los tokens dentro de una capa para una entrada fija.
2. Marco HeadKV: Un marco de compresión novedoso y sin entrenamiento que asigna presupuestos asimétricos entre las cabezas de atención. Asigna presupuestos compactos a las cabezas locales y presupuestos más grandes a las cabezas globales, guiado por un mecanismo de clasificación en etapa temprana.
3. Expulsión de Tokens Estratificada (STE): Una estrategia específica para cabezas globales que divide los tokens por distancia para asegurar la retención de dependencias de largo alcance informativas, abordando la limitación de la puntuación de importancia estándar que a menudo favorece los tokens locales.
4. Validación Empírica: Experimentos extensos demuestran que HeadKV reduce significativamente el uso de memoria de la caché KV y mejora el rendimiento de inferencia en múltiples modelos (Janus-Pro, Lumina-mGPT, LlamaGen-XL) mientras mantiene una alta calidad de imagen.

Resultados

Los autores evaluaron HeadKV en cuatro modelos base de generación de imágenes autoregresivas: Janus-Pro-1B, Janus-Pro-7B, Lumina-mGPT-768 y LlamaGen-XL.

• Calidad de Imagen: HeadKV logra un rendimiento comparable a la línea base de caché completa en benchmarks como GenEval y DPG. Por ejemplo, en Janus-Pro-7B con una relación de compresión ($\rho$) de 1/8, HeadKV mantiene una puntuación de GenEval de 0.99 (vs. 0.99 para Completo) y una puntuación de DPG de 83.78 (vs. 84.20 para Completo).
• Comparación con LineAR: En comparación con el método de asignación uniforme LineAR, HeadKV muestra una degradación menor del rendimiento a medida que aumentan las relaciones de compresión (por ejemplo, en $\rho=1/8$). En LlamaGen-XL, HeadKV logra una puntuación FID-50K mejor (2.75) en comparación con LineAR (2.95) bajo compresión extrema.
• Eficiencia: HeadKV reduce significativamente el uso de memoria y mejora el rendimiento.
  • En LlamaGen-XL ($\rho=1/6$), HeadKV reduce la memoria en un 63.60% y logra una aceleración de 3.00× en comparación con la configuración completa, superando a LineAR (2.93×).
  • En Janus-Pro-7B ($\rho=1/6$), logra una aceleración de 3.68× con un ahorro de memoria del 52.08%.
  • En Lumina-mGPT-768, aunque la latencia es ligeramente superior a la de LineAR debido a sobrecargas de implementación en la agrupación, aún logra una aceleración de 1.35× con un ahorro de memoria significativo.

Significado y Afirmaciones

El artículo afirma que HeadKV ofrece una solución práctica al cuello de botella de memoria en la generación de imágenes autoregresiva sin requerir entrenamiento adicional ni supervisión externa. Al reconocer y explotar la heterogeneidad inherente de las cabezas de atención, el método va más allá de la asignación de caché "talla única" de trabajos anteriores.

Los autores enfatizan que su enfoque es generalizable a través de diferentes entradas y arquitecturas de modelos porque se basa en patrones de atención intrínsecos observados durante las etapas tempranas de generación. La estrategia de Expulsión de Tokens Estratificada se destaca como un componente crítico para mantener la calidad de las imágenes generadas bajo presupuestos de memoria ajustados, específicamente al preservar la información estructural de largo alcance que a menudo se pierde en los métodos de expulsión estándar. El trabajo sugiere que la asignación de recursos adaptativa y específica por cabeza es un paradigma más eficiente para la generación visual que la retención uniforme de tokens.