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

TL;DR - CoMem exploits transformers’ layer-wise division of labor to provide bounded-compute, unbounded-context memory by caching intermediate residual states and recomputing upper layers. It substantially improves long-context benchmarks while reducing 128k-context memory use and prefill time.

• CoMem retrieves a fixed number of cached states, making read compute and memory independent of total stored-context length.
• A frozen Qwen3-8B with rank-32 self-distillation LoRA scores 97.05 on RULER and 38.27 on LoCoMo.
• At 128k context on an NVIDIA H20, an adapter-free variant uses 18.26 GB versus 89.36 GB and delivers 7.83× faster prefill.
• Caching at deeper layers reduces recomputation but loses fidelity; self-distillation substantially repairs this tradeoff.