Resumen Técnico: HeatKV - Compresión de Caché KV Ajustada por Cabeza para Modelado Autoregresivo Visual

Enunciado del Problema

Los modelos Autoregresivos Visuales (VAR) han surgido como una alternativa competitiva a los modelos de difusión para la generación de imágenes, ofreciendo velocidades de inferencia rápidas y propiedades de escalado favorables al generar imágenes como una secuencia de mapas de tokens multiescala. Sin embargo, estos modelos enfrentan restricciones de memoria severas debido a la naturaleza de su caché KV. A diferencia de los Modelos de Lenguaje Grande (LLM), que típicamente generan un token por paso, los modelos VAR generan mapas completos de tokens a resoluciones crecientes. En consecuencia, el número de tokens en la caché KV crece de forma superlineal con la escala de generación.

Los modelos VAR más avanzados, como Infinity-2B e Infinity-8B, pueden requerir hasta 10 GB de memoria de caché KV para generar una sola imagen de 1024 × 1024. Esta alta huella de memoria limita el tamaño de lote máximo y el número de usuarios concurrentes que un proveedor de servicios puede soportar en aceleradores de hardware especializados. Los métodos de compresión existentes para VARs, como ScaleKV y HACK, a menudo dependen de estrategias de granularidad gruesa, como clasificar las cabezas de atención o las capas en categorías binarias (por ejemplo, "presupuesto alto" frente a "presupuesto bajo"). Los autores argumentan que estas clasificaciones binarias son excesivamente restrictivas, ya que los patrones de atención dentro de la misma capa varían significativamente entre diferentes cabezas, lo que conduce a un desperdicio de memoria o a un podado subóptimo.

Metodología: HeatKV

Los autores proponen HeatKV (Compresión de Caché KV Ajustada por Cabeza), una estrategia novedosa que adapta la asignación de caché a la granularidad de las cabezas de atención individuales y escalas de imagen específicas. El método consta de tres componentes principales:

1. Importancia Cabeza-Escala (Calibración Offline)

A diferencia de los métodos anteriores que utilizan métricas sensibles a la varianza, HeatKV clasifica las cabezas de atención basándose en su Puntuación de Atención Acumulada Dependiente de la Escala (S-CAS).

• Mecanismo: Utilizando un pequeño conjunto de datos de calibración offline, el método agrega las probabilidades de atención desde los tokens en escalas futuras hasta los tokens en escalas anteriores.
• Métrica: Para una capa $\ell$ dada, una cabeza $h$ y una escala de origen $k$, el S-CAS mide la masa de atención promedio desde todas las escalas futuras ($\tau > k$) hacia los tokens en la escala $k$.
• Resultado: Esto genera una clasificación específica de pares "cabeza-escala" para cada escala de origen, indicando cuán crítico es el caché de una cabeza específica de una escala anterior específica para la generación de tokens futuros. Esto permite al sistema identificar que una cabeza específica podría necesitar retener la Escala 1 pero puede descartar la Escala 2, un nivel de granularidad que los métodos binarios no pueden lograr.

2. Asignación de Presupuesto

HeatKV calcula el número de cabezas ($N_k$) que deben podarse después de generar cada escala $k$ para asegurar que el tamaño total de la caché permanezca dentro de un presupuesto de memoria definido por el usuario $b$.

• Restricción: El método se adhiere al principio de "Streaming-LLM" de retener las primeras $s$ escalas como un "sumidero de atención" (esencial para la estabilidad) y nunca almacenar en caché la escala actual que se está generando.
• Cálculo: El algoritmo determina el número mínimo de cabezas a podar para satisfacer la restricción del presupuesto, permitiendo una reducción dinámica de la caché a medida que avanza la generación.

3. Optimización del Orden de Eliminación (Podado Temprano Codicioso)

Un desafío crítico en los VARs es que el tamaño de la caché crece de forma superlineal dentro de un solo paso de generación (a través de las capas del transformador). Los métodos estándar que podan solo entre escalas pueden causar sobreasignaciones de memoria transitorias que exceden los límites del hardware.

