Context Memorization for Efficient Long Context Generation

Technische Zusammenfassung: Attention-State Memory für eine effiziente Generierung langer Kontexte

Problemstellung
Moderne Anwendungen großer Sprachmodelle (LLM) stützen sich zunehmend auf lange Konditionierungspräfixe (z. B. In-Context-Beispiele, abgerufene Dokumente oder agentische Anweisungen), um die Inferenz zu steuern. Die Inferenz mit präfixaugmentierten Eingaben stößt jedoch auf zwei primäre strukturelle Grenzen:

1. Präfix-Decay: Während der Generierung verteilt sich die Aufmerksamkeit des Modells auf neu generierte Tokens, wodurch der Einfluss des Präfixes nachlässt, insbesondere in Szenarien mit langem Kontext.
2. Ineffizienz der Inferenz: Die Berechnung der Aufmerksamkeit über das Präfix skaliert linear mit seiner Länge ($O(L)$) und verursacht erhebliche Latenz- und Speicherkosten sowohl während der Vorbelegung (Prefill) als auch bei jedem Decodierschritt. Obwohl das Caching von Präfixen die Vorbelegungskosten amortisiert, entstehen dennoch erhebliche Speicherverbräuche.

Bestehende Lösungen lassen sich in zwei Kategorien einteilen, die beide Nachteile aufweisen:

• Präfix-Internalisierung: Methoden, die Modellparameter feinabstimmen oder Hypernetzwerke verwenden, um das Präfixverhalten zu internalisieren. Diese erfordern gradientenbasiertes Training, was sie ressourcenintensiv, langsam zu aktualisieren und für dynamische Präfixänderungen ungeeignet macht.
• Präfix-Kompression: Methoden, die das Präfix im Aufmerksamkeitsmechanismus belassen, aber komprimieren (z. B. durch Token-Pruning oder KV-Cache-Eviction). Diese erfordern weiterhin eine Aufmerksamkeit auf das komprimierte Präfix während der Inferenz, sodass die Latenz mit der komprimierten Länge skaliert und der Einfluss des Präfixes weiterhin dem Decay unterliegt.

Methodik: Attention-State Memory (ASM)
Die Autoren schlagen Attention-State Memory vor, einen trainingsfreien Ansatz, der das Präfix in ein leichtgewichtiges, nachschlagebasiertes Gedächtnis aus vorkalkulierten Aufmerksamkeitszuständen externalisiert. Diese Methode nutzt die Online-Softmax-Identität, um die Aufmerksamkeit über einen verketteten Schlüsselblock in Aufmerksamkeiten über Teilblöcke zu zerlegen.

Die zentrale Erkenntnis besteht darin, dass für eine gegebene Abfrage der Aufmerksamkeitszustand $(a, Z)$ – bestehend aus dem Aufmerksamkeitsausgang und dem Normalisierungsterm – ausreicht, um den Beitrag des Präfixes zur Aufmerksamkeit wiederherzustellen, ohne erneut auf die ursprünglichen Schlüssel und Werte zu achten.

Das Framework arbeitet in zwei Phasen:

1. Offline-Kalibrierungsphase (Speicherkonstruktion):

   • Sammlung: Das Modell führt Vorwärtsdurchläufe über eine Reihe repräsentativer Abfrage-Antwort-Spuren durch, die mit dem Präfix verkettet sind. Für jeden Abfrage-Token in der Antwort zeichnet das System den Abfragevektor und den entsprechenden Aufmerksamkeitszustand $(a, Z)$ über das Präfix auf.
   • Clustering: Die gesammelten Abfragevektoren werden (z. B. via K-Means) in $K$ Zentren (Centroids) gruppiert.
   • Aggregation: Für jedes Cluster wird ein einzelner repräsentativer Eintrag erstellt. Das Centroid-Abfragevektor ist der Mittelwert der Abfragen des Clusters. Der Aufmerksamkeitszustand $(\bar{a}, \bar{Z})$ wird mittels aufmerksamkeitsbewusster Aggregation berechnet, wobei die Zustände aller Mitglieder des Clusters unter Verwendung der Online-Softmax-Aktualisierungsregel zusammengeführt werden. Dies stellt sicher, dass das Centroid einen verlustfreien Durchschnitt der Präfix-Aufmerksamkeit des Clusters repräsentiert.
   • Effizienz: Um sehr lange Präfixe (z. B. 16K+ Tokens) zu verarbeiten, ohne den GPU-Speicher zu überschreiten, kann das Präfix in Blöcke partitioniert, unabhängig kodiert und innerhalb des Speichers unter Ausnutzung der Komposabilitätseigenschaft der Online-Softmax-Identität zusammengeführt werden.

