Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Résumé Technique : Tensor Cache

Énoncé du Problème
L'inférence des Transformers autoregressifs repose sur la mise en cache des paires Clé-Valeur (KV) pour éviter de recalculer les préfixes. Cependant, l'état KV conservé croît linéairement avec la longueur du contexte, la profondeur et la concurrence, créant un goulot d'étranglement mémoire. Le cache à fenêtre glissante limite l'usage mémoire en ne conservant que les $W$ tokens les plus récents, mais cette approche souffre d'un « oubli brutal » : une fois un token expulsé de la fenêtre, ses informations sont entièrement rejetées, rendant les preuves pertinentes en dehors de la fenêtre inaccessibles. Les méthodes à état borné existantes soit rejettent entièrement les tokens expulsés, soit tentent de compresser chaque token en une moyenne courante, ce qui peut introduire du bruit ou échouer à distinguer efficacement le contexte récent du contexte lointain.

Méthodologie : Tensor Cache (TC)
Tensor Cache introduit une architecture à état borné à deux niveaux qui transforme l'expulsion KV d'un événement de suppression en une opération d'écriture en mémoire associative.

1. Architecture à Deux Niveaux :

   • L1 (Cache de Premier Niveau) : Un tampon circulaire KV de taille fixe contenant les $W$ tokens les plus récents. Cela prend en charge l'attention softmax exacte pour le contexte local.
   • L2 (Cache de Deuxième Niveau) : Une matrice de produit extérieur $A$ de taille fixe (mémoire à poids rapides) qui stocke exclusivement les informations des tokens expulsés de la fenêtre L1.
   • Mécanisme d'Écriture : Lorsque le tampon circulaire écrase une entrée, la paire $(k, v)$ déplacée n'est pas rejetée. Elle est plutôt écrite dans la matrice L2 $A$ via une mise à jour de produit extérieur : $A \leftarrow \lambda A + \eta (k \otimes v)$. Ici, $\lambda$ (décroissance) et $\eta$ (taux d'écriture) sont des scalaires appris par tête.
   • Mécanisme de Lecture : Les requêtes futures $q_t$ accèdent au passé compressé en calculant $r_t = q_t A$. En exploitant l'identité de l'attention linéaire $q(k \otimes v) = \langle q, k \rangle v$, cette opération récupère efficacement les informations de tous les tokens expulsés sans les stocker individuellement.
   • Fusion : Une porte scalaire apprise $g$ fusionne les sorties de l'attention locale (L1) et de la lecture mémoire (L2) : $y_t = y_{local} + \sigma(g) m_t$.

2. Optimisation de l'Entraînement (Balayage Parallèle) :
   L'article identifie un problème critique avec les raccourcis d'entraînement par blocs standards. Une approche naïve résume un bloc par sa moyenne ($\bar{k} \otimes \bar{v}$), ce qui introduit $C^2 - C$ produits extérieurs croisés de tokens spuriaires ($k_i \otimes v_j$ où $i \neq j$) qui n'existent pas lors de l'inférence en flux continu.

   • Solution : Les auteurs implémentent un balayage pondéré parallèle mathématiquement équivalent aux écritures séquentielles par token. Cela comble l'écart entre l'entraînement et l'inférence dans l'epsilon float32 (erreur relative < $10^{-7}$), garantissant que le modèle apprend les dynamiques associatives correctes sans le bruit des interférences croisées de tokens.

Contributions Clés

• Mémoire Associative Conditionnée par l'Expulsion : Contrairement aux mémoires à produit extérieur précédentes (par exemple, Infini-attention, mLSTM) qui ingèrent chaque token, TC utilise la matrice de produit extérieur strictement comme un cache L2 alimenté uniquement par les expulsions de la fenêtre glissante. Cela partitionne proprement les tokens en « résidents dans L1 » et « compressés dans L2 ».
• Correction des Artifacts d'Entraînement : L'article identifie et résout un raccourci d'entraînement (moyenne par bloc) précédemment négligé qui introduit des produits croisés de tokens spuriaires. Le balayage pondéré parallèle proposé élimine cette discordance au même coût computationnel.
• Efficacité à État Borné : TC fournit un mécanisme où le contexte plus ancien influence les prédictions sans conserver les entrées individuelles expulsées, offrant une alternative pratique à la conservation complète des KV.

Résultats Expérimentaux
Les auteurs ont évalué TC sur des systèmes de mise à l'échelle, un rappel associatif synthétique et un modèle de langage à long contexte sur texte réel (OpenWebText et Shakespeare).

• Rappel Associatif Synthétique : Dans des tâches nécessitant la récupération d'informations au-delà de la fenêtre glissante, TC a atteint une précision de 100 %, égalant les performances du Full KV. Il a réduit l'état d'inférence conservé de 72 à 84 % par rapport au Full KV tout en surpassant significativement Infini-attention (qui se dégradait à mesure que les écarts augmentaient) et les bases de référence à fenêtre glissante (qui performaient au niveau du hasard).
• Modélisation du Langage à Long Contexte : Sur OpenWebText, TC a obtenu la plus faible NLL (Négative Log-Likelihood) moyenne à chaque longueur de contexte d'évaluation de 1 024 à 32 768 tokens, avec la seule exception de $L=2\,048$ où il était statistiquement indiscernable d'Infini-attention. À $L=32\,768$, TC a atteint une NLL de 5,14, surpassant le Full KV (6,00) et Infini-attention (5,42), tout en utilisant 2,4 fois moins de mémoire GPU de pointe que le Full KV.
• Débit : TC a engendré une surcharge de débit d'environ 30 % par rapport aux méthodes purement à fenêtre glissante (due à la lecture/mise à jour L2) mais est resté constamment plus rapide qu'Infini-attention sur toutes les longueurs de contexte testées.
• Diagnostics de Capacité : Les expériences sur la capacité brute de mémoire associative ont montré que les performances de récupération se dégradent à mesure que davantage d'associations non liées sont compressées dans l'état fixe, confirmant la nature bornée de la mémoire.

Signification et Revendications
L'article revendique que Tensor Cache comble avec succès le fossé entre l'efficacité mémoire de l'attention à fenêtre glissante et les capacités de rappel de l'attention à contexte complet. En traitant l'expulsion KV comme une opération mémoire structurée plutôt qu'une suppression destructive, TC permet aux Transformers d'accéder aux preuves pertinentes en dehors de la fenêtre immédiate sans le coût mémoire linéaire des caches KV complets. Les auteurs soulignent que leur contribution principale n'est pas l'invention de la mémoire à produit extérieur elle-même, mais son application spécifique en tant que cache L2 conditionné par l'expulsion et la correction rigoureuse des artifacts d'entraînement qui entravaient précédemment de telles méthodes. TC est présenté comme une alternative pratique à état borné qui améliore la frontière mémoire-qualité pour l'inférence à long contexte.