Skip to content
Preprint

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

Jul 2026 · 0 citations · 37 references
Computer Science

TL;DR

These results show that long-context memory can be organized along the layer axis, not only the token axis, and expose both the benefits of bounded retrieval and its in-window compression tax.

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.

View source

Similar papers

Preprint Jul 2026

Convolution for Large Language Models

These results support depthwise convolution as a lightweight complement to self-attention for modeling short-range token interactions and suggest that the convolution makes repeated token IDs more sensitive to their immediate context.

Yuchuan Tian, Yingte Shu, Wei He et al. · 0 citations
Preprint Jul 2026

Reproducing LightMem: Naive RAG Is Just as Good for Memory Management

Long-term conversational agents require access to information from earlier interactions, such as a user's preferences, past requests, or previously mentioned facts. Repeatedly providing the full dialogue history can be expensive as conversations grow, so many memory approaches instead transform past interactions into compact entries that can be retrieved when needed. LightMem is a recent lightweight memory-management approach that reports strong effectiveness while maintaining relatively low construction cost. However, it still relies on a separate constructed memory representation and is evaluated with only one retriever, leaving unclear how sensitive its results are to retriever choice and whether memory construction discards answer-relevant information. In this study, we reproduce LightMem and compare it with Naive RAG, which retrieves directly from raw user turns. We recover LightMem's main configuration trend, but find that retriever choice is a major source of performance variation: changing only the retriever over a fixed LightMem store shifts answer accuracy from 58.1% to 75.5%. Constructed memories also do not consistently outperform raw-turn retrieval. Naive RAG generally performs better at matched retrieval depths, whereas LightMem performs better mainly under tight answering-token budgets. Oracle evaluation further shows that memory construction removes some answer-relevant information. Overall, LightMem offers a context-efficiency trade-off rather than a general advantage over Naive RAG. Its value depends on the retriever and available token budget, motivating future work on retrieval, reranking, query formulation, and their interaction with raw and constructed memory representations.

Yong Zhou, Shuai Wang, B. Koopman et al. · 0 citations
Preprint Jul 2026

Frontier Language Models Struggle to Copy: Text Can Be Better Viewed in 2D

While large language models (LLMs) can solve advanced reasoning problems in seconds, we show that even frontier models fail to perform a much simpler operation: exactly copying an input string that lies well within their context windows. We attribute this failure to positional encodings in Transformer architectures, whose inductive bias favors copying through a shortcut based on matching local contexts rather than carefully locating the corresponding input positions. To address this issue, we introduce 2D-RoPE, which organizes text into a 2D grid rather than a 1D sequence and assigns each token a row ID and a column ID. Under this view, copying becomes simply retrieving input tokens at a fixed column offset, which makes the task easy to learn. In synthetic copy experiments, shallow Transformers with 2D-RoPE achieve perfect copying at input lengths hundreds of times longer than those seen during training, whereas standard positional encodings fall far behind. We further show that the advantage of 2D-RoPE language models on copy tasks consistently holds in large-scale pretraining on DCLM with model sizes up to 1.4B parameters. Overall, our results suggest that viewing text in 2D can benefit language modeling, and we hope this encourages future work to further explore the potential of 2D positional encodings.

Haodong Wen, Yiran Zhang, Yingfa Chen et al. · 0 citations
#artificial intelligence Preprint Aug 2026

Cross-Model Memory Transfer via Target-Side Reader Adaptation

The results suggest that Engram can serve as a reusable external knowledge artifact, provided that the target has access to a compatible reader interface and target-side adaptation can further improve alignment when direct reader reuse is insufficient.

Mingyuan Li, Guangsheng Yu, Xu Wang et al. · 0 citations
Preprint Jul 2026

Co-LMLM: Continuous-Query Limited Memory Language Models

This work proposes continuous-query LMLM (CO-LMLM), where the KB pairs continuous keys with textual knowledge values, a significant departure from prior reliance on relational KB and queries, which outperforms prior LMLMs and vanilla LLMs in both perplexity and factual precision.

Yair Feldman, Lin-Xi Zhao, Nathan Godey et al. · 0 citations
Conference Open access 2026

Octopus: Gated Selective Attention for Memory-Bounded Long-Context Inference in Large Language Models

O CTOPUS is proposed, a framework that confers fixed-memory inference onto pretrained Transform-ers without the information loss of linearization and outperforms state-of-the-art linearized baselines on the GSM8K benchmark, demonstrating that learned sparse retention serves as an effective regular-izer for long-horizon reasoning.

C. Nguyen, Ryan A. Rossi, L. Van et al. · 0 citations