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

1. Probleemstelling

Hoewel Large Language Models (LLM's) state-of-the-art prestaties behalen, belemmeren hun enorme omvang en rekenkosten de implementatie in omgevingen met beperkte middelen. Hoewel neuron-pruning een veelbelovende compressietechniek is, richt het merendeel van het bestaande onderzoek zich op algemene modellen. Het blijft onduidelijk of neuronen in task-specifieke modellen (bijvoorbeeld gespecialiseerd in wiskunde of code) uniform bijdragen aan de prestaties, of dat er onderscheidende "task-specifieke neuronen" bestaan. Bovendien zijn de grenzen van pruning voordat modelinstorting optreedt, de aard van structurele instabiliteit (bijvoorbeeld generatielussen) en de effectiviteit van herstel na pruning via fine-tuning niet goed begrepen.

2. Methodologie

Modellen en Datasets

• Modellen: De studie evalueerde Qwen2.5-modellen in twee schalen: 1.5B en 7B.
  • Wiskundemodellen: Qwen2.5-Math-1.5B/7B-Instruct.
  • Codemodellen: Qwen2.5-Coder-1.5B/7B-Instruct.
• Datasets:
  • Doeltaken: GSM8K (Wiskunde) en CodeFeedback-Filtered-Instruction (Code).
  • Afstrepende taken: SQuAD (Vraagbeantwoording) en everyday-Conversational-cleaned (Dialoog).
  • Opmerking: Wiskunde en code werden niet behandeld als wederzijdse afstrepende taken vanwege hun gedeelde onderliggende structurele logica.

Activeringsgebaseerde Selectiviteitsmetriek

De kerninnovatie is een activeringsgebaseerde selectiviteitsscore om neuronen te identificeren die kritiek zijn voor de doeltaak versus die welke geassocieerd zijn met afstrepende taken.

1. Activeringsopname: Neuronactiveringen ($h^{(\ell)}$) uit de intermediaire SwiGLU-laag van Feed-Forward Networks (FFN) worden geëxtraheerd voor zowel doelprompts als afstrepende prompts.
2. Berekening van selectiviteit: Voor elke neuron $n$ in laag $\ell$ wordt een directionele selectiviteitsscore $S_n^{(\ell)}$ berekend:
   $$S_n^{(\ell)} = \frac{\mu_{target,n}^{(\ell)} - \mu_{distractor,n}^{(\ell)}}{\sigma_n^{(\ell)} + \epsilon}$$
   Waar $\mu$ de gemiddelde activatie voorstelt en $\sigma$ de standaardafwijking.
   • Positieve score: Doelrelevant.
   • Negatieve score: Geassocieerd met afstrepende taken (kandidaten voor pruning).

Pruning-strategieën

• Selectieve Pruning: Verwijdert neuronen met de laagste selectiviteitsscores (vriendelijk voor afstrepende taken).
• Willekeurige Pruning: Baseline waarbij neuronen willekeurig worden verwijderd.
• Reverse Pruning: Verwijdert neuronen met de hoogste selectiviteitsscores (task-kritiek) om de noodzaak van specifieke neuronen te testen.
• Implementatie: Structurele pruning wordt fysiek toegepast door rijen/kolommen uit gewichtsmatrices ($W_{gate}, W_{up}, W_{down}$) in MLP-blokken te snijden. Pruning is uniform over de lagen om hardware-efficiëntie te behouden (deelbaar door 128).

Herstelmechanisme

Na pruning ondergaan de modellen LoRA (Low-Rank Adaptation) fine-tuning om prestaties te herstellen.

• Wiskundemodellen: $r=8$, LR=$2\times10^{-5}$.
• Codemodellen: $r=16$, LR=$2\times10^{-4}$.
• Getraind gedurende 2 epochs op de trainingssets van de doeltaak.

Evaluatiemetrieken

• Nauwkeurigheid: Exact Match (EM) voor Wiskunde; Pass@1 op HumanEval voor Code.
• Stabiliteit: "Trap"-detectie (degeneratie-lussen, ontbrekende EOS-tokens, repetitieve tekst).
• Prestaties op afstrepende taken: Semantische gelijkenis (BERTScore + SBERT) om taakspecificiteit te waarborgen.
• Ressourcen: VRAM-gebruik en inferentiethroughput (tokens/sec).

---

3. Belangrijkste Bijdragen

1. Task-specifieke neuron-specialisatie: Empirisch bewijs dat task-specifieke LLM's een kleine, highly gespecialiseerde subset van neuronen bevatten (ongeveer 10%) die kritiek zijn voor prestaties, terwijl de rest overbodig of gerelateerd is aan afstrepende taken.
2. Robuustheidsthorshouding: Identificatie van een kritieke pruning-drempel (15–20%). Daarboven daalt de nauwkeurigheid scherp en neemt structurele instabiliteit (generatielussen) exponentieel toe.
3. Validatie van Reverse Pruning: Aangetoond dat het verwijderen van slechts ongeveer 10% van de meest task-relevante neuronen volledige modelinstorting veroorzaakt, terwijl het verwijderen van 30–35% van de minst relevante neuronen aanzienlijke functionaliteit behoudt.
4. Herstel via Fine-tuning: Aangetoond dat LoRA fine-tuning de prestaties aanzienlijk kan herstellen, zelfs bij agressief gepreinde modellen, waardoor zowel verlies aan nauwkeurigheid als generatie-instabiliteit wordt gemitigeerd.

---

4. Belangrijkste Resultaten

A. Selectieve versus Willekeurige versus Reverse Pruning

• Selectief versus Willekeurig: Selectieve pruning presteerde consistent beter dan willekeurige pruning over alle ratios, wat bevestigt dat activeringsgebaseerde selectie een systematisch voordeel biedt.
• Reverse Pruning: Het verwijderen van de top 10% van task-specifieke neuronen veroorzaakte volledige prestatie-instorting (nagenoeg nul nauwkeurigheid) en frequente generatiefouten. Dit bewijst dat kritieke taakinformatie geconcentreerd is in een tiny fractie van het netwerk.
• Selectieve Pruning: Staarde tot 35% pruning toe met een geleidelijke daling in nauwkeurigheid, hoewel de stabiliteit bij hogere ratios leed.

B. Modelgrootte en Robuustheid

• 7B versus 1.5B: Grotere modellen (7B) vertoonden hogere redundantie en konden hogere pruning-ratio's tolereren met minder relatief verlies aan nauwkeurigheid in vergelijking met kleinere modellen (1.5B).
• Drempel: Beide schalen vertoonden een scherpe degradatie in prestaties en een piek in "traps" (generatielussen) wanneer pruning 15–20% overschreed.

C. Herstel na Pruning (LoRA)

• Fine-tuning herstelde succesvol de prestaties over alle pruning-niveaus.
• Tegenintuïtieve bevinding: In sommige gevallen (bijvoorbeeld Qwen 1.5B Code) presteerden matig gepreinde modellen (10–15%) na fine-tuning beter dan licht gepreinde modellen (5%), wat suggereert dat het verwijderen van enige redundantie het fine-tuning-proces soms kan helpen.
• Stabiliteit: Fine-tuning verlaagde de frequentie van generatielussen aanzienlijk en herstelde structurele stabiliteit.

D. Resource-efficiëntie

• VRAM & Throughput: Pruning leidde tot lineaire reducties in VRAM-gebruik en toenames in inferentiethroughput.
  • Voorbeeld (Qwen-7B): 35% pruning verlaagde VRAM van 14,26 GB naar 10,56 GB en verhoogde de throughput van 35,83 naar 42,70 tokens/sec.

E. Laag-voor-laag Analyse

• Vroege lagen: Overwegend gericht op afstrepende taken (algemene verwerking).
• Middenlagen: Overgang naar doelrelevante activeringen.
• Late lagen: Wiskundemodellen bleven sterk doel-dominant; codemodellen vertoonden een mix, waarbij ze soms terugkeerden naar dominantie van afstrepende taken.

---

5. Betekenis en Implicaties

• Neuron-specialisatie: De studie levert sterk bewijs dat task-specifieke LLM's niet uniform overbodig zijn; ze vertrouwen op een "kern" van gespecialiseerde neuronen. Dit daagt de aanname uit dat alle neuronen evenveel bijdragen.
• Pruning-grenzen: Het definieert een praktische "veilige zone" voor pruning (tot ongeveer 15%) voordat structurele instabiliteit onbeheersbaar wordt zonder herstel.
• Herstelstrategie: Het valideert LoRA fine-tuning als een essentiële, low-cost rehabilitatiestrategie voor gepreinde modellen, waardoor agressieve pruning haalbaar wordt voor implementatie.
• Hardware-optimalisatie: Door neuronen fysiek te verwijderen (structurele pruning) in plaats van alleen gewichten te maskeren, biedt de aanpak tastbare hardwarevoordelen (verminderde VRAM, hogere throughput) die geschikt zijn voor implementatie aan de rand (edge).

Conclusie

Het paper concludeert dat task-specifieke taalmodellen een distincte architectuur bezitten waarbij een kleine subset van neuronen de taakprestaties drijft. Hoewel agressieve pruning risico's op modelinstorting en instabiliteit met zich meebrengt, biedt een combinatie van activeringsgebaseerde selectieve pruning en post-pruning LoRA fine-tuning een robuust pad om efficiënte, gespecialiseerde modellen te implementeren zonder kritieke capaciteiten te offeren.