Search papers, labs, and topics across Lattice.
This paper explores the non-uniform utilization of transformer depth in large language models (LLMs), revealing that lower and middle layers focus on building semantic representations while upper layers specialize in prediction. The authors introduce CoMem (Comprehension Memory), a method that optimizes memory retrieval by writing context chunks through intermediate layers and recomputing upper layers, achieving significant performance improvements on long-context tasks. CoMem demonstrates a remarkable efficiency, using substantially less memory and achieving faster prefill speeds compared to traditional methods, while maintaining high fidelity in predictions.
CoMem achieves a 7.83x prefill speedup and drastically reduces memory usage while maintaining high performance on long-context tasks, challenging conventional memory management in LLMs.
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.
Tsinghua AI