💬 NLP

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

This paper introduces CoMem, a memory mechanism that exploits the non-uniform utilization of Transformer layers by caching intermediate residual states and recomputing upper layers, thereby achieving unbounded-context memory with constant compute and memory costs while significantly outperforming full-context baselines on long-context benchmarks.

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

Published 2026-07-31

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a super-smart robot that can read entire libraries of books in a split second. This robot is a type of Artificial Intelligence called a "Large Language Model." To understand a story, the robot has to remember every word it has read so far. Think of its memory like a giant whiteboard. In older robots, every time they read a new word, they had to write a fresh note on the board for every single word they had ever seen. If the story got really long, the whiteboard would run out of space, or the robot would get so tired trying to read its own notes that it would move incredibly slowly. This is the big problem scientists are trying to solve: how to let these robots remember huge amounts of information without running out of space or time.

The secret to how these robots think lies in layers, kind of like a multi-story building. The lower floors are great at understanding the basic meaning of words (like knowing that "apple" is a fruit). The upper floors are where the robot gets really smart, using that basic meaning to answer specific questions or predict what happens next. Usually, the robot has to process the entire story from the bottom floor all the way to the top floor every time it wants to answer a question. This paper suggests a clever shortcut: what if we only let the robot read the whole story up to the middle floor, save that "half-baked" understanding, and then only finish the job on the top floor when we actually need an answer?

This is exactly what the researchers behind a new method called CoMem (Comprehension Memory) have discovered. They found that instead of trying to remember every single detail of a massive document all at once, the robot can be much smarter about how it stores and retrieves information. They realized that the robot's "understanding" of a text chunk is mostly ready by the time it reaches the middle of its brain (the middle layers). So, CoMem works like a librarian who doesn't keep the whole book on the desk. Instead, the librarian reads the first half of a chapter, writes down a quick summary on a sticky note, and sticks that note on a shelf. When you ask a question, the librarian doesn't pull out the whole book. They just grab the few relevant sticky notes, finish reading the rest of the chapter in their head, and give you the answer.

The paper shows that this "depth division of labor" works incredibly well. By only processing the first half of the text and caching (saving) that intermediate result, the robot saves a massive amount of space. In their tests, using a model called Qwen3-8B, CoMem used only 18.26 GB of memory to handle a 128k token context, whereas the traditional method needed 89.36 GB. That's a huge difference! It also made the robot much faster, achieving a 7.83× speedup when preparing to answer a question.

However, the researchers are careful to point out that this isn't magic. You can't just stop reading halfway and expect a perfect answer. If you stop too early, the robot forgets the details. If you stop too late, you lose the speed benefit. The team found a "sweet spot" in the middle layers (specifically layer 12 out of 36) where the robot understands the meaning well enough to save it, but hasn't yet specialized it for a specific question. They also found that simply saving the notes isn't enough; the robot still needs to be able to look at all the relevant notes at once to connect the dots.

The results are impressive but specific. On a test called RULER, which checks if a robot can find a "needle in a haystack" (a specific fact in a huge text), CoMem scored 97.05, beating the standard method which scored 78.80. On a long conversation test called LoCoMo, CoMem scored 38.27 compared to 34.59 for the standard method. The paper suggests that this approach is a strong way to organize memory, but it's not a perfect fix for every single task. For example, the robot still struggles a bit with certain types of complex questions (like some "qa2" tasks in the BABILong test), showing that while this method is a huge step forward for efficiency, it doesn't make the robot perfect at everything.

In short, CoMem suggests that we don't need to force the robot to re-read the whole library every time we ask a question. By splitting the work—letting the lower layers do the heavy lifting of understanding and the upper layers do the specific answering—we can build robots that remember more, think faster, and use less energy. It's a bit like realizing you don't need to carry the whole encyclopedia in your backpack; you just need to know which pages to pull out when you need them.

Technical Summary: CoMem (Comprehension Memory)

Problem Statement

Long-context language modeling faces two primary bottlenecks: the quadratic prefill cost of self-attention and the linear memory growth of the Key-Value (KV) cache. Existing efficient methods typically address these by evicting or compressing KV states along the token axis (retaining a subset of tokens) or by retrieving external text units. However, these approaches often struggle to balance memory efficiency with the fidelity required for complex reasoning over long contexts. Furthermore, many methods rely on training specialized memory networks or compressing the entire context, which can introduce significant overhead or fidelity loss.

The authors posit that Transformer layers exhibit a division of labor: semantic information becomes accessible in lower and middle layers, while upper layers increasingly specialize representations for the current query and prediction. Current methods do not fully exploit this depth-based structure to organize memory.

Methodology: CoMem

