TF-Engram：一种基于SSD存储的无训练记忆印记系统，用于大型语言模型

# TF-Engram：一种基于SSD备份记忆的无训练脑痕系统，适用于大语言模型  
来源：https://arxiv.org/html/2607.07388  

###### 摘要  

大语言模型（LLM）将事实知识和领域特定模式隐式存储在密集的Transformer参数中，使得通过预训练、微调、检索增强或更长上下文进行知识扩展的成本高昂。脑痕式（Engram）记忆提供了一种紧凑的隐藏状态注入路径，但现有驻留GPU的设计通常依赖基于哈希的压缩，导致不相关的短语在共享槽中发生冲突，削弱了短语级别的语义保真度。我们提出TF-Engram，一种无训练的脑痕系统，它从外部语料库离线构建短语特定的语义记忆，将大型记忆表分布在GPU、DRAM和SSD层级中，并采用早退引导预测预取（Early-Exit Guided Predictive Prefetching）技术，在自回归解码期间隐藏外部内存延迟。在Qwen3-0.6B上，TF-Engram将下游平均得分从57.6提升至59.4，优于冻结骨干网络和参数匹配的LoRA基线。系统评估表明，大型TF-Engram表可以以适度的离线成本构建，SSD备份存储显著降低了GPU内存需求，预测预取在很大程度上恢复了因外部内存访问导致的吞吐量损失。这些结果表明，静态短语记忆可以作为可扩展、无训练且低开销的系统组件集成到LLM推理中。  

## 1. 引言  

大语言模型（LLM）已成为自然语言理解、推理、代码生成和知识密集型问答的主流架构（Vaswani等人，2017；Devlin等人，2019；Brown等人，2020；Chowdhery等人，2022；Touvron等人，2023a，b）。它们的成功主要归功于密集Transformer参数和训练数据的规模化（Kaplan等人，2020；Hoffmann等人，2022）。然而，这种设计也将知识容量与模型尺寸紧密耦合：事实知识、领域特定术语以及频繁出现的语义模式被隐式存储在数十亿个神经参数中。改进此类知识通常需要更多的预训练、监督微调、检索增强或更长的上下文窗口（Brown等人，2020；Lewis等人，2020；Guu等人，2020；Borgeaud等人，2022）。这些解决方案增加了训练成本、推理延迟、内存消耗或部署复杂度。随着LLM越来越多地用于领域特定和知识密集型应用，将静态知识存储与密集Transformer计算解耦变得日益重要。  

近期关于记忆增强型LLM的研究从多个方向探索了这种解耦。检索增强生成（RAG）系统维护一个外部文档语料库，在推理期间检索相关段落，通常将检索到的文本附加到提示中（Lewis等人，2020；Guu等人，2020；Borgeaud等人，2022；Izacard等人，2023）。长上下文和基于KV缓存的方法通过保留更多历史token或中间状态来扩展模型可访问的信息量（Dao等人，2022；Dao，2024；Kwon等人，2023）。基于适配器和参数高效的记忆方法向模型注入额外的可训练模块，使得新知识可以存储在比全微调更少的参数中（Houlsby等人，2019；Li和Liang，2021；Lester等人，2021；Hu等人，2022）。另一类工作以脑痕式（Engram）记忆为代表，引入了一种轻量级侧边记忆路径，存储token级或短语级模式，并将检索到的记忆信号注入骨干模型的隐藏状态（Cheng等人，2026a）。这些类别在记忆粒度、训练成本、推理开销和系统可扩展性方面存在显著差异。  

