Résumé Technique : TF-Engram

Problématique

Les grands modèles de langage (LLM) stockent implicitement des connaissances factuelles et des motifs spécifiques à un domaine au sein des paramètres denses des Transformers. L'expansion de ces connaissances nécessite généralement un pré-entraînement coûteux, un ajustement fin supervisé (SFT), l'augmentation par récupération (RAG) ou des fenêtres de contexte étendues, autant de méthodes qui augmentent les coûts d'entraînement, la latence d'inférence ou la complexité de déploiement.

Les approches actuelles de mémoire augmentée, particulièrement les mémoires de type Engram, tentent de découpler le stockage des connaissances des paramètres du modèle en injectant des signaux récupérés dans les états cachés. Cependant, les conceptions d'Engram actuelles font face à trois limitations critiques lors du passage à l'échelle vers de vastes bases de connaissances statiques :

1. Dépendance à l'entraînement : Les Engrams conventionnels nécessitent souvent d'apprendre conjointement les entrées de mémoire, les couches de projection ou les portes avec le modèle de base (backbone), ce qui les rend difficiles à déployer sur différents LLM ou à adapter à de nouveaux domaines sans réentraînement.
2. Compromis Fidélité Sémantique vs Capacité : Pour faire tenir de vastes espaces de phrases dans une mémoire GPU limitée, les systèmes existants s'appuient sur une compression par hachage. Cela force des phrases sans rapport à entrer en collision dans des créneaux de mémoire partagés, créant des représentations ambiguës et mixtes qui dégradent la fidélité sémantique des connaissances spécifiques à une phrase (ex: « New York City » ou « théorie quantique des champs »).
3. Latence dans l'Accès Externe : Le déplacement de la mémoire vers un stockage externe (DRAM/SSD) pour éviter les collisions de hachage introduit une latence qui bloque le décodage autorégressif. Le préchargement (prefetching) naïf basé sur la tête de modélisation du langage finale est trop tardif pour chevaucher le transfert de stockage, car la distribution finale n'est disponible qu'après le passage complet (forward pass).

Méthodologie

Le papier propose TF-Engram, un système Engram sans entraînement (train-free) qui utilise une mémoire supportée par SSD et un préchargage prédictif guidé par sortie précoce (early-exit). Le système est conçu pour construire une mémoire sémantique spécifique aux phrases hors ligne et l'intégrer à l'inférence du LLM sans modifier le modèle de base.

1. Construction de Mémoire Sémantique Sans Entraînement

Au lieu d'apprendre des paramètres de mémoire, TF-Engram construit une table de mémoire de phrases statique hors ligne :

• Minage : Les phrases candidates (bigrammes, trigrammes, entités nommées, termes techniques) sont extraites de corpus à grande échelle (ex: texte web, Wikipédia, articles scientifiques).
• Encodage : Les phrases sont tokenisées à l'aide du tokenizer du LLM cible et encodées via un encodeur sémantique externe gelé (ex: Qwen3-Embedding).
• Indexation : Un index à deux niveaux est construit : préfixe de token → ID de phrase → adresse mémoire. Cela permet une recherche efficace sans scanner l'intégralité de la table.
• Sélection : Lors de l'inférence, le système utilise une politique déterministe (ex: sélection par correspondance la plus longue) pour récupérer des vecteurs de phrases spécifiques, évitant ainsi l'ambiguïté des collisions de hachage.

2. Hiérarchie de Mémoire Supportée par SSD

Pour supporter de vastes tables de mémoire sans épuiser la mémoire GPU, TF-Engram organise le stockage en une hiérarchie à trois niveaux :

• Cache GPU (Chaud) : Stocke les entrées les plus fréquemment accédées pour une injection à faible latence.
• DRAM Hôte (Tiède) : Agit comme un cache de transit pour les entrées qui ne sont pas actuellement dans la mémoire GPU.
• NVMe SSD (Froid) : Stocke l'intégralité de la table de mémoire des phrases, absorbant la demande de capacité.
  Cette hiérarchie permet au système de préserver des entrées logiques précises et spécifiques aux phrases plutôt que de les compresser dans des créneaux de hachage ambigus.

3. Préchargage Prédictif Guidé par Sortie Précoce (Early-Exit)

