Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Technische Zusammenfassung: Tensor Cache

Problemstellung
Die Inferenz autoregressiver Transformer basiert auf dem Cachen von Key-Value (KV)-Paaren, um die Neuberechnung von Präfixen zu vermeiden. Der beibehaltene KV-Zustand wächst jedoch linear mit der Kontextlänge, der Tiefe und der Parallelität und erzeugt so einen Engpass im Speicherbedarf. Das Cachen mit gleitendem Fenster begrenzt die Speichernutzung, indem nur die neuesten $W$ Tokens behalten werden, doch dieser Ansatz leidet unter „hartem Vergessen": Sobald ein Token aus dem Fenster verdrängt wird, gehen seine Informationen vollständig verloren, wodurch relevante Beweise außerhalb des Fensters unzugänglich werden. Bestehende Methoden mit begrenztem Zustand verwerfen entweder verdrängte Tokens vollständig oder versuchen, jeden Token in einen laufenden Durchschnitt zu komprimieren, was Rauschen einführen oder die Unterscheidung zwischen neuem und fernem Kontext unwirksam gestalten kann.

Methodik: Tensor Cache (TC)
Tensor Cache führt eine zweistufige Architektur mit begrenztem Zustand ein, die KV-Verdrängung von einem Löschvorgang in einen assoziativen Speicher-Schreibvorgang transformiert.

1. Zweistufige Architektur:

   • L1 (Erststufiger Cache): Ein KV-Ringpuffer fester Größe, der die neuesten $W$ Tokens hält. Dies unterstützt eine exakte Softmax-Aufmerksamkeit für den lokalen Kontext.
   • L2 (Zweitstufiger Cache): Eine Matrix $A$ fester Größe für äußere Produkte (Fast-Weight-Speicher), die ausschließlich Informationen von Tokens speichert, die aus dem L1-Fenster verdrängt wurden.
   • Schreibmechanismus: Wenn der Ringpuffer einen Eintrag überschreibt, wird das verdrängte $(k, v)$-Paar nicht verworfen. Stattdessen wird es über eine Aktualisierung des äußeren Produkts in die L2-Matrix $A$ geschrieben: $A \leftarrow \lambda A + \eta (k \otimes v)$. Hier sind $\lambda$ (Abklingfaktor) und $\eta$ (Schreibrate) pro-Kopf gelernte Skalare.
   • Lese-Mechanismus: Zukünftige Abfragen $q_t$ greifen auf die komprimierte Vergangenheit zu, indem sie $r_t = q_t A$ berechnen. Unter Ausnutzung der Identität der linearen Aufmerksamkeit $q(k \otimes v) = \langle q, k \rangle v$ holt dieser Vorgang Informationen aus allen verdrängten Tokens ab, ohne diese einzeln speichern zu müssen.
   • Fusion: Ein gelernter skalare Gate $g$ fusioniert die Ausgaben der lokalen Aufmerksamkeit (L1) und des Speicherzugriffs (L2): $y_t = y_{local} + \sigma(g) m_t$.

2. Optimierung des Trainings (Paralleler Scan):
   Die Arbeit identifiziert ein kritisches Problem bei Standard-Abkürzungen für das Chunking im Training. Ein naiver Ansatz fasst ein Chunk durch seinen Mittelwert ($\bar{k} \otimes \bar{v}$) zusammen, was $C^2 - C$ spurartige kreuz-Token-Außere-Produkte ($k_i \otimes v_j$ mit $i \neq j$) einführt, die während der Streaming-Inferenz nicht existieren.

   • Lösung: Die Autoren implementieren einen parallelen gewichteten Summen-Scan, der mathematisch äquivalent zu sequenziellen Schreibvorgängen pro Token ist. Dies schließt die Lücke zwischen Training und Inferenz innerhalb von Float32-Epsilon ($< 10^{-7}$ relativer Fehler) und stellt sicher, dass das Modell die korrekten assoziativen Dynamiken ohne das Rauschen kreuz-Token-Interferenzen lernt.

Hauptbeiträge

