CATP: Cross-Attention Token Pruning for Accuracy Preserved Multimodal Model Inference

Resumen Técnico: CATP (Poda de Tokens por Atención Cruzada)

1. El Problema

El modelo multimodal de última generación BLIP-2 enfrenta un desafío crítico de eficiencia computacional durante la inferencia.

• Costo Computacional: El modelo tiene 3.1 mil millones de parámetros, de los cuales el Decodificador de LLM (Modelo de Lenguaje Grande) representa más del 87% (2.7 mil millones). Este componente domina el tiempo de inferencia.
• Limitación de las Estrategias Actuales: Las técnicas de poda (pruning) existentes, diseñadas principalmente para modelos unimodales (como SpAtten basado en auto-atención o poda por magnitud L2), fallan al aplicarse a BLIP-2. Estas estrategias provocan una degradación severa de la precisión, especialmente en tareas de Respuesta a Preguntas Visuales (VQA), al eliminar tokens de consulta (query tokens) sin considerar su relación con la información visual.
• Complejidad: La complejidad de la atención en los transformadores es cuadrática respecto a la longitud de la secuencia. Reducir la longitud de los tokens de entrada es necesario para la eficiencia, pero hacerlo sin perder precisión es el reto principal.

2. Metodología: CATP

Los autores proponen CATP (Cross-Attention Token Pruning), un método de poda post-entrenamiento diseñado específicamente para modelos multimodales como BLIP-2. La premisa central es que la importancia de un token de consulta (query token) en el LLM depende de su relevancia con los tokens de imagen del módulo Q-Former.

Mecanismo Clave:

1. Uso de Atención Cruzada: A diferencia de los métodos que usan solo la auto-atención, CATP explota las capas de atención cruzada dentro del Q-Former (que conecta el codificador de imágenes con el decodificador de texto).
2. Estrategia de Votación (Voting Strategy):
   • Dado que las probabilidades de atención cruzada se normalizan con softmax (sumando a 1), no se pueden simplemente acumular para medir importancia.
   • CATP introduce un sistema de votación ponderada: Para cada token de consulta, cada token de imagen "vota" asignando puntos basados en la probabilidad de atención cruzada.
   • Algoritmo de Puntuación: Para cada mapa de atención cruzada (por cabeza y por capa), el token de imagen con la $n$-ésima probabilidad más alta asigna $L_0 - n$ puntos al token de consulta correspondiente (donde $L_0$ es el número total de tokens de consulta).
   • La puntuación final de importancia de un token de consulta es la suma de los puntos recibidos a través de todas las capas y cabezas de atención.
3. Proceso de Poda: Se eliminan los tokens de consulta con las puntuaciones de importancia más bajas, manteniendo una proporción $p$ de los tokens más relevantes.

Variantes Exploradas:

• Variantes de Capas: Se evaluó el uso de todas las capas de atención cruzada frente a solo la primera capa.
• Ponderación de Tokens de Imagen: Se propuso una mejora donde los tokens de imagen tienen pesos variables basados en sus propias puntuaciones de auto-atención (tokens de imagen más importantes tienen más peso en la votación).

3. Contribuciones Clave

• Primera Poda Específica para Multimodales: CATP es una solución diseñada desde cero para abordar la poda en modelos multimodales, reconociendo que los métodos unimodales no son transferibles directamente.
• Mecanismo de Votación Innovador: Transforma las probabilidades de atención cruzada en puntuaciones de importancia discriminativas mediante un sistema de votación jerárquica, resolviendo el problema de la normalización softmax.
• Integración de Información Multimodal: Al basarse en la atención cruzada, el método considera simultáneamente la información de texto e imagen para tomar decisiones de poda, asegurando que los tokens de texto críticos para la comprensión visual no se eliminen.
• Análisis de Importancia de Capas: Descubrimiento de que las capas inicial y final de atención cruzada en el Q-Former aportan la información más valiosa para la poda, superando a las capas intermedias.

4. Resultados Experimentales

Los experimentos se realizaron sobre el modelo BLIP-2 (blip2-opt-2.7b) utilizando un subconjunto del dataset VQA (Visual Question Answering).

• Comparativa con Baselines:
  • Baseline L2-norm: Al mantener solo 1/8 de los tokens, la precisión cayó al 8.91%.
  • Baseline Auto-atención (SpAtten): Al mantener 1/8 de los tokens, la precisión cayó drásticamente al 1.92%.
  • CATP (Todas las capas): Al mantener 1/8 de los tokens, logró una precisión del 23.27%.
• Mejoras Cuantitativas:
  • CATP supera a la línea base de norma L2 en un factor de 6.6x.
  • CATP supera a la línea base de auto-atención en un factor de 12.1x en términos de precisión relativa bajo poda agresiva.
• Impacto de las Mejoras:
  • La inclusión de la ponderación de tokens de imagen mejoró aún más la precisión (ej. 29.30% vs 23.27% al mantener 1/8 de tokens).
  • El uso de capas específicas (inicial y final) demostró ser más efectivo que el uso de capas intermedias.

5. Significado e Impacto

Este trabajo es significativo porque resuelve el dilema fundamental entre eficiencia computacional y precisión del modelo en la inferencia de grandes modelos multimodales.

• Viabilidad de Despliegue: CATP permite reducir drásticamente el costo de inferencia (al reducir la longitud de secuencia de los tokens de consulta) sin sacrificar la capacidad del modelo para entender imágenes y responder preguntas complejas.
• Nueva Dirección de Investigación: Establece que para los modelos multimodales, la poda debe basarse en la interacción entre modalidades (atención cruzada) y no solo en la importancia interna de los tokens de texto (auto-atención).
• Potencial de Escalabilidad: La metodología propuesta es aplicable a otros modelos grandes multimodales (como GPT-4 o Gemini, aunque el estudio se centró en BLIP-2), ofreciendo una ruta para hacer estos modelos más accesibles y rápidos sin reentrenamiento costoso.