2. Online-Inferenzphase:

   • Abruf: Für eine eingehende Abfrage identifiziert das System das nächstgelegene Speicher-Centroid mittels Kosinus-Ähnlichkeit (optional mit Whitening-Transformationen und RoPE-Behandlung).
   • Zusammenführung: Der abgerufene vorkalkulierte Aufmerksamkeitszustand $(\bar{a}, \bar{Z})$ wird mit dem Selbst-Aufmerksamkeitsausgang der Abfrage (berechnet über Nicht-Präfix-Tokens) unter Verwendung des durch die Online-Softmax-Identität definierten Zusammenführungsoperators verschmolzen. Dies stellt den exakten Aufmerksamkeitsausgang wieder her, der entstanden wäre, wenn das Präfix vorhanden gewesen wäre, ohne dass auf die Präfix-Tokens geachtet werden muss.
   • Komplexität: Die Suchkosten skalieren logarithmisch ($O(\log K)$) mit der Speichergröße durch hierarchische Indizierung, unabhängig von der ursprünglichen Präfixlänge.

Hauptbeiträge

1. Trainingsfreie Externalisierung: Die Arbeit schlägt eine Methode vor, lange Präfixe in ein kompaktes Speicherwörterbuch zu externalisieren, ohne gradientenbasiertes Training oder zusätzliche Modelle zu erfordern.
2. Erweiterung der Online-Softmax: Die Autoren erweitern die Online-Softmax-Identität von der effizienten Aufmerksamkeitsberechnung auf den Bereich der wiederverwendbaren Präfixe über verschiedene Abfragen hinweg, was eine verlustfreie Rekonstruktion der Präfix-Aufmerksamkeit durch Abruf ermöglicht.
3. Leistung und Effizienz: Experimente zeigen, dass ASM die Leistung der vollen Aufmerksamkeit erreicht oder übertrifft, während die Rechenkosten erheblich reduziert werden.

Experimentelle Ergebnisse
Die Methode wurde an ManyICLBench (In-Context-Lernen) und RuleArena/NBA (Retrieval-Augmented Generation) unter Verwendung von LLaMA-3.1-8B evaluiert.

• In-Context-Lernen (ICL): Auf ManyICLBench verbesserte ASM die Genauigkeit gegenüber dem Standard-In-Context-Lernen über Speicherbudgets von 1K bis 8K Einträgen. Bei einem Speicherbudget von 8K reduzierte sich die Aufmerksamkeits-Latenz im Vergleich zur vollen Aufmerksamkeit um den Faktor 1,36. ASM übertraf Baselines zur KV-Cache-Kompression (KVzip), die oft kritische Informationen wie Label-Tokens kollabieren ließen.
• Retrieval-Augmented Generation (RAG): Auf dem NBA-Benchmark (mit ca. 20K Tokens an Regeln) übertraf ASM die Leistung der RAG mit voller Aufmerksamkeit unter Verwendung von nur 20 % des Speicheraufwands. Die optimale Leistung wurde bei einer intermediären Codebook-Größe (4K Einträge) beobachtet, was darauf hindeutet, dass die Anzahl der Einträge ein aufgaben spezifischer Hyperparameter ist und nicht einfach maximiert werden sollte.
• Latenz: Während der hierarchische Abruf bei sehr kleinen Speicherbudgets einen höheren Overhead verursacht, ist das Latenzwachstum sublinear. ASM wird bei etwa 4K Speichereinträgen schneller als die volle Aufmerksamkeit und erreicht bei 16K Einträgen eine 1,8-fache Beschleunigung.

Bedeutung und Behauptungen
Die Arbeit behauptet, dass Attention-State Memory die strukturellen Kosten der Generierung langer Kontexte adressiert, indem es Präfixwissen von der pro-Abfrage-Aufmerksamkeitsberechnung entkoppelt. Durch die Externalisierung des Präfixes in ein vorkalkuliertes, nachschlagebasiertes Gedächtnis eliminiert die Methode die lineare Skalierung der Inferenz-Latenz mit der Präfixlänge und beseitigt das Problem des Präfix-Decay, das mit herkömmlichen Aufmerksamkeitsmechanismen verbunden ist.

Die Autoren positionieren diese Arbeit als einen Schritt in Richtung einer breiteren Forschungsrichtung: die Externalisierung von LLM-Wissen in kompakte, wiederverwendbare Darstellungen, die über Text-Prompts und Modellparameter hinaus existieren. Sie betonen, dass ihr Ansatz einen günstigen Kompromiss zwischen Genauigkeit und Effizienz bietet und eine trainingsfreie Alternative zur Parameter-Internalisierung sowie eine effizientere Alternative zur Präfix-Kompression darstellt.