Memorize Step by Step：一种处理 LLM 长上下文的更智能方法

“百万级 token 上下文窗口”的时代已经到来。随着 Claude 3 和 Gemini 1.5 等模型的出现，我们正从处理简短的提示词，迈向一次性处理整本书籍、代码库和法律档案的领域。

但这里有一个陷阱: 硬件的发展并没有跟上我们野心的步伐。处理 100 万个 token 的序列需要巨大的计算资源和 GPU 内存。如果你尝试将这么长的文本输入到标准 GPU 中，几乎肯定会在模型生成第一个字之前就遇到“内存溢出” (Out of Memory，OOM) 错误。

这种瓶颈主要发生在预填充阶段 (prefilling stage) ——即模型读取并处理输入提示词 (上下文) 的阶段。为了解决这个问题，复旦大学和华为诺亚方舟实验室的研究人员在他们的论文 “Memorize Step by Step” 中提出了一种巧妙且直观的方法。他们的方法利用增量记忆 (Incremental Memory, IM) 和减量分块 (Decremental Chunks, IMDC) , 证明了我们不需要同等对待序列的每一部分。通过动态调整记忆和输入的大小，我们可以用更少的硬件更快地处理无限长度的上下文。

在这篇文章中，我们将拆解这篇论文，解释为什么传统方法也是一种浪费，以及这种新的“循序渐进”策略如何优化大语言模型 (LLM) 的推理过程。

问题所在: 预填充的沉重代价

LLM 推理包含两个截然不同的阶段:

1. 预填充 (Prefilling) : 模型处理输入提示词 (上下文) 以计算键值 (KV) 缓存。
2. 解码 (Decoding) : 模型逐个 token 生成输出。

虽然许多研究侧重于优化解码阶段，但对于长上下文来说，预填充往往是瓶颈所在。Transformer 中的注意力机制具有二次时间复杂度 \(O(L^2)\)，这意味着输入长度每翻一倍，计算时间就会增加四倍。此外，存储海量输入的 KV 缓存所需的显存超过了大多数 GPU 的容量。

标准解决方案: 固定大小记忆的分块处理

当输入太长而无法放入 GPU 内存时，标准的变通方法是迭代压缩 (Iterative Compression) 。 序列被分成较小的段，或称“块 (chunks) ”。模型一次处理一个块，将信息压缩到传递给下一个块的“记忆”缓冲区 (KV 缓存) 中。

传统上，这个记忆缓冲区具有固定大小 (Fixed-Size) 。 无论模型是在读取第 1 个块还是第 100 个块，它都会分配相同数量的内存来存储历史上下文。

如图 3 第一行 (FM) 所示，固定大小记忆方法从一开始就分配了一个静态的内存块 (黄色方框) 。

但这真的高效吗? 研究人员提出了一个根本性的问题: 我们真的需要在序列的一开始就需要一个大的记忆缓冲区吗？

洞察: 早期记忆没那么重要

为了回答这个问题，作者分析了模型在不同处理阶段实际上对记忆给予了多少“关注 (attention) ”。他们使用了两种流行的压缩方法 SnapKV 和 StreamingLLM 来可视化注意力分数。

图 1(a) 的结果具有启示性。红色和青色的线显示了记忆的平均注意力分数。请注意，在开始时 (低 Chunk ID) ，注意力分数非常低。随着处理的进行，它们逐渐增加。

这表明早期阶段的记忆对下一步计算的影响微乎其微。

此外，图 1(b) 显示了过程结束时的记忆内容分布。绝大多数保留的信息来自最近的块 (右侧的高条形图) 。当模型到达末尾时，大部分早期阶段的记忆已经被丢弃或忽略。

结论: 在预填充的早期阶段维持一个巨大的、固定大小的记忆缓冲区是对 GPU 内存和计算能力的浪费。

解决方案: 增量记忆 (IM)

基于这一洞察，作者提出了增量记忆 (Incremental Memory, IM) 。 记忆缓冲区不是立即分配最大内存大小，而是从很小开始，随着模型处理更多的块而逐步增长。

再次参考 图 3 的中间行，你可以看到“增量记忆”策略。记忆 (橙色块) 起初很小，然后一步步扩展。

从数学上讲，如果 \(n\) 是块的数量，\(m_{max}\) 是最大记忆限制，\(m_0\) 是初始记忆，则第 \(i\) 步的记忆大小呈线性增加:

为什么这有帮助？

通过在早期步骤保持较小的记忆，依赖于记忆大小的注意力计算变得显著更快。时间复杂度下降了，因为整个过程中的平均记忆大小大约是最大大小的一半 (\(m_{max}/2\)) 。

自适应增长

研究人员更进一步。通过分析 LLM 的不同层，他们发现高层 (靠近输出) 往往比低层保留更多的记忆。