• Solución: HeatKV introduce un algoritmo de Podado Temprano Codicioso (Algoritmo 1). Hace cumplir la restricción del presupuesto después de cada capa del transformador, no solo entre escalas.
• Estrategia: Si el tamaño de la caché excede el presupuesto después de una capa específica, el algoritmo selecciona codiciosamente pares cabeza-escala adicionales para podar desde esa capa o capas anteriores. Prioriza el podado de cabezas con rangos S-CAS más bajos y, en caso de empate, prioriza el podado de capas posteriores o escalas de origen posteriores para maximizar el alivio de memoria.
• Optimalidad: Los autores demuestran (Proposición 1) que este enfoque codicioso produce un conjunto de cabezas podadas tempranamente de cardinalidad mínima, asegurando que el presupuesto nunca se viole mientras se maximiza la utilización de la memoria disponible.

Contribuciones Clave

1. Observación de Granularidad Fina: Los autores demuestran que los patrones de atención en los modelos VAR varían significativamente entre diferentes escalas dentro de la misma cabeza de atención, pero permanecen relativamente estables a través de diferentes indicaciones de entrada. Esto motiva un cambio desde la clasificación binaria de cabezas hacia una asignación de presupuesto continua y dependiente de la escala.
2. Métrica S-CAS: La introducción de la Puntuación de Atención Acumulada Dependiente de la Escala permite la construcción de un programa de podado estático adaptado a un presupuesto de memoria específico, permitiendo que cada cabeza descarte escalas no esenciales mientras retiene las críticas.
3. Algoritmo de Podado Optimizado: El desarrollo de un algoritmo de expulsión codicioso que hace cumplir las restricciones de presupuesto a nivel de capa (en lugar de solo a nivel de escala) previene desbordamientos de memoria transitorios, lo cual es crucial para el crecimiento superlineal característico de los VARs.
4. Rendimiento de Vanguardia: El método logra un nuevo estado del arte para la compresión de caché KV de modelos VAR, demostrando que la asignación de granularidad fina supera significativamente a los enfoques de granularidad gruesa.

Resultados Experimentales

Los autores evaluaron HeatKV en los modelos Infinity-2B e Infinity-8B utilizando el conjunto de validación MS-COCO 2017. Compararon HeatKV contra líneas base que incluyen StreamingLLM, SnapKV, HACK y ScaleKV en presupuestos de memoria del 4%, 10% y 20%.

• Fidelidad: HeatKV superó consistentemente a todas las líneas base. Notablemente, en el modelo Infinity-2B, HeatKV con un presupuesto del 10% logró puntuaciones FID, LPIPS y PSNR mejores que ScaleKV con un presupuesto del 20%. Con un presupuesto del 4%, HeatKV mantuvo un FID de 2.70 (Infinity-2B), significativamente más bajo que el siguiente mejor método (ScaleKV con 4.92).
• Alineación con la Indicación: Utilizando GenEval, HeatKV mantuvo una fuerte alineación con la indicación incluso con presupuestos bajos, con puntuaciones comparables a la línea base de caché completa.
• Calidad Perceptiva: Los resultados de la Puntuación de Preferencia Humana (HPSv2.1) indicaron que HeatKV preservó una alta calidad perceptiva, superando a otros métodos con un presupuesto del 4% en Infinity-2B.
• Eficiencia: HeatKV logró una relación de compresión 2× mayor en la asignación de memoria en comparación con los métodos existentes, manteniendo una fidelidad de imagen similar o mejor. El método también permitió tamaños de lote más grandes, reduciendo efectivamente la latencia de generación en comparación con la inferencia de caché completa.

Significado y Afirmaciones

El artículo afirma que HeatKV representa un avance significativo para hacer viables los modelos VAR para su implementación en hardware con capacidad de memoria limitada, como GPUs de nivel de consumo, o para servir a un mayor número de usuarios concurrentes en infraestructura existente.

Los autores enfatizan que su método no requiere cambios en la arquitectura del modelo ni en el proceso de entrenamiento; se basa en un pequeño conjunto de calibración offline (tan pocos como 1–10 indicaciones) para generar un programa de podado estático. Los resultados sugieren que la asignación de caché de granularidad fina, específica por cabeza es esencial para preservar la estructura semántica y la calidad de generación de los modelos VAR bajo restricciones de memoria agresivas, superando a los métodos que dependen de clasificaciones binarias de granularidad gruesa de cabezas o capas de atención.

El artículo concluye señalando que, aunque HeatKV optimiza la memoria, el trabajo futuro podría explorar combinarlo con otras técnicas de optimización, como la explotación de la dispersión de atención o la cuantización del modelo, para mejorar aún más la eficiencia.