理解在早期完成：大型语言模型中的深度分工及其用于无界上下文记忆

# 大型语言模型中的深度分工及其在无界上下文记忆中的应用
来源：https://arxiv.org/html/2607.28263
韩卓刘†‡ 琦轩†∗ 春宇刘†∗ 昊天钟†∗ 玉龙王‡ Rayying‡ Key‡ Alex Lamb† 明宇高†
†清华大学‡腾讯
\{lhz24@mails\., gaomy@\}tsinghua\.edu\.cn
https://github.com/liuhanzuo/COMem

###### 摘要

Transformer 的各层似乎扮演着不同的角色：语义信息在较底层和中间层变得可访问，而较高层则越来越专门化用于预测的表征。我们将这种分工转化为 CoMem（Comprehension Memory，理解记忆），它仅通过一个中间层写入每个上下文分块，检索固定数量的缓存残差张量，并在由此得到的分组上重新计算以查询为条件的较高层。对于固定的检索预算，模型侧的读取计算量和内存使用与已存储的上下文长度无关。我们在统一的、无聊天模板的协议下评估了一个持续训练的 Qwen3-8B 基础语言模型。骨干网络被冻结；主模型仅在纯 PG-19 上训练一个秩为 32 的自蒸馏 LoRA，我们还单独报告了一个无适配器的分支。CoMem 在 RULER 上达到 97.05，在 LoCoMo 上达到 38.27，而全上下文 KV-Direct 为 34.59；对话记忆优势在对话簇重采样和独立评判下仍然存在。在额外的长上下文和长文档任务上的结果既展示了有界检索的优势，也暴露了其窗口内压缩代价。受控的深度扫描表明，更深的缓存会降低每次查询的重计算量，但会导致保真度损失，而自蒸馏可以大幅修复这种损失。在一项单独的无适配器效率控制实验中，NVIDIA H20 在 128k 下，CoMem 使用 18.26 GB 而非 89.36 GB，并实现了 7.83× 的预填充加速。这些结果表明，长上下文记忆不仅可以沿着 token 轴组织，也可以沿着层轴组织。

理解早已发生：大型语言模型中的深度分工及其在无界上下文记忆中的应用

预印本

韩卓刘†‡ 琦轩†∗ 春宇刘†∗ 昊天钟†∗ 玉龙王‡ Rayying‡ Key‡ Alex Lamb† 明宇高†
†清华大学‡腾讯
\{lhz24@mails\., gaomy@\}tsinghua\.edu\.cn
https://github.com/liuhanzuo/COMem
††脚注：匿名仓库：https://anonymous.4open.science/r/COMem-Anonymous/

## 1 引言

参见标题图1：CoMem 一图概览。语义信息在中间深度已可被使用，因此 CoMem 在深度 \(j\) 处缓存每个 token 的残差张量，并在一个固定的检索分组上重新计算较高层。对于固定的 \(k\)、分块大小和查询长度，模型侧的在线读取工作集与已存储的上下文长度无关。准确率曲线使用蒸馏后的主模型，而 7.83×/18.26 GB 的插图是单独的无适配器效率对照。CoMem 在未扩展骨干网络的原生窗口之外仍然准确；我们不与 YaRN 扩展的全上下文模型进行比较。

长上下文语言建模面临自注意力的二次方预填充成本，以及键值（KV）缓存的线性内存增长。（Vaswani 等人，2017；Kwon 等人，2023）。它还继承了骨干网络的位置范围：Qwen3-8B（Yang 等人，2025）拥有 40,960 token 的原生窗口（配置的 131,072 限制需要未启用的 YaRN 扩展，Peng 等人，2023），因此我们将更长的全上下文运行仅视为未扩展的压力参考。

大多数内存高效方法会驱逐或压缩 KV 状态（Xiao 等人，2024b；Zhang 等人，2023；Li 等人，2024；Cai 等人，2024；Liu 等人，2024a）、检索相关文本或记忆单元（Lewis 等人，2020；Xiao 等人，2024a；Tang 等人，2024；Qiu 等人，2025b），或者写入经过训练的记忆（Wang 等人，2024；Wang 等人，2023；Qiu 等人，2025a）。残差流重用可以在不独立存储 KV 状态的情况下重建它们（Qasim 等人，2026）。CoMem 将有界 token 选择与*深度划分的计算重用*相结合：它既要回答检索哪些分块，也要回答它们有多少计算可以被重用。其前提是，较底层和中间层使语义信息可访问，而较高层则越来越多地使其针对当前查询和预测进行专门化（Lad 等人，2024）。

