💬 NLP

Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

Die Arbeit stellt ASL vor, eine trainingsfreie Methode, die die Schicht für die Token-Pruning in LLM-Inferenz basierend auf der Varianz der Aufmerksamkeitswerte dynamisch anpasst, um die Genauigkeit bei schwierigen Aufgaben zu verbessern und gleichzeitig die KV-Cache-Budgetanforderungen einzuhalten.

Ursprüngliche Autoren: Rei Taniguchi, Yuyang Dong, Makoto Onizuka, Chuan Xiao

Veröffentlicht 2026-04-17

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Rei Taniguchi, Yuyang Dong, Makoto Onizuka, Chuan Xiao

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

Stell dir vor, ein großes Sprachmodell (LLM) ist wie ein genialer Bibliothekar, der gerade einen riesigen Roman liest, um dir eine Frage zu beantworten. Je länger der Roman ist, desto mehr Notizen muss der Bibliothekar machen, um sich an alles zu erinnern. Diese Notizen nennt man im Fachjargon „KV-Cache".

Das Problem: Wenn der Roman 100.000 Seiten hat, werden die Notizen so riesig, dass der Bibliothekar (der Computer) vor lauter Papierstapeln erstickt. Er braucht zu viel Speicherplatz und wird langsam.

Bisherige Lösungen waren wie ein starrer Schalter: „Ab Seite 500 werfen wir alle Notizen weg, außer den wichtigsten." Das Problem dabei ist: Manchmal ist Seite 500 noch viel zu früh. Vielleicht versteckt sich die entscheidende Information erst auf Seite 800. Wenn man zu früh schneidet, vergisst der Bibliothekar die Antwort. Wenn man zu spät schneidet, ist der Speicher immer noch voll.

ASL (Adaptive Selection Layer) ist die neue, clevere Lösung. Hier ist die Erklärung in einfachen Bildern:

1. Das Problem: Der starre Schalter

Stell dir vor, du hast einen Fotografen, der eine lange Reise dokumentiert.

