MemoryLLM : Mémoire interprétable feed-forward plug-and-play pour les transformeurs
MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers
January 30, 2026
papers.authors: Ajay Jaiswal, Lauren Hannah, Han-Byul Kim, Duc Hoang, Arnav Kundu, Mehrdad Farajtabar, Minsik Cho
cs.AI
papers.abstract
Comprendre le fonctionnement des composants des transformateurs dans les LLMs est essentiel, car il est au cœur des récentes avancées technologiques en intelligence artificielle. Dans ce travail, nous revisitons les défis associés à l'interprétabilité des modules feed-forward (FFN) et proposons MemoryLLM, qui vise à découpler les FFN de l'auto-attention et nous permet d'étudier les FFN découplés comme une mémoire neuronale de récupération token-wise et indépendante du contexte. En détail, nous étudions comment les tokens d'entrée accèdent aux emplacements mémoire au sein des paramètres des FFN et l'importance de la mémoire des FFN pour différentes tâches en aval. MemoryLLM permet d'obtenir des FFN indépendants du contexte en les entraînant de manière isolée de l'auto-attention, directement à partir des embeddings de tokens. Cette approche permet de pré-calculer les FFN sous forme de tables de consultation token-wise (ToLs), autorisant un transfert à la demande entre la VRAM et le stockage, améliorant ainsi l'efficacité de l'inférence. Nous présentons également Flex-MemoryLLM, que nous positionnons entre une architecture de transformateur conventionnelle et MemoryLLM. Cette architecture comble l'écart de performance induit par l'entraînement des FFN avec des embeddings token-wise indépendants du contexte.
English
Understanding how transformer components operate in LLMs is important, as it is at the core of recent technological advances in artificial intelligence. In this work, we revisit the challenges associated with interpretability of feed-forward modules (FFNs) and propose MemoryLLM, which aims to decouple FFNs from self-attention and enables us to study the decoupled FFNs as context-free token-wise neural retrieval memory. In detail, we investigate how input tokens access memory locations within FFN parameters and the importance of FFN memory across different downstream tasks. MemoryLLM achieves context-free FFNs by training them in isolation from self-attention directly using the token embeddings. This approach allows FFNs to be pre-computed as token-wise lookups (ToLs), enabling on-demand transfer between VRAM and storage, additionally enhancing inference efficiency. We also introduce Flex-MemoryLLM, positioning it between a conventional transformer design and MemoryLLM. This architecture bridges the performance gap caused by training FFNs with context-free token-wise embeddings.