我们用探针、因果截断和分裂深度扫描验证了这一深度分工（§3），然后将其转化为 CoMem（图1）。在写入阶段，每个分块仅处理到深度 \(j\)，并缓存其残差张量。在读取阶段，一个有界的相关张量集合与查询打包在一起，只重新计算 \([j:L]\) 层。外部存储仍随已存储上下文线性增长，但对于固定的检索广度，模型侧的在线读取并非如此。主模型在冻结的骨干网络上使用秩为 32 的自蒸馏 LoRA（Hinton 等人，2015；Hu 等人，2022）；一个无适配器分支单独隔离架构效果。

##### 贡献。

- **深度是一个可测量的重用轴。** 语义内容在中间深度附近变得可访问，而零样本可读边界取决于模型规模，并且可以通过轻量级适配来移动。
- **一种有界的深度划分记忆。** CoMem 为每个 token 缓存一个残差状态，精确地恢复上层 Transformer，并将有界检索与完整的分块间注意力相结合。在 Qwen3-8B 上，缓存状态仅为每个已存储 token 的完整 bf16 KV 字节数的 1/18。
- **受控的质量和效率证据。** 在统一的、无聊天模板的 Qwen3-8B 评估协议下，CoMem 在 RULER 上达到 97.05，在 LoCoMo 上达到 38.27，而 KV-Direct 为 34.59；对话收益在簇重采样和独立评判下仍然存在。一项单独的无适配器 H20 控制在 128k 下使用 18.26 GB 而非 89.36 GB，并实现了 7.83× 的全写入包含预填充加速。消融研究分别隔离了检索、分裂深度、分块间交互和蒸馏。

## 2 相关工作

我们按照设计结构的三个相同轴来组织先前工作（§3）：缓存*什么*、训练*多深*/训练多少，以及*如何*读回。

##### 缓存什么：全深度 KV 与一个中间残差张量。

长上下文的主流方法沿着 *token* 轴限制键值（KV）缓存，保留*全深度* KV 的一个压缩子集。StreamingLLM（Xiao 等人，2024b）保留最近窗口加上注意力汇聚 token；H2O（Zhang 等人，2023）和 SnapKV（Li 等人，2024）逐出除重击者或预测注意力位置之外的所有内容；PyramidKV（Cai 等人，2024）和 LCKV（Wu 和 Tu，2024）跨层重新分配或共享预算；MiniCache（Liu 等人，2024a）合并相邻层的冗余 KV。CompressKV（Lin 等人，2026）增加了检索头逐出标准。与我们最接近的是 HCache（Gao 等人，2025），它缓存一个中间层激活并重新计算其上的层——但它是事后、未训练且无检索的；KV-Direct（Qasim 等人，2026）表明 KV 缓存是残差流的确定性投影，并精确地重新计算*全深度* KV，但它保留每个 token，因此其工作集仍随上下文增长。大多数 token 轴方法保留选定的全深度 KV，而 HCache 是最接近的中间深度重计算设计，但缺少相关性选择。遵循深度分工（§3，*缓存什么*）和 Lad 等人（2024）的深度语义原理，CoMem 改为为每个分块缓存一个深度 \(j\) 的残差张量（每个 token 一个向量），并按需重建较高层。

##### 多深：固定或训练记忆与可调深度分裂。

第二条线通过架构或训练来固定记忆。MemoryLLM（Wang 等人，2024）在 Transformer 内部维护一个固定大小的潜在池，自动更新（我们将它发布的 Llama-3 检查点作为骨干外参考）；LongMem（Wang 等人，2023）通过一个训练过的侧网络检索外部状态；RecursiveSummarizing（Wang 等人，2023）保留自然语言摘要。YOCO（Sun 等人，2024）在重新设计的解码器-解码器中共享一个全局 KV，而 CEPE（Yen 等人，2024）添加一个并行上下文编码器和交叉注意力。其他方法通过持续训练引入可压缩性：Activation Beacon（Zhang 等人，2024）学习将每层激活压缩为信标 token；KV-CAT（Gelberg 等人，2026）在 token/槽轴上训练缓存可压缩性。CoMem 改为在*深度*轴上划分，并保持骨干网络冻结。其推理时记忆操作不需要参数更新，而评估的主模型只训练一个轻量级自蒸馏 LoRA；无适配器变体单独报告。CoMem 将分裂 \(j\) 暴露为一个单一的 RAG↔闭卷旋钮（§3，*多深*）。

