面向大语言模型持续适应的记忆库压缩方法 Memory Bank Compression for Continual Adaptation of Large Language Models

MBC的完整流程分为训练阶段和在线适应阶段。训练阶段（Algorithm 1）包含以下步骤：首先，初始化码本 $E$ 中的 $N_c = 512$ 个向量为均匀分布，所有码字的使用计数 $u_j$ 初始化为0。然后进入50个epoch的训练循环：(1) 从训练语料 $D_{train}$ 中采样一个mini-batch的 $K$ 个文档 $\{d_1, ..., d_K\}$；(2) 对每个文档采样对应的QA对 $(q_i, y_i)$；(3) 将每个文档 $d_i$ 输入amortization网络（T5编码器-解码器），得到连续上下文向量 $\phi_i = g_{\theta_{amort}}(d_i) \in \mathbb{R}^{T \times D}$；(4) 对每个 $\phi_i$ 执行VQ量化，找到最近的码本条目 $c_i = \arg\min_j \|\phi_i - E_j\|_2^2$，得到硬量化向量 $\hat{\phi}_i^{hard} = E_{c_i}$；(5) 使用直通估计器（STE）定义可微分的前向传播 $\hat{\phi}_i = \phi_i + sg[\hat{\phi}_i^{hard} - \phi_i]$；(6) 更新码字使用的EMA $u_j \leftarrow \gamma u_j + (1-\gamma) n_j$，其中 $\gamma = 0.99$，$n_j$ 是当前批次中使用码字 $j$ 的样本数；(7) 识别死码 $I_{dead} = \{j | u_j < \epsilon\}$，并用当前批次随机采样的编码输出重新初始化死码对应的码本向量，同时将重置后的使用计数设为全局均值 $\bar{u}$；(8) 用aggregation网络（4层cross-attention）将查询表示和量化后的记忆向量融合，生成调制信号 $\hat{\phi}_i^*$；(9) 将调制信号注入KV-LoRA增强的base LLM，计算QA损失 $\mathcal{L}_{QA}$ 和VQ损失 $\mathcal{L}_{VQ}$，总损失为 $\mathcal{L}_{total} = \mathcal{L}_{QA} + \lambda_{VQ} \mathcal{L}_{VQ}$；(10) 通过Adam优化器更新除base LLM参数 $\theta_b$ 以外的所有参数。在线适应阶段（Algorithm 2）无需梯度更新：对测试流中每个新文档 $d_k$，用训练好的amortization网络编码并量化，将离散索引 $c_k$ 存入压缩记忆库 $M_{VQ}$；推理时对每个查询 $q_i$，从码本 $E$ 中检索对应的量化向量，通过aggregation网络和KV-LoRA增强的base LLM生成答案。

MBC的技术新颖性体现在三个层面。第一，首次将VQ码本压缩引入LLM持续学习的记忆库场景。此前的记忆增强方法（如MAC、CAMELoT、MEMORYLLM等）均存储连续向量，MBC证明了离散化索引+小规模码本足以保留关键信息，且能将存储压缩至原来的0.3%。这是一个重要的工程和理论发现——意味着记忆库中的信息存在大量冗余，可以通过有限码本有效建模。第二，EMA在线重置机制解决了VQ在流式学习场景中的特有挑战。传统VQ-VAE在固定数据集上训练，码本坍塌可以通过重启或重初始化解决；但在在线适应场景中，数据分布持续变化，MBC的重置机制通过追踪平滑使用频率并及时替换死码，实现了对非平稳分布的鲁棒适应。第三，KV-LoRA的引入解决了「压缩-利用」的权衡问题。压缩不可避免地损失部分信息，而MAC完全冻结base LLM的做法限制了模型对压缩信号的利用能力。MBC仅在注意力层的K和V投影上添加少量可训练参数（增加不到0.5%），就显著提升了QA准确率，说明base LLM本身需要一定的适应能力来匹配压缩后的记忆质量。这种「压缩外部记忆+轻量内部适配」的双轨策略是一个有启发性的设计模式。

第一，码本大小的选择缺乏理论依据和系统实验。$N_c = 512$ 是固定的，但不同数据规模和不同backbone可能需要不同的码本大小。在文档数量极大时，512个码字可能不足以覆盖文档表示的多样性；而在文档数量较小时，512个码字又可能过多导致冗余。改进方向是设计自适应码本大小策略，例如根据码字使用率的统计特性动态增减码本条目，或者采用分层码本（hierarchical codebook）来捕获不同粒度的语义模式。第二，VQ的最近邻搜索在码本较大时会成为计算瓶颈。当前 $N_c = 512$ 时影响不大，但如果扩展到更大码本（如 $N_c = 4096$），每次量化都需要计算 $T \times N_c$ 次距离，开销会显著增加。改进方向包括使用产品量化（Product Quantization）将码本分解为多个子空间，或使用近似最近邻搜索算法。第三，aggregation网络的复杂度没有因为记忆库压缩而降低——它仍然需要对所有存储的量化索引对应的码字向量做cross-attention。虽然存储量下降了，但推理时仍需加载完整的码本并对每个查询进行attention计算。改进方向是引入稀疏检索或聚类机制，只对与当前查询最相关的部分记忆进行聚合，而不是处理整个记忆库。第四，论文未讨论码本的可迁移性——如果换一个数据集或任务，是否需要重新训练码本？这关系到方法在实际部署中的灵活性。

作者在结论中提出了两个明确的未来方向：一是引入强化学习信号来自适应地指导记忆使用。当前MBC的aggregation网络对所有记忆条目一视同仁地做attention，但理想情况下应该根据查询内容选择性地关注相关记忆，强化学习可以帮助模型学会这种选择策略。二是探索分布式记忆库，实现联邦持续学习——多个设备各自维护压缩的记忆库，定期聚合码本和索引，避免中心化存储的瓶颈。基于论文的成果还可以延伸出更多方向：(1) 将MBC扩展到多模态场景，压缩图像、音频等模态的记忆表示；(2) 探索码本的层次化或语义化组织，使码字不仅代表最近邻关系，还编码文档的领域、主题等高层语义；(3) 结合检索增强生成（RAG），将压缩记忆库作为「快速路径」处理高频查询，将原始文本检索作为「慢速路径」处理低频长尾查询；(4) 研究码本与知识图谱的融合，使离散化的索引能够关联到结构化的知识表示，提升推理能力。

论文的代码已开源在 https://github.com/Thomkat/MBC，这为复现提供了良好基础。数据集方面，StreamingQA、SQuAD和ArchivalQA均为公开可用的QA数据集，下载和预处理流程在论文和MAC的代码库中有详细说明。算力需求方面，论文明确说明所有实验在单张NVIDIA A100 80GB GPU上完成，这对于大多数研究机构来说是可获取的。但需要注意的是，LLaMA-2-7B的训练使用了4-bit量化（通过QLoRA库），且对于超过1B参数的模型使用了概率为0.75的反向传播dropout来降低显存占用，这些技巧虽然降低了硬件门槛，但也可能影响最终性能的完全复现。训练配置方面，论文详细报告了所有超参数（码本大小512、VQ承诺代价0.25、EMA衰减率0.99、重置阈值10^-4、KV-LoRA秩16-32、训练50个epoch等），以及不同backbone的学习率选择和梯度累积策略。总体而言，复现难度中等偏低，预计一个有经验的ML工程师可以在1-2周内完成全部实验。