MemoryLLM: Plug-n-Play interpretierbarer Feed-Forward-Speicher für Transformer
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
Das Verständnis der Funktionsweise von Transformer-Komponenten in LLMs ist von zentraler Bedeutung, da sie den Kern der jüngsten technologischen Fortschritte im Bereich der künstlichen Intelligenz darstellen. In dieser Arbeit beleuchten wir erneut die Herausforderungen im Zusammenhang mit der Interpretierbarkeit von Feed-Forward-Modulen (FFNs) und schlagen MemoryLLM vor, das darauf abzielt, FFNs von der Self-Attention zu entkoppeln. Dies ermöglicht es uns, die entkoppelten FFNs als kontextfreien, token-basierten neuronalen Abrufspeicher zu untersuchen. Im Einzelnen untersuchen wir, wie Eingabe-Tokens auf Speicherstellen innerhalb der FFN-Parameter zugreifen und welche Bedeutung dem FFN-Speicher in verschiedenen Downstream-Aufgaben zukommt. MemoryLLM erreicht kontextfreie FFNs, indem diese isoliert von der Self-Attention direkt mit Token-Embeddings trainiert werden. Dieser Ansatz ermöglicht es, FFNs als token-basierte Lookup-Tabellen (ToLs) vorzuberechnen, was einen bedarfsgesteuerten Transfer zwischen VRAM und Speicher erlaubt und zusätzlich die Inferenzeffizienz steigert. Wir stellen zudem Flex-MemoryLLM vor, das sich zwischen einem konventionellen Transformer-Design und MemoryLLM positioniert. Diese Architektur überbrückt die Leistungslücke, die durch das Training von FFNs mit kontextfreien Token-Embeddings entsteht.
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.