TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Technische Samenvatting: TAPIOCA – Waarom Taakbewuste Pruning de OOD-Capaciteit van Modellen Verbetert

Probleemstelling
Recent werk heeft aangetoond dat het verwijderen van lagen uit Large Language Models (LLM's) tijdens de inferentie – specifiek via taakbewuste pruning-methoden zoals TALE (Task-Aware Layer Elimination) – de prestaties op specifieke taken kan verbeteren zonder opnieuw te trainen. De mechanismen die deze winsten aandrijven, blijven echter onduidelijk. Een centrale vraag is of deze verbeteringen voortkomen uit algemene compressievoordelen, regularisatie-effecten of specifieke interacties tussen de interne representaties van het model en de invoerverdeling. Bovendien is het onbekend of deze winsten universeel zijn of specifiek voor bepaalde soorten data-verschuivingen.

Methodologie
De auteurs onderzoeken dit fenomeen met een combinatie van gecontroleerde synthetische experimenten en grootschalige evaluaties op fijngefineerde LLM's.

1. Gecontroleerde Regressie-Setting: De studie begint met een 12-laags transformer dat vanaf nul is getraind op contextgebonden lineaire regressie. De trainingsverdeling is expliciet gedefinieerd (coëfficiënten bemonsterd uit $U(-1, 1)$). De auteurs evalueren het model op zowel in-distribution (ID) data (dezelfde verdeling) als out-of-distribution (OOD) data (coëfficiënten bemonsterd uit bredere intervallen $U(-\sigma, \sigma)$ waarbij $\sigma > 1$). Zij passen TALE toe, een greedige pruning-algoritme dat lagen verwijdert om de validatie-accuraatheid op een specifieke verdeling te maximaliseren.
2. Grootschalige LLM-Evaluatie: De bevindingen worden uitgebreid naar fijngefineerde LLM's (Llama 3.1 8B en GPT-OSS 120B) die gespecialiseerd zijn in wiskundig redeneren en codegeneratie. Hier is de "aangepaste" verdeling de fijnfineer-taak (bijvoorbeeld MATH500), en OOD-taken omvatten algemene benchmarks zoals MMLU en BoolQ.
3. Geometrische en Causale Analyse: Om te verklaren waarom pruning werkt, analyseren de auteurs de geometrie van verborgen toestanden. Zij meten laagniveau representatienormen en paarwise token-afstanden. Zij maken ook gebruik van lokale lineaire surrogaten om laagtransformaties te benaderen en voeren causale interventies uit, zoals het invoegen van inverse surrogaatkaarten en het herschalen van residu-updates, om het effect van geometrische vervorming te isoleren.

Belangrijkste Bevindingen

• OOD-specifieke Winsten: De primaire empirische bevinding is dat taakbewuste pruning prestatiewinsten oplevert alleen onder verdelingsverschuiving. Bij evaluatie op de verdeling waarop het model is aangepast (ID), faalt TALE om de prestaties te verbeteren of verslechtert deze actief door lagen te verwijderen die noodzakelijk zijn voor optimale precisie. Daarentegen identificeert en verwijdert TALE op OOD-data consistent lagen die de nauwkeurigheid verbeteren.
• Mismatch in Representatiegeometrie: De auteurs tonen aan dat OOD-invoeren een vervormde representatiegeometrie induceert in vergelijking met ID-invoeren. Specifiek zorgen OOD-invoeren ervoor dat laagniveau verborgen-toestandnormen en paarwise afstanden systematisch afwijken van de profielen die op ID-data worden waargenomen. In de regressiesetting zorgen OOD-invoeren voor een inflatie van afstanden over de lagen heen; bij LLM's divergeren OOD-trajecten in de middelste tot late lagen van ID-trajecten.
• Versterkingsmechanisme: De studie identificeert dat bepaalde lagen fungeren als "verdelingsafhankelijke versterkers". Op ID-invoeren voeren deze lagen nuttige verfijningen uit. Echter, op OOD-invoeren vertonen dezelfde lagen hoog-gain, low-rank versterking, waardoor representaties verder worden geduwd van de taak-aangepaste geometrie. TALE verbetert de prestaties door deze specifieke lagen te verwijderen, waardoor de versterking van geometrische vervorming wordt voorkomen.
• Alignering, Niet Alleen Contractie: Pruning onderdrukt niet simpelweg activatienormen. In experimenten waarbij de richting van de verdelingsverschuiving wordt omgekeerd (trainen op brede verdelingen, testen op smalle verdelingen), vergroten pruning de representatieafstanden om te aligneren met het ID-profiel. Het kernmechanisme is de heralignering van OOD-representaties naar de geometrie die tijdens de aanpassing is geleerd, ongeacht of dit contractie of expansie vereist.
• Causaal Bewijs: Causale interventies bevestigen dat geometrische vervorming de drijvende kracht is van prestatieverlies. Het invoegen van een inverse surrogaatkaart na een verwijderde laag (om het geometrische effect van de laag ongedaan te maken) herstelt de OOD-prestaties bijna tot het niveau van het gepreunde model. Evenzo verbetert het herschalen van de residu-update van een geselecteerde laag continu de nauwkeurigheid naarmate de bijdrage van de laag wordt verzwakt, waardoor topologische artefacten of eenvoudige redundantie als enige oorzaak worden uitgesloten.

Betekenis en Aanspraken
Het artikel claimt een nieuwe geometrische verklaring te bieden voor het succes van pruning tijdens de inferentie. Het betoogt dat de belangrijkheid van lagen niet vaststaat, maar conditioneel is voor zowel de taak als de invoerverdeling. De auteurs stellen dat voorgetrainde modellen meerdere taakgespecialiseerde functies coderen, elk geassocieerd met een lokale representatiegeometrie. ID-invoeren passen bij deze geometrie, terwijl OOD-invoeren deze schenden.

De betekenis van dit werk ligt in:

1. Verduidelijking van de Grenzen van Pruning: Het stelt vast dat taakbewuste pruning geen universele prestatieboost is, maar een corrigerend mechanisme voor verdelingsverschuiving.
2. Geometrische Interpretatie: Het verschuift het begrip van lagenredundantie van statische parameteraantallen naar dynamische, verdelingsafhankelijke geometrische vervormingen.
3. Mechanistisch Inzicht: Het biedt causaal bewijs dat prestatiewinsten voortkomen uit het corrigeren van het traject van interne representaties om te matchen met de aangepaste geometrie van het model, en niet uit het vergroten van de modelcapaciteit of het verminderen van ruis.

De auteurs concluderen dat pruning de OOD-prestaties verbetert door lagen te verzwakken die de mismatch tussen de invoerverdeling en de geleerde taakgeometrie van het model versterken, waardoor de interne toestand van het model effectief wordt heraligneerd met de omstandigheden waaronder het is geoptimaliseerd.