Technische Zusammenfassung: RATEQUANT

Problemstellung

Large Language Models (LLMs) erfordern während der Generierung das Caching aller zuvor berechneten Schlüssel-Wert-(KV-)Paare, was einen Speicherengpass schafft, der sich linear mit der Sequenzlänge, der Batch-Größe und der Modelltiefe skaliert. Obwohl die Quantisierung des KV-Caches diesen Speicherbedarf reduziert, weisen bestehende Methoden eine uniforme Zuweisung der gleichen Bit-Breite für jeden Attention-Head auf. Dieser Ansatz ignoriert die gut dokumentierte Heterogenität in der Bedeutung der einzelnen Heads.

Neuere Versuche einer gemischten Präzisionszuweisung auf Ebene der Schichten oder Kanäle verlassen sich auf heuristische Regeln, die an spezifische Quantisierer gebunden sind, und fehlen theoretische Garantien sowie Übertragbarkeit. Die Autoren identifizieren ein tieferes, bisher nicht erkannte Hindernis: Fehlanpassung des Verzerrungsmodells. Unterschiedliche Quantisierer folgen unterschiedlichen Verzerrungs-Raten-Kurven, modelliert als $D(b) = \alpha \beta^{-b}$. Die Abklingrate $\beta$ variiert erheblich zwischen verschiedenen Designs (z. B. von 3,6 bei TURBOQUANT bis 5,3 bei QuaRot). Eine naive Anwendung einer gemischten Präzisionszuweisungsstrategie, die aus dem Verzerrungsmodell eines Quantisierers abgeleitet wurde, auf einen anderen, kehrt die Reihenfolge des marginalen Gewinns der Attention-Heads um und führt häufig zu Ergebnissen, die schlechter sind als eine uniforme Quantisierung.

Methodik

RATEQUANT adressiert diese Herausforderungen, indem es die Bit-Zuweisung für den KV-Cache pro Head als ein Rate-Distortion-Optimierungsproblem formalisiert. Das Framework besteht aus vier Phasen:

1. Problemformulierung und optimale Zuweisung

Die Autoren modellieren den mittleren quadratischen Fehler (MSE) der Quantisierung pro Head als exponentielle Funktion der Bit-Breite: $D(b) = \alpha \beta^{-b}$. Das Ziel ist es, die gewichtete Summe der Verzerrungen über $N$ Quantisierungsgruppen (Heads) unter Einhaltung eines Gesamtbudgets an Bits $B$ zu minimieren:
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{s.t.} \quad \sum b_i = B$$
wobei $w_i$ die Sensitivität des Heads $i$ darstellt.

Unter Verwendung der Lagrange-Stationarität leiten die Autoren eine geschlossene Lösung mittels Reverse Waterfilling ab:
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
wobei $\bar{b}$ die durchschnittliche Bit-Breite und $\bar{w}$ das geometrische Mittel der Sensitivitäten ist. Diese Lösung weist Heads mit höherer Sensitivität mehr Bits zu. Theoretisch ist die erreichbare Reduktion der Verzerrung durch das Verhältnis des arithmetischen Mittels zum geometrischen Mittel (AM/GM) der Head-Sensitivitäten begrenzt.

2. Schätzung der Sensitivität

Um die Gewichte $w_i$ zu bestimmen, verwendet RATEQUANT eine gradientenbasierte Sensitivität. Es berechnet die quadrierten Gradientennormen der KV-Projektionsausgaben in Bezug auf den Verlust des kausalen Sprachmodells:
$$w_{l,h} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial K_{l,h,t}} \right\|^2 \right]$$
Die Autoren zeigen, dass gradientenbasierte Proxy-Metriken aktivierungsbasierten Metriken (z. B. Kanal-Magnitude) überlegen sind, die eine Fehlerverstärkung statt einer direkten Verlustauswirkung messen. Die gradientenbasierte Sensitivität offenbart ein U-förmiges Muster der Bedeutung (hohe Sensitivität in frühen und späten Schichten), während aktivierungsbasierte Metriken einen monotonen Anstieg mit der Tiefe zeigen.

3. Ganzzahlige Zuweisung und Erweiterungen

Da Bit-Breiten ganzzahlig sein müssen, wird die kontinuierliche Lösung mit einem greedy marginal gain Algorithmus (Algorithmus 1) verfeinert. Dieser Algorithmus weist iterativ Bits an den Head zu, der die größte Reduktion der gewichteten Verzerrung bietet, bis das Budget erschöpft ist.

Zwei kritische Erweiterungen machen RATEQUANT quantisiereragnostisch:

