Demystifying When Pruning Works via Representation Hierarchies

Probleemstelling

Netwerkpruning (het verwijderen van minder belangrijke parameters of architectuurelementen) wordt vaak gebruikt om de efficiëntie van Large Language Models (LLMs) te verbeteren. Echter, de prestaties van gepreunde modellen zijn inconsistent afhankelijk van het type taak:

• Niet-generatieve taken (zoals retrieval, tekstclassificatie en meerkeuzevragen): Gepreunde modellen behouden vaak een sterke prestatie.
• Generatieve taken (zoals het genereren van tekst, code of het oplossen van wiskundeproblemen stap-voor-stap): Deze modellen lijden vaak onder catastrofale prestatieverlies, zelfs bij matige pruning.

De vraag die dit artikel beantwoordt is: Waarom faalt pruning bij generatieve taken terwijl het werkt bij niet-generatieve taken? Bestaande literatuur legt dit vaak uit door te verwijzen naar foutpropagatie, maar de onderliggende mechanismen op representatieniveau blijven onduidelijk.

Methodologie

De auteurs analyseren netwerkpruning vanuit een representatie-hiërarchie perspectief. Ze decomponeren de interne berekening van een taalmodel in drie opeenvolgende ruimtes:

1. Embedding-ruimte: De verborgen representaties (hidden states) binnen het model.
2. Logit-ruimte: De lineaire uitkomsten vóór de softmax-functie (pre-softmax).
3. Kansruimte (Probability): De post-softmax verdeling over het vocabulaire.

Analyse-aanpak:

• Empirische Visualisatie: Ze meten de cosinus-ähnelijkheid en de KL-divergentie tussen de oorspronkelijke en gepreunde modellen op elk van deze drie niveaus tijdens het generatieproces.
• Theoretische Analyse: Ze gebruiken Taylor-reeks-expansies (lokaal) om af te leiden hoe perturbaties (verstoringen) veroorzaakt door pruning zich door de verschillende ruimtes voortplanten en versterken.
• Experimenten: Ze testen verschillende pruning-strategieën (zoals het verwijderen van volledige lagen, "Layer Drop", en fijner granulaire pruning zoals "Wanda") op modellen zoals Mistral-7B, Llama-3 en Qwen-2.5.

Belangrijkste Bevindingen en Resultaten

1. Hiërarchische Robuustheid

De studie onthult een duidelijke hiërarchie in de gevoeligheid voor pruning:

• Embedding-ruimte: Blijft grotendeels robuust. Zelfs bij het verwijderen van een aanzienlijk deel van de parameters vertonen de verborgen representaties slechts kleine afwijkingen.
• Logit-ruimte: De lineaire transformatie van embedding naar logit (via de LM-head) verzwakt de perturbaties zelfs verder. De cosinus-ähnelijkheid tussen de oorspronkelijke en gepreunde logits is vaak zelfs hoger dan in de embedding-ruimte.
• Kansruimte (Probability): Dit is het kritieke punt. De niet-lineaire transformatie (softmax) van logits naar kansen versterkt de kleine afwijkingen enorm. Zelfs kleine verschillen in de logits leiden tot grote verschuivingen in de kansverdeling, vooral omdat de softmax-functie gevoelig is voor variatie in de input.

2. Het Mechanisme van Falen bij Generatieve Taken

Bij generatieve taken (autoregressieve generatie) treden twee versterkende factoren op:

• Niet-lineaire Amplificatie: De softmax-functie versterkt de perturbaties in de kansruimte.
• Foutpropagatie: Omdat generatieve taken iteratief zijn, worden fouten die in een vroeg tijdstap worden geïntroduceerd, doorgegeven aan de volgende stappen via de context (key-value cache). Omdat de gegenereerde tokens van het gepreunde model afwijken, verandert de input voor de volgende stap, wat leidt tot een cumulatief en exponentieel groeiend afwijking (error accumulation). Dit resulteert vaak in onzin, herhaling of het volledig instorten van de generatie na slechts een paar stappen.

3. Succes bij Niet-Generatieve Taken

Niet-generatieve taken (zoals retrieval of meerkeuzevragen) zijn robuust tegen pruning omdat:

• Ze vaak slechts één beslissingstap vereisen (geen iteratieve feedbacklus).
• Ze vaak werken op een beperkte subset van het vocabulaire (bijv. alleen de A/B/C/D opties in een meerkeuzevraag). Zelfs als de globale kansverdeling verschuift, blijven de relatieve kansen van de relevante categorie-tokens stabiel genoeg om de juiste keuze te maken.
• Ze vertrouwen direct op de robuuste embedding- of logit-ruimtes, waar de perturbaties klein blijven.

Bijdragen

1. Identificatie van de Discrepantie: Het werk bevestigt en kwantificeert de fundamentele ongelijkheid in de effectiviteit van pruning tussen generatieve en niet-generatieve taken.
2. Mechanistische Uitleg: Het identificeert de niet-lineaire mapping van logits naar kansen als de sleutelfactor die perturbaties versterkt, in plaats van alleen te verwijzen naar foutpropagatie of modelgrootte.
3. Theoretisch Kader: Het biedt wiskundige theorema's (gebaseerd op Taylor-expansies) die de versterking van fouten in de kansruimte verklaren en de stabiliteit in de embedding/logit-ruimtes kwantificeren.
4. Praktische Richtlijnen: Het biedt richtlijnen voor het toepassen van pruning:
   • Veilig toepasbaar voor retrieval, classificatie en meerkeuze-taken.
   • Risicovol en vaak ongeschikt voor open-ended generatie zonder aanvullende fine-tuning of herstelstrategieën.

Significantie

Deze studie is cruciaal voor de efficiënte implementatie van LLMs. Het weerlegt de aanname dat een model dat goed presteert op classificatie-benchmarks (niet-generatief) ook goed zal presteren bij generatieve taken na compressie. Het biedt ontwikkelaars een dieper inzicht in waarom generatieve modellen instorten bij pruning en benadrukt dat compressiestrategieën moeten worden afgestemd op het specifieke gebruiksscenario (generatief vs. niet-generatief). Het stelt ook een theoretische basis voor toekomstig onderzoek naar het stabiliseren van generatieve modellen na compressie, bijvoorbeeld door de niet-lineaire transformaties te regulariseren of door post-pruning fine-tuning.