2602.00398

MemoryLLM：可插拔的可解释前馈记忆Transformer架构

一句话定位

提出MemoryLLM，通过将FFN模块与自注意力解耦并直接基于token嵌入训练，使其成为可解释的上下文无关神经检索记忆，支持预计算为Token-wise Lookups（ToLs）以提升推理效率。

小学生也能听懂

这篇论文把大模型里负责记忆的部分（FFN）变成像字典一样的“查找表”，让它能直接记住单词对应的知识，还能提前算好存起来，用的时候快速查，既省内存又好懂，特别适合查资料类任务。

为什么值得记录

• FFN占LLM约三分之二参数但长期缺乏可解释性研究，MemoryLLM首次实现FFN作为离散、人类可理解的token级键值记忆；
• 支持FFN预计算为静态查找表（ToLs），显著降低VRAM占用与推理计算开销，适用于资源受限场景；
• 引入Flex-MemoryLLM架构，在保持部分记忆能力的同时弥合与传统LLM的性能差距；
• 实证发现FFN记忆对检索类任务影响远大于推理类任务，深化了对LLM内部知识存储机制的理解。

核心方法

• 设计MemoryLLM架构：所有FFN模块独立于残差流，直接接收来自嵌入层的上下文无关token嵌入向量 $ X_0 $ 作为输入；
• 采用TKV框架（Token-Key-Value）解释FFN：$ W_{Up} $ 为Key矩阵，$ W_{Down} $ 为Value矩阵，$ W_{Gate} $ 为可学习重加权函数；
• FFN输出通过静态Token-wise Lookups（ToLs）预计算并存储，推理时按需加载至VRAM；
• 提出Flex-MemoryLLM：将FFN参数划分为FFN-C（参与残差计算）和FFN-M（作为上下文无关记忆），实现性能与效率平衡；
• 在不同规模（250M/750M/1B参数）上从头训练模型，使用C4数据集与Llama-3.1 tokenizer。

关键结果

• MemoryLLM在相同总参数量下性能低于传统Base模型，但按有效活跃参数（不含ToLs）计算则优于Base模型（如1B MemoryLLM活跃参数402M，C4 PPL 20.933 vs Base-750M活跃737M，PPL 19.730）；
• Flex-MemoryLLM-$3h^2$（1B总参，704M活跃）性能接近Base-1B（1208M活跃），且优于Base-737M模型；
• 降低FFN贡献（α从1.0→0.5）导致Wikitext-2困惑度上升120.57%，而逻辑任务如HellaSwag仅下降38.52%，表明FFN主导检索类任务；
• ToLs支持8-bit量化（存储减半）与低秩SVD压缩（如20%秩减，存储降18.74%），性能损失极小；
• 中间层ToLs冗余度高，移除后对性能影响小，支持分层压缩策略。

局限与风险

• MemoryLLM因完全解耦FFN导致表达能力受限，需依赖Flex-MemoryLLM补足性能；
• ToLs总存储需求仍较大（1B模型约12.6GB@F16），尽管支持压缩与按需加载；
• 实验未涵盖更大规模模型（如7B以上）及真实部署环境下的端到端延迟测试；
• TKV框架依赖SwiGLU结构，其他FFN变体（如ReLU）是否适用尚不明确。

适合沉淀的概念

feed-forward-network-interpretability, token-wise-lookup-tol, memory-augmented-transformer, flex-memoryllm-architecture, neural-key-value-memory, llm-model-compression, context-free-ffn-training, tkv-framework

阅读注意

• 关注图1与图2：理解MemoryLLM与传统Transformer在残差流处理上的根本差异；
• 细读表1：对比不同α值下检索类与推理类任务的性能变化趋势；
• 查看附录D：了解ToLs的存储成本估算与三种压缩策略（量化、低秩、分层）的具体效果；
• 注意Flex-MemoryLLM参数划分方式（附录表4）及其对活跃参数量的影响。

质量说明

• 采用来源：clean，papers/2026/01/2602.00398.md
• 生成模型：LongCat-Flash-Chat
• 源材料判断：论文材料完整，包含方法设计、实验设置、结果分析与附录细节，数据充分支持结论。