• Verzerrungskalibrierung: Anstatt ein universelles $\beta$ anzunehmen, passt RATEQUANT ein verzerrungsmodell pro Quantisierer ($\alpha, \beta$) unter Verwendung eines kleinen Kalibrierungsdatensatzes (16 Sequenzen) an. Dies verhindert den Fehlermodus der „Fehlanpassung des Verzerrungsmodells".
• Trennung von K/V: Da Keys und Values oft unterschiedliche Quantisierungsschemata verwenden (z. B. pro-Kanal-symmetrisch für Keys vs. pro-Token-asymmetrisch für Values bei KIVI), behandelt das Framework sie als $2N$ unabhängige Komponenten mit unterschiedlichen Verzerrungskurven. Dies ermöglicht dem Algorithmus, asymmetrische Bit-Budgets zu entdecken (z. B. mehr Bits für Keys als für Values).

4. Pipeline

Der Prozess umfasst einmalige Offline-Kosten:

1. Sensitivitätskalibrierung: Vorwärts- und Rückwärtsdurchläufe auf einem kleinen Datensatz (~1,6 s für ein 8B-Modell).
2. Verzerrungsmodellierung: Anpassung von $\alpha$ und $\beta$ bei 5 Bit-Breiten (<0,1 s).
3. Zuweisung: Ausführung des Greedy-Algorithmus (<0,01 s).
4. Inferenz: Anwendung der statischen Bit-Breiten mit null Laufzeit-Overhead.

Wichtige Ergebnisse

Experimente wurden an Qwen3-Modellen (4B, 8B, 32B) unter Verwendung der Perplexität (PPL) auf WikiText-2 und Downstream-Aufgaben (ARC-Challenge, HellaSwag usw.) durchgeführt.

• Behebung der Fehlanpassung: Die Anwendung eines fehlangepassten Verzerrungsmodells (z. B. Verwendung von TURBOQUANTs $\beta$ für KIVI) verschlechterte die PPL von KIVI bei 2,5 durchschnittlichen Bits von 49,3 auf 87,0. Die kalibrierte RATEQUANT-Lösung behob dies und reduzierte die PPL auf 14,9 (eine Reduktion um 70 %).
• Leistungssteigerungen:
  • KIVI (2,5 Bits): RATEQUANT verbesserte die PPL von 49,3 auf 14,9.
  • QuaRot (2,5 Bits): Verbesserung der PPL um 6,6 Punkte.
  • TURBOQUANT (4,0 Bits): Wiederherstellung von 66 % der durch Quantisierung verursachten Degradation für Qwen3-8B.
• Trennung von K/V: Für KIVI bei 2,5 Bits betrug die von RATEQUANT entdeckte optimale Zuweisung 2,85 Bits für Keys und 2,15 Bits für Values, was die gemeinsame Zuweisung deutlich übertraf.
• Effizienz: Die Methode fügt keinen Laufzeit-Overhead hinzu. Die Kalibrierung dauert auf einer einzelnen GPU weniger als 2 Sekunden.
• Downstream-Genauigkeit: Bei 4,0 Bits erreichte RATEQUANT eine durchschnittliche Genauigkeit von 64,2 % bei Downstream-Aufgaben, was nahezu der FP16-Baseline (64,4 %) entspricht und 89,8 % der Genauigkeitslücke schließt.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass RATEQUANT das erste Framework ist, das die Rate-Distortion-Theorie auf die Quantisierung des KV-Caches auf Granularität pro Head anwendet und gleichzeitig die strukturellen Unterschiede zwischen Keys und Values adressiert.

Die zentrale Erkenntnis ist, dass die Fehlanpassung des Verzerrungsmodells ein kritischer Fehlermodus bei der gemischten Präzisionszuweisung ist. Ohne Kalibrierung der Verzerrungskurve ($\beta$) an den spezifischen Basis-Quantisierer können gemischte Präzisionsstrategien kontraproduktiv sein. Durch die Kombination von:

1. Gradientenbasierter Sensitivität (überlegen gegenüber aktivierungsbasierten Proxies),
2. Empirischer Verzerrungskalibrierung (zur Anpassung an den Quantisierer) und
3. Trennung von K/V (zum Umgang mit asymmetrischen Fehlercharakteristika),

verwandelt RATEQUANT die gemischte Präzisions-KV-Quantisierung in eine allgemeine, quantisiereragnostische Zuweisungsschicht. Die Autoren gehen davon aus, dass das AM/GM-Verhältnis der Head-Sensitivitäten als praktischer Prädiktor dient, wann eine gemischte Präzisionszuweisung signifikante Gewinne bringt. Die Arbeit legt nahe, dass die Rate-Distortion-Perspektive auf die Ressourcenallokation in anderen heterogenen neuronalen Netzwerkkomponenten jenseits von KV-Caches verallgemeinert werden kann.