MemoryLLM: Plug-n-Play interpreteerbaar feedforward-geheugen voor Transformers
MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers
January 30, 2026
Auteurs: Ajay Jaiswal, Lauren Hannah, Han-Byul Kim, Duc Hoang, Arnav Kundu, Mehrdad Farajtabar, Minsik Cho
cs.AI
Samenvatting
Het begrijpen van de werking van transformatorcomponenten in grote taalmodel(len) is belangrijk, aangezien dit de kern vormt van recente technologische vooruitgang in kunstmatige intelligentie. In dit werk herbezien we de uitdagingen die gepaard gaan met de interpreteerbaarheid van feed-forward modules (FFN's) en stellen we MemoryLLM voor, dat als doel heeft FFN's los te koppelen van zelf-attentie en ons in staat stelt de losgekoppelde FFN's te bestuderen als contextvrij, token-gewijs neurale ophaalgeheugen. In detail onderzoeken we hoe invoertokens toegang krijgen tot geheugenlocaties binnen de FFN-parameters en het belang van FFN-geheugen bij verschillende downstreamtaken. MemoryLLM bereikt contextvrije FFN's door ze geïsoleerd van zelf-attentie te trainen, rechtstreeks met behulp van de token-embeddingen. Deze aanpak maakt het mogelijk FFN's vooraf te berekenen als token-gewijze opzoekingen (ToL's), wat een on-demand overdracht tussen VRAM en opslag mogelijk maakt en bovendien de inferentie-efficiëntie verbetert. We introduceren ook Flex-MemoryLLM en positioneren dit tussen een conventionele transformatorarchitectuur en MemoryLLM. Deze architectuur overbrugt het prestatieverschil dat wordt veroorzaakt door het trainen van FFN's met contextvrije token-gewijze embeddingen.
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.