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

1. Il Problema

I Modelli Linguistici Multimodali (MLLM) hanno rivoluzionato l'interazione tra testo e immagini, ma affrontano una sfida significativa legata all'efficienza computazionale.

• Sovraccarico di Token Visivi: Gli input visivi (immagini o video) vengono codificati in un numero elevato di "token visivi" (spesso centinaia o migliaia), che superano di gran lunga il numero di token testuali.
• Complessità Quadratica: La complessità computazionale dei modelli basati su attenzione (Transformer) cresce quadraticamente rispetto alla lunghezza della sequenza di token. L'accumulo di token visivi ridondanti aumenta drasticamente il tempo di inferenza e l'uso della memoria GPU, rendendo difficile l'applicazione in scenari real-time o su dispositivi con risorse limitate.
• Limitazioni delle Soluzioni Esistenti: Le tecniche di pruning (potatura) dei token visivi attuali si basano su tre approcci principali:
  1. Basati sulla similarità: Rimuovono regioni ridondanti ignorando il contesto dell'istruzione utente.
  2. Basati sul fine-tuning: Richiedono un addestramento costoso e hanno scarsa generalizzazione.
  3. Basati sull'attenzione: Stimano l'importanza dei token usando i punteggi di attenzione, ma assumono erroneamente che tutte le "teste" (heads) dell'attenzione contribuiscano in modo uguale all'interpretazione visiva, calcolando una media aritmetica semplice.

2. Metodologia: HAWK

Gli autori propongono HAWK (Head Importance-Aware Visual Token Pruning), un metodo che rivoluziona l'approccio al pruning riconoscendo che diverse teste di attenzione catturano semantiche visive distinte e hanno ruoli diversi.

Il metodo è completamente free-training (non richiede ri-addestramento) e si articola in tre fasi principali:

A. Pesatura Statica dell'Importanza delle Teste (Static Head Importance Weights)

Gli autori hanno scoperto che abilitare o disabilitare singole teste di attenzione ha impatti drasticamente diversi sulle prestazioni del modello.

• Analisi Offline: Viene eseguita un'analisi una tantum su dataset di benchmark standard.
• Calcolo del Peso: Per ogni testa $i$, si calcola il calo di prestazioni ($\Delta S$) quando quella testa viene "ablata" (disabilitata) rispetto al modello baseline.
• Vettore di Peso: Si ottiene un vettore di importanza $W$ normalizzando e mediando questi cali di prestazioni su tutti i dataset. Questo assegna un peso statico a ogni testa, riflettendo la sua rilevanza intrinseca per la comprensione visiva.

B. Punteggi di Attenzione Guidati dal Testo (Dynamic Text-Guided Attention Scores)

L'importanza di un token visivo non è fissa, ma dipende dall'istruzione testuale dell'utente.

• Meccanismo Dinamico: Viene utilizzata una proiezione delle query ($W_q$) e delle chiavi ($W_k$) dalla prima layer di attenzione del LLM per calcolare l'attenzione tra i token testuali e quelli visivi.
• Rimozione del Bias Posizionale: Per garantire che l'attenzione dipenda solo dalla corrispondenza semantica e non dalla posizione, l'operazione di Rotary Position Embedding (RoPE) viene omessa in questo passaggio specifico.
• Punteggio Aggregato: Si calcola un punteggio di rilevanza media per ogni token visivo rispetto all'intero input testuale.

C. Pruning Consapevole dell'Importanza delle Teste

La decisione finale su quali token conservare combina le due componenti precedenti:

• Punteggio di Importanza Finale ($I_k$): Per ogni token visivo $k$, il punteggio è la somma pesata dei punteggi di attenzione guidati dal testo, dove i pesi sono quelli delle teste calcolati nella fase statica:
  $$I_k = \sum_{i=1}^{N_h} w_i \cdot c_k^i$$
  Dove $w_i$ è l'importanza della testa $i$ e $c_k^i$ è l'attenzione della testa $i$ verso il token $k$.
• Selezione: Vengono mantenuti i token con i punteggi $I_k$ più alti, scartando quelli ridondanti.

3. Contributi Chiave

1. Scoperta Fondamentale: Dimostrazione empirica che le diverse teste di attenzione nei MLLM svolgono ruoli unici e non intercambiabili nella elaborazione visiva, invalidando l'assunzione di uguale contributo fatta dalle ricerche precedenti.
2. Proposta HAWK: Un nuovo metodo di pruning che integra dinamicamente l'importanza delle teste con la rilevanza semantica guidata dal testo.
3. Efficienza e Generalizzazione: Il metodo è training-free, richiede solo un calcolo offline una tantum e si applica seamless a diverse architetture (es. Qwen2.5-VL, InternVL3) senza costi di calibrazione aggiuntivi.

4. Risultati Sperimentali

Il paper presenta valutazioni estensive su benchmark visivi e video (HallBench, MME, TextVQA, VideoMME, ecc.) utilizzando modelli come Qwen2.5-VL-7B e InternVL3-8B.

• Prestazioni di Accuratezza:
  • Su Qwen2.5-VL, rimuovendo l'80.2% dei token visivi, HAWK mantiene il 96.0% dell'accuratezza originale.
  • Con un pruning del 90.1%, mantiene il 90.0% delle prestazioni, superando di 4.1 punti percentuali i migliori metodi esistenti (come FastV, DivPrune, CDPruner).
  • Su InternVL3-8B, con un pruning dell'80%, mantiene il 94.1% delle prestazioni, superando i baselines di 7 punti percentuali.
• Efficienza Computazionale:
  • Latency: Riduce la latenza end-to-end al 74.4% del tempo originale (un speedup di 1.34x).
  • Memoria: Riduce significativamente l'uso della memoria GPU e la dimensione della KV cache (es. da 668 MB a 148 MB con pruning 0.8).
• Robustezza: Le prestazioni rimangono superiori sia in scenari a risoluzione fissa che nativa (variabile), e anche in compiti di comprensione video con sequenze lunghe.

5. Significato e Impatto

HAWK rappresenta un passo avanti significativo verso l'efficienza dei modelli multimodali su larga scala.

• Paradigma Shift: Sposta il focus dal semplice "rimuovere token simili" al "comprendere quale parte del meccanismo di attenzione è cruciale per il compito specifico".
• Praticità: Essendo un metodo training-free, può essere adottato immediatamente da sviluppatori e ricercatori per ottimizzare modelli esistenti senza i costi proibitivi del ri-addestramento.
• Scalabilità: Abilita l'uso di MLLM in applicazioni real-time e su hardware con risorse limitate, aprendo la strada a un'interazione più fluida e veloce tra uomo e macchina in contesti visivi complessi.

In sintesi, HAWK dimostra che una comprensione più profonda dell'architettura interna del modello (l'eterogeneità delle teste di attenzione) è la chiave per ottenere un pruning aggressivo senza sacrificare l'intelligenza del modello.