Do LLMs Encode Functional Importance of Reasoning Tokens?

1. Problemstellung

Große Sprachmodelle (LLMs) lösen komplexe Aufgaben oft durch die Generierung langer, expliziter Denkketten (Reasoning Chains), was die Genauigkeit erhöht, aber erhebliche Kosten verursacht (höhere Latenz, Speicherbedarf und Trainingskosten). Ein zentrales Problem besteht darin, dass bestehende Methoden zur Komprimierung dieser Ketten (z. B. durch Temperatur-Sampling, Heuristiken oder Supervision durch stärkere Modelle) wenig Einblick darin geben, ob das Modell selbst eine funktionale Bedeutung einzelner Tokens für die Antwortgenerierung intern kodiert.

Die Autoren fragen sich: Kodieren LLMs die funktionale Wichtigkeit von Reasoning-Tokens intern, und kann diese Struktur diagnostisch genutzt werden, um effizientere Reasoning-Ketten zu erzeugen?

2. Methodik: Greedy Pruning

Die Autoren schlagen Greedy Pruning vor, ein diagnostisches Verfahren zur likelihood-erhaltenden Löschung von Tokens. Im Gegensatz zu bisherigen Ansätzen, die externe Heuristiken nutzen, basiert diese Methode auf den internen Wahrscheinlichkeitsverteilungen des Modells.

• Prinzip: Ausgehend von einer vollständigen Reasoning-Kette werden Tokens iterativ entfernt.
• Bewertungsfunktion: Für jeden Kandidaten-Token wird berechnet, wie stark die Likelihood des Modells sinkt, wenn dieser Token entfernt wird (Post-deletion Likelihood, $L_{del}$).
• Greedy-Strategie: Der Token, dessen Entfernung die Likelihood am wenigsten verschlechtert, wird gelöscht. Dieser Prozess wird wiederholt, bis eine gewünschte Länge (Keep-Fraction $\rho$) erreicht ist.
• Zielobjektive:
  1. $L_{ANS}$: Erhaltung der Likelihood der finalen Antwort.
  2. $L_{JOINT}$: Erhaltung der gemeinsamen Likelihood von Reasoning-Kette und Antwort (strenger, da die Struktur der Denkweise erhalten bleiben muss).
• Ergebnis: Das Verfahren erzeugt eine Rangliste (Pruning Ranks) der Tokens nach ihrer funktionalen Wichtigkeit für das Modell.

3. Experimentelles Setup

Die Methode wurde im Rahmen eines Teacher-Pruner-Student Distillation-Frameworks evaluiert:

• Teacher/Pruner: Stärkere Modelle (z. B. Llama-3.1-8B, Qwen-2.5-7B) generieren die Ketten und führen das Pruning durch.
• Student: Schwächere Modelle (z. B. Llama-2-7B, Mistral-7B) werden auf den gekürzten Ketten trainiert (Supervised Fine-Tuning).
• Benchmarks: GSM8K (Mathematik), MATH (Olympiade-Mathematik) und MMLU-Pro (Multi-Domain).
• Benchmarks: Vergleich mit TokenSkip (gelernte semantische Wichtigkeit via GPT-4), H2O (Attention-basiert), Surprisal und Uniform Random.

4. Wichtige Ergebnisse

A. Distillationsleistung

• Überlegenheit: Schülermodelle, die auf mittels Greedy Pruning gekürzten Ketten trainiert wurden, übertreffen bei gleicher Länge alle Baselines (inklusive TokenSkip) signifikant.
• Stabilität: Die Genauigkeit nimmt mit sinkender Keep-Fraction ($\rho$) sanft ab, was auf eine effektive Kompression ohne katastrophalen Informationsverlust hindeutet.
• Schlussfolgerung: Greedy Pruning erhält die funktional kritischen Tokens für die Antwortgenerierung besser als externe Heuristiken.

B. Funktionale Struktur unter Pruning

Durch Annotation der Tokens in funktionale Kategorien (z. B. SYMBMATH für symbolische Berechnungen, COREF für Referenzen, GRAMMAR für Grammatik) wurden folgende Muster identifiziert:

• Symbolische Berechnung (SYMBMATH): Wird überproportional stark erhalten, selbst bei aggressivem Pruning.
• Referenzen & Scaffolding (COREF, METADISC, GRAMMAR): Werden frühzeitig entfernt. Das Modell scheint diese als weniger funktional notwendig für die finale Berechnung zu erachten.
• Einfluss des Ziels: Bei Verwendung von $L_{ANS}$ (nur Antwort) wird die funktionale Struktur weniger scharf; bei $L_{JOINT}$ (Kette + Antwort) bleibt die Unterscheidung zwischen rechnerischen und narrativen Tokens klarer.
• Einfluss der Pruner-Stärke: Ein schwächerer Pruner (Llama-2-7B) erhält zwar symbolische Teile, stört aber das Gleichgewicht der nicht-symbolischen Struktur, was zu schlechteren Distillationsergebnissen führt.

C. Dynamik und Vorhersagbarkeit

• Dynamische Wichtigkeit: Die Wichtigkeit eines Tokens ist nicht statisch. Die Rangliste ändert sich dynamisch, wenn sich der Kontext durch das Entfernen vorheriger Tokens verkürzt. Ein statischer „Frozen"-Ranking-Ansatz ist weniger effektiv.
• Attention als Proxy: Ein einfacher MLP-Surrogat-Modell konnte die Pruning-Ränge basierend auf Attention-Scores des Lehrers mit einer Korrelation von 0,88 vorhersagen. Dies zeigt, dass die interne Struktur der Token-Wichtigkeit über Attention-Mechanismen zugänglich ist.

5. Bedeutung und Beiträge

• Diagnostisches Werkzeug: Greedy Pruning dient nicht nur der Kompression, sondern als „Sonde" (Probe), um die interne Organisation von Reasoning in LLMs zu entschlüsseln. Es zeigt, dass Modelle eine nicht-triviale funktionale Hierarchie über Tokens kodieren.
• Effiziente Kompression: Es bietet einen prinzipiellen Weg, Reasoning-Ketten zu komprimieren, ohne die Leistungsfähigkeit zu beeinträchtigen, indem redundante narrative Scaffolding-Elemente entfernt werden, während der rechnerische Kern erhalten bleibt.
• Interpretierbarkeit: Die Ergebnisse deuten darauf hin, dass LLMs interne Signale für die Token-Wichtigkeit besitzen, die für effizientere Inferenz und Training genutzt werden können, ohne auf externe, teure Modelle (wie GPT-4) angewiesen zu sein.

6. Limitationen

• Rechenkosten: Das naive Greedy Pruning skaliert superlinear ($O(n^2)$ oder höher) mit der Sequenzlänge, da für jeden Schritt viele Likelihood-Berechnungen nötig sind. Dies ist derzeit ein Offline-Prozess, aber für sehr lange Ketten (z. B. >1000 Tokens) rechenintensiv.
• Lesbarkeit: Bei starker Kompression (z. B. $\rho=0.7$) leiden Oberflächenmerkmale (Punktuation, Grammatik), während die mathematische Integrität erhalten bleibt.

Fazit: Das Paper belegt, dass LLMs eine interne, funktionale Hierarchie ihrer Reasoning-Tokens besitzen. Durch likelihood-erhaltendes Pruning kann diese Struktur genutzt werden, um effizientere, kompakte Reasoning-Ketten zu erzeugen, die für das Training leistungsfähigerer, aber ressourcenschonenderer Modelle überlegen sind.