Stacked from One: Multi-Scale Self-Injection for Context Window Extension

Wei Han; Pan Zhou; Soujanya Poria; Shuicheng Yan

cs.CLMar 5, 2026

Stacked from One: Multi-Scale Self-Injection for Context Window Extension

Authors: Wei Han, Pan Zhou, Soujanya Poria, Shuicheng Yan

Organizations: Singapore University of Technology and Design (SUTD), Singapore · Singapore Management University (SMU), Singapore · Nanyang Technological University (NTU), Singapore · National University of Singapore (NUS), Singapore

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~\modelname, 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 query-aware retrieval of contextual information. Despite being trained on sequences of only 8K tokens, \modelname~~effectively generalizes to inputs exceeding 128K tokens. Across a comprehensive suite of long-context modeling and understanding benchmarks, \modelname~~achieves performance superior or comparable to strong baselines, striking an optimal balance between efficiency and accuracy. Furthermore, these design choices allow \modelname~to substantially reduce the memory footprint and yield notable inference speedups ( $2\times$ over streaming and $3\times$ over encoder-decoder architectures).

Figures & tables

Figure 1: Overview of SharedLLM . The architecture resembles general encoder-decoder architecture like T5 ( Raffel et al., 2020 ) , but the interaction occurs at the first $M$ layers between lower and upper model through shared key-values which are encoded and compressed from the text chunk into a sequence of trees (top-left).

Figure 2: An running example of our tree (depth=3). Each box indexed by $i$ represents the $i$ th iteration of node split and selection.

Base Model	Arxiv				PG19				ProofPile
Base Model	4K	8K	32K	128K	4K	8K	32K	128K	4K	8K	32K	128K
LLaMA-2-32K ( Together, 2023 )	3.58	3.34	2.96	OOM	6.93	6.81	7.04	OOM	2.87	2.58	2.47	OOM
PI ( Chen et al., 2023 )	3.49	3.21	2.77	OOM	6.97	6.77	6.89	OOM	2.77	2.64	2.51	OOM
YaRN ( Peng et al., 2023 )	3.35	3.09	2.58	OOM	6.85	6.62	6.91	OOM	2.82	2.56	2.47	OOM
CEPE ( Yen et al., 2024 )	3.03	3.02	2.51	2.97	6.69	6.40	6.80	6.10	2.38	2.43	2.45	2.39
SharedLLM	2.99	2.97	2.46	2.91	6.55	6.28	6.65	5.96	2.33	2.34	2.38	2.40

Table 1: Language modeling results (perplexity) of the continual pretraining setting on downsampled RedPajama. Best results on context-extended models are marked in bold. Perplexity higher than

10^{2}

are denoted by dash (”-”). LLaMA-3.1 has the declared 128K context-length since release, and we list the direct inference results separately for reference only.

Base Model	Method	PG19				ProofPile				CodeParrot
Base Model	Method	4K	16K	32K	100K	4K	16K	32K	100K	4K	16K	32K	100K
LLaMA-2	StreamingLLM	9.21	9.25	9.24	9.32	3.47	3.51	3.50	3.55	2.55	2.60	2.54	2.56
	LongAlpaca-16K	9.96	9.83	-	OOM	3.82	3.37	-	OOM	2.81	2.54	-	OOM
	Activation Beacon	9.21	8.34	8.27	8.50	3.47	3.34	3.32	3.31	2.55	2.43	2.41	2.62
	SharedLLM	8.68	8.01	7.96	8.24	3.36	3.24	3.21	3.19	2.33	2.25	2.23	2.36
Mistral-7B	StreamingLLM	9.58	9.63	9.52	9.55	4.08	4.19	4.16	4.23	2.99	3.05	3.13	3.02

Table 2: Langauge modeling results of the supervised fine-tuning setting. “OOM” denotes the out-of-memory exception is raised during inference. Excessively large perplexities (

>10^{2}

) are hidden with a dash (“-”).