尽管RAG、长上下文和基于适配器的方法在许多场景中有效，但它们不太适用于细粒度的静态短语记忆。RAG检索的是粗粒度的文本块，而非紧凑的短语级记忆条目，且检索到的段落会增加上下文长度和注意力成本（Lewis等人，2020；Borgeaud等人，2022；Izacard等人，2023）。长上下文方法允许模型关注更多信息，但即使有用知识对应于简短且固定的短语，它们仍然需要Transformer处理额外的token（Dao等人，2022；Dao，2024）。基于适配器的记忆方法降低了全微调的成本，但仍然需要训练，并且通常绑定到特定的骨干模型（Houlsby等人，2019；Li和Liang，2021；Lester等人，2021；Hu等人，2022）。相比之下，脑痕式记忆是面向系统设计的更合适起点。它提供了一条紧凑的侧边路径，在token或短语粒度上运行，并且无需重构主要Transformer计算即可集成到LLM推理中（Cheng等人，2026a）。因此，本文专注于将脑痕式记忆扩展为LLM的可扩展外部记忆系统。  

尽管前景广阔，脑痕式记忆尚未成为实用的规模化静态记忆系统。现有设计主要优化为紧凑的驻留GPU神经记忆模块。为了将大量token模式适配到有限的加速器内存中，它们通常依赖基于哈希的压缩，将许多token序列映射到固定大小的表中。这种设计减少了内存占用，但也引入了记忆覆盖范围与记忆保真度之间的紧张关系。为了使脑痕适用于大规模静态短语知识，必须解决三个关键问题：如何在不额外训练的情况下构建短语特定的记忆；如何在不依赖模糊的哈希冲突槽的情况下扩展记忆容量；以及如何在不让自回归解码停滞的情况下访问外部内存。  

第一个关键问题是记忆构建仍然依赖于训练且针对特定模型。传统的脑痕设计通常从数据中学习记忆条目、投影层、门控或辅助参数，要么与骨干模型联合学习，要么通过额外的后训练（Cheng等人，2026a；Houlsby等人，2019；Hu等人，2022）。这使得记忆难以跨不同LLM部署，或在不进行额外优化的情况下适应新领域。此外，当记忆条目以token或哈希槽粒度构建时，所得值并不显式保留短语级语义。因此，可扩展的脑痕系统应该能够离线构建短语级记忆，保留语义信息，并避免重新训练骨干模型。  

第二个关键问题是容量与语义保真度之间的权衡。现有的脑痕式记忆可以通过将许多token模式映射到固定大小的哈希表来保持紧凑。然而，当候选短语空间远大于表格时，不相关的短语不可避免地会冲突在相同记忆槽中。单个槽可能接收到来自多个不相关短语的信号，导致其记忆值变成模糊的混合体，而非清晰的短语特定表示。这削弱了静态短语记忆的原始目标：诸如*New York City*、*quantum field theory*或*BGP route reflector*这样的短语应该理想地检索到绑定于该短语的稳定记忆信号，而非与许多不相关模式共享的值。因此，简单地将更多短语压缩到小型GPU表中会增加表观覆盖率，但可能降低检索记忆的实用性。扩展脑痕记忆需要保留精确的短语条目，并将容量负担转移到存储层次结构中。  

第三个关键问题是在将精确记忆外部化后出现的访问延迟。将记忆条目从GPU内存移动到主机DRAM或NVMe SSD可以大幅扩展容量并减少哈希冲突，但朴素的外部查找可能破坏自回归解码的延迟要求。在生成过程中，每个token依赖于先前输出，模型几乎无法等待随机的存储访问。先前的LLM服务与卸载系统表明，内存层级、调度和数据移动是高效推理的核心（Yu等人，2022；Kwon等人，2023；Sheng等人，2023；Aminabadi等人，2022）。直接的预取设计可能使用原始LLM头部产生的最终下一个token分布。然而，这个信号只有在完整前向传播完成后才可用，剩余的计算时间不足以与SSD或CPU-GPU传输重叠。因此，SSD备份的脑痕记忆需要一种更早的预测机制，能够在最终层完成之前预测未来的内存请求。  

