Understanding Is Done Early: A Depth Division of Labor in Large Language Models and Its Use for Unbounded-Context Memory

📄 arXiv: 2607.28263v1 📥 PDF

作者: Hanzuo Liu, Xuan Qi, Chunyu Liu, Haotian Zhong, Yulong Wang, Rayying, Key, Alex Lamb, Mingyu Gao

分类: cs.CL

发布日期: 2026-07-30

备注: 19 pages, 4 figures, 27 tables. Submitted to ACL Rolling Review


💡 一句话要点

提出CoMem以优化长上下文记忆的处理效率

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 长上下文记忆 Transformer模型 自蒸馏训练 计算效率 信息检索

📋 核心要点

  1. 现有的Transformer模型在不同层次的功能分配不均,导致上下文处理效率低下。
  2. 论文提出的CoMem方法通过中间层处理上下文,优化了记忆检索和计算过程。
  3. 实验结果显示,CoMem在多个任务上超越了传统方法,并显著降低了内存使用和计算时间。

📝 摘要(中文)

本文探讨了Transformer模型不同层次的功能分工,提出了CoMem(Comprehension Memory)方法。该方法通过中间层写入上下文块,仅检索固定数量的缓存残差状态,并在此基础上重新计算查询条件的上层。实验表明,CoMem在RULER和LoCoMo任务上分别达到了97.05和38.27的性能,相较于传统的全上下文KV-Direct方法有显著提升。此外,CoMem在内存使用和计算速度上也表现出色,显示了长上下文记忆的层次化组织潜力。

🔬 方法详解

问题定义:本文旨在解决Transformer模型在处理长上下文时的效率问题,现有方法在上下文长度增加时,计算和内存需求急剧上升,影响模型性能。

核心思路:CoMem方法通过将上下文块的写入限制在中间层,减少了对上层的重复计算,从而提高了处理效率。此设计使得模型在固定的检索预算下,读取计算和存储上下文的需求相互独立。

技术框架:CoMem的整体架构包括上下文块的写入、缓存残差状态的检索和基于查询条件的上层重新计算。模型的主干保持冻结,仅在特定层进行自蒸馏训练。

关键创新:CoMem的最大创新在于其通过层次化的方式组织长上下文记忆,而非传统的基于token的方式,这一设计显著提升了模型的效率和性能。

关键设计:在实验中,CoMem使用了固定的缓存数量和层次化的记忆结构,参数设置上采用了rank-32的自蒸馏LoRA,并在NVIDIA H20平台上实现了18.26 GB的内存使用,相较于89.36 GB有显著降低。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,CoMem在RULER任务上达到了97.05的性能,在LoCoMo任务上达到了38.27,相较于全上下文KV-Direct方法的34.59有显著提升。此外,CoMem在内存使用上仅需18.26 GB,相比传统方法节省了71%,并实现了7.83倍的预填充速度提升。

🎯 应用场景

该研究的潜在应用领域包括对话系统、长文本理解和信息检索等场景。通过优化长上下文的处理效率,CoMem可以提升模型在实际应用中的响应速度和准确性,具有广泛的实际价值和未来影响。

📄 摘要(原文)

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.