Die alte Methode (z. B. FastKV): Der Fotograf hat einen festen Plan: „Ab Bild Nr. 150 mache ich nur noch Fotos von den wichtigsten Personen und ignoriere den Rest."
Das Problem: Bei einer entspannten Wanderung (einfache Aufgabe) ist das super. Aber bei einer komplexen Jagd nach einem versteckten Schatz (schwere Aufgabe wie „KV Retrieval") ist das Bild Nr. 150 noch voller verwirrender Details. Wenn der Fotograf jetzt aufhört, alle Details zu speichern, verpasst er den Schatz, weil er ihn erst auf Bild 200 sieht.

2. Die Lösung: Der adaptive Beobachter (ASL)

ASL ist wie ein intelligenter Assistent, der dem Fotografen nicht sagt, wann er aufhören soll, sondern ihm sagt: „Schau mal, wann sich die Dinge beruhigen."

Der Assistent beobachtet die Aufmerksamkeit des Bibliothekars (die „Aufmerksamkeits-Scores"):

Am Anfang (die ersten Seiten): Der Bibliothekar schaut sich alles an. Die Aufmerksamkeit ist wie ein wilder Wirbelwind – sie springt von Wort zu Wort, von Person zu Person. Niemand steht im Fokus. Das ist wie ein Sturm, bei dem man noch nicht weiß, wohin der Wind weht.
Im Verlauf: Langsam beruhigt sich der Sturm. Der Bibliothekar beginnt, sich auf bestimmte Schlüsselwörter zu konzentrieren. Die „Wahrscheinlichkeiten" stabilisieren sich.
Der Moment der Entscheidung: ASL misst die Unruhe (Varianz) in den Notizen. Solange die Unruhe hoch ist, macht der Assistent weiter: „Noch nicht schneiden! Wir brauchen noch mehr Kontext." Sobald die Unruhe unter einen bestimmten Wert fällt (der Wind sich beruhigt hat und nur noch auf ein paar Bäume bläst), sagt er: „Jetzt! Hier ist der Punkt, an dem wir wissen, was wichtig ist."

Erst an diesem Punkt werden die unwichtigen Notizen weggeworfen.

3. Warum ist das besser?

Bei einfachen Aufgaben: Der Sturm legt sich schnell. ASL schneidet früh ab (vielleicht bei Seite 150). Das spart viel Zeit und Speicher.
Bei schwierigen Aufgaben: Der Sturm dauert länger. ASL wartet geduldig bis Seite 250, bevor er schneidet. Der Bibliothekar hat genug Zeit, den Schatz zu finden, bevor er anfängt, Notizen zu vernichten.

Das Ergebnis: Der Bibliothekar ist nicht mehr starr an einen festen Punkt gebunden. Er passt sich der Aufgabe an.

Einfache Aufgaben? Schnell und effizient.
Schwere Aufgaben? Genau und zuverlässig, ohne den Speicher zu sprengen.

Zusammenfassung in einem Satz

ASL ist wie ein dynamischer Schiedsrichter, der nicht nach einer festen Uhrzeit pfeift, sondern wartet, bis das Spiel klar genug ist, um die wichtigsten Spieler auf dem Feld zu behalten und den Rest nach Hause zu schicken – egal, ob das Spiel schnell oder langsam verläuft.

Dadurch wird die KI schneller, braucht weniger Speicher, macht aber trotzdem weniger Fehler, besonders bei den kniffligen Aufgaben, bei denen andere Methoden scheitern.

1. Problemstellung

Große Sprachmodelle (LLMs) stoßen bei der Verarbeitung langer Kontexte (z. B. 100k+ Token) auf ein kritisches Problem: Der hohe Speicherbedarf des Key-Value (KV) Caches. Um diesen zu reduzieren, wurden Methoden zur Token-Pruning (Beschneiden von Tokens) entwickelt, die eine Teilmenge wichtiger Tokens auswählen und den Rest verwerfen.

Ein etablierter Ansatz ist das layer-wise token pruning (schichtweises Token-Pruning), bei dem Tokens in einer bestimmten, vordefinierten Schicht des Transformer-Modells ausgewählt werden. Ab dieser Schicht werden nur noch die ausgewählten Tokens weiterverarbeitet.

Die Schwachstelle: Bestehende Methoden (wie FastKV, GemFilter) verwenden starre, vordefinierte Schichten für die Token-Auswahl, unabhängig von der spezifischen Aufgabe.
Die Konsequenz: Dies führt zu einer erheblichen Leistungsschwankung. Bei einfachen Aufgaben (z. B. einfaches QA) kann die Auswahl in mittleren Schichten gut funktionieren. Bei komplexen Aufgaben (z. B. KV-Retrieval, wo semantische Ähnlichkeit zwischen Frage und Kontext hoch ist) führt eine zu frühe Auswahl zu einem massiven Genauigkeitsverlust, da die relevanten Informationen noch nicht stabil genug sind. Umgekehrt führt eine zu späte Auswahl bei einfachen Aufgaben zu unnötigem Speicherverbrauch.

2. Methodik: Adaptive Selection Layer (ASL)

Die Autoren schlagen ASL vor, eine training-freie Methode, die die Auswahl-Schicht dynamisch und aufgabenabhängig bestimmt.

Kernidee:
Die Methode nutzt die Beobachtung, dass sich die Rangfolge der Tokens basierend auf ihren Aufmerksamkeits-Scores (Attention Scores) über die Schichten hinweg stabilisiert.

In frühen Schichten sind die Scores oft gleichmäßig verteilt.
In tieferen Schichten konzentrieren sich die Scores auf eine stabile Teilmenge von Tokens.

Algorithmus:

Beobachtungsphase: Während der Prefilling-Phase (Verarbeitung des Eingabe-Prompts) beginnt ASL ab einer minimalen Schicht $L_{min}$ , die Aufmerksamkeits-Scores zu überwachen.
Rang-Varianz: Für eine Fenstergröße von $L_{obs}$ $L_{o b s}$ aufeinanderfolgenden Schichten wird die Varianz der Token-Ränge (sortiert nach Attention Score) berechnet.
- Zuerst werden die Scores über die Attention-Heads gemittelt und gepoolt (ähnlich wie bei SnapKV), um Rauschen zu reduzieren.
- Es wird eine Vereinigung der Top- $k$ Tokens über die letzten $L_{obs}$ Schichten gebildet.
- Die Varianz der Ränge dieser Tokens innerhalb des Fensters wird berechnet.
Adaptive Entscheidung: Die berechnete Varianz wird durch die initiale Varianz bei $L_{min}$ $L_{min}$ geteilt, um eine relative Varianz zu erhalten.
- Sobald diese relative Varianz einen benutzerdefinierten Schwellenwert $\tau$ unterschreitet, wird die aktuelle Schicht als Auswahl-Schicht (Selection Layer) festgelegt.
- Dies signalisiert, dass die Aufmerksamkeit sich nun konsistent auf eine stabile Teilmenge von Tokens fokussiert hat.
Integration:
- One-Shot: Sobald die Schicht gewählt ist, werden die Tokens ausgewählt und für alle folgenden Schichten weitergegeben.
- Budget-Management: Um den KV-Budget-Requirement (z. B. 2048 Tokens) zu erfüllen, kann ASL mit bestehenden Methoden kombiniert werden (z. B. SnapKV für die Schichten vor der Auswahl oder GemFilter für einen Zwei-Pass-Ansatz).

3. Hauptbeiträge

Identifikation der Starrheit: Die Arbeit zeigt auf, dass feste Auswahl-Schichten in aktuellen Methoden ineffizient sind und die Leistung bei schwierigen Aufgaben (wie Retrieval) drastisch sinken lässt.
Entwicklung von ASL: Ein neuer, aufgabenbewusster Algorithmus, der die Auswahl-Schicht basierend auf der Stabilität der Token-Ränge (Varianz der Attention-Scores) dynamisch bestimmt.
Flexibilität und Kompatibilität: ASL ist training-frei, fügt sich nahtlos in bestehende KV-Cache-Reduktionsmethoden ein und ermöglicht einen Trade-off zwischen Geschwindigkeit und Genauigkeit basierend auf dem Schwierigkeitsgrad der Aufgabe.
Umfassende Evaluation: Die Methode wurde auf drei großen Benchmarks (InfiniteBench, RULER, NIAH) mit Kontextlängen bis zu 256k Token evaluiert.

4. Ergebnisse

Die Evaluation erfolgte an den Modellen Llama-3.1-8B-UL und Qwen2.5-7B.

Genauigkeit (Accuracy):
- ASL übertrifft State-of-the-Art-Methoden (FastKV, GemFilter, PyramidInfer) signifikant in schwierigen Aufgaben, insbesondere beim KV-Retrieval und bei langen Kontexten (z. B. 128k+).
- Bei einfachen Aufgaben (z. B. QA) bleibt die Leistung vergleichbar oder leicht verbessert.
- Im Vergleich zu FastKV zeigt ASL bei schwierigen Aufgaben eine Genauigkeit, die nahe an der eines vollen KV-Caches liegt, während FastKV hier stark einbricht.
Effizienz (Geschwindigkeit & Speicher):
- TTFT (Time to First Token): Da ASL die Auswahl oft in tieferen Schichten trifft als FastKV (um Genauigkeit zu sichern), ist der TTFT bei ASL etwas höher als bei FastKV, aber immer noch deutlich besser als bei vollen KV-Caches.
- TPOT (Time per Output Token): Die Decodier-Geschwindigkeit ist mit anderen Pruning-Methoden vergleichbar, da nach der Auswahl nur noch wenige Tokens verarbeitet werden.
- Speicher: Der zusätzliche Overhead für die Berechnung der Varianz ist vernachlässigbar gering (ca. 0,2–0,3 GB zusätzlich), und der finale KV-Speicherbedarf entspricht dem der anderen Pruning-Methoden.
Trade-off: ASL bietet einen besseren Kompromiss zwischen Genauigkeit und Geschwindigkeit, insbesondere wenn die Aufgabe komplex ist.

5. Bedeutung und Fazit

Die Arbeit zeigt, dass die Annahme einer „ein Schicht passt für alle"-Strategie bei der Token-Auswahl für LLMs nicht haltbar ist. Durch die Einführung einer adaptiven, datengetriebenen Auswahl kann die Leistung von LLMs bei langen Kontexten erheblich verbessert werden, ohne dass ein aufwendiges Fine-Tuning notwendig ist.

ASL ermöglicht es, die inhärente Dynamik von Transformer-Modellen zu nutzen: Je schwieriger die Aufgabe, desto später wird die Auswahl getroffen, um sicherzustellen, dass die relevanten Informationen tatsächlich identifiziert wurden. Dies macht ASL zu einer vielversprechenden Lösung für die effiziente Inferenz von LLMs in realen Szenarien mit variierenden Aufgabenkomplexitäten und langen Kontexten.