为了解决这些问题，我们提出了TF-Engram，一种基于SSD备份记忆的无训练脑痕系统，适用于大语言模型。其核心思想是将脑痕从一个紧凑的哈希压缩学习表转变为一个存储管理的外部记忆系统。TF-Engram没有将大型短语空间强制映射到模糊的共享槽中，而是离线构建短语特定的语义记忆，将其分布在GPU、DRAM和SSD层级中，并利用早期层的预测信号在需要之前预取未来的记忆条目。  

针对第一个问题，TF-Engram完全离线构建语义短语记忆。它首先从大规模语料库（如通用网络文本、百科全书文本和领域特定的科学语料库）中挖掘候选短语。然后使用目标LLM分词器对这些短语进行分词，使用外部语义编码器进行编码，并作为静态键值记忆条目存储。与从模型训练中学习记忆参数的传统脑痕记忆不同，TF-Engram将记忆构建与骨干优化分离。这种设计允许记忆表一次性构建，在多次推理运行中重复使用，并且可以扩展新语料库而无需重新训练LLM本身。  

针对第二个问题，TF-Engram将记忆组织为SSD备份的层级结构。频繁访问的条目保留在GPU或主机内存缓存中以实现低延迟查找，而冷条目则卸载到普通NVMe SSD。这种层级结构允许记忆容量独立于加速器内存进行扩展。TF-Engram没有将哈希压缩作为将记忆适配到GPU的主要机制，而是保留短语特定的逻辑条目，并依靠层级存储来吸收容量需求。这种设计使得支持更大的短语表成为可能，同时将稀缺的GPU内存留给模型权重、激活值、KV缓存和最热门的记忆条目（Kwon等人，2023；Sheng等人，2023；Rajbhandari等人，2021）。  

针对第三个问题，TF-Engram引入了早退引导预测预取（Early-Exit Guided Predictive Prefetching）。TF-Engram没有等待最终语言建模头部产生下一个token分布，而是在LLM的中间层（靠近末尾，如层L−r）附加一个辅助预取头部。这个早退头部在全前向传播完成之前产生近似的下一个token分布。TF-Engram将预测出的前K个token扩展为可能的未来短语访问，并在剩余Transformer层继续计算的同时，从SSD或主机DRAM异步预取这些短语的记忆条目。通过将剩余层用作延迟隐藏窗口，TF-Engram减少了阻塞性外部内存访问，使基于SSD的短语记忆对自回归推理变得实用。  

我们在Qwen3-0.6B上的评估表明，TF-Engram在保持骨干网络冻结的同时提升了端到端模型质量。在十个下游基准测试中，TF-Engram将平均得分从57.6提升到59.4，优于原始骨干网络和参数匹配的LoRA基线。系统分析进一步表明，大型TF-Engram表可以以适度的离线成本构建，而SSD备份的层级结构相比驻留GPU记忆显著降低了GPU内存需求。借助早退引导预取，TF-Engram在很大程度上恢复了因外部内存访问导致的吞吐量损失，并将解码延迟开销保持在较低水平。这些结果表明，TF-Engram为LLM推理中的无训练、可扩展且低开销的记忆增强提供了一条实用路径。  

## 2. 背景  

### 2.1. 记忆增强型LLM  

大语言模型（LLM）通常实现为仅解码器Transformer（Vaswani等人，2017；Brown等人，2020；Touvron等人，2023a，b）。给定输入序列x={x_1, x_2, ..., x_t}，模型首先将token ID映射为连续嵌入，然后通过一堆Transformer层进行处理。在每一层，模型通过应用自注意力和前馈变换来细化每个token的隐藏表示（Vaswani等人，2017）。在最终层之后，位置t的隐藏状态由语言建模头部投影，生成下一个token的概率分布：  

(1) p(x_{t+1} | x_{≤t}) = softmax(W_lm h_t^L)  

其中h_t^L表示最终层的隐藏状态，L是Transformer层的数量。在自回归解码过程中，这一过程逐token重复（Brown等人，2020；Touvron等人，2023b）。由于每个生成的token依赖于……