🤖 machine learning

Mixture of Chapters: Scaling Learnt Memory in Transformers

Ce papier propose une architecture de transformateurs enrichie par des banques de mémoire latentes parcourues via un routage par chapitres, permettant d'augmenter la capacité de stockage des connaissances et de réduire l'oubli lors du fine-tuning sans alourdir les coûts de calcul.

Auteurs originaux : Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

Publié 2026-03-24

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

🧠 Le Problème : Le Cerveau qui Oublie

Imaginez un étudiant très brillant (un modèle d'intelligence artificielle de type "Transformer") qui a lu des millions de livres pendant ses études (l'entraînement). Il est très intelligent, mais il a un gros défaut : il n'a pas de carnet de notes.

Pour répondre à une question, il doit tout retenir dans sa tête (ses "paramètres"). Si on lui demande de réviser pour un examen très spécifique (l'apprentissage par instruction), il risque d'effacer accidentellement ce qu'il savait sur l'histoire ou la science pour faire de la place aux nouvelles règles. C'est ce qu'on appelle l'oubli catastrophique. De plus, s'il veut se souvenir de quelque chose, il doit "rejouer" tout son cerveau, ce qui est lent et énergivore.

💡 La Solution : La Bibliothèque Intelligente (Mixture of Chapters)

Les auteurs de cet article proposent d'ajouter à ce cerveau une bibliothèque interne qu'il peut consulter.

La Bibliothèque (La Mémoire Apprise) :
Au lieu d'écrire des livres entiers, le modèle crée une immense bibliothèque de "cartes mémoires" (des jetons latents). Ces cartes contiennent des faits, des concepts et des connaissances acquises pendant l'entraînement. C'est comme si l'étudiant avait un index mental géant où chaque fait est rangé sur une petite carte.
Le Problème de la Taille :
Cette bibliothèque est énorme (262 000 cartes !). Si l'étudiant devait feuilleter toutes les cartes pour chaque question, il mettrait des heures à répondre. C'est trop lent et trop cher en énergie.
Le Système de Tri (Le "Mélange de Chapitres") :
C'est ici que vient l'idée géniale : le "Mélange de Chapitres" (Mixture of Chapters).
Imaginez que la bibliothèque est divisée en 4 000 étagères (les chapitres).
- Quand l'étudiant reçoit une question, un libraire intelligent (le "routeur") regarde la question.
- Au lieu de parcourir toute la bibliothèque, le libraire dit : "Attends, cette question concerne la cuisine. Oublie l'histoire et la physique. Va directement aux étagères 12, 45 et 89."
- L'étudiant ne consulte alors que quelques étagères spécifiques.

🚀 Pourquoi c'est génial ?

Économie d'énergie : Comme il ne lit que quelques étagères au lieu de toute la bibliothèque, il répond très vite, même avec une mémoire gigantesque.
Pas d'oubli : Puisque les faits sont rangés sur des cartes physiques dans la bibliothèque, l'étudiant n'a pas besoin de les "garder en tête" en permanence. Même s'il apprend de nouvelles choses (comme des règles de politesse pour un chatbot), il ne perd pas ses connaissances sur la cuisine ou l'histoire, car elles sont bien rangées dans leur étagère.
Flexibilité : On peut ajouter plus d'étagères (plus de mémoire) sans avoir à reconstruire tout le cerveau de l'étudiant.

📊 Les Résultats de l'expérience

Les chercheurs ont testé cette idée :

Avant : Le modèle classique (sans bibliothèque) apprenait de nouvelles choses mais oubliait ses anciennes connaissances. Ses réponses devenaient aléatoires sur des questions de culture générale.
Après : Le modèle avec la "Bibliothèque à Chapitres" a appris les nouvelles choses sans oublier les anciennes. Il est resté stable et précis.
Le petit secret : Ils ont même découvert qu'on pouvait geler la bibliothèque (ne plus la modifier) pendant l'apprentissage des nouvelles règles, et le modèle fonctionnait tout aussi bien ! Cela signifie que la bibliothèque est une base solide et fiable.

En résumé

C'est comme passer d'un cerveau qui essaie de tout retenir par cœur (et qui finit par mélanger les choses) à un cerveau qui possède une bibliothèque organisée avec un système de tri ultra-rapide. Cela permet à l'intelligence artificielle d'être plus grande, plus intelligente, et surtout, de ne jamais oublier ce qu'elle a appris, tout en restant rapide et efficace.

1. Problématique

Les architectures Transformer actuelles stockent les connaissances acquises lors de l'entraînement de manière implicite dans leurs paramètres denses. Cette approche présente plusieurs limites majeures :

Absence de mémoire explicite : Il n'existe pas de mécanisme architectural natif pour stocker, organiser ou interroger des connaissances de manière persistante et adressable.
Difficulté de mise à l'échelle : Augmenter la capacité de mémoire via des mécanismes d'attention denses sur de grands ensembles de tokens est prohibitif en termes de coût computationnel (complexité quadratique).
Oubli catastrophique : Lors de l'entraînement continu (par exemple, le passage du pré-entraînement au fine-tuning d'instructions), les modèles souffrent souvent d'un oubli des connaissances factuelles précédemment apprises, car les nouveaux gradients interfèrent avec les paramètres existants.

Les solutions existantes (comme le cache KV, la RAG ou la récurrence) gèrent principalement le contexte d'inférence ou récupèrent des données externes, mais ne fournissent pas de véritable "mémoire interne apprise" pour les faits.

