HeadRouter: Dynamic Head-Weight Routing for Task-Adaptive Audio Token Pruning in Large Audio Language Models

1. Enunciato del Problema

I Large Audio-Language Models (LALM) hanno dimostrato capacità notevoli nell'elaborazione di sequenze multi-modali, ma soffrono di costi di inferenza proibitivi, in particolare quando gestiscono input audio a contesto lungo.

• Il Collo di Bottiglia: L'elaborazione di audio della durata di minuti genera enormi sequenze di token (ad esempio, 1.500 token al minuto), portando a latenze eccessive di prefilling e ingenti footprint di memoria KV-cache.
• Limitazioni dei Metodi Esistenti: Le attuali tecniche di pruning dei token (basate sulla similarità, sul tempo e sull'attenzione) generalmente falliscono nei contesti audio a causa di tre aspetti trascurati:
  1. Eterogeneità dei Compiti: I compiti audio non sono uniformi; variano da compiti semantici (ad esempio, ASR, trascrizione) a compiti acustici (ad esempio, genere del parlante, rilevamento di eventi sonori). I metodi esistenti presuppongono che una singola strategia di pruning sia adatta a tutti.
  2. Eterogeneità del Comportamento delle Teste: Le teste di attenzione nei LALM non contribuiscono in modo uguale. Il paper rivela che i compiti semantici e acustici attivano sottoinsiemi distinti di teste di attenzione con distribuzioni di importanza diverse.
  3. Mancanza di Consapevolezza della Natura Audio: Metodi come FastV soffrono di bias posizionale (mantenendo i token finali), mentre il sottocampionamento Frame è cieco al contenuto (ignorando i segmenti vocali rispetto a quelli silenziosi).

2. Metodologia: HeadRouter

Gli autori propongono HeadRouter, un metodo di pruning dinamico dei token senza necessità di addestramento che si adatta alla natura specifica dell'input audio instradando i pesi dell'attenzione.

Meccanismo Principale

HeadRouter opera in tre fasi principali:

1. Sondaggio QK Posizionalmente Agnostico:

   • Per evitare il bias posizionale, il metodo sonde le interazioni testo-audio utilizzando le proiezioni Query (Q) e Key (K) dal livello precedente al livello di pruning ($M-1$).
   • Crucialmente, durante questo passaggio, le Rotary Positional Embeddings (RoPE) vengono rimosse. Ciò garantisce che l'importanza del token sia basata puramente sull'allineamento semantico tra contenuto testuale e audio, e non sulla posizione assoluta.

2. Stima della Selettività e Segnale di Instradamento:

   • Per ogni testa di attenzione $h$, viene calcolato un punteggio di selettività ($sel_h$) basato sull'entropia della distribuzione di attenzione sui token audio. Un'alta selettività implica che la testa si concentra su token specifici; una bassa selettività implica una distribuzione diffusa.
   • Un segnale di instradamento ($s_{pr}$) viene calcolato come la deviazione standard dei punteggi di selettività su tutte le teste.
     • Bassa $s_{pr}$: Indica che le teste si comportano in modo simile (tipico dei compiti semantici).
     • Alta $s_{pr}$: Indica un sottoinsieme sparso di teste altamente selettive (tipico dei compiti acustici).

3. Instradamento Soft Gaussiano:

   • Invece di un commutazione rigida, HeadRouter utilizza un meccanismo di miscelazione soft Gaussiana per fondere tre profili di pesi delle teste pre-calibrati:
     • Profilo Semantico ($w_{sem}$): Derivato da compiti semantici.
     • Profilo Acustico ($w_{aco}$): Derivato da compiti acustici.
     • Profilo Uniforme ($w_{uni}$): Pesi uguali.
   • Il segnale di instradamento $s_{pr}$ determina i coefficienti di miscelazione ($\alpha$) tramite un kernel Gaussiano, producendo un vettore di pesi delle teste ($w^*$) continuo e specifico per il campione.
   • L'importanza del token viene quindi calcolata ponderando i punteggi di attenzione di ogni testa utilizzando $w^*$.

3. Contributi Chiave

1. Divergenza delle Teste Semantiche-Acustiche: Il paper fornisce prove empiriche che le teste di attenzione mostrano comportamenti fondamentalmente diversi per compiti semantici rispetto a quelli acustici, sfidando l'assunzione che tutte le teste contribuiscano in modo uguale.
2. Instradamento Adattivo all'Audio: Introduzione di HeadRouter, un meccanismo innovativo senza addestramento che instrada dinamicamente i pesi delle teste in base alle caratteristiche dell'input (dispersione della selettività) senza richiedere fine-tuning.
3. Compressione che Migliora le Prestazioni: Il metodo dimostra che un pruning aggressivo (rimozione del 30–60% dei token) può effettivamente migliorare le prestazioni rispetto alla baseline non prunata in certi scenari, filtrando rumore e token ridondanti.

4. Risultati Sperimentali

Il metodo è stato valutato sui benchmark AudioMarathon e MMAU-Pro utilizzando modelli come Qwen2.5-Omni (3B/7B) e Phi-4-Multimodal.

• Prestazioni Superiori:
  • Su Qwen2.5-Omni-3B con il 70% di ritenzione dei token (30% di pruning), HeadRouter ha raggiunto il 101,8% delle prestazioni della baseline vanilla, superando tutti i concorrenti (FastV, DART, Frame).
  • Su Qwen2.5-Omni-7B, ha raggiunto il 103,0% della baseline.
  • Ha costantemente superato le baseline attraverso compiti semantici (ASR, SER) e acustici (SED, SGR), mentre altri metodi hanno subito un degrado significativo, specialmente nell'ASR.
• Robustezza al Pruning Aggressivo: Anche con un pruning del 90% (mantenendo solo il 10% dei token), HeadRouter ha mantenuto il miglior equilibrio tra prestazioni semantiche e acustiche, mentre altri metodi sono collassati.
• Efficienza: HeadRouter occupa la frontiera di Pareto nel compromesso prestazioni-memoria. Riduce significativamente l'uso di memoria GPU di picco (simile a Frame/DART) mantenendo una maggiore accuratezza.
• Generalizzazione: Il metodo si trasferisce efficacemente tra diverse architetture di modelli (famiglie Qwen e Phi) e benchmark, dimostrando che il segnale di instradamento cattura proprietà audio generali piuttosto che adattarsi eccessivamente a un dataset specifico.

5. Significato

• Cambiamento di Paradigma: Questo lavoro sfida l'approccio "taglia unica" al pruning dei token negli LLM. Stabilisce che l'instradamento adattivo al compito è essenziale per l'elaborazione audio multi-modale.
• Deploy Pratico: Essendo senza addestramento e aggiungendo un sovraccarico computazionale trascurabile (<1% del tempo di prefilling), HeadRouter offre una soluzione immediata e deployabile al collo di bottiglia dell'efficienza dei LLM Audio a Contesto Lungo.
• Approfondimento sui LALM: Le scoperte rivelano che il pruning può essere benefico non solo per la velocità, ma anche per il denoising, poiché il meccanismo adattivo scarta efficacemente i token irrilevanti per il compito specifico in esame.

In conclusione, HeadRouter rappresenta un avanzamento significativo nel rendere i Large Audio-Language Models più efficienti e robusti, sfruttando l'eterogeneità intrinseca delle teste di attenzione per ottimizzare dinamicamente la ritenzione dei token.