TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Technische Zusammenfassung: TAPIOCA – Warum tasksensitives Pruning die OOD-Fähigkeit von Modellen verbessert

Problemstellung
Neuere Arbeiten haben gezeigt, dass das Entfernen von Schichten aus Large Language Models (LLMs) zur Inferenzzeit – spezifisch durch tasksensitive Pruning-Methoden wie TALE (Task-Aware Layer Elimination) – die Leistung bei bestimmten Aufgaben ohne Nachtraining verbessern kann. Die Mechanismen, die diese Gewinne antreiben, bleiben jedoch unklar. Eine zentrale Frage ist, ob diese Verbesserungen aus allgemeinen Kompressionsvorteilen, Regularisierungseffekten oder spezifischen Interaktionen zwischen den internen Repräsentationen des Modells und der Eingabeverteilung resultieren. Ferner ist unbekannt, ob diese Gewinne universell sind oder nur für bestimmte Arten von Datenverschiebungen gelten.

Methodik
Die Autoren untersuchen dieses Phänomen mittels einer Kombination aus kontrollierten synthetischen Experimenten und groß angelegten Evaluierungen an feinabgestimmten LLMs.

1. Kontrolliertes Regressions-Setting: Die Studie beginnt mit einem 12-Schichten-Transformer, der von Grund auf für kontextbezogene lineare Regression trainiert wurde. Die Trainingsverteilung ist explizit definiert (Koeffizienten, die aus $U(-1, 1)$ gezogen werden). Die Autoren evaluieren das Modell sowohl auf in-distribution (ID)-Daten (gleiche Verteilung) als auch auf out-of-distribution (OOD)-Daten (Koeffizienten, die aus breiteren Intervallen $U(-\sigma, \sigma)$ mit $\sigma > 1$ gezogen werden). Sie wenden TALE an, einen gierigen Pruning-Algorithmus, der Schichten entfernt, um die Validierungsgenauigkeit auf einer spezifischen Verteilung zu maximieren.
2. Groß angelegte LLM-Evaluierung: Die Erkenntnisse werden auf feinabgestimmte LLMs (Llama 3.1 8B und GPT-OSS 120B) erweitert, die auf mathematisches Schlussfolgern und Code-Generierung spezialisiert sind. Hier ist die „angepasste" Verteilung die Feinabstimmungsaufgabe (z. B. MATH500), und OOD-Aufgaben umfassen allgemeine Benchmarks wie MMLU und BoolQ.
3. Geometrische und kausale Analyse: Um zu erklären, warum Pruning funktioniert, analysieren die Autoren die Geometrie der versteckten Zustände. Sie messen schichtweise Repräsentationsnormen und paarweise Token-Abstände. Zudem setzen sie lokale lineare Surrogate ein, um Schichttransformationen zu approximieren, und führen kausale Interventionen durch, wie das Einfügen inverser Surrogat-Karten und das Skalieren von Residual-Updates, um den Effekt geometrischer Verzerrung zu isolieren.

Hauptergebnisse

