ASAP: Attention-Shift-Aware Pruning for Efficient LVLM Inference

Resumen Técnico: ASAP

1. El Problema: Costo Computacional y Sesgo de Atención

Los Modelos de Lenguaje y Visión Grandes (LVLMs) como LLaVA-NeXT han demostrado capacidades excepcionales, pero enfrentan un cuello de botella crítico: el costo computacional cuadrático ($O(N^2)$) asociado al procesamiento de secuencias de tokens visuales de alta resolución.

• Redundancia de Tokens: Las imágenes de alta resolución se dividen en múltiples "baldosas" (tiles), generando miles de tokens visuales (ej. ~2,304 en LLaVA-NeXT), muchos de los cuales son redundantes (cielo, océano).
• Fallo de los Métodos Actuales: Las estrategias de poda (pruning) existentes fallan por dos razones principales:
  1. Ignoran la redundancia semántica: A menudo eliminan tokens basándose solo en la atención cruda sin consolidar tokens similares.
  2. El Fenómeno de "Desplazamiento de Atención" (Attention Shift): Debido a la propiedad de decaimiento por distancia inherente a las Posiciones de Codificación Rotatoria (RoPE), los tokens visuales posteriores (generalmente en la parte inferior de la imagen) reciben puntuaciones de atención artificialmente altas, no por su contenido semántico, sino por su posición relativa al texto. Esto provoca que los métodos de poda basados en atención seleccionen incorrectamente las regiones inferiores de la imagen, ignorando detalles críticos en el fondo o en la parte superior.

2. Metodología: El Marco ASAP

ASAP es un marco de poda sin entrenamiento (training-free) y compatible con la caché de clave-valor (KV-Cache), diseñado para mitigar el desplazamiento de atención y eliminar redundancia. Se compone de tres etapas principales:

A. Poda Consciente del Desplazamiento (Salience-Guided Bidirectional Attention)

• Mecanismo: En lugar de usar una máscara de atención causal estricta (unidireccional), ASAP introduce una Máscara Bidireccional Guiada por la Saliencia (SG-BiMask).
• Funcionamiento: Permite que los tokens visuales "miren hacia adelante" (bidireccionalidad) para agregar contexto global, pero de forma controlada.
• Cálculo de Saliencia: Utiliza la matriz de alineación pre-softmax ($W$) para calcular la "masa de atención" ($s_j$) de cada token. Esta métrica se normaliza para determinar un factor de descuento ($\lambda_j$).
• Resultado: Los tokens estructuralmente importantes (anclajes visuales) pueden propagar sus características hacia atrás en la secuencia con poca atenuación, mientras que los fondos no informativos son suprimidos. Esto corrige el sesgo espacial inducido por RoPE.

B. Consolidación Ponderada por Saliencia (Salience-Weighted Consolidation)

• Objetivo: Eliminar tokens redundantes (ej. parches de un cielo uniforme) sin perder información densa.
• Proceso: Calcula la similitud coseno entre los estados ocultos de los tokens. A diferencia de promedios simples, fusiona los tokens redundantes utilizando una combinación convexa ponderada por la masa de atención normalizada ($\hat{s}$).
• Ventaja: Asegura que al fusionar tokens, la representación latente resultante esté anclada al token más estructuralmente crítico, evitando la dilución de características.

C. Reasignación de Presupuesto (Budget Reallocation)

• Mecanismo: Cuando la fusión de tokens libera espacios en el presupuesto de tokens (k), estos espacios vacantes no se dejan ociosos.
• Estrategia: Se reasignan a los tokens "descartados" previamente que tenían la mayor masa de atención, asegurando que la secuencia final maximice la diversidad semántica y no solo recorte la representación espacial.

3. Contribuciones Clave

1. Identificación del Sesgo RoPE: Se atribuye sistemáticamente el "desplazamiento de atención" en LVLMs al mecanismo de decaimiento por distancia de RoPE, demostrando cómo penaliza injustamente las características visuales tempranas.
2. Mitigación Estructural: Primera propuesta para mitigar estructuralmente este artefacto mediante un mecanismo de atención bidireccional suave y dinámico, sin eliminar RoPE (lo cual degradaría el razonamiento espacial).
3. Marco de Poda Sin Entrenamiento: Presentación de ASAP, que combina la poda consciente del desplazamiento con la eliminación de redundancia basada en características, manteniendo compatibilidad total con la inferencia de múltiples turnos y la caché KV estándar.
4. Eficiencia sin Pérdida: Demostración de que es posible lograr una compresión casi sin pérdidas manteniendo >99% del rendimiento original.

4. Resultados Experimentales

Los experimentos se realizaron en modelos como LLaVA-1.5, LLaVA-NeXT-7B e InternVL 2.5 en múltiples benchmarks (VQAv2, GQA, MME, MMBench, etc.).

• Rendimiento General:
  • En LLaVA-NeXT-7B, ASAP retiene el 99.02% del rendimiento del modelo original (vanilla) mientras reduce los FLOPs en un ~80% (manteniendo solo ~22% de los tokens).
  • En algunos benchmarks (como MMBench), supera incluso al modelo original ("supra-vanilla").
• Comparación con Baselines:
  • Supera consistentemente a métodos de vanguardia como FastV, LLaVA-PruMerge+, SparseVLM y DivPrune.
  • En el escenario de 128 tokens (LLaVA-1.5-13B), ASAP alcanza un 97.10% de retención de rendimiento promedio, superando a DivPrune (96.02%).
• Eficiencia del Sistema:
  • Reduce drásticamente el uso de memoria GPU y el almacenamiento de la caché KV.
  • En el escenario de 128 tokens, reduce el almacenamiento de la caché en 15% comparado con FastV, con una sobrecarga de inferencia mínima (2 ms).

5. Significado e Impacto

ASAP representa un avance significativo en la inferencia eficiente de modelos multimodales. Al abordar la raíz estructural del problema (el sesgo de RoPE y la redundancia semántica) en lugar de aplicar heurísticas superficiales, permite:

• Desplegar LVLMs de alta resolución en hardware con recursos limitados.
• Mantener la precisión en tareas de razonamiento visual complejo que requieren comprensión global de la escena.
• Ofrecer una solución "plug-and-play" que no requiere reentrenamiento costoso, facilitando su adopción inmediata en la industria y la investigación.

En resumen, ASAP demuestra que una poda inteligente, consciente de la arquitectura y la semántica, puede reducir drásticamente la carga computacional sin sacrificar la inteligencia del modelo.