POP: Prefill-Only Pruning for Efficient Large Model Inference

1. Problemstellung

Große Sprachmodelle (LLMs) und Vision-Language-Modelle (VLMs) sind rechenintensiv, was ihre Bereitstellung erschwert. Bestehende Methoden zur strukturierten Beschleunigung (Pruning) entfernen oft redundante Komponenten wie ganze Schichten oder Kanäle.

• Herausforderung: Herkömmliche strukturierte Pruning-Methoden sind meist „stufenagnostisch" (stage-agnostic), d. h., sie wenden das gleiche reduzierte Modell sowohl auf die Prefill-Phase (Verarbeitung des Eingabe-Kontextes) als auch auf die Decode-Phase (generative Next-Token-Vorhersage) an.
• Nachteil: Dies führt zu erheblichen Genauigkeitsverlusten, insbesondere bei offenen Generierungsaufgaben. Die Autoren argumentieren, dass dies daran liegt, dass die funktionale Asymmetrie zwischen den beiden Phasen ignoriert wird: Die Prefill-Phase dient nur der Kodierung des Kontextes, während die Decode-Phase sowohl kodiert als auch die Wahrscheinlichkeitsverteilung für das nächste Token modelliert.

2. Methodik: Prefill-Only Pruning (POP)

Die Kernidee von POP ist die Erkenntnis, dass tiefe Schichten für die Generierung (Decode) kritisch, für die Kontextkodierung (Prefill) jedoch weitgehend redundant sind.

A. Analyse der Schichtwichtigkeit (Virtual Gate Mechanism)

Um diese Asymmetrie zu quantifizieren, führen die Autoren einen virtuellen Gate-Mechanismus ein:

• Sie modellieren die Wichtigkeit einer Schicht $l$ als den Anstieg des Verlusts, wenn diese Schicht entfernt wird.
• Statt jede Schicht physisch zu entfernen (was rechenintensiv wäre), multiplizieren sie die Residualzweige (Attention und FFN) mit einem skalaren Parameter $g_l$.
• Durch eine Taylor-Entwicklung und die Nutzung der Fisher-Information (unter der Annahme, dass das Modell auf seiner eigenen Verteilung trainiert wird) wird die Wichtigkeit basierend auf dem Gradientenquadrat geschätzt: $\tilde{I}_l = E[(\partial L / \partial g_l)^2]$.
• Ergebnis der Analyse: Die Wichtigkeitskurven zeigen eine deutliche Diskrepanz:
  • Decode: Tiefe Schichten sind extrem wichtig für die Next-Token-Vorhersage.
  • Prefill: Tiefe Schichten haben eine sehr geringe Wichtigkeit (nahe Null), während flache Schichten für die initiale Merkmalsextraktion entscheidend sind.

B. Die POP-Strategie

Basierend auf dieser Analyse schlägt POP eine stufenbewusste (stage-aware) Pruning-Strategie vor:

1. Prefill-Phase: Die tiefen Schichten (z. B. die letzten 1/3 des Modells) werden übersprungen. Dies reduziert die Rechenlast (FLOPs) erheblich, da diese Phase bei langen Eingaben rechenintensiv ist.
2. Decode-Phase: Das vollständige Modell wird verwendet, um die Genauigkeit der Generierung zu gewährleisten.

C. Technische Implementierung der Übergänge

Da die Decode-Phase KV-Caches (Key-Value Caches) für alle Schichten benötigt, auch für die im Prefill übersprungenen, müssen zwei Probleme gelöst werden:

• Unabhängige KV-Projektionen: Für die übersprungenen Schichten werden die KV-Caches nicht durch den vollen Attention/FFN-Block berechnet, sondern durch eine separate, leichte Projektion der Eingabe ($x_l$) direkt in Keys und Values. Dies ist rechnerisch vernachlässigbar (< 5% Overhead), stellt aber sicher, dass der Decoder über vollständige KV-Daten verfügt.
• Boundary Handling (Grenzbearbeitung): Um sicherzustellen, dass das erste generierte Token korrekt ist, wird die letzte Eingabe-Token ($x_N$) nicht mit dem prunierten Prefill-Modell verarbeitet, sondern als erster Schritt der Decode-Phase mit dem vollen Modell behandelt. Nur die vorherigen Tokens ($x_{1:N-1}$) werden im prunierten Modus verarbeitet.

3. Wichtige Beiträge

1. Entdeckung der Asymmetrie: Nachweis, dass tiefe Schichten für die Kontextkodierung redundant, aber für die Generierung essenziell sind.
2. POP-Algorithmus: Entwicklung einer Methode, die tiefen Schichten nur während der Prefill-Phase entfernt, während die Decode-Phase intakt bleibt.
3. Mechanismen für Integrität: Einführung unabhängiger KV-Projektionen und einer speziellen Grenzbearbeitung, um den Übergang zwischen den Phasen nahtlos und genau zu gestalten.
4. Breite Evaluierung: Validierung auf verschiedenen Modellen (Llama-3.1, Qwen3-VL, Gemma-3) und Modalitäten (Text, Vision).

4. Ergebnisse

Die Experimente wurden auf NVIDIA A100 GPUs durchgeführt und umfassten Benchmarks wie MMLU, GSM8K, HumanEval und multimodale Aufgaben.

• Genauigkeit:
  • Herkömmliche strukturierte Pruning-Methoden (wie SliceGPT, ShortGPT) führten zu katastrophalen Einbrüchen bei Generierungsaufgaben (z. B. Rückgang von 79% auf <1% bei GSM8K).
  • POP behält die Genauigkeit des Vollmodells nahezu vollständig bei (z. B. 97% der Leistung auf GSM8K für Llama-3.1), selbst bei einem Pruning-Anteil von ca. 33%.
  • Die Genauigkeit ist vergleichbar mit unstrukturierten Methoden (wie Wanda), bietet aber bessere Hardware-Kompatibilität.
• Geschwindigkeit (Inference Speedup):
  • Prefill-Latenz (TTFT): POP erzielt bis zu 1,37-fache Beschleunigung bei langen Kontexten (2048 Tokens) und hohen Auflösungen.
  • Im Gegensatz dazu zeigen unstrukturierte Methoden (Wanda) auf Standard-Hardware keine Beschleunigung (1,0x), da sie spezielle Sparse-Kernels benötigen.
  • Der Geschwindigkeitsgewinn steigt mit der Eingabelänge, da der rechenintensive Prefill-Teil dominiert.

5. Bedeutung und Fazit

POP stellt einen Paradigmenwechsel in der Effizienzoptimierung von LLMs dar. Anstatt das Modell global zu vereinfachen, nutzt es die inhärente Asymmetrie der Inferenzphasen aus.

• Vorteil: Es bietet eine „Plug-and-Play"-Lösung, die keine Neukalibrierung des gesamten Modells erfordert und auf Standard-Hardware läuft.
• Einschränkung: Da das volle Modell für die Decode-Phase geladen werden muss, reduziert POP nicht den VRAM-Speicherbedarf (Peak Memory), sondern optimiert primär die Rechenzeit (Compute-bound).
• Zukunftspotenzial: Die Arbeit legt nahe, dass stufenbewusste Optimierungen (für Quantisierung oder Architekturdesign) ein vielversprechender Weg für effiziente Inferenz sind, insbesondere in Kombination mit verteilten Systemen, die Prefill und Decode auf separaten Hardware-Ressourcen ausführen.