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

Resumo Técnico: HAWK

1. O Problema

Os Modelos de Linguagem Multimodal (MLLMs) modernos, como o Qwen2.5-VL e o LLaVA, enfrentam um gargalo computacional significativo devido à quantidade massiva de tokens visuais gerados a partir de imagens e vídeos.

• Custo Computacional: A complexidade do mecanismo de atenção nos modelos baseados em transformers cresce quadraticamente com o comprimento da sequência de tokens. Como as imagens (especialmente em alta resolução ou vídeos longos) geram centenas ou milhares de tokens visuais, o tempo de inferência e o uso de memória (VRAM) tornam-se proibitivos para aplicações em tempo real ou com recursos limitados.
• Limitação das Soluções Atuais: Métodos existentes de poda de tokens visuais (token pruning) geralmente assumem que todos os cabeçalhos de atenção (attention heads) contribuem igualmente para a interpretação visual. Eles tendem a usar médias aritméticas simples ou similaridade de características, ignorando que diferentes cabeçalhos capturam semânticas visuais distintas e possuem níveis variados de importância para tarefas específicas. Isso pode levar à remoção acidental de tokens semanticamente valiosos ou à retenção de redundâncias.

2. Metodologia: HAWK

O HAWK (Head Importance-Aware Visual Token Pruning) é uma abordagem de poda de tokens visuais que é consciente da importância dos cabeçalhos de atenção. O método é inteiramente livre de treinamento (training-free) e não requer ajuste fino (fine-tuning) do modelo.

O processo funciona em três etapas principais:

1. Pesos de Importância Estáticos dos Cabeçalhos (Static Head Importance Weights):

   • Os autores realizam uma análise offline única para quantificar a contribuição intrínseca de cada cabeçalho de atenção para a compreensão visual.
   • Isso é feito através de um processo de ablação em conjuntos de dados de referência: cada cabeçalho é desativado individualmente e a queda de desempenho do modelo é medida.
   • Os pesos de importância ($w_i$) são calculados com base na média normalizada da queda de desempenho causada pela ablação de cada cabeçalho em múltiplos conjuntos de dados. Cabeçalhos que causam grandes quedas de desempenho quando removidos recebem pesos mais altos.

2. Pontuações de Atenção Guiadas por Texto Dinâmicas (Dynamic Text-Guided Attention Scores):

   • Para capturar a dependência da tarefa, o HAWK utiliza um mecanismo de atenção guiado pelo texto para avaliar a relevância semântica entre os tokens visuais e as instruções de texto do usuário.
   • Utiliza as matrizes de projeção de query e key da primeira camada de atenção do LLM.
   • Nota Técnica: O método deliberadamente omite o Rotary Positional Embedding (RoPE) nesta etapa para eliminar viés posicional e focar puramente na correspondência semântica entre texto e visão.
   • Calcula-se uma pontuação de relevância agregada para cada token visual baseada na média das atenções de todos os tokens de texto.

3. Poda Consciente da Importância (Head Importance-aware Pruning):

   • Durante a inferência, o HAWK combina os pesos estáticos dos cabeçalhos com as pontuações dinâmicas de atenção guiada por texto.
   • A importância final de cada token visual ($I_k$) é calculada como a soma ponderada das suas pontuações de atenção em todos os cabeçalhos, multiplicada pelo peso de importância daquele cabeçalho:
     $$I_k = \sum_{i=1}^{N_h} w_i \cdot c^i_k$$
   • Com base nessas pontuações, os tokens visuais são classificados e os top-$k$ (os mais importantes) são retidos, enquanto os redundantes são descartados antes de entrarem nas camadas subsequentes do LLM.

3. Principais Contribuições

• Descoberta Fundamental: Evidencia que diferentes cabeçalhos de atenção em MLLMs desempenham papéis distintos e desiguais no processamento visual, desafiando a suposição comum de que todos contribuem igualmente.
• Método HAWK: Propõe um algoritmo de poda que integra pesos de importância de cabeçalhos com atenção guiada por texto, permitindo uma seleção de tokens mais precisa e adaptada à tarefa.
• Eficiência e Generalização: O método é training-free, exigindo apenas uma computação offline para calcular os pesos, o que o torna aplicável a diversas arquiteturas de MLLMs (como LLaVA-NeXT e Qwen2.5-VL) sem custo adicional de treinamento.
• Desempenho SOTA: Demonstra resultados superiores em múltiplos benchmarks de visão e linguagem, superando métodos baseados em similaridade, atenção e híbridos.

4. Resultados Experimentais

Os experimentos foram realizados em modelos de ponta (Qwen2.5-VL-7B e InternVL3-8B) em diversos benchmarks (MME, TextVQA, ChartQA, Video-MME, etc.).

• Precisão:
  • No modelo Qwen2.5-VL, ao podar 80,2% dos tokens visuais, o HAWK manteve 96,0% da precisão original.
  • Em cenários de poda extrema (90,1% de tokens removidos), o método ainda manteve 90,0% da performance, superando os métodos concorrentes em mais de 4 pontos percentuais.
  • Em tarefas de vídeo (Video-MME), manteve 96,6% da performance com 80% de poda.
• Eficiência (Latência e Memória):
  • Redução da latência de ponta a ponta para 74,4% do original (aceleração de ~1,34x) com 80% de poda.
  • Redução significativa no uso de memória GPU (KV Cache e memória total), tornando a execução mais viável em hardware limitado.
• Estudos de Ablação: Confirmaram que cada componente (pesos de cabeçalho, estratégia de guia de texto e ausência de RoPE na etapa de seleção) é essencial para o desempenho superior.

5. Significado e Impacto

O HAWK representa um avanço significativo na eficiência de MLLMs. Ao reconhecer e explorar a heterogeneidade dos cabeçalhos de atenção, ele resolve o dilema entre compressão agressiva e perda de informação crítica.

• Viabilidade Prática: Torna modelos multimodais grandes viáveis para aplicações em tempo real e dispositivos com recursos limitados, sem a necessidade de re-treinamento custoso.
• Direção Futura: Estabelece uma nova linha de pesquisa onde a estrutura interna do mecanismo de atenção (importância dos cabeçalhos) é utilizada diretamente para otimização de inferência, em vez de tratar o modelo como uma caixa preta.

Em resumo, o HAWK oferece uma solução elegante e altamente eficaz para o gargalo de tokens visuais, permitindo que modelos multimodais operem com maior velocidade e menor consumo de memória, mantendo quase toda a sua capacidade de raciocínio original.