💬 NLP

TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

TraceNAS ist ein hocheffizientes, trainingsfreies Neural Architecture Search-Framework, das die Korrelation von Gradientenspuren nutzt, um optimale nicht-uniform beschnittene Large Language Models zu identifizieren, indem es deren Verlustlandschaften mit den ursprünglichen vortrainierten Modellen in Einklang bringt und so eine konkurrenzfähige Leistung im Vergleich zu trainingsbewussten Methoden zu einem Bruchteil der Rechenkosten erzielt.

Ursprüngliche Autoren: Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

Veröffentlicht 2026-02-04

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie besitzen eine riesige, unglaublich kluge Bibliothek (ein Large Language Model oder LLM), die die Summe des menschlichen Wissens enthält. Sie ist so groß und schwer, dass es unmöglich ist, sie mit sich herumzutragen oder auf einem normalen Telefon zu nutzen. Sie möchten diese Bibliothek auf eine taschengroße Version schrumpfen, ohne dabei ihre Fähigkeit zu verlieren, Geschichten zu erzählen, Rätsel zu lösen oder Witze zu verstehen.

Dies ist das Problem, das die Arbeit TraceNAS zu lösen versucht.

Das Problem: Die falschen Dinge wegschneiden

Normalerweise verhalten sich Menschen, wenn sie versuchen, diese riesigen Modelle zu verkleinern, wie ein ungeschickter Gärtner. Sie sagen vielleicht: „Lass uns 50 % der Äste von jedem Baum abschneiden“, oder „Lass uns die schwersten Blätter entfernen“. Dies nennt man uniformes Pruning (gleichmäßiges Beschneiden).

Aber hier liegt der Haken: Nicht alle Teile des Gehirns (oder der Bibliothek) sind gleich. Einige Teile sind die „lebenswichtigen Organe“, die die komplexe Logik halten, während andere nur „Fett“ sind, das leicht entfernt werden kann. Wenn Sie das falsche lebenswichtige Organ abschneiden, vergisst das Modell alles. Wenn Sie nur das Fett abschneiden, bleibt es zu schwer.

Bestehende Methoden versuchen herauszufinden, was zu schneiden ist, indem sie entweder:

Einen Teil nach dem anderen betrachten: Sie prüfen, ob ein bestimmtes Neuron wichtig ist, aber sie übersehen, wie dieses Neuron mit dem Rest des Gehirns kommuniziert.
Das Modell während des Schneidens trainieren: Sie versuchen, das Modell zu verkleinern und ihm gleichzeitig wieder beizubringen, wie man denkt. Das ist so, als würde man versuchen, Autofahren zu lernen, während man gleichzeitig den Motor des Autos umbaut. Das dauert ewig und benötigt eine enorme Menge an Treibstoff (Rechenleistung).

Die Lösung: TraceNAS (Der „Gradient Trace“-Detektiv)

Die Autoren schlagen eine neue Methode namens TraceNAS vor. Betrachten Sie dies als einen hochmodernen Röntgenscanner, der in der Lage ist, die „Seele“ des Modells zu betrachten, ohne es tatsächlich zu berühren oder neu zu lehren.

So funktioniert es, unter Verwendung einfacher Analogien:

1. Das „Echo“ des ursprünglichen Verstandes

Stellen Sie sich vor, das ursprüngliche, riesige Modell ist ein Meisterkoch, der eine perfekte Mahlzeit zubereitet hat. Wenn Sie einen Bissen nehmen, hinterlässt der Geschmack eine spezifische „Spur“ auf Ihrer Zunge.

Der alte Weg: Um zu sehen, ob ein neues, kleineres Rezept funktioniert, müssten Sie das ganze Gericht kochen, probieren und – falls es schlecht schmeckt – wieder von vorne anfangen.
Der TraceNAS-Weg: Anstatt das ganze Gericht zu kochen, schaut TraceNAS auf die Zutaten und die Rezeptschritte (die Gradienten), um vorherzusagen, ob der Geschmack stimmen wird. Es prüft, ob die „Geschmackspur“ der kleinen, verkleinerten Version mit der „Geschmackspur“ des ursprünglichen Meisterkochs übereinstimmt.

2. Der „Schatten“-Test (Gradient Trace Correlation)

Das Paper verwendet ein mathematisches Konzept namens Gradient Trace Correlation.

