Exploring the Limits of Pruning: Task-Specific Neurons, Model Collapse, and Recovery in Task-Specific Large Language Models

1. Enunciato del Problema

Sebbene i Large Language Models (LLM) raggiungano prestazioni all'avanguardia, le loro dimensioni massive e i costi computazionali ostacolano il deployment in ambienti con risorse limitate. Sebbene il pruning dei neuroni sia una tecnica di compressione promettente, la maggior parte della ricerca esistente si concentra su modelli generici. Rimane poco chiaro se i neuroni nei modelli specifici per compito (ad esempio, specializzati in matematica o codice) contribuiscano uniformemente alle prestazioni o se esistano distinti "neuroni specifici per compito". Inoltre, non sono ben compresi i limiti del pruning prima del collasso del modello, la natura dell'instabilità strutturale (ad esempio, loop di generazione) e l'efficacia del recupero post-pruning tramite fine-tuning.

2. Metodologia

Modelli e Dataset

• Modelli: Lo studio ha valutato i modelli Qwen2.5 in due scale: 1.5B e 7B.
  • Modelli Matematici: Qwen2.5-Math-1.5B/7B-Instruct.
  • Modelli Codice: Qwen2.5-Coder-1.5B/7B-Instruct.
• Dataset:
  • Compiti Target: GSM8K (Matematica) e CodeFeedback-Filtered-Instruction (Codice).
  • Compiti Distrattori: SQuAD (Risposta alle domande) e everyday-Conversational-cleaned (Dialogo).
  • Nota: Matematica e Codice non sono stati trattati come distrattori reciproci a causa della loro logica strutturale sottostante condivisa.

Metrica di Selettività Basata sull'Attivazione

L'innovazione centrale è un punteggio di selettività basato sull'attivazione per identificare i neuroni critici per il compito target rispetto a quelli associati ai distrattori.

1. Cattura dell'Attivazione: Vengono estratte le attivazioni dei neuroni ($h^{(\ell)}$) dal layer intermedio SwiGLU delle Reti Feed-Forward (FFN) sia per i prompt target che per quelli distrattori.
2. Calcolo della Selettività: Per ogni neurone $n$ nel layer $\ell$, viene calcolato un punteggio di selettività direzionale $S_n^{(\ell)}$:
   $$S_n^{(\ell)} = \frac{\mu_{target,n}^{(\ell)} - \mu_{distractor,n}^{(\ell)}}{\sigma_n^{(\ell)} + \epsilon}$$
   Dove $\mu$ rappresenta l'attivazione media e $\sigma$ è la deviazione standard.
   • Punteggio Positivo: Rilevante per il target.
   • Punteggio Negativo: Associato al distrattore (candidati per il pruning).

Strategie di Pruning

• Pruning Selettivo: Rimuove i neuroni con i punteggi di selettività più bassi (favorevoli ai distrattori).
• Pruning Casuale: Baseline in cui i neuroni vengono rimossi casualmente.
• Pruning Inverso: Rimuove i neuroni con i punteggi di selettività più alti (critici per il compito) per testare la necessità di neuroni specifici.
• Implementazione: Il pruning strutturale viene applicato fisicamente tagliando righe/colonne dalle matrici dei pesi ($W_{gate}, W_{up}, W_{down}$) nei blocchi MLP. Il pruning è uniforme su tutti i layer per mantenere l'efficienza hardware (divisibile per 128).

Meccanismo di Recupero

Dopo il pruning, i modelli subiscono un fine-tuning LoRA (Low-Rank Adaptation) per recuperare le prestazioni.

• Modelli Matematici: $r=8$, LR=$2\times10^{-5}$.
• Modelli Codice: $r=16$, LR=$2\times10^{-4}$.
• Addestrati per 2 epoche sui set di addestramento del compito target.

Metriche di Valutazione

• Accuratezza: Exact Match (EM) per la Matematica; Pass@1 su HumanEval per il Codice.
• Stabilità: Rilevazione di "trappole" (loop di degenerazione, token EOS mancanti, testo ripetitivo).
• Prestazioni sui Distrattori: Similarità semantica (BERTScore + SBERT) per garantire la specificità del compito.
• Risorse: Utilizzo della VRAM e throughput di inferenza (token/sec).

---

3. Contributi Chiave

