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

2026-07-30

AI Summary

This work addresses the bottleneck in conventional large language models, where memory and computational costs grow linearly with context length. The authors propose CoMem, a novel mechanism that organizes long-term memory along the model’s depth axis: comprehension tasks are concentrated in lower-to-mid layers, with context blocks written and residual states cached only at intermediate layers. During inference, recomputation is restricted to upper layers conditioned on queries, yielding context-length-independent overhead. To mitigate semantic fidelity loss, the method integrates self-distilled LoRA while employing a frozen backbone, fixed-budget retrieval, and adapter-free deployment. Evaluated on Qwen3-8B, CoMem achieves 97.05 on RULER and 38.27 on LoCoMo—substantially outperforming KV-Direct—reduces GPU memory from 89.36 GB to 18.26 GB for 128k-context inputs, accelerates prefilling by 7.83×, and demonstrates robust multi-turn dialogue memory retention.

Abstract

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.

Problem

• long-context memory
• transformer depth
• unbounded-context
• memory efficiency
• large language models

Innovation

• depth division of labor
• comprehension memory
• long-context memory
• layer-wise caching
• self-distillation LoRA