• Verdrängungsbedingte assoziative Speicher: Im Gegensatz zu früheren Speichern für äußere Produkte (z. B. Infini-Attention, mLSTM), die jeden Token aufnehmen, verwendet TC die Matrix für äußere Produkte strikt als L2-Cache, der ausschließlich durch Verdrängungen aus dem gleitenden Fenster gespeist wird. Dies partitioniert Tokens sauber in „im L1 resident" und „in L2 komprimiert".
• Korrektur von Trainingsartefakten: Die Arbeit identifiziert und löst eine zuvor übersehene Trainings-Abkürzung (Chunk-Mittelwert), die spurartige kreuz-Token-Produkte einführt. Der vorgeschlagene parallele gewichtete Summen-Scan eliminiert diese Diskrepanz bei gleichen Rechenkosten.
• Effizienz des begrenzten Zustands: TC bietet einen Mechanismus, bei dem älterer Kontext Vorhersagen beeinflusst, ohne einzelne verdrängte Einträge zu behalten, und stellt eine praktische Alternative zur vollständigen KV-Behaltung dar.

Experimentelle Ergebnisse
Die Autoren bewerteten TC über Skalierungssysteme, synthetisches assoziatives Abrufen und Sprachmodellierung mit echtem Text für lange Kontexte (OpenWebText und Shakespeare).

• Synthetisches assoziatives Abrufen: In Aufgaben, die den Abruf von Informationen jenseits des gleitenden Fensters erfordern, erreichte TC eine Genauigkeit von 100 % und entsprach damit der Leistung von Full KV. Es reduzierte den beibehaltenen Inferenz-Zustand um 72–84 % im Vergleich zu Full KV, übertraf dabei deutlich Infini-Attention (die bei zunehmenden Lücken an Leistung verlor) und Baselines mit gleitendem Fenster (die auf Zufallsniveau performten).
• Sprachmodellierung mit langem Kontext: Auf OpenWebText erreichte TC bei jeder Evaluierungskontextlänge von 1.024 bis 32.768 Tokens den niedrigsten mittleren negativen Log-Likelihood (NLL), mit der einzigen Ausnahme von $L=2.048$, wo es statistisch nicht von Infini-Attention zu unterscheiden war. Bei $L=32.768$ erreichte TC einen NLL von 5,14 und übertraf dabei Full KV (6,00) und Infini-Attention (5,42), während es 2,4$\times$ weniger Peak-GPU-Speicher als Full KV benötigte.
• Durchsatz: TC verursachte im Vergleich zu reinen Methoden mit gleitendem Fenster einen Durchsatz-Overhead von ca. 30 % (aufgrund von L2-Lese-/Schreibvorgängen), blieb aber über alle getesteten Kontextlängen hinweg konsistent schneller als Infini-Attention.
• Kapazitätsdiagnostik: Experimente zur rohen Kapazität des assoziativen Speichers zeigten, dass die Abringleistung abnimmt, wenn mehr nicht zusammenhängende Assoziationen in den festen Zustand komprimiert werden, was die begrenzte Natur des Speichers bestätigt.

Bedeutung und Behauptungen
Die Arbeit behauptet, dass Tensor Cache erfolgreich die Lücke zwischen der Speichereffizienz der Aufmerksamkeit mit gleitendem Fenster und den Abrufkapazitäten der Aufmerksamkeit mit vollem Kontext schließt. Indem KV-Verdrängung als strukturierter Speicheroperation und nicht als destruktive Löschung behandelt wird, ermöglicht TC Transformern den Zugriff auf relevante Beweise außerhalb des unmittelbaren Fensters ohne die linearen Speicherkosten vollständiger KV-Caches. Die Autoren betonen, dass ihr Hauptbeitrag nicht die Erfindung des Speichers für äußere Produkte an sich ist, sondern dessen spezifische Anwendung als verdrängungsbedingter L2-Cache sowie die rigorose Korrektur von Trainingsartefakten, die solche Methoden zuvor behindert hatten. TC wird als praktische, zustandsbegrenzte Alternative präsentiert, die die Grenze zwischen Speicher und Qualität für die Inferenz mit langem Kontext verbessert.