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
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, \modelnameeffectively generalizes to inputs exceeding 128K tokens. Across a comprehensive suite of long-context modeling and understanding benchmarks, \modelnameachieves 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× over streaming and 3× 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
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 102 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
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 ( >102 ) are hidden with a dash (“-”).
Base Model
Base Model
LongBench
InfBench
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.
Figure 7
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
σ
l/5
l/10
chunk size
1024
512
α
1/16,1/8,1/4
AdamW ( β1,β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 ∼ 4K
4 ∼ 6K
6 ∼ 8K
Proportion
47%
29%
8%
16%
Appendix
Table 7: Proportion of samples within each length interval.
Model
Arxiv
PG19
ProofPile
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
Books3 involved in training
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. † represents the values in corresponding rows are reproduced from open-sourced code.
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 β=8 . Our default setting is highlighted in bold .
β
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 β with constant h=3 and M=4 . Our default setting is highlighted in bold . For β∈{1,2} , we are not able to set levelwise compression ratios and thus we set the compression ratio same as the β 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 β=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 β 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.
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.
Gabriele Berton, Jayakrishnan Unnikrishnan, Son Tran +1
Amazon · Center For Research in Computer Vision, University of Central Florida
Long-context understanding remains challenging for LLMs due to limited context windows. This paper introduces Long Input Fine-Tuning (LIFT), a framework that improves the long-context performance of arbitrary short-context LLMs by dynamically adapting their parameters to each long input. Instead of endlessly extending context windows to fit longer inputs in context, LIFT stores and absorbs the input in parameters. By fine-tuning long inputs into parameters, LIFT enables short-context LLMs to answer questions even when required information is absent from the inference context, avoiding the quadratic input-length complexity of standard long-context models. Rather than simple continued pretraining on new long contexts, LIFT uses carefully designed LLM-generated synthetic tasks to enhance comprehension beyond memorization. To offset fine-tuning overhead, we design a highly optimized pipeline that reduces Time to First Token (TTFT) to under 10 seconds for 8k context. We further analyze LIFT's strengths and limitations, discuss large-scale deployment feasibility, and highlight future research directions. Implementation is open-sourced at https://github.com/MuLabPKU/LIFT.
Yansheng Mao, Yufei Xu, Jiaqi Li +5
Institute for Artificial Intelligence, Peking University · School of Electronics Engineering and Computer Science, Peking University · State Key Laboratory of General Artificial Intelligence, BIGAI
Long-context language model inference is bottlenecked by memory, as the KV cache grows with context length. Recent techniques to compress the KV cache fall short: they either degrade model quality substantially or require considerable time and compute to compress a single long prompt. Furthermore, many methods require the input to fit within the target model's context window, and are generally incompatible with modern production inference engines. Encoder-decoder compressors, which map a long token sequence to a shorter sequence of latent embeddings consumed by a decoder, are an appealing alternative in principle. However, existing approaches are not competitive with KV cache compression on the accuracy-efficiency frontier. In this work, we revisit encoder-decoder compression and close this gap. We first perform an architecture search, pre-training many variants from scratch to determine how best to design and train encoder-decoder compressors. Guided by our findings, we continually pre-train a family of 0.6B-encoder, 4B-decoder models on nearly 400B tokens each, at compression ratios of 1:4, 1:8, and 1:16. We introduce Latent Context Language Models (LCLMs), a family of compressors that improve the Pareto frontier across general-task performance, compression speed, and peak memory usage. LCLMs also unlock new efficient downstream applications: they serve as backbones for long-horizon agents that skim a compressed long context and adaptively expand relevant segments on demand, and they can compact the model's own generated reasoning into latent memory during inference, which is nearly lossless on the full MATH test set while being highly efficient, reducing peak KV usage by up to 7.9×.
Ang Li, Sean McLeish, Haozhe Chen +12
New York University · Modal Labs · Harvard University +5