Pour masquer la latence de l'accès à la mémoire externe pendant le décodage autorégressif, TF-Engram emploie un mécanisme de préchargage prédictif :

• Tête de Sortie Précoce (Early-Exit Head) : Une tête de projection auxiliaire est attachée à une couche intermédiaire (ex: couche $L-r$) proche de la fin du Transformer.
• Prédiction : Cette tête génère une distribution approximative du prochain token avant que le passage complet (forward pass) ne soit terminé.
• Préchargage (Prefetching) : Basé sur les top-$K$ tokens prédits, le système infère les accès futurs probables aux phrases et précharge de manière asynchrone les entrées de mémoire correspondantes du SSD ou de la DRAM vers le cache GPU.
• Masquage de Latence : Le transfert de mémoire se produit pendant que les couches restantes du Transformer ($L-r+1$ à $L$) poursuivent leur calcul, permettant de chevaucher efficacement les E/S avec le traitement. La génération du token final repose toujours sur l'originale tête de LM une fois toutes les couches terminées.

Contributions Clés

• Construction Sans Entraînement : TF-Engram découple la construction de la mémoire de l'optimisation du modèle, permettant à la mémoire spécifique aux phrases d'être construite hors ligne à partir de corpus externes et réutilisée à travers les exécutions d'inférence sans réentraînement du backbone.
• Hiérarchie de Stockage Évolutive : En déplaçant la charge de capacité vers une hiérarchie GPU–DRAM–SSD, le système évite la dégradation sémantique causée par la compression par hachage, permettant le stockage de millions d'entrées de phrases précises.
• Préchargage Masquant la Latence : L'introduction du préchargage prédictif guidé par sortie précoce permet d'anticiper les requêtes de mémoire suffisamment tôt pour chevaucher l'accès au stockage externe avec le calcul des couches restantes du Transformer, atténuant ainsi la pénalité de latence du stockage par SSD.

Résultats

Le système a été évalué sur Qwen3-0.6B sur dix benchmarks en aval (incluant MMLU, ARC-Challenge et TruthfulQA).

• Performance : TF-Engram a amélioré le score moyen en aval de 57,6 (backbone gelé) à 59,4. Il a surpassé à la fois le backbone gelé et une ligne de base LoRA à paramètres équivalents (58,7).
• Efficacité :
  • Coût Hors Ligne : La construction d'une table de 100 millions d'entrées a pris environ 2,0 heures, principalement dominée par l'extraction de n-grammes plutôt que par la génération de vecteurs.
  • Utilisation de la Mémoire : La conception supportée par SSD a réduit l'utilisation de la mémoire GPU de 24,03 Go (pour une table de 10 millions d'entrées résidente en GPU) à 3,06 Go, déplaçant la charge de stockage vers le CPU/SSD.
  • Débit et Latence : Bien qu'un accès externe naïf réduise le débit, le préchargage prédictif (Top-128) a permis de récupérer une grande partie de la perte, atteignant 462,4 tokens/s (contre 481,2 tokens/s pour la ligne de base) et maintenant un surcoût de latence par token modeste (2,16 ms/token contre 2,08 ms/token pour la ligne de base).
• Sensibilité : La performance a progressé avec la taille de la mémoire (jusqu'à 100 millions d'entrées), et une couche de sortie précoce à la couche 8 a fourni le compromis optimal entre précision de prédiction et masquage de latence.

Signification

Le papier affirme que TF-Engram démontre une voie pratique vers une augmentation de mémoire évolutive, sans entraînement et à faible surcharge pour les LLM. En découplant le stockage de la connaissance statique des paramètres denses du modèle, TF-Engram offre une solution orientée système qui :

1. Évite les coûts de fine-tuning ou de réentraînement pour les mises à jour de connaissances.
2. Préserve la fidélité sémantique en éliminant l'ambiguïté des collisions de hachage grâce au stockage externe.
3. Maintient l'efficacité de l'inférence en masquant la latence de la mémoire externe par la prédiction de sortie précoce.

Les auteurs positionnent ce travail comme une alternative viable au réglage efficace des paramètres (comme LoRA) et à la génération augmentée par récupération (RAG) pour les scénarios nécessitant une injection de connaissances fines au niveau des phrases statiques.