• OOD-spezifische Gewinne: Das primäre empirische Ergebnis ist, dass tasksensitives Pruning Leistungsgewinne nur unter Verteilungsverschiebung erzielt. Bei der Evaluierung auf der Verteilung, an die das Modell angepasst ist (ID), versagt TALE entweder darin, die Leistung zu verbessern, oder verschlechtert sie aktiv, indem es Schichten entfernt, die für optimale Präzision notwendig sind. Im Gegensatz dazu identifiziert und entfernt TALE auf OOD-Daten konsistent Schichten, die die Genauigkeit verbessern.
• Fehlanpassung der Repräsentationsgeometrie: Die Autoren zeigen, dass OOD-Eingaben im Vergleich zu ID-Eingaben eine verzerrte Repräsentationsgeometrie hervorrufen. Konkret führen OOD-Eingaben dazu, dass schichtweise Normen versteckter Zustände und paarweise Abstände systematisch von den Profilen abweichen, die bei ID-Daten beobachtet werden. Im Regressions-Setting führen OOD-Eingaben zu einer Aufblähung der Abstände über die Schichten hinweg; bei LLMs weichen OOD-Trajektorien in den mittleren bis späten Schichten von ID-Trajektorien ab.
• Verstärkungsmechanismus: Die Studie identifiziert, dass bestimmte Schichten als „verteilungsabhängige Verstärker" wirken. Bei ID-Eingaben führen diese Schichten nützliche Verfeinerungen durch. Bei OOD-Eingaben zeigen dieselben Schichten jedoch eine Verstärkung mit hohem Gewinn und niedrigem Rang, die Repräsentationen weiter von der an die Aufgabe angepassten Geometrie wegtreiben. TALE verbessert die Leistung, indem es diese spezifischen Schichten entfernt und damit die Verstärkung geometrischer Verzerrung verhindert.
• Ausrichtung, nicht nur Kontraktion: Pruning unterdrückt nicht einfach Aktivierungsnormen. In Experimenten, bei denen die Richtung der Verteilungsverschiebung umgekehrt wurde (Training auf breiten Verteilungen, Testen auf schmalen), erweitert Pruning Repräsentationsabstände, um sie mit dem ID-Profil in Einklang zu bringen. Der Kernmechanismus ist die Neuausrichtung von OOD-Repräsentationen hin zur Geometrie, die während der Anpassung gelernt wurde, unabhängig davon, ob dies Kontraktion oder Expansion erfordert.
• Kausale Evidenz: Kausale Interventionen bestätigen, dass geometrische Verzerrung der Treiber für Leistungsverlust ist. Das Einfügen einer inversen Surrogat-Karte nach einer entfernten Schicht (um den geometrischen Effekt der Schicht rückgängig zu machen) stellt die OOD-Leistung fast auf das Niveau des pruned Modells wieder her. Ebenso verbessert das Skalieren des Residual-Updates einer ausgewählten Schicht die Genauigkeit kontinuierlich, wenn der Beitrag der Schicht abgeschwächt wird, was topologische Artefakte oder einfache Redundanz als alleinige Ursache ausschließt.

Bedeutung und Behauptungen
Die Arbeit behauptet, eine neue geometrische Erklärung für den Erfolg von Pruning zur Inferenzzeit zu liefern. Sie argumentiert, dass die Wichtigkeit von Schichten nicht festgelegt ist, sondern sowohl von der Aufgabe als auch von der Eingabeverteilung abhängt. Die Autoren gehen davon aus, dass vortrainierte Modelle mehrere aufgaben-spezifische Funktionen kodieren, die jeweils mit einer lokalen Repräsentationsgeometrie assoziiert sind. ID-Eingaben passen zu dieser Geometrie, während OOD-Eingaben sie verletzen.

Die Bedeutung dieser Arbeit liegt in:

1. Klärung der Grenzen von Pruning: Sie stellt fest, dass tasksensitives Pruning kein universeller Leistungsbooster ist, sondern ein Korrekturmechanismus für Verteilungsverschiebung.
2. Geometrische Interpretation: Sie verlagert das Verständnis von Schichtredundanz von statischen Parameterzahlen hin zu dynamischen, verteilungsabhängigen geometrischen Verzerrungen.
3. Mechanistische Einsicht: Sie liefert kausale Evidenz, dass Leistungsgewinne aus der Korrektur der Trajektorie interner Repräsentationen resultieren, um sie mit der angepassten Geometrie des Modells in Einklang zu bringen, und nicht aus der Erhöhung der Modellkapazität oder der Reduktion von Rauschen.

Die Autoren schließen, dass Pruning die OOD-Leistung verbessert, indem es Schichten abschwächt, die die Diskrepanz zwischen der Eingabeverteilung und der gelernten Aufgabengeometrie des Modells verstärken, und damit den internen Zustand des Modells effektiv mit den Bedingungen neu ausrichtet, unter denen es optimiert wurde.