BOSCH: Black-Box Binary Optimization for Short-Context Attention-Head Selection in LLMs

1. Il Problema

I Large Language Models (LLM) basati su Transformer utilizzano l'attenzione self-attention, che ha una complessità quadratica rispetto alla lunghezza della sequenza. Per gestire contesti lunghi in modo efficiente, una strategia comune è l'ibridazione post-training, che sostituisce l'attenzione self-attention completa con l'attenzione a finestra scorrevole (Sliding Window Attention - SWA) per ridurre l'uso della cache KV e migliorare la latenza.

Le sfide principali identificate dagli autori sono:

• Granularità a livello di layer: I metodi esistenti spesso decidono quali layer interi convertire in SWA (es. interleaving o begin-middle-end). Tuttavia, all'interno dello stesso layer, diversi attention heads gestiscono dipendenze locali e globali in modo diverso. Convertire un intero layer può rimuovere informazioni globali critiche.
• Problema di "Entanglement" (Intreccio): I metodi statici a livello di head classificano le teste da locali a globali prima dell'ibridazione e convertono quelle più locali in base a un rapporto fisso. Il problema è che il comportamento locale/globale di una testa può cambiare dopo l'ibridazione a causa delle interazioni con le altre teste modificate, rendendo le selezioni statiche subottimali.
• Complessità Computazionale: Una ricerca diretta a livello di testa (head-level) su modelli moderni (centinaia o migliaia di teste) è computazionalmente proibitiva per gli algoritmi di ottimizzazione black-box standard, che tendono a bloccarsi o richiedere un numero di valutazioni ingestibile.

2. Metodologia: BOSCH

Gli autori propongono BOSCH (Black-Box Binary Optimization for Short-Context Head Selection), un metodo training-free che formula la selezione delle teste come un problema di Ricerca nel Vicinato di Grande Dimensione (Large Neighborhood Search - LNS). Il processo è decomposto in tre fasi sequenziali per rendere la ricerca fattibile entro un budget di valutazione realistico:

1. Rilevamento dell'Importanza del Layer (Layer-Importance Detection):

   • Si valuta la sensibilità di ogni layer alla localizzazione dell'attenzione.
   • Si esegue una ricerca black-box a budget ridotto su un layer alla volta (dall'alto verso il basso), convertendo un numero fisso di teste in SWA per quel layer specifico, mantenendo fissi gli altri.
   • Si registra il punteggio di performance migliore per ogni layer per guidare le fasi successive.

2. Assegnazione Adattiva del Rapporto SWA per Layer:

   • Basandosi sui punteggi di sensibilità ottenuti nella fase 1, si calcola il calo di performance relativo per ogni layer.
   • I layer vengono raggruppati in "bucket" in base alla loro tolleranza alla localizzazione (layer più facili da convertire vs. layer più critici).
   • Si assegna un rapporto SWA specifico per ogni layer ($r_\ell$) che rispetti il budget globale totale ($\rho$), massimizzando la performance complessiva.

3. Selezione delle Teste a Livello di Gruppo (Grouped Head-Level Optimization):

   • I layer con lo stesso rapporto assegnato vengono raggruppati.
   • Si esegue un'ottimizzazione black-box congiunta sulle teste all'interno di ciascun gruppo, convertendo esattamente il numero di teste necessario per soddisfare il rapporto del gruppo.
   • Questo passaggio risolve il problema dell'entanglement ottimizzando le interazioni tra le teste all'interno del contesto ibrido specifico.

3. Contributi Chiave

• Superamento dell'approccio statico: BOSCH dimostra che la selezione delle teste non può essere basata su ranking statici pre-ibridazione a causa dell'entanglement, ma richiede un'ottimizzazione congiunta.
• Decomposizione LNS: La trasformazione di un problema di ottimizzazione binaria su migliaia di variabili in tre sottoproblemi gestibili (rilevamento layer, assegnazione budget, ottimizzazione di gruppo) rende possibile la ricerca senza addestramento.
• Metodo Training-Free: Non richiede ri-addestramento del modello, rendendolo applicabile immediatamente a modelli LLM pre-addestrati.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su 4 modelli della famiglia Qwen3 (da 1.7B a 30B parametri) con 4 diversi rapporti SWA ($\rho \in \{0.25, 0.5, 0.75, 0.875\}$).

• Prestazioni Zero-Shot: BOSCH supera sistematicamente sia le euristiche a livello di layer (come INTR, BME, RAND) che i 6 metodi statici a livello di testa (DCAM, APL, FISHER, RAZOR, ecc.) sui benchmark NIAH (Needle-in-a-Haystack) e LongBench.
  • Il vantaggio di BOSCH aumenta significativamente a rapporti SWA più alti (es. $\rho=0.75$ o $0.875$), dove i metodi statici falliscono drasticamente.
• Analisi dell'Entanglement: L'analisi mostra che l'insieme delle teste selezionate da BOSCH cambia sostanzialmente al variare del rapporto SWA (alto tasso di turnover). Questo conferma che una selezione fissa non è adatta e che l'ottimizzazione specifica per il target è cruciale.
• Recupero con Pre-training Continuo: Anche quando si permette un breve pre-training continuo (2.5B token) per recuperare le prestazioni, BOSCH recupera più velocemente e raggiunge un livello di performance superiore rispetto ai metodi basati su layer o statici.
• Scalabilità: Le prestazioni sono coerenti attraverso diverse dimensioni di modello (da 1.7B a 30B) e si mantengono robuste anche con estrazione di lunghezza (length extrapolation) fino a 128k token.

5. Significato e Impatto

BOSCH rappresenta un avanzamento significativo nell'efficienza dei LLM per contesti lunghi. Dimostra che:

1. La granularità a livello di testa è superiore a quella a livello di layer per l'ibridazione SWA.
2. L'ottimizzazione dinamica e congiunta delle teste è necessaria per mitigare gli effetti di entanglement che i metodi statici ignorano.
3. È possibile ottenere modelli ibridi ad alte prestazioni e a basso costo computazionale (in fase di selezione) senza richiedere costosi cicli di addestramento da zero.

In sintesi, BOSCH fornisce un framework pratico e robusto per trasformare modelli LLM esistenti in versioni efficienti per contesti lunghi, superando i limiti delle euristiche strutturali semplici e dei metodi di selezione statica.