Base Model	Base Model	LongBench					InfBench
Base Model	Base Model	SDQA	MDQA	Summ.	FS	Code	Math.F	En.MC	Ret.N
LLaMA-2	Base	24.90	22.60	24.70	60.00	48.10	2.85	22.79	1.85
	StreamingLLM	21.47	22.22	22.20	50.05	48.00	6.00	32.31	5.23
	LongAlpaca-16K	28.70	28.10	27.80	63.70	56.00	6.23	25.74	4.87
	SnapKV	24.05	22.98	17.25	16.11	58.87	9.95	28.83	2.31
	OmniKV	23.86	22.77	21.09	35.74	49.37	8.81	26.25	3.66

Table 3: Evaluation results of different SFT methods on two benchmarks from LongBench and Infini-Bench. Note that for some baselines we follow their default settings to truncate the input below their window length, which may cast positive effects on their performance.

Figure 3: Comparison of memory usage (left) and total inference time on 100 examples (right) between SharedLLM and other training time baseline methods. The data is collected by running a tiny experiment on 100 examples in corresponding lengths. “OOM” means out-of-memory exception triggered during test time.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Item	Continual Pretraining	Supervised Fine-tuning
training epoch	1	2
warmup ratio	0.01	0.001
$\sigma$	l/5	l/10
chunk size	1024	512
$\alpha$	$1/16,1/8,1/4$
AdamW ( $\beta_{1},\beta_{2})$	0.9, 0.999

Appendix

Table 5: Configurations for training on both tasks.

Domain	Proportion (%)
Arxiv ( Clement et al., 2019 )	2.5
Books (w/o S3) ( Rae et al., 2020 )	4.5
C4 ( Roberts et al., 2019 )	15.0
CommonCrawl ( Luccioni and Viviano, 2021 )	67.0
Github	4.5
StackExchange	2.0

Appendix

Table 6: Dataset composition in our downsampled Redpajama (10B) tokens.

Length	$<$ 2K	2 $\sim$ 4K	4 $\sim$ 6K	6 $\sim$ 8K
Proportion	47%	29%	8%	16%

Appendix

Table 7: Proportion of samples within each length interval.

Books3 involved in training
Model	Arxiv				PG19				ProofPile
Model	4K	8K	32K	128K	4K	8K	32K	128K	4K	8K	32K	128K
LLaMA-2-7B (4K)	2.60	-	-	OOM	6.49	-	-	OOM	2.28	-	-	OOM
YaRN-2-128K	3.13	2.96	2.34	OOM	6.15	6.02	6.32	OOM	2.70	2.47	2.41	OOM
CEPE	2.86	2.84	2.34	2.91	6.60	6.24	6.66	5.99	2.22	2.33	2.26	2.23
Books3 not involved in training

Appendix

Table 8: Perplexity increment as a negative effect from the lack of books3.

\dagger

represents the values in corresponding rows are reproduced from open-sourced code.

Figure 5: Accuracy comparison on passkey retrieval (single key-value pair) task.

Setting	Arxiv	MD-QA
Default	2.46 ( $\pm$ 0.01)	30.93 ( $\pm$ 0.16)
Random	2.61 ( $\pm$ 0.13)	28.85 ( $\pm$ 0.45)
Fixed Left	2.49 ( $\pm$ 0.02)	29.62 ( $\pm$ 0.15)
Query-aware	2.48 ( $\pm$ 0.02)	29.27 ( $\pm$ 0.12 )
$\Lambda$ -shape	2.52 ( $\pm$ 0.04)	29.48 ( $\pm$ 0.18)

Appendix

Table 9: Pilot studies of branch-selection policies.

$M$	1	2	4	8	16
Time (s)	6.78	9.35	11.81	16.81	25.85
Memory (GB)	21.04	21.50	22.39	24.08	27.82

Appendix

Table 10: Inference time under various

M

with constant

h=3

and

\beta=8

. Our default setting is highlighted in bold .

