MemoryLLM: Memória Feed-Forward Interpretável Plug-and-Play para 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
Resumo
Compreender o funcionamento dos componentes dos transformadores em LLMs é crucial, pois está no cerne dos recentes avanços tecnológicos em inteligência artificial. Neste trabalho, revisitamos os desafios associados à interpretabilidade dos módulos *feed-forward* (FFNs) e propomos o MemoryLLM, que visa desacoplar os FFNs da auto-atenção e nos permite estudar os FFNs desacoplados como uma memória neural de recuperação *token-wise* e livre de contexto. Detalhadamente, investigamos como os *tokens* de entrada acessam locais de memória dentro dos parâmetros do FFN e a importância da memória do FFN em diferentes tarefas *downstream*. O MemoryLLM alcança FFNs livres de contexto ao treiná-los isoladamente da auto-atenção, utilizando diretamente os *embeddings* de *tokens*. Esta abordagem permite que os FFNs sejam pré-computados como pesquisas *token-wise* (ToLs), possibilitando a transferência sob demanda entre a VRAM e o armazenamento, adicionalmente melhorando a eficiência da inferência. Também introduzimos o Flex-MemoryLLM, posicionando-o entre um projeto convencional de transformador e o MemoryLLM. Esta arquitetura preenche a lacuna de desempenho causada pelo treinamento de FFNs com *embeddings token-wise* livres 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.