MemoryLLM: Интерпретируемая память прямого распространения для трансформаторов по принципу plug-n-play

MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers

January 30, 2026

Авторы: Ajay Jaiswal, Lauren Hannah, Han-Byul Kim, Duc Hoang, Arnav Kundu, Mehrdad Farajtabar, Minsik Cho

cs.AI

Аннотация

Понимание принципов работы компонентов трансформаторов в больших языковых моделях (LLM) имеет важное значение, поскольку они лежат в основе последних технологических достижений в области искусственного интеллекта. В данной работе мы вновь обращаемся к проблемам интерпретируемости прямых (feed-forward) модулей (FFN) и предлагаем MemoryLLM — подход, который позволяет отделить FFN от механизма самовнимания (self-attention) и изучать декомпозированные FFN как контекстно-независимую поэлементную нейросетевую память выборки. В частности, мы исследуем, как входные токены обращаются к ячейкам памяти внутри параметров FFN, и оцениваем важность памяти FFN для различных прикладных задач. MemoryLLM обеспечивает контекстную независимость FFN за счёт их изолированного обучения напрямую на эмбеддингах токенов, без использования механизма самовнимания. Этот подход позволяет предварительно вычислять FFN в виде таблиц поиска по токенам (ToL), что даёт возможность гибко перемещать их между видеопамятью (VRAM) и долговременным хранилищем, дополнительно повышая эффективность вывода. Мы также представляем Flex-MemoryLLM — архитектуру, занимающую промежуточное положение между классической архитектурой трансформатора и MemoryLLM. Данная архитектура сокращает разрыв в производительности, вызванный обучением FFN на контекстно-независимых эмбеддингах токенов.

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.