Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

Sintesi Tecnica: Causal Attribution Pruning (CAP)

Definizione del Problema
I Large Language Models (LLM) dimostrano forti capacità di ragionamento multi-step, ma comportano costi di inferenza sostanziali, ostacolando l'implementazione nel mondo reale. Sebbene il pruning dei modelli offra una soluzione rimuovendo i parametri ridondanti senza riaddestramento, i metodi esistenti spesso degradano le prestazioni nei compiti di ragionamento complesso. Gli approcci attuali si basano su segnali correlazionali, come la magnitudo dei pesi (ad es. Wanda) o statistiche locali del gradiente. Questi metodi assumono che i pesi piccoli siano poco importanti; tuttavia, i parametri a bassa magnitudo possono essere critici per i circuiti di ragionamento multi-step. Inoltre, i criteri puramente correlazionali non riescono a catturare le interazioni cross-layer essenziali per l'inferenza di tipo chain-of-thought (CoT), portando frequentemente al "collasso dello strato" (layer collapse), dove interi moduli perdono la propria funzione.

Metodologia: Causal Attribution Pruning (CAP)
Gli autori introducono il Causal Attribution Pruning (CAP), un metodo di pruning one-shot, privo di addestramento (training-free), progettato per preservare le prestazioni di ragionamento. A differenza dei metodi che classificano i pesi in base a proprietà statiche, CAP misura l'impatto causale di specifici componenti del modello sulle prestazioni del task. Il framework opera in tre fasi:

1. Calibrazione e Punteggio Causale (Fase 1):

   • Un piccolo dataset di calibrazione ($D_{cal}$) viene costruito campionando da benchmark orientati al ragionamento: GSM8K (matematica), StrategyQA (senso comune multi-hop) e ARC-Challenge (scienza).
   • Per ogni testa di attenzione $h_\ell$ nel modello, CAP esegue un esperimento di masking interventistico. La testa viene temporaneamente mascherata durante il forward pass, e l'incremento risultante della cross-entropy loss a livello di token ($\Delta(h_\ell)$) viene misurato rispetto al modello integro.
   • Per garantire la robustezza contro la varianza derivante da campioni di piccole dimensioni, il dataset è partizionato in $K$ sottoinsiemi, e viene calcolato il valore mediano dell'incremento della loss.
   • Questo processo identifica le teste che sono funzionalmente critiche per il ragionamento, distinguendole da quelle che sono semplicemente numericamente grandi ma ridondanti.

2. Conversione dell'Importanza a Livello di Peso (Fase 2):

   • I punteggi causali a livello di testa vengono convertiti in fattori di importanza normalizzati $I(h_\ell)$, mappati in un intervallo $[1.0, 10.0]$. Le teste ad alto impatto ricevono fattori più elevati.
   • Questi fattori vengono propagati ai singoli pesi all'interno delle corrispondenti matrici di proiezione ($W_q, W_k, W_v, W_o$). Il punteggio di importanza per un peso specifico $W_{ij}$ è calcolato come il prodotto della sua magnitudo e del fattore di importanza della regione di testa associata: $S_{ij} = |W_{ij}| \cdot I(h_{region})$.
   • Per gli strati MLP, che mancano di una struttura naturale partizionata per teste, il metodo assegna un'importanza media per strato derivata dalle teste di attenzione co-locate. Questo viene indicato come un'approssimazione.

3. Pruning Globale (Fase 3):

   • Tutti i pesi attraverso il modello vengono classificati in base ai loro punteggi di magnitudo pesati per l'importanza ($S_{ij}$).
   • Una soglia globale viene determinata tramite ricerca binaria per raggiungere la sparsità target (ad es. 10%, 20%, 50%). I pesi al di sotto della soglia vengono azzerati, creando un pattern di sparsità non strutturata compatibile con le infrastrutture di inferenza standard.

Contributi Chiave

• Formalizzazione dell'Attribuzione Causale per il Pruning: Gli autori formalizzano un metodo per misurare l'incremento atteso della loss specifica del task quando le singole teste di attenzione vengono mascherate, andando oltre le proxy correlazionali per quantificare la criticità funzionale.
• Pruning a Livello di Peso Consapevole del Budget: Introducono una procedura che converte i punteggi causali a livello di testa in valori di importanza a livello di peso. Ciò consente un pruning della magnitudo fine-grained sotto un vincolo di sparsità globale, preservando la struttura dei circuiti di ragionamento critici.
• Validazione Empirica: Il metodo è valutato su Llama-3-8B-Instruct e Mistral-7B-Instruct attraverso i benchmark GSM8K, StrategyQA e ARC-Challenge.

Risultati

• Sparsità Moderata (10–20%): CAP generalmente supera il baseline Wanda nel preservare le prestazioni di ragionamento.
  • Su Llama-3-8B-Instruct al 20% di sparsità, CAP ottiene un miglioramento relativo del 61% rispetto a Wanda su ARC-Challenge (70.8% vs 43.9% di accuratezza).
  • Su GSM8K al 10% di sparsità, CAP mostra un miglioramento relativo del 20% rispetto a Wanda (71.5% vs 59.5%).
  • Gli autori notano che, sebbene CAP a volte resti indietro rispetto a Wanda su task specifici (ad es. Llama-3 GSM8K al 20%), esso preserva costantemente le prestazioni nei compiti di inferenza multi-step dove i criteri correlazionali di Wanda falliscono.
• Alta Sparsità (50%): Le prestazioni degradano significativamente per entrambi i metodi. CAP subisce un fallimento catastrofico su Llama-3 (la perplexity balza a 428.2), attribuito al trattamento grossolano degli strati MLP. A questo livello, la soglia globale rimuove pesi MLP essenziali che sono stati assegnati una bassa importanza in base all'approssimazione delle teste di attenzione.
• Perplexity vs. Ragionamento: A sparsità moderata, CAP spesso ottiene un'accuratezza di ragionamento migliore nonostante una perplexity leggermente superiore rispetto a Wanda, suggerendo un disaccoppiamento tra la qualità del modellamento del linguaggio generale e la preservazione del ragionamento specifico del task.
• Limitazioni dell'Architettura: Esperimenti preliminari su modelli Mixture-of-Experts (MoE) suggeriscono che CAP sia meno efficace, poiché il routing dinamico rende l'importanza delle teste non stazionaria.

Significatività e Rivendicazioni
Il paper sostiene che l'attribuzione causale a livello di testa di attenzione sia un criterio superiore per preservare le prestazioni di ragionamento rispetto ai metodi di pruning correlazionali a parità di livelli di sparsità. Misurando esplicitamente il contributo causale delle teste al compito di ragionamento, CAP può identificare e proteggere parametri "numericamente piccoli ma funzionalmente critici" che i metodi basati sulla magnitudo scartano.

Gli autori sottolineano che il loro lavoro è di portata modesta:

• È un metodo training-free, il che lo rende computazionalmente efficiente rispetto agli approcci che richiedono aggiornamenti dei pesi (come SparseGPT).
• È task-aligned, calibrato specificamente su dati di ragionamento piuttosto che su un modellamento del linguaggio generico.
• Riconosce le limitazioni, in particolare la mancanza di attribuzione causale fine-grained per i neuroni MLP e la difficoltà di applicare il metodo ad architetture con routing dinamico (MoE).

Il contributo principale è la dimostrazione che il pruning interventistico e allineato al task può preservare meglio capacità specifiche (ragionamento) rispetto ai metodi ottimizzati per la perplexity generale o le statistiche statiche dei pesi, a condizione che il livello di sparsità rimanga moderato.