$\beta$	64	32	16	8	4	2	1
Time (s)	11.68	11.73	11.78	11.81	11.87	12.04	12.47
Memory (GB)	22.20	22.20	22.20	22.39	22.40	22.35	22.97

Appendix

Table 11: Inference time under various

\beta

with constant

h=3

and

M=4

. Our default setting is highlighted in bold . For

\beta\in\{1,2\}

, we are not able to set levelwise compression ratios and thus we set the compression ratio same as the

\beta

for every level of the tree.

$h$	1	2	3	4
Time (s)	11.16	11.55	11.81	11.86
Memory (GB)	19.72	22.42	22.39	22.41

Appendix

Table 12: Inference time under various

h

with constant

\beta=8

and

M=4

. Our default setting is highlighted in bold .

Setting	Time	Memory
w/o query-aware retrieval	11.81	22.39
w query-aware retrieval	13.18	25.44

Appendix

Table 13: Comparison of time and memory consumption when query-based retrieval is incorporated/not incorporated in SharedLLM .

h

M

and

\beta

are fixed at the default values.

Figure 6: A living example of tree growth and split on the passkey retrieval task. The numbers above the text boxes are the correlation score between the text chunk and user query. Green/Red boxes indicate the chunk is selected/not selected.

Long-Context Language Model Inference

Long-context language model inference makes processing and generating over lengthy prompts faster and more memory-efficient.

See all 227 papers

CompLLM: Compression for Long Context Q&A
Sep 23, 2025Gabriele Berton, Jayakrishnan Unnikrishnan, Son Tran +1LLM Compression Long-Context Language Model Inference
Amazon · Center For Research in Computer Vision, University of Central Florida
LIFT: A Novel Framework for Enhancing Long-Context Understanding of LLMs via Long Input Fine-Tuning
Feb 20, 2025Yansheng Mao, Yufei Xu, Jiaqi Li +5Long-Context QA LLM Fine-Tuning
Institute for Artificial Intelligence, Peking University · School of Electronics Engineering and Computer Science, Peking University · State Key Laboratory of General Artificial Intelligence, BIGAI
End-to-End Context Compression at Scale
Jun 8, 2026Ang Li, Sean McLeish, Haozhe Chen +12LLM Compression Long-Context Language Model Inference
New York University · Modal Labs · Harvard University +5

Sep 23, 2025J/K move · Enter open · S save

CompLLM: Compression for Long Context Q&A

Gabriele Berton, Jayakrishnan Unnikrishnan, Son Tran, Mubarak Shah

Amazon · Center For Research in Computer Vision, University of Central Florida

LLM Compression Long-Context Language Model Inference Long-Context QA Context Compression LLM Context Compression

Open paper AskPDF arXiv

Large Language Models (LLMs) face significant computational challenges when processing long contexts due to the quadratic complexity of self-attention. While soft context compression methods, which map input text to smaller latent representations, have shown promise, their real-world adoption is limited. Existing techniques typically compress the context as a single unit, which leads to quadratic compression complexity and an inability to reuse computations across queries with overlapping contexts. In this work, we introduce CompLLM, a soft compression technique designed for practical deployment. Instead of processing the context holistically, CompLLM divides it into segments and compresses each one independently. This simple design choice yields three critical properties: efficiency, as the compression step scales linearly with the context length; scalability, enabling models trained on short sequences (e.g., 1k tokens) to generalize to contexts of 100k tokens; and reusability, allowing compressed segments to be cached and reused across different queries. Our experiments show that with a 2x compression rate, at high context lengths CompLLM speeds up Time To First Token (TTFT) by up to 4x and reduces the KV cache size by 50%. Furthermore, CompLLM achieves performance comparable to that obtained with the uncompressed context, and even surpasses it on very long sequences, demonstrating its effectiveness and practical utility.

Abstract

Long-Context Language Model Inference

Explore similar work

CompLLM: Compression for Long Context Q&A

LIFT: A Novel Framework for Enhancing Long-Context Understanding of LLMs via Long Input Fine-Tuning

End-to-End Context Compression at Scale