2. Méthodologie : Mixture of Chapters (MoC)

Les auteurs proposent une architecture augmentée intégrant une banque de mémoire apprise (un ensemble de tokens latents) accessible via une attention croisée (cross-attention), combinée à une stratégie de routage sparse inspirée des Mixture of Experts (MoE).

A. Banque de Mémoire Apprise

Structure : Une banque de mémoire $M$ composée de $N_m$ tokens latents (initialisés aléatoirement et appris de bout en bout).
Accès : À chaque couche du Transformer, le flux de tokens génère des requêtes ( $Q$ ), tandis que la banque de mémoire fournit les clés ( $K$ ) et les valeurs ( $V$ ).
Intégration : Cette opération d'attention croisée est ajoutée aux blocs Transformer standards via une connexion résiduelle.

B. Mise à l'échelle via le "Mélange de Chapitres" (MoC)

Pour éviter le coût prohibitif d'une attention dense sur une banque massive (262 000 tokens), les auteurs partitionnent la banque en chapitres ( $C$ ).

Partitionnement : La banque est divisée en $C$ chapitres, chacun contenant $T$ tokens.
Routage : Un routeur léger, basé sur les états cachés du modèle (via un pooling), sélectionne un sous-ensemble restreint de chapitres ( $k$ ) pour chaque séquence d'entrée.
Calcul Sparse : Le modèle n'effectue l'attention croisée que sur les $k$ chapitres sélectionnés. Cela réduit la complexité de $O(L \cdot N_m)$ à $O(L \cdot k \cdot T)$ , rendant la mise à l'échelle de la mémoire computationnellement gérable.

C. Configuration Expérimentale

Modèle de base : Transformer en mode decoder-only (16 couches, taille cachée 768).
Mémoire : 262 208 tokens de mémoire répartis en 4 097 chapitres.
Routage : Sélection de $k=64$ chapitres par séquence (plus un chapitre partagé).
Entraînement : Pré-entraînement sur 9,6 milliards de tokens, suivi d'un Instruction Fine-Tuning (IFT) intensif sur 230 millions de tokens.

3. Contributions Clés

Banque de mémoire associative apprise : Introduction d'un stockage persistant de tokens latents accessibles via attention croisée, offrant une capacité de mémoire interne complémentaire aux paramètres du modèle.
Mixture of Chapters (MoC) : Une méthode de routage sparse efficace permettant de mettre à l'échelle la capacité de mémoire à des centaines de milliers de tokens tout en maintenant un coût de calcul constant et prévisible.
Évaluation iso-FLOP : Comparaison rigoureuse avec des baselines denses ayant le même budget de calcul (FLOPs), démontrant que l'ajout de mémoire apporte une valeur ajoutée au-delà de la simple augmentation de la taille du modèle.
Résilience à l'oubli : Démonstration empirique que la mémoire explicite améliore la rétention des connaissances factuelles lors de phases d'entraînement continues (transition pré-entraînement vers IFT).

4. Résultats Expérimentaux

A. Performance en Pré-entraînement

Le modèle MoC atteint une perte de validation plus faible (2,79) que le modèle dense iso-FLOP (2,86) et le modèle sans mémoire (2,92).
Cela suggère que la mémoire apporte une capacité supplémentaire pour l'apprentissage des motifs linguistiques et factuels.

B. Rétention des Connaissances et Oubli (Fine-Tuning)

L'expérience clé consiste à mesurer la dégradation des performances sur des benchmarks de connaissances (MMLU, ARC-Challenge, BoolQ) après un fine-tuning intensif :

Modèle Dense (Iso-FLOP) : Subit un oubli catastrophique significatif.
- ARC-Challenge : Chute de 6,69 points (31,77% $\to$ 25,08%).
- BoolQ : Chute de 6,24 points (57,13% $\to$ 50,89%), se rapprochant du hasard.
Modèle MoC : Reste remarquablement stable.
- ARC-Challenge : Chute limitée à 2,68 points.
- BoolQ : Légère amélioration (+0,24 points).
Conclusion : La banque de mémoire agit comme une "ancre" pour les faits, réduisant l'interférence lors de l'adaptation aux nouvelles tâches.

C. Gel de la Mémoire

Les auteurs ont testé le gel de la banque de mémoire pendant le fine-tuning. Les résultats montrent que geler la banque donne des performances équivalentes à la mettre à jour, indiquant que la mémoire pré-entraînée est déjà suffisante et robuste, ce qui réduit les coûts de calcul et de mémoire (VRAM) lors du post-entraînement.

5. Signification et Perspectives

Ce travail ouvre une nouvelle axe de mise à l'échelle pour les modèles de langage : au lieu d'augmenter uniquement la profondeur ou la largeur des paramètres denses, on peut augmenter la capacité de mémoire externe apprise.

Complémentarité : La mémoire explicite ne remplace pas les paramètres, mais complète la capacité implicite du modèle, offrant un mécanisme plus efficace pour stocker les faits.
Efficacité : La méthode MoC permet d'intégrer des mémoires massives sans exploser le coût computationnel, grâce au routage sparse.
Apprentissage Continu : La robustesse observée face à l'oubli catastrophique suggère que cette architecture est particulièrement adaptée aux scénarios d'apprentissage continu et de mise à jour de connaissances sans ré-entraînement complet.

En résumé, le papier démontre que l'intégration d'une mémoire latente structurée et accessible via un routage intelligent est une voie prometteuse pour construire des modèles plus intelligents, plus stables et capables de retenir des connaissances à long terme.