Demystifying When Pruning Works via Representation Hierarchies

1. Il Problema

Il pruning (potatura) delle reti neurali è una tecnica consolidata per migliorare l'efficienza computazionale e ridurre l'uso di memoria nei Large Language Models (LLM), rimuovendo parametri o architetture meno importanti. Tuttavia, il paper evidenzia una discrepanza fondamentale nelle prestazioni dei modelli potati:

• Compiti non generativi: I modelli potati mantengono spesso prestazioni elevate in compiti come la classificazione, il retrieval (recupero informazioni) e la selezione a scelta multipla.
• Compiti generativi: Gli stessi modelli subiscono un crollo catastrofico delle prestazioni in compiti di generazione di testo (es. scrittura creativa, risoluzione di problemi matematici passo-passo, coding).

La comunità scientifica non aveva ancora una spiegazione teorica chiara del perché questa divergenza si verifichi, nonostante l'uso di tecniche di pruning simili.

2. Metodologia

Gli autori analizzano il pruning non solo a livello di parametri, ma da una prospettiva di gerarchia delle rappresentazioni interne del modello. Decompongono il processo di inferenza di un LLM in tre spazi sequenziali distinti:

1. Spazio di Embedding: Le rappresentazioni nascoste (hidden states) all'interno del modello.
2. Spazio Logit: Gli output pre-softmax (valori lineari prima della normalizzazione).
3. Spazio di Probabilità: Le distribuzioni finali post-softmax (dove avviene la selezione del token).

Approccio Analitico:

• Analisi Empirica: Hanno misurato la similarità (cosine similarity) e la divergenza (KL Divergence) tra le rappresentazioni di un modello originale e uno potato attraverso i diversi strati e passaggi di generazione.
• Analisi Teorica: Hanno utilizzato espansioni di Taylor del secondo ordine per derivare teoremi che quantificano come le perturbazioni indotte dal pruning si propagano e vengono amplificate attraverso le trasformazioni lineari e non lineari del modello.
• Esperimenti: Hanno testato diverse strategie di pruning (rimozione di layer interi, pruning non strutturato, pruning strutturato) su modelli come Mistral, Llama e Qwen, valutandoli su benchmark generativi (es. GSM8K, HumanEval) e non generativi (es. MMLU, HellaSwag, Retrieval).

3. Contributi Chiave e Risultati Principali

A. Robustezza degli Spazi di Embedding e Logit

L'analisi rivela che le perturbazioni introdotte dal pruning rimangono minime e stabili negli spazi di embedding e logit.

• La trasformazione lineare dalla rappresentazione nascosta agli logit (tramite l'LM Head) tende ad attenuare le deviazioni.
• Di conseguenza, per compiti che dipendono direttamente da questi spazi (come il retrieval basato su similarità di embedding o la classificazione basata su un piccolo set di logit), il pruning è efficace e le prestazioni non ne risentono significativamente.

B. Amplificazione Non Lineare nello Spazio di Probabilità

Il punto critico è la trasformazione non lineare dagli logit alle probabilità tramite la funzione Softmax.

• Teorema 2: Gli autori dimostrano che la funzione softmax agisce come un amplificatore delle perturbazioni. Anche piccole variazioni negli logit ($\Delta z$) possono causare grandi deviazioni nella distribuzione di probabilità ($\Delta p$), specialmente quando la varianza degli errori è alta.
• La divergenza KL (Kullback-Leibler) tra la distribuzione originale e quella potata cresce esponenzialmente a causa di questa non linearità.

C. Propagazione dell'Errore nella Generazione Autoregressiva

Per i compiti generativi, il modello opera in un ciclo di feedback:

1. L'errore iniziale nella distribuzione di probabilità porta alla selezione di un token errato.
2. Questo token errato diventa parte del contesto storico per il passo successivo.
3. Il modello genera basandosi su un contesto "inquinato", amplificando l'errore a ogni passo successivo.

• Risultato: Una piccola deviazione iniziale, amplificata dal softmax e propagata dal ciclo autoregressivo, porta al collasso della coerenza del testo generato (es. ripetizioni, allucinazioni, perdita di senso).

D. Stabilità dei Sottospazi di Probabilità nei Compiti Non Generativi

Nei compiti non generativi (es. scelta multipla A/B/C/D), il modello deve scegliere tra un numero molto ridotto di token candidati rispetto all'intero vocabolario.

• Anche se la distribuzione globale delle probabilità cambia drasticamente, i sottospazi relativi ai token candidati specifici tendono a rimanere stabili.
• Di conseguenza, il token selezionato (argmax) rimane spesso corretto, spiegando la robustezza del pruning in questi scenari.

4. Significato e Implicazioni

Questo lavoro offre una spiegazione teorica unificata per un fenomeno empirico osservato ma non compreso:

1. Guida Pratica: Sconsiglia l'uso del pruning "senza riaddestramento" (training-free) per compiti generativi complessi, a meno che non siano adottate strategie di compensazione avanzate. Al contrario, valida l'uso del pruning per ottimizzare modelli per compiti di retrieval e classificazione.
2. Nuova Prospettiva: Sposta il focus dalla semplice rimozione dei parametri all'analisi di come le perturbazioni viaggiano attraverso le diverse fasi di trasformazione del modello (lineare vs non lineare).
3. Impatto sulla Ricerca: Suggerisce che le future tecniche di compressione per la generazione devono concentrarsi specificamente sulla stabilizzazione dello spazio di probabilità e sulla mitigazione della propagazione dell'errore nel ciclo autoregressivo, piuttosto che solo sulla riduzione della dimensione del modello.

In sintesi, il paper dimostra che il pruning fallisce nella generazione non perché il modello perde "conoscenza" generale, ma perché la non linearità del softmax trasforma piccole imperfezioni strutturali in errori catastrofici che si auto-alimentano durante la generazione sequenziale.