让大模型分层记忆,用中间层存信息,上层重算推理,省显存还快。
Understanding Is Done Early: A Depth Division of Labor in Large Language Models and Its Use for Unbounded-Context Memory

- 分层设计:用中间层存上下文,上层按需重算,实现高效记忆
- 97.05分(RULER)和38.27分(LoCoMo),远超全上下文直接存储
- 不依赖适配器,128k长度仅用18.26GB显存,提速7.83倍
Transformer的深度并非均匀使用:低中层构建语义表征,高层逐步专业化用于预测。我们据此提出CoMem(Comprehension Memory),将每个上下文块仅通过中间层写入,检索固定数量缓存的残差状态,并在结果组合上重新计算查询相关的高层。在固定检索预算下,模型读取计算与内存开销独立于存储上下文长度。我们在统一无模板协议下评估持续训练的Qwen3-8B基础语言模型,主干冻结;旗舰模型仅在PG19上训练一个秩为32的自蒸馏LoRA,另报告无适配器版本。CoMem在RULER上达97.05,在LoCoMo上达38.27,远超全上下文KV-Direct的34.59;对话记忆优势在对话簇重采样和独立评价中仍成立。额外长上下文与长文档任务揭示了有限检索的优势及其窗口内压缩代价。深度消融实验表明更深缓存降低每查询重算成本,但牺牲保真度,而自蒸馏可显著修复。独立无适配器效率对比显示,在NVIDIA H20上处理128k长度时,CoMem仅需18.26GB而非89.36GB显存,预填充速度提升7.83倍。结果表明,长上下文记忆可沿层轴组织,而不仅是标记轴。
原文摘要 · Abstract (English)
Transformer depth is not used uniformly: lower and middle layers build semantic representations, while upper layers increasingly specialize them for prediction. We turn this division of labor into CoMem (Comprehension Memory), which writes each context chunk only through an intermediate layer, retrieves a fixed number of cached residual states, and recomputes the query-conditioned upper layers over the resulting pack. For a fixed retrieval budget, model-side read compute and memory are independent of stored-context length. We evaluate a continued-trained Qwen3-8B base LM under a unified chat-template-free protocol. The backbone is frozen; the flagship trains only a rank-32 self-distillation LoRA on plain PG19, and we report an adapter-free arm separately. CoMem reaches 97.05 on RULER and 38.27 on LoCoMo versus 34.59 for full-context KV-Direct; the dialogue-memory advantage survives conversation-cluster resampling and an independent judge. Results on additional long-context and long-document tasks expose both the benefits of bounded retrieval and its in-window compression tax. Controlled depth sweeps show that deeper caching lowers per-query recomputation but incurs a fidelity loss that self-distillation substantially repairs. In a separate adapter-free efficiency control on an NVIDIA H20 at 128k, CoMem uses 18.26 GB rather than 89.36 GB and achieves a 7.83x prefill speedup. These results show that long-context memory can be organized along the layer axis, not only the token axis.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。