Stellen Sie sich das ursprüngliche Modell als ein riesiges Schiff vor, das auf einem ruhigen Ozean segelt. Es hinterlässt ein spezifisches Kielwasser (eine Spur aus Wellen) hinter sich.
Wenn Sie versuchen, ein kleineres Boot zu bauen (ein gepruntes Modell), fragt TraceNAS: „Hinterlässt dieses kleinere Boot ein Kielwasser, das exakt wie das Kielwasser des großen Schiffes aussieht?“
Wenn die Kielwasser übereinstimmen, bedeutet das, dass das kleine Boot demselben Pfad folgt und wahrscheinlich dasselbe Ziel erreichen wird (gut performen wird), sobald es ein wenig Übung erhält. Wenn die Kielwasser chaotisch sind, wird das Boot Schiffbruch erleiden.

3. Die „Low-Rank“-Abkürzung

Das Berechnen dieser Kielwasser für ein riesiges Schiff erfordert normalerweise einen Supercomputer. TraceNAS ist clever: Es erkennt, dass die Bewegung des Schiffes hauptsächlich in ein paar Hauptrichtungen stattfindet. Es nutzt einen Low-Rank-Trick, um nur die wichtigsten Richtungen der Spur zu betrachten.

Analogie: Anstatt jede einzelne Kräuselung im Ozean zu messen, misst es nur die drei größten Wellen. Dies ermöglicht es ihnen, den Test auf einer einzigen Grafikkarte (GPU in nur 8,5 Stunden durchzuführen, während die alten Methoden Tage oder Wochen auf einem ganzen Cluster von Computern benötigt hätten.

Die Ergebnisse: Schnell, günstig und klug

Das Paper hat dies an berühmten Modellen wie Llama und Qwen getestet.

Geschwindigkeit: Sie fanden den besten „Schnitt“ in 8,5 Stunden auf einem Computer. Die alten Methoden dauerten 10-mal länger und verbrauchten 10-mal mehr Energie.
Genauigkeit: Die resultierenden kleinen Modelle schnitten genauso gut ab wie Modelle, die über Wochen hinweg trainiert wurden, um die besten Schnitte zu finden.
Nicht-Uniformität: Im Gegensatz zum „ungeschickten Gärtner“, der alles gleichmäßig schneidet, fand TraceNAS eine „maßgeschneiderte Passform“. Es hielt die schweren, komplexen Teile des Gehirns intakt und schnitt nur das Fett weg, wodurch ein Modell entstand, das effizient, aber immer noch sehr klug ist.

Das Fazsit

TraceNAS ist ein Werkzeug, mit dem Sie ein riesiges KI-Gehirn in eine taschengroße Version schrumpfen können, ohne es neu lehren zu müssen oder ein Vermögen an Elektrizität auszugeben. Es tut dies, indem es prüft, ob der „Schatten“ des kleinen Gehirns mit dem „Schatten“ des großen Gehirns übereinstimmt. Wenn die Schatten übereinstimmen, ist das kleine Gehirn bereit für den Einsatz.

Dies macht es möglich, leistungsstarke KI auf gewöhnlichen Geräten laufen zu lassen, ohne ein riesiges Rechenzentrum zu benötigen, und das bei gleichzeitiger Einsparung einer enormen Menge an Zeit und Energie.

Technisches Resümee: TraceNAS

Problemdefinition

Die Bereitstellung von Large Language Models (LLMs) wird durch massive Parameterzahlen und hohe Rechenanforderungen eingeschränkt. Während strukturiertes Pruning einen Weg zur effizienten Bereitstellung bietet, indem architektonische Einheiten (z. B. Attention-Heads, Transformer-Blöcke) entfernt werden, stehen bestehende Methoden vor signifikanten Einschränkungen:

Lokale vs. globale Abhängigkeiten: Die meisten Pruning-Methoden bewerten die Wichtigkeit von Layern, Heads oder Kanälen isoliert anhand lokaler Heuristiken (z. B. aktivierungsgewichtete Magnituden oder Hessian-basierte Rekonstruktion). Diese Ansätze ignorieren komplexe globale strukturelle Abhängigkeiten und die Ausbreitung von Pruning-bedingten Fehlern über das Modell hinweg.
Rechenaufwand von trainingsbewusstem Suchprozess: Methoden, die versuchen, globale Abhängigkeiten zu erfassen, verlassen sich oft auf ein trainingsbewusstes Neural Architecture Search (NAS). Diese Ansätze erfordern iterative Gradienten-Updates und kontinuierliche Ströme ungesehener Token, um Overfitting zu vermeiden, wodurch der Suchprozess so rectechenintensiv wird wie das eigentliche Recovery-Training nach dem Pruning.
Unzulänglichkeit bestehender Zero-Shot-Proxies: Aktuelle trainingsfreie Proxies sind darauf ausgelegt, trainierbare, zufällig initialisierte Architekturen zu identifizieren, indem sie sich auf die Trainierbarkeit (z. B. Gradientenstabilität, Glätte der Loss-Landschaft) konzentrieren. Sie versagen jedoch bei der Bewertung der „funktionalen Vererbung“ – also der Fähigkeit eines geprunten Modells, die Performance wiederherzustellen, indem die vortrainierte Loss-Landschaft bewahrt wird.

Methodik

TraceNAS wird als trainingsfreies NAS-Framework vorgeschlagen, das darauf ausgelegt ist, Tiefe und Breite von LLMs ohne Training während der Suche gemeinsam zu optimieren.

1. Suchraum und Kodierung

Das Framework behandelt das vortrainierte LLM als Super-Netzwerk. Der Suchraum $G$ ist definiert als eine gemeinsame Verteilung über:

Tiefe ( $d$ ): Eine diskrete Maske über $L$ Transformer-Blöcke, die eine nicht-uniforme Tiefenreduktion ermöglicht und gleichzeitig die Kontinuität des Residual-Streams aufrechterhält.
Breite ( $\kappa$ ): Parameter-Retention-Raten für Attention- und MLP-Sub-Blöcke, was eine heterogene Sparsität über das Modell hinweg ermöglicht.

2. In-Place Architektonische Realisierung

Um den Speicheroverhead durch die Instanziierung tausender Kandidatenmodelle zu vermeiden, verwendet TraceNAS eine In-Place-Maskierungsstrategie:

Breiten-Maskierung: Verwendet ein aktivierungsgewichtetes Produkt (ähnlich wie Wanda), um binäre Masken für Gewichtskanäle zu generieren. Diese $O(d^2)$ -Heuristik identifiziert funktionell kritische Kanäle und bewahrt gleichzeitig Ausreißer-Features (Outlier Features).
Tiefen-Maskierung: Umgeht ganze Blöcke, falls diese deaktiviert werden, und leitet die Aktivierungen über Residual-Verbindungen weiter.
Ausführung: Das Framework modifiziert die Gewichtmatrizen des Basismodells temporär, um Gradienten auf einem minimalen Kalibrierungsdatensatz zu berechnen, und stellt anschließend die Originalgewichte wieder her, um das Basismodell für nachfolgende Iterationen zu bewahren.

3. Der Zero-Shot-Proxy: Gradient Trace Correlation

Die Kerninnovation ist ein skala-invarianter Proxy $\Phi$ , der die funktionale Ausrichtung zwischen einem geprunten Kandidaten ( $M_{sub}$ ) und dem vortrainierten Basismodell ( $M_{base}$ ) misst.

Gradienten-Spuren (Gradient Traces): Die Gradienten werden auf einem Kalibrierungsdatensatz berechnet. Um die Skalierbarkeit zu gewährleisten, werden diese in einen Low-Rank (LoRA) Subraum projiziert, wobei die intrinsische Niedrigdimensionalität von LLMs genutzt wird, um den Speicherbedarf um Größenordnungen zu senken.
Ausrichtungsmetrik: Der Proxy berechnet den Pearson-Korrelationskoeffizienten zwischen den standardisierten Gradienten-Spuren des geprunten und des Basismodells für jede Schicht. Dies gewährleistet Skaleninvarianz und konzentriert sich auf die Richtungs-Ausrichtung statt auf Magnituden-Verschiebungen.
Sparsity-gewichtete Aggregation: Der endgültige Score $\Phi$ ist eine gewichtete Summe der Schicht-Korrelationen, wobei die Gewichte den Parameter-Retention-Raten entsprechen. Dies verankert den globalen Score in hochkapazitiven Regionen und verhindert, dass Rauschen aus aggressiv geprunten Sub-Blöcken das Ranking verzerrt.

4. Evolutionäre Suche

TraceNAS verwendet eine evolutionäre Strategie, um den diskreten Suchraum zu navigieren:

Initialisierung: Eine Population wird basierend auf globalen Wichtigkeits-Priors generiert.
Evolution: Die Population durchläuft Crossover und Mutation. Die Tiefen-Evolution nutzt Bit-Flip-Mutationen und kontinuierliche Bereichs-Crossover, während die Breiten-Evolution Interpolation und Gaußsche Jitter verwendet.
Selektion: Kandidaten werden ausschließlich durch den Proxy $\Phi$ bewertet, wobei Modelle mit dem höchsten Potenzial für die Performance-Wiederherstellung während des Post-Pruning Continued Pretraining (CPT) ausgewählt werden.

Wichtigste Beiträge

Skalierbares Trainingsfreies NAS: TraceNAS bietet ein einheitliches Framework für das gemeinsame Tiefe- und Breiten-Pruning, das die Entdeckung nicht-uniformer Architekturen ohne Training während der Suche ermöglicht.
Funktionaler Vererbungs-Proxy: Führt einen neuartigen, skala-invarianten gradientenbasierten Proxy ein, der die Ausrichtung der Gradienten-Spuren zwischen geprunten und vortrainierten Modellen bewertet. Er erreicht eine hohe Korrelation mit der Downstream-Genauigkeit (Spearman $\rho = 0,94$ , Kendall $\tau = 0,82$ ).
Effizienz: Das Framework reduziert den Suchaufwand durch die Nutzung von Low-Rank-Gradientensubräumen und In-Place-Maskierung, was die hochpräzise Modellentdeckung auf einer einzelnen GPU ermöglicht.

Experimentelle Ergebnisse

TraceNAS wurde auf Llama-2-7B, Llama-3.1-8B und Qwen-2.5-14B evaluiert, wobei diese auf 2,7B, 4,6B bzw. 8,4B Parameter reduziert wurden.

Sucheffizienz: TraceNAS identifizierte optimale Architekturen in 8,5 GPU-Stunden (auf einer A100), was eine 10-fache Reduktion der Suchzeit gegenüber trainingsbewussten Baselines wie DarwinLM darstellt. Es benötigte nur 98 Mio. Such-Token, eine 10-fache Reduktion im Vergleich zu DarwinLM.
Performance:
- Llama-2-7B: TraceNAS erreichte die höchste durchschnittliche Genauigkeit (62,81 %) über acht Reasoning-Benchmarks hinweg und übertraf damit DarwinLM (62,57 %) und ShearedLLaMA (62,63 %). Bemerkenswert ist, dass die Performance von ShearedLLaMA (auf 60,86 %) signifikant sank, wenn es auf das gleiche Token-Budget beschränkt wurde, das von TraceNAS verwendet wurde.
- Llama-3.1-8B & Qwen-2.5-14B: TraceNAS übertraf Uniform-Pruning-Baselines und trainingsbewusste Methoden (DarwinLM, E3-Pruner) bei komplexen Reasoning-Aufgaben (z. B. ARC-Challenge, BoolQ), während es deutlich weniger Recovery-Token benötigte (20B vs. 94B für Minitron).
Proxy-Validierung: Der Proxy zeigte eine hohe Stabilität über variierende Hyperparameter hinweg (Stichprobengröße, Kontextlänge und LoRA-Rang), was bestätigt, dass die Ausrichtung der Gradienten-Spuren die Modellvererbung selbst unter beschränkten Bedingungen effektiv erfasst.

Bedeutung und Ansprüche

Das Paper behauptet, dass TraceNAS den kritischen Engpass der Rechenkosten im strukturierten Pruning adressiert. Durch die Verschiebung des Fokus von der „Trainierbarkeit“ (Lernen von Grund auf) hin zur „Vererbung“ (Bewahrung der vortrainierten Loss-Landschaft) ermöglicht TraceNAS die Entdeckung leistungsstarker, nicht-uniformer geprunter Modelle mit minimalem Ressourcenaufwand.

Die Autoren betonen, dass ihre Methode:

Den prohibitiven Kosten des Trainings während der Suche entgeht, wodurch groß angelegte Modellkompression zugänglich wird.
Globale strukturelle Abhängigkeiten erfasst, die lokale Heuristiken übersehen, und so die Erhaltung komplexer Denkfähigkeiten sicherstellt.
Den CO2-Fußabdruck der Modellkompression reduziert, indem sie die benötigten GPU-Stunden für die Architektur-Suche drastisch senkt.

Die Arbeit positioniert TraceNAS als skalierbares Fundament für den Einsatz effizienter LLMs in ressourcenbeschränkten Umgebungen bei gleichzeitiger Wahrung der funktionalen Integrität dichter vortrainierter Modelle.