ICLR2026
Stacked from One: Multi-Scale Self-Injection for Context Window Extension
Wei Han, Pan Zhou, Shuicheng YAN
Abstract
The limited context window of contemporary large language models (LLMs) remains a primary bottleneck for their broader application across diverse domains. Although continual pre-training on long-context data offers a straightforward solution, it incurs prohibitive data acquisition and computational costs. To address this challenge, we propose SHAREDLLM, a novel framework based on multi-grained context compression and query-aware information acquisition. SHAREDLLM comprises two stacked short-context LLMs: a lower model serving as a compressor and an upper model acting as a decoder. The lower model compresses long inputs into compact, multi-grained representations, which are then forwarded to the upper model for context-aware processing. To maximize efficiency, this information transfer occurs exclusively at the lowest layers, bypassing lengthy forward passes and redundant cross-attention operations. This entire process, wherein the upper and lower models are derived from the same underlying LLM layers, is termed self-injection. To support this architecture, a specialized tree-based data structure enables the efficient encoding and queryaware retrieval of contextual information. Despite being trained on sequences of only 8K tokens, SHAREDLLM effectively generalizes to inputs exceeding 128K tokens. Across a comprehensive suite of long-context modeling and understanding benchmarks, SHAREDLLM achieves performance superior or comparable to strong baselines, striking an optimal balance between efficiency and accuracy. Furthermore, these design choices allow SHAREDLLM to substantially reduce the memory footprint and yield notable inference speedups (2× over streaming and 3× over encoder-decoder architectures). The core implementation code, along with training and evaluation details, is open-sourced at https://github.com/Clement25/SharedLLM . More detailes are provided in the appendix and supplementary materials.