HAWK: Head Importance-Aware Visual Token Pruning in Multimodal Models

1. El Problema

Los Modelos de Lenguaje Multimodal Grandes (MLLMs) han avanzado rápidamente, pero enfrentan un cuello de botella computacional crítico: la explosión de tokens visuales.

• Carga Computacional: Las entradas visuales (imágenes de alta resolución o videos largos) se convierten en cientos o miles de tokens, superando con creces a los tokens de texto. Dado que la complejidad de los modelos basados en atención crece cuadráticamente con la longitud de los tokens, esto incrementa drásticamente el tiempo de inferencia y el uso de memoria, haciéndolos poco prácticos para aplicaciones en tiempo real o con recursos limitados.
• Limitaciones de Métodos Existentes: Las estrategias actuales de poda (pruning) de tokens visuales se dividen en métodos basados en similitud, ajuste fino (fine-tuning) y atención. Sin embargo, la mayoría asume erróneamente que todos los cabezales de atención (attention heads) contribuyen por igual a la interpretación visual. Esto lleva a decisiones de poda subóptimas, donde se eliminan tokens semánticamente valiosos o se conservan redundantes, ya que no consideran que diferentes cabezales capturan distintas semánticas visuales.

2. Metodología: HAWK

Los autores proponen HAWK (Head Importance-Aware Visual Token Pruning), un método libre de entrenamiento (training-free) que adapta la poda de tokens visuales basándose en la importancia específica de cada cabezal de atención.

El enfoque se basa en tres pilares principales:

• Descubrimiento de la Importancia de los Cabezales:
  Mediante un proceso de ablación sistemático, los autores demostraron que desactivar diferentes cabezales de atención afecta la capacidad de comprensión visual del modelo de manera desigual y consistente entre conjuntos de datos. Esto confirma que cada cabezal juega un rol único.

• Pesos de Importancia Estáticos de los Cabezales (Static Head Importance Weights):
  HAWK calcula una vez (offline) un vector de pesos que cuantifica la contribución intrínseca de cada cabezal de atención a la comprensión visual.

  • Se utiliza un análisis de ablación en varios conjuntos de datos de referencia.
  • Se mide la caída de rendimiento ($\Delta S$) al eliminar cada cabezal.
  • Estos valores se normalizan y promedian para obtener un peso de importancia ($w_i$) para cada cabezal $i$.

• Poda Consciente de la Importancia (Head Importance-Aware Pruning):
  Durante la inferencia, HAWK combina dos componentes para decidir qué tokens conservar:

  1. Atención Guiada por Texto (Dynamic Text-Guided Attention): Evalúa dinámicamente la relevancia semántica entre los tokens visuales y las instrucciones de texto del usuario. Se utiliza un mecanismo de atención en la primera capa del LLM, omitiendo intencionalmente las codificaciones de posición (RoPE) para centrarse puramente en la correspondencia semántica.
  2. Fusión de Pesos: Se calcula una puntuación de importancia final ($I_k$) para cada token visual $k$ combinando la relevancia dinámica con los pesos estáticos de los cabezales:
     $$I_k = \sum_{i=1}^{N_h} w_i \cdot c_{k}^i$$
     Donde $w_i$ es el peso del cabezal $i$ y $c_{k}^i$ es la puntuación de atención del token $k$ en ese cabezal.
  • Finalmente, se retienen los tokens con las puntuaciones $I_k$ más altas, eliminando los redundantes.

3. Contribuciones Clave

• Identificación de Roles Distintos: Demostraron empíricamente que los cabezales de atención en MLLMs tienen roles distintos en el procesamiento visual y cuantificaron su importancia individual.
• Método HAWK: Propusieron un algoritmo de poda de tokens visuales consciente de la importancia de los cabezales, que es completamente libre de entrenamiento y no requiere ajuste fino.
• Generalización y Eficiencia: El método se aplica de manera transparente a diversas arquitecturas de MLLM (como LLaVA-NeXT y Qwen2.5-VL) y funciona tanto en resoluciones fijas como nativas, así como en tareas de video.

4. Resultados Experimentales

Los experimentos se realizaron en múltiples benchmarks de visión-lenguaje (MME, TextVQA, ChartQA, AI2D, Video-MME, etc.) utilizando modelos de vanguardia como Qwen2.5-VL-7B e InternVL3-8B.

• Precisión (SOTA): HAWK supera consistentemente a los métodos baselines (FastV, DivPrune, CDPruner).
  • En Qwen2.5-VL, al podar el 80.2% de los tokens visuales, HAWK retiene el 96.0% de la precisión original.
  • En condiciones extremas (90% de poda), mantiene un 90.0% de rendimiento, superando a los métodos existentes en más de 4 puntos porcentuales.
  • En tareas de video, mantiene un 96.6% de rendimiento con una poda del 80%.
• Eficiencia Computacional:
  • Latencia: Reduce la latencia de extremo a extremo a un 74.4% del tiempo original (aceleración de 1.34x).
  • Memoria: Disminuye significativamente el uso de memoria de la GPU y el tamaño de la caché KV (Key-Value cache). Por ejemplo, en Qwen2.5-VL, reduce el uso de memoria de 16.9 GB a 15.7 GB con una poda del 80%.

5. Significado e Impacto

El trabajo de HAWK es significativo porque aborda una limitación fundamental en la optimización de MLLMs: la homogeneización de los cabezales de atención. Al reconocer y explotar la heterogeneidad de los cabezales, HAWK logra una compresión de tokens mucho más inteligente y eficiente.

• Viabilidad en Producción: Al ser un método libre de entrenamiento y compatible con modelos existentes, facilita la implementación de MLLMs en entornos con restricciones de hardware o requisitos de tiempo real.
• Escalabilidad: Permite procesar imágenes de mayor resolución y videos más largos sin colapsar los recursos computacionales, abriendo la puerta a aplicaciones multimodales más complejas y accesibles.

En resumen, HAWK representa un avance hacia la inferencia multimodal eficiente, logrando un equilibrio superior entre la reducción de costos computacionales y la preservación de la inteligencia del modelo.