MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915764612628480 |
|---|---|
| author | Jaiswal, Ajay Hannah, Lauren Kim, Han-Byul Hoang, Duc Kundu, Arnav Farajtabar, Mehrdad Cho, Minsik |
| author_facet | Jaiswal, Ajay Hannah, Lauren Kim, Han-Byul Hoang, Duc Kundu, Arnav Farajtabar, Mehrdad Cho, Minsik |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_00398 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers Jaiswal, Ajay Hannah, Lauren Kim, Han-Byul Hoang, Duc Kundu, Arnav Farajtabar, Mehrdad Cho, Minsik Machine Learning 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. |
| title | MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2602.00398 |