CoMem (Comprehension Memory) reorganizes long-context memory along the depth axis rather than solely the token axis. It combines bounded token selection with depth-partitioned computation reuse.

Core Mechanism

Write Phase (Depth Partitioning):
- The input context is streamed in disjoint chunks (e.g., 512 tokens).
- Each chunk is processed only through a specific intermediate layer $j$ (the "split depth").
- Instead of storing full-depth KV states, CoMem caches the residual tensor $h_j$ (the hidden state) at layer $j$ for each token in the chunk.
- These tensors are stored in an external index keyed by the chunk text.
- Storage Efficiency: For Qwen3-8B, storing $h_j$ requires only 1/18th the bytes of a full bf16 KV pair per token.
Read Phase (Bounded Retrieval & Recomputation):
- Upon receiving a query, a selector (using BM25) retrieves a bounded set of top- $k$ relevant chunks.
- The cached residual tensors $h_j$ for these chunks are packed together with a "sink" token and the current query.
- The model recomputes only the upper layers $[j:L]$ over this fixed-size pack using full cross-chunk attention.
- Complexity: For a fixed retrieval budget ( $k$ ) and chunk size, the model-side online read compute and memory working set are independent of the total stored context length.

Training and Adaptation

Backbone: The base model (Qwen3-8B) remains frozen.
Self-Distillation: To mitigate the fidelity loss caused by skipping lower layers during the read phase, a lightweight LoRA (Rank-32) is trained via self-distillation. The student model (at split $j=12$ ) learns to match the logits of a frozen teacher model (at $j=0$ ) on the top-64 logit support.
Training Data: The LoRA is trained on plain PG-19 text without task-specific supervision or retrieval data, allowing the model to generalize the "depth split" capability.

Key Contributions

Depth as a Reuse Axis: The paper establishes that semantic accessibility peaks near mid-depth, while the "zero-shot readable boundary" (where cached states can be directly used without recomputation) depends on model scale. CoMem treats this depth split ( $j$ ) as a tunable knob between precomputed work and query-time cost.
Bounded Depth-Partitioned Memory: CoMem caches one residual state per token and exactly resumes the upper transformer. It combines bounded retrieval with full cross-chunk attention, achieving a storage footprint of 1/18th of full KV caches per token on Qwen3-8B.
Controlled Quality-Efficiency Trade-off: The authors provide evidence that deeper caching reduces per-query computation but incurs a fidelity loss, which self-distillation substantially repairs. The method isolates the benefits of retrieval from the costs of depth reuse.

Experimental Results

Evaluated on a unified, chat-template-free protocol using Qwen3-8B:

Benchmark Performance:
- RULER: CoMem achieves 97.05, outperforming the full-context KV-Direct baseline (78.80).
- LoCoMo (Dialogue Memory): CoMem scores 38.27 vs. 34.59 for KV-Direct. This advantage persists under conversation-cluster resampling and independent judge audits.
- LongEval & LongBench: CoMem matches or exceeds baselines (69.0 on LongEval; 12.15 macro F1 on LongBench).
Efficiency (NVIDIA H20 at 128k context):
- Memory: CoMem uses 18.26 GB peak memory compared to 89.36 GB for full-context KV-Direct.
- Speed: CoMem achieves a 7.83× speedup in full-write-inclusive prefill time.
- Decode: Resumed-band decode remains within 11% of dense inference up to 64k and becomes 1.07× faster at 128k.
Ablations:
- Retrieval vs. Depth: Bounded selection alone (without depth reuse) improves dialogue memory, but depth reuse is critical for efficiency.
- Distillation: Self-distillation recovers ~13.75 LoCoMo points lost by the frozen split at $j=12$ .
- Cross-Chunk Attention: Full cross-chunk attention is essential for multi-fact disambiguation; block-diagonal reuse causes significant performance drops.

Significance and Claims

The paper claims that long-context memory can be effectively organized along the layer axis, not just the token axis. CoMem demonstrates that:

Semantic information is accessible at intermediate depths, allowing for the caching of query-blind states that are later specialized by recomputing upper layers.
A bounded retrieval strategy combined with depth-partitioned computation can decouple online inference cost from stored context length.
Lightweight self-distillation can repair the fidelity loss inherent in skipping lower layers, enabling high accuracy without retraining the backbone.

The authors emphasize that CoMem is an efficiency method offering a quality-memory-compute trade-off, rather than a state-of-the-art accuracy solution for every in-window task. They explicitly state that their results on contexts exceeding the native window (40,960 tokens) are "unextended stress references" and do not claim superiority over models using position-extended techniques like YaRN. The method is presented as a complementary approach to existing token-axis compression and retrieval methods.