TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Sintesi Tecnica: TAPIOCA – Perché la Potatura Consapevole del Compito Migliora le Capacità del Modello OOD

Enunciato del Problema
Recenti lavori hanno dimostrato che la rimozione di livelli dai Large Language Models (LLM) al momento dell'inferenza—specificamente attraverso metodi di potatura consapevoli del compito come TALE (Task-Aware Layer Elimination)—può migliorare le prestazioni su compiti specifici senza riaddestramento. Tuttavia, i meccanismi che guidano questi guadagni rimangono poco chiari. Una domanda centrale è se questi miglioramenti derivino da benefici generali di compressione, effetti di regolarizzazione o interazioni specifiche tra le rappresentazioni interne del modello e la distribuzione degli input. Inoltre, è sconosciuto se questi guadagni siano universali o specifici per certi tipi di spostamenti della distribuzione.

Metodologia
Gli autori indagano questo fenomeno utilizzando una combinazione di esperimenti sintetici controllati e valutazioni su larga scala su LLM finetunati.

1. Setting di Regressione Controllato: Lo studio inizia con un transformer a 12 livelli addestrato da zero sulla regressione lineare in contesto. La distribuzione di addestramento è definita esplicitamente (coefficienti campionati da $U(-1, 1)$). Gli autori valutano il modello sia su dati in-distribution (ID) (stessa distribuzione) sia su dati out-of-distribution (OOD) (coefficienti campionati da intervalli più ampi $U(-\sigma, \sigma)$ dove $\sigma > 1$). Applicano TALE, un algoritmo di potatura greedy che rimuove livelli per massimizzare l'accuratezza di validazione su una specifica distribuzione.
2. Valutazione su LLM su Larga Scala: I risultati sono estesi a LLM finetunati (Llama 3.1 8B e GPT-OSS 120B) specializzati nel ragionamento matematico e nella generazione di codice. Qui, la distribuzione "adattata" è il compito di finetuning (ad esempio, MATH500), e i compiti OOD includono benchmark generali come MMLU e BoolQ.
3. Analisi Geometrica e Causale: Per spiegare perché la potatura funziona, gli autori analizzano la geometria degli stati nascosti. Misurano le norme rappresentazionali livello per livello e le distanze tra token a coppie. Utilizzano anche surrogati lineari locali per approssimare le trasformazioni dei livelli ed eseguono interventi causali, come l'inserimento di mappe surrogate inverse e il ridimensionamento degli aggiornamenti residui, per isolare l'effetto della distorsione geometrica.

Risultati Chiave

• Guadagni Specifici per OOD: Il principale risultato empirico è che la potatura consapevole del compito produce guadagni di prestazioni solo sotto spostamento della distribuzione. Quando valutato sulla distribuzione alla quale il modello è adattato (ID), TALE non riesce a migliorare le prestazioni o le degrada attivamente rimuovendo livelli necessari per una precisione ottimale. Al contrario, su dati OOD, TALE identifica e rimuove consistentemente livelli che migliorano l'accuratezza.
• Disallineamento della Geometria delle Rappresentazioni: Gli autori dimostrano che gli input OOD inducono una geometria rappresentazionale distorta rispetto agli input ID. Nello specifico, gli input OOD causano norme degli stati nascosti livello per livello e distanze a coppie che si discostano sistematicamente dai profili osservati sui dati ID. Nel setting di regressione, gli input OOD causano un'espansione delle distanze attraverso i livelli; negli LLM, le traiettorie OOD divergono dalle traiettorie ID nei livelli centrali e tardivi.
• Meccanismo di Amplificazione: Lo studio identifica che certi livelli agiscono come "amplificatori dipendenti dalla distribuzione". Sui input ID, questi livelli eseguono raffinamenti utili. Tuttavia, sui input OOD, gli stessi livelli esibiscono un'amplificazione ad alto guadagno e basso rango, spingendo le rappresentazioni ulteriormente lontano dalla geometria adattata al compito. TALE migliora le prestazioni rimuovendo questi specifici livelli, prevenendo così l'amplificazione della distorsione geometrica.
• Allineamento, non solo Contrazione: La potatura non sopprime semplicemente le norme di attivazione. Negli esperimenti in cui la direzione dello spostamento della distribuzione è invertita (addestramento su distribuzioni ampie, test su distribuzioni strette), la potatura espande le distanze rappresentazionali per allinearle al profilo ID. Il meccanismo fondamentale è il riallineamento delle rappresentazioni OOD verso la geometria appresa durante l'adattamento, indipendentemente dal fatto che ciò richieda contrazione o espansione.
• Evidenza Causale: Gli interventi causali confermano che la distorsione geometrica è il motore della perdita di prestazioni. Inserire una mappa surrogate inversa dopo un livello potato (per annullare l'effetto geometrico del livello) ripristina le prestazioni OOD quasi al livello del modello potato. Analogamente, ridimensionare l'aggiornamento residuo di un livello selezionato migliora continuamente l'accuratezza man mano che il contributo del livello viene attenuato, escludendo artefatti topologici o semplice ridondanza come unica causa.

Significato e Affermazioni
Il paper afferma di fornire una spiegazione geometrica innovativa per il successo della potatura al momento dell'inferenza. Sostiene che l'importanza dei livelli non è fissa ma è condizionata sia dal compito sia dalla distribuzione degli input. Gli autori ipotizzano che i modelli preaddestrati codifichino molteplici funzioni specializzate per il compito, ciascuna associata a una geometria rappresentazionale locale. Gli input ID si adattano a questa geometria, mentre gli input OOD la violano.

Il significato di questo lavoro risiede in:

1. Chiarezza sui Limiti della Potatura: Stabilisce che la potatura consapevole del compito non è un potenziatore universale delle prestazioni, ma un meccanismo correttivo per lo spostamento della distribuzione.
2. Interpretazione Geometrica: Sposta la comprensione della ridondanza dei livelli dai conteggi statici dei parametri alle distorsioni geometriche dinamiche e dipendenti dalla distribuzione.
3. Insight Meccanicistico: Fornisce evidenze causali che i guadagni di prestazioni derivano dalla correzione della traiettoria delle rappresentazioni interne per farle corrispondere alla geometria adattata del modello, piuttosto che dall'aumento della capacità del modello o dalla riduzione del rumore.

Gli autori concludono che la potatura migliora le prestazioni OOD attenuando i livelli che amplificano il disallineamento tra la distribuzione degli input e la geometria del compito appresa dal modello, riallineando efficacemente lo stato interno del modello con le condizioni in cui è stato ottimizzato.