让Transformer的前馈模块变成可解释的查表式记忆，提升推理效率。
MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers
- 将前馈模块独立训练为无上下文的词元查找表，实现可解释性
- 支持按需加载内存，降低显存占用并加速推理
- 适合需要高效推理与模型可解释性的部署场景
理解Transformer组件在大模型中的运作机制至关重要，是当前人工智能技术进步的核心。本文重新审视前馈网络（FFN）的可解释性挑战，提出MemoryLLM，旨在将FFN从自注意力机制中解耦，使其作为与上下文无关的词元级神经检索记忆进行研究。具体而言，我们分析输入词元如何访问FFN参数中的记忆位置，并评估其在不同下游任务中的重要性。MemoryLLM通过仅使用词元嵌入独立训练FFN，实现上下文无关的前馈结构，使FFN可预计算为词元级查找表（ToL），支持按需在显存与存储间转移，进一步提升推理效率。此外，我们提出Flex-MemoryLLM，在传统Transformer与MemoryLLM之间建立桥梁，缓解因使用无上下文词元嵌入训练FFN导致的性能差距。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。