1. Specializzazione dei Neuroni Specifici per Compito: Evidenza empirica che gli LLM specifici per compito contengono un piccolo sottoinsieme altamente specializzato di neuroni (circa il 10%) critico per le prestazioni, mentre il resto è ridondante o correlato ai distrattori.
2. Soglia di Robustezza: Identificazione di una soglia critica di pruning (15–20%). Oltre questo limite, l'accuratezza crolla bruscamente e l'instabilità strutturale (trappole di generazione) aumenta esponenzialmente.
3. Validazione del Pruning Inverso: Dimostrato che rimuovere solo ~10% dei neuroni più rilevanti per il compito causa un completo collasso del modello, mentre la rimozione del 30–35% dei neuroni meno rilevanti preserva funzionalità significative.
4. Recupero tramite Fine-Tuning: Mostrato che il fine-tuning LoRA può recuperare sostanzialmente le prestazioni anche nei modelli aggressivamente prunati, mitigando sia la perdita di accuratezza che l'instabilità di generazione.

---

4. Risultati Chiave

A. Pruning Selettivo vs. Casuale vs. Inverso

• Selettivo vs. Casuale: Il pruning selettivo ha costantemente superato il pruning casuale in tutti i rapporti, confermando che la selezione basata sull'attivazione fornisce un vantaggio sistematico.
• Pruning Inverso: La rimozione del 10% superiore dei neuroni specifici per il compito ha causato un completo collasso delle prestazioni (accuratezza quasi zero) e frequenti fallimenti di generazione. Ciò dimostra che le informazioni critiche del compito sono concentrate in una minuscola frazione della rete.
• Pruning Selettivo: Ha permesso fino al 35% di pruning con un declino graduale dell'accuratezza, sebbene la stabilità ne abbia risentito a rapporti più elevati.

B. Scala del Modello e Robustezza

• 7B vs. 1.5B: I modelli più grandi (7B) hanno mostrato una maggiore ridondanza e hanno potuto tollerare rapporti di pruning più elevati con una perdita di accuratezza relativa inferiore rispetto ai modelli più piccoli (1.5B).
• Soglia: Entrambe le scale hanno mostrato un degrado netto delle prestazioni e un picco di "trappole" (loop di generazione) quando il pruning ha superato il 15–20%.

C. Recupero Post-Pruning (LoRA)

• Il fine-tuning ha recuperato con successo le prestazioni a tutti i livelli di pruning.
• Riscontro Controintuitivo: In alcuni casi (ad esempio, Qwen 1.5B Codice), i modelli moderatamente prunati (10–15%) hanno superato i modelli leggermente prunati (5%) dopo il fine-tuning, suggerendo che la rimozione di una certa ridondanza può talvolta aiutare il processo di fine-tuning.
• Stabilità: Il fine-tuning ha ridotto significativamente la frequenza delle trappole di generazione, ripristinando la stabilità strutturale.

D. Efficienza delle Risorse

• VRAM e Throughput: Il pruning ha portato a riduzioni lineari nell'uso della VRAM e aumenti nel throughput di inferenza.
  • Esempio (Qwen-7B): Il 35% di pruning ha ridotto la VRAM da 14,26 GB a 10,56 GB e aumentato il throughput da 35,83 a 42,70 token/sec.

E. Analisi per Livelli

• Livelli Iniziali: Prevalentemente orientati ai distrattori (elaborazione generale).
• Livelli Medi: Transizione verso attivazioni rilevanti per il target.
• Livelli Finali: I modelli matematici sono rimasti fortemente dominanti per il target; i modelli codice hanno mostrato un mix, occasionalmente revertando alla dominanza dei distrattori.

---

5. Significato e Implicazioni

• Specializzazione dei Neuroni: Lo studio fornisce prove solide che gli LLM specifici per compito non sono uniformemente ridondanti; si basano su un "nucleo" di neuroni specializzati. Ciò sfida l'assunzione che tutti i neuroni contribuiscano in modo uguale.
• Limiti del Pruning: Definisce una "zona sicura" pratica per il pruning (fino a ~15%) prima che l'instabilità strutturale diventi ingestibile senza recupero.
• Strategia di Recupero: Convalida il fine-tuning LoRA come strategia di riabilitazione essenziale e a basso costo per i modelli prunati, rendendo il pruning aggressivo fattibile per il deployment.
• Ottimizzazione Hardware: Rimuovendo fisicamente i neuroni (pruning strutturale) invece di mascherare semplicemente i pesi, l'approccio offre benefici hardware tangibili (riduzione VRAM, throughput più elevato) adatti al deployment edge.

Conclusione

Il paper conclude che i modelli linguistici specifici per compito possiedono un'architettura distinta in cui un piccolo sottoinsieme di neuroni guida le prestazioni del compito. Sebbene il pruning aggressivo rischi il collasso e l'instabilità del modello, una combinazione di pruning selettivo basato sull'attivazione e fine-tuning LoRA post-pruning offre un percorso robusto per deployare modelli specializzati ed efficienti senza sacrificare le capacità critiche.