图 6 说明了这一现象。在 LLaMA2-7b (左) 中，较高层 (顶部行) 显示出更均匀的记忆使用分布，而较低层则主要集中在最近的 token 上。这引出了自适应函数 (Adaptive Function) , 即根据该层实际保留的信息量，为每一层定制记忆大小。

优化: 减量分块 (IMDC)

增量记忆使过程变得更快，但在峰值内存使用量方面存在缺陷。

即使你从极小的记忆开始，在最后一步你最终还是会达到最大记忆大小 (\(m_{max}\)) 。由于你的 GPU 内存限制取决于峰值使用量，IM 并不一定能让你比固定方法运行更大的模型或批处理大小；它只是完成得更快而已。

为了解决这个问题，作者引入了减量分块 (Decremental Chunk, IMDC) 。

这个想法是一种平衡之术:

• 随着 记忆 (Memory) 变大…
• 输入 分块 (Chunk) 应该变小。

这确保了在任何单一步骤中处理的 token 总数 (分块长度 + 记忆长度) 保持大致恒定。

图 2 完美地展示了这一点:

• 固定大小记忆 (左) : 总高度 (记忆 + 分块) 从一开始就很高。
• 增量记忆 (中) : 分块大小恒定，但记忆像阶梯一样增长。峰值使用量在最后很高。
• 减量分块 (右) : 随着黄色记忆阶梯上升，蓝色分块方块缩小。总高度保持稳定且低于其他两者。

IMDC背后的数学原理

目标是保持计算负载稳定。第 \(i\) 步的分块大小 \(c_i\) 是动态计算的，使得当前分块与先前记忆之和等于平均负载:

这里，\(\hat{m}\) 代表平均记忆大小。由于分块大小随着记忆增加而减小，峰值 GPU 内存需求显著降低。这使得模型能够使用更少的显存处理长上下文，甚至可能在固定方法会崩溃的小型消费级 GPU (如 RTX 3090) 上运行。

实验结果

理论听起来很扎实，但在实践中表现如何？研究人员在 NVIDIA A800 和 RTX 3090 GPU 上使用 LLaMA-2 和 InternLM2 对比测试了 IM、IMDC 和固定大小记忆 (FM) 。

1. 速度和内存效率

图 4 展示了效率的提升。

• 第一行 (时间) : 绿色 (IM) 和蓝色 (IMDC) 的条形图始终低于红色 (FM) 。在使用 SnapKV 的 A800 GPU 上，IMDC 比基线快了 1.45 倍 。
• 第二行 (内存) : 看 IMDC 的蓝色条形图。它们显示 GPU 内存使用量显著下降——高达 23.3% 的缩减 。

这证实了 IMDC 不仅更快，而且对硬件的要求也低得多。

2. 模型性能 (困惑度)

这种优化是否损害了模型的智能？大幅压缩记忆或缩小分块理论上可能导致模型“遗忘”或混淆。

图 5 显示了困惑度 (PPL) ，这是衡量模型预测文本好坏的指标 (越低越好) 。IM (绿色) 和 IMDC (蓝色) 的性能与固定大小记忆 (红色) 几乎完全相同，有时甚至略好。这种“循序渐进”的方法并没有降低语言建模的质量。

3. 长上下文基准测试

研究人员还在 LongBench 上运行了模型，这是一套旨在测试长上下文理解 (如摘要和问答) 的任务。

如表 2 所示，不同方法之间的差异微乎其微 (通常小于 0.5 分) 。IMDC 取得了与资源密集型的固定大小记忆相当的结果，证明它是“免费的午餐”——在不牺牲准确性的情况下提高了效率。

4. 扩展到 100 万 Token

最后是终极测试: 它能处理 100 万个 token 吗？

表 4 显示了在 1M token 序列上的结果。标准的“全注意力 (Full Attention) ”立即导致了内存溢出 (OOM) 错误。然而，IMDC 成功处理了该序列，且困惑度很低 (约 2.11) ，延迟也大幅降低 (328 秒 vs FM 的 505 秒) 。

结论

“Memorize Step by Step” 这篇论文为我们如何思考处理长文档提供了一个引人注目的转变。通过认识到早期阶段的记忆不那么关键以及计算负载应该平衡 , 作者创造了一种方法，它:

1. 更快速: 通过从小记忆开始。
2. 更轻量: 通过随记忆增长缩小分块。
3. 能力相当: 在基准测试中保持高准确性。

对于学生和从业者来说，这凸显了系统设计中的一个重要教训: 静态分配 (如固定缓冲区) 通常不是最优的。能够根据处理状态进行反应的动态、自适应系统 (如增量记忆) 往往能带来显著的效率提升。随着我们推向无限上下文 LLM，像 IMDC 这样的策略对于在可获得的硬件上运行强大的模型将至关重要。