##### 如何读：检索集成记忆与外部重计算。

检索增强生成向模型提供选定的文本（Lewis 等人，2020）。在长上下文架构中，Landmark Attention（Mohtashami 和 Jaggi，2023）通过学习的地标 token 选择块；FoT/LongLLaMA（Tworkowski 等人，2023）使用外部对比训练的 \((k,v)\) 记忆；InfLLM（Xiao 等人，2024a）在注意力内部查找 token 相关的记忆单元。Quest（Tang 等人，2024）和 RetrievalAttention（Liu 等人，2024b）转而在注意力期间检索查询相关的 KV 页或向量。服务系统如 CacheBlend（Yao 等人，2024）和 RAGCache（Jin 等人，2024）为检索到的分块预计算或重用 KV 状态。这些方法在存储单元和训练要求上有所不同，但检索与 token/KV 记忆或注意力路径耦合。CoMem 将选择保持为外部：迭代 BM25（Robertson 和 Zaragoza，2009）选择分块残差张量，之后以完整的分块间注意力重新计算较高层。一个检索到的 \(j=0\) 对照单独衡量没有深度重用的选择效果。

## 3 动机

CoMem 源自一个经验属性——Transformer 的深度并不是均匀使用的——以及由此得出的三个记忆设计推论。我们在这里只保留激发这些推论所需的证据；受控消融在 §5 和附录 A.1 中。

##### （I1）缓存一个中间残差张量。

逐层探针将语义可访问性置于中间深度附近，而因果截断表明较浅的前缀保留了大量的下游信号。较高层状态更多地以查询为条件（Lad 等人，2024），因此缓存的无查询状态不应被视为最终结果。CoMem 因此存储每个 token 的完整残差 \(h_j\)，并在查询存在的情况下重新计算 \([j:L]\)。在 Qwen3-8B 上，\(h_j\) 仅使用每个已存储 token 的完整 bf16 KV 字节数的 1/18。

##### （I2）将深度视为一个质量-成本旋钮。

对于较浅的 \(j\)，召回率保持较高，然后随着缓存状态跨越零样本可读边界而下降（图2a）。这个边界不同于语义*内容*峰值，并随着模型规模增大而加深（图2b）。我们将 \(j\) 暴露为一个部署旋钮，并使用自蒸馏将工作点移向更深；主模型使用 \(j=12/36\)。精确的 \(j=0\) 等价性和受控深度/适配结果见 §5。

参见标题图2：深度分工设置了一个有界、可调的缓存深度。（a）在 Qwen3-8B 上，零样本针检索在分裂 \(j\) 较浅时保持完美，一旦缓存进入查询条件带则崩溃；一个轻量自蒸馏适配器将可读深度推得更深。面板（a）是一个独立的趋势扫描；表10给出了受控的固定 \(j\) 队列。（b）跨规模来看，语义*内容*深度接近规模不变（≈0.45L），而零样本*可读*深度随模型大小加深；阴影间隙是探针提示的适配余量，到 32B 时消失。完整表格见附录 A.1。

##### （I3）读取一个有界的相关分组。

当金标准分块存在时，一个 oracle 选择器在报告的 RULER 针扫描上达到 100%，而新近度和读者注意力选择随着长度增长而失败（图3）。因此，在这个诊断任务上观察到的瓶颈是选择，而不是缓存状态的保真度。CoMem 使用有界的相关性检索和完整的分块间上层注意力；§5 直接测试检索广度和分块交互。

参见标题图3：选择，而非保真度，是观察到的长距离瓶颈。在无聊天模板的 RULER 单针扫描中，Oracle 和 BM25 保持在 100，而新近度和读者注意力在 32k 处下降。完整选择器结果见附录表12。

## 4 方法

CoMem 沿深度划分解码器（图4）：写入阶段为每个分块缓存一个中间残差张量；选择阶段检索一个有界子集；读取阶段在查询存在的情况下重新计算较高层。恢复的计算与 KV-Direct（Qasim 等人，2026）相关，但 CoMem 将其与浅层写入、可调分裂和有界外部检