CoMem: Reusing Transformer Depth across Queries with Persistent Intermediate Residuals
cs.CL
Submitted: 2026-07-30
Updated: 2026-10-07
Comments: 19 pages, 4 figures, 27 tables. Submitted to ACL Rolling Review
Code: https://github.com/liuhanzuo/COMem
License: http://creativecommons.org/licenses/by/4.0/
The gist: Transformer depth is not used uniformly: lower and middle layers build semantic representations, while upper layers increasingly specialize them for prediction.
Terminology
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.
Sources
- LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding
- PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
- Training Transformers for KV Cache Compressibility
- The Llama 3 Herd of Models
- Distilling the Knowledge in a Neural Network
- RULER: What's the Real Context Size of Your Long-Context Language Models?
- RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation
- The Remarkable Robustness of LLMs: Stages of Inference?
- SnapKV: LLM Knows What You are Looking for Before Generation
- CompressKV: Semantic-Retrieval-Guided KV-Cache Compression for Resource-Efficient Long-Context LLM Inference
- MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
- RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval
- Evaluating Very Long-Term Conversational Memory of LLM Agents
- Landmark Attention: Random-Access Infinite Context Length for Transformers
- YaRN: Efficient Context Window Extension of Large Language Models
- The Residual Stream Is All You Need: On the Redundancy of the KV Cache in Transformer Inference
- Alita-G: Self-Evolving Generative Agent for Agent Generation
- Alita: Generalist Agent Enabling Scalable Agentic Reasoning with Minimal Predefinition and Maximal Self-Evolution
- Compressive Transformers for Long-Range Sequence Modelling
- RoFormer: Enhanced Transformer with Rotary Position Embedding
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering