MemoryLLM: Memoria de retroalimentación interpretable lista para usar en Transformers
MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers
January 30, 2026
Autores: Ajay Jaiswal, Lauren Hannah, Han-Byul Kim, Duc Hoang, Arnav Kundu, Mehrdad Farajtabar, Minsik Cho
cs.AI
Resumen
Comprender el funcionamiento de los componentes del transformador en los LLM es fundamental, ya que constituye el núcleo de los recientes avances tecnológicos en inteligencia artificial. En este trabajo, revisamos los desafíos asociados con la interpretabilidad de los módulos de retroalimentación (FFN) y proponemos MemoryLLM, que busca desacoplar los FFN de la autoatención y nos permite estudiar los FFN desacoplados como una memoria de recuperación neuronal libre de contexto a nivel de token. En detalle, investigamos cómo los tokens de entrada acceden a las ubicaciones de memoria dentro de los parámetros del FFN y la importancia de la memoria del FFN en diferentes tareas posteriores. MemoryLLM logra FFN libres de contexto entrenándolos de forma aislada de la autoatención, utilizando directamente las incrustaciones de tokens. Este enfoque permite precalcular los FFN como búsquedas por token (ToL), posibilitando la transferencia bajo demanda entre la VRAM y el almacenamiento, mejorando además la eficiencia de la inferencia. También presentamos Flex-MemoryLLM, posicionándolo entre un diseño de transformador convencional y MemoryLLM. Esta arquitectura salva la brecha de rendimiento causada por el entrenamiento de FFN con incrustaciones de tokens libres de contexto.
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.