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

1. O Problema

O artigo aborda os desafios de eficiência computacional e preservação de precisão em Modelos Multimodais de Grande Escala (LMMs), especificamente focando na arquitetura BLIP-2.

• Custo Computacional: O BLIP-2 possui 3,1 bilhões de parâmetros, dos quais o Decodificador de LLM (Large Language Model) representa mais de 87% (2,7 bilhões). Durante a inferência, este componente domina o tempo de processamento.
• Complexidade Quadrática: O mecanismo de atenção nos transformadores tem complexidade quadrática em relação ao comprimento da sequência de entrada. Reduzir o número de tokens de consulta (query tokens) pode diminuir drasticamente o custo computacional.
• Falha das Métodos Atuais: Estratégias de poda (pruning) de última geração (SOTA), como poda baseada em magnitude ou probabilidade de auto-atenção (SpAtten), funcionam bem em modelos unimodais, mas resultam em uma degradação severa da precisão quando aplicadas a modelos multimodais como o BLIP-2 em tarefas de Resposta a Perguntas Visuais (VQA).

2. Metodologia: CATP (Cross-Attention Token Pruning)

Os autores propõem o CATP, um método de poda pós-treinamento que utiliza as camadas de atenção cruzada (cross-attention) do módulo Q-Former do BLIP-2 para determinar a importância dos tokens de consulta.

Princípio Central

A importância de cada token de consulta do LLM depende de sua relação com os tokens de imagem. A atenção cruzada indica a relevância entre os tokens de imagem e os de texto. O objetivo é manter os tokens de consulta que são altamente relevantes para os tokens de imagem, evitando a perda de informações cruciais.

O Algoritmo de Votação

Como as probabilidades de atenção cruzada são normalizadas (somam 1 via softmax), não é possível simplesmente somá-las para obter uma pontuação de importância. O CATP introduz uma estratégia de votação refinada:

1. Extração: O método extrai mapas de probabilidade de atenção cruzada de todas as cabeças (heads) e camadas do Q-Former.
2. Votação por Token de Imagem: Para cada token de consulta, cada token de imagem "vota" pontos baseados na probabilidade de atenção cruzada entre eles.
   • Especificamente, para cada mapa de atenção, o token de imagem com a $n$-ésima maior probabilidade de atenção para um token de consulta atribui $L_0 - n$ pontos a esse token de consulta (onde $L_0$ é o número total de tokens de consulta).
3. Agregação: A pontuação de importância final de cada token de consulta é a soma dos pontos recebidos através de todas as cabeças e camadas.
4. Poda: Os tokens de consulta com as pontuações de importância mais baixas são removidos, mantendo apenas uma fração definida pela taxa de poda ($p$).

Variações Exploradas

• Variação 1: Cálculo de importância utilizando probabilidades de todas as camadas de atenção do Q-Former.
• Variação 2: Cálculo baseado exclusivamente na primeira camada de atenção.
• Votação Ponderada (Otimização): Uma versão avançada onde os tokens de imagem recebem pesos baseados em suas próprias pontuações de auto-atenção (tokens de imagem mais importantes "votam" mais), utilizando a última camada do codificador visual.

3. Resultados Principais

Os experimentos foram conduzidos no modelo BLIP-2 (blip2-opt-2.7b) utilizando o conjunto de dados VQA (Visual Question Answering).

• Comparação com Baselines:
  • Baseline L2-norm: Seleciona tokens com maior magnitude. Resultou em precisão muito baixa (ex: 9,50% com 50% dos tokens mantidos).
  • Baseline Auto-atenção (SpAtten): Baseada em auto-atenção. Também apresentou queda drástica (ex: 44,22% com 50% mantidos, caindo para 1,92% com 12,5% mantidos).
  • CATP: Demonstrou superioridade significativa.
    • Com 50% dos tokens mantidos: 41,14% (todas as camadas) vs 44,22% (baseline auto-atenção).
    • Com 12,5% (1/8) dos tokens mantidos: 23,27% (todas as camadas) vs 1,92% (baseline auto-atenção).
• Ganhos de Precisão: O CATP alcançou até 6,6x mais precisão que a baseline baseada em L2-norm e até 12,1x mais precisão que a baseline baseada em auto-atenção em taxas de poda agressivas.
• Análise de Camadas: A pesquisa mostrou que as camadas de atenção cruzada inicial e final fornecem informações mais valiosas para a poda do que as camadas intermediárias.
• Impacto do Ponderamento: A adição de pesos baseados na importância dos tokens de imagem (votação ponderada) melhorou ainda mais a precisão, especialmente em taxas de poda mais agressivas (ex: 29,30% vs 23,27% com 1/8 mantido).

4. Contribuições Chave

1. Método CATP: Introdução de uma nova técnica de poda de tokens específica para modelos multimodais, utilizando a atenção cruzada entre modalidades (imagem-texto) em vez de apenas auto-atenção.
2. Estratégia de Votação: Proposta de um mecanismo inovador para transformar probabilidades de atenção (normalizadas) em pontuações de importância discriminativas através de um sistema de votação distribuída entre tokens.
3. Validação de Eficiência vs. Precisão: Demonstração de que é possível reduzir drasticamente o custo computacional (poda de tokens) sem sacrificar a precisão do modelo, resolvendo o dilema comum em compressão de modelos multimodais.
4. Insights Arquiteturais: Identificação de que as camadas iniciais e finais do Q-Former são críticas para a determinação da importância dos tokens, e que a ponderação baseada na importância dos tokens de imagem é benéfica.

5. Significado

O trabalho é significativo porque preenche uma lacuna na literatura de compressão de modelos, onde métodos existentes falham em cenários multimodais. O CATP oferece uma solução prática para viabilizar a inferência de modelos multimodais de grande escala em ambientes com recursos limitados, mantendo a utilidade do modelo em tarefas complexas como VQA. A abordagem sugere que a interação entre modalidades (atenção cruzada) é um sinal muito mais forte para a seleção de informações relevantes do que a estrutura interna de uma única modalidade.