Stacked from One: Multi-Scale Self-Injection for Context Window Extension
Authors: Wei Han
Pan Zhou
Shuicheng YAN
Venue: ICLR 2026 Poster (2026)
Keywords: long-context modeling; continual pretraining; extrapolation
Primary area: applications to computer vision, audio, language, and other modalities
TL;DR: An efficient method to fast extend the context window length of small language models

Abstract: The limited context window of contemporary large language models (LLMs) hinders broader application. In this work, we present SharedLLM, a novel approach grounded in the design philosophy of multi-grained context compression and query-aware information retrieval. SharedLLM is composed of two short-context LLMs: a lower moel (compressor) and an upper model (decoder). The lower model compresses context information, while the upper model processes compressed, context information from the lower model and performs context-aware modeling. Information transfer between the compressor and decoder occurs only at the lowest layers to reduce redundant computation. Based on this architecture, we introduce a specialized tree-style data structure to efficiently encode, store and retrieve multi-grained contextual information from text chunks. This entire process, wherein the sender and receiver are derived from the same LLM layer, is referred to as self-injection. In our evaluation on long-context modeling and understanding tasks, SharedLLM achieves superior or comparable results to several strong baselines, striking an effective balance between efficiency and performance. Meanwhile, with the aforementioned design choices, SharedLLM can greatly reduce memory consumption, and demonstrates substantial speed-ups over other advanced baselines. The core code of our implementation along with training and evaluation is available in appendix and supplementary.

Published as a conference paper at ICLR 2026
STACKED FROM ONE: MULTI-SCALE SELF-INJECTION FOR CONTEXT WINDOW EXTENSION
Wei Han'; Pan Zhou?, Yan
1 Singapore University of Technology and Design (SUTD), Singapore
2 Singapore Management University (SMU), Singapore
3 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 do- mains. Although continual pre-training on long-context data offers a straight- forward solution, it incurs prohibitive data acquisition and computational costs. To address this challenge, we propose SHAREDLLM, novel framework based on multi culti-grraned context compression and query-awart information acquisition. SHAREDLLM comprises two stacked short-context LLMs: a lower model serv- ing compressor and 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, by- passing 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. SHAREDLLM effectively generalizes to inputs exceeding 128K tokens. Across comprehensive suite of long-context modeling and un- derstanding benchmarks, SHAREDLLM achieves performance superior or com- parable to strong baselines, striking an optimal balance between efficiency and accuracy. Furthermore, these design choices allow SHAREDLLM to substan- tially reduce the memory footprint and yield notable inference speedups (2x over streaming and 3x over deccder architectures). The core imple- mentation code, along with training and evaluation details, is open-sourced at thu.. com/ C1 ement 2ete225/Shaardiiie More detailes pro- https: //gi pittps///gil
vided in the appendix and supplementary materials.
INTRODUCTION
Since the release of GPT-3 Brown 2020), the rapid advancement of large language models (LLMS) (Chowdhery et al. 2022 Achiam et al. 2023, Touvron et al. 2023a b Dubey et al. 2024 Ma al. 2024 Guo et al. 2025) has revolutionized the NLP research community and transformed
various workflows. Pretrained on trillions of tokens, LLMs exhibit remarkable abilities, such as completing unfinished text code and following human instructions perform designated tasks after simple supervised fine-tuning (Wii et al. 2021 Chung et al. 2024] Wang et al. 2025). Despite their impressive capabilities, several factors limit their broader application. One major constraint the context window size (Hsich et al. 2024 Liu et al. (2025). which refers to the maximum number of tokens on which an I.I.M can behave normally. When the input text exceeds this limit, L.MSs may suffer from severe performance degradation or hallucination during inference.
Many researchers have attempted to extend the context window of LLMs with minimal training costs (Peng et al. 2023 Together 2023 Xiong et al. 2024). One common approach involves post-
"Email: wei han @ sutd.edu.sg corresponding author: panzhou3 @ gmail.com

---

Published as a conference paper at ICLR 2026
pretraining LLMs long-context corpora with substantial computational resources (TogetherA] 2023 Xiong et al. 2024: Ma et al. 2024). Advanced positional encoding methods, which usually extend RoPE to rescale attention scores in a more principled manner, are integrated to minimize the size of training corpus (Chen et al. 2023 Peng et al. 2023). Although they achieve extrapolation- "train short, test long" ", the efficiency relatively low. For example, to reach the context length of 128K tokens, using YaRN Peng et al. (2023) one has to pretrain an LLM on 64K tokens. Prompt compression (Ge al. 2023 Gao et al. 2025 accelerates the inference process by replacing long prompts with -.M-eenertted semantic tokens, but fails to extend the context window of LLMs or only applies on limited scenarios. Other approaches upgrade the conventional transformer architec-
tures to enable streaming processing of long context (Xiao et al., 2024b) Yen et al. 2024 Zhang ct al. 2025), which maintain a sliding window of constant-sized memory. Although these designs significantly alleviate the memory-bout -ouuor issue of matrix multiplication, their specialized attention patterns may cause incompatibility with nigh-performance attention implementations (c.g., FlashAt- Dao 2023)), potentially leading to slower inference speeds.
tention (Dao et al. 2022
To strike a balance between efficiency and performance, we propose SHAREDLLM, a lightweight architecture which consists of one upper model and one lower model. The lower model compresses text chunks into multi- culli-grained representations, while the upper model integrates the encoded infor- mation and generates the final output. This multi-grained setting helps LLM focus on task-related fine-grained information while regulating other auxiliary coarse-grained information to a secondary role. Both models are initialized from the same off-the-shelf checkpoint of a short-context LLM, either in full or in part. Since there is no disparity between the hidden spaces of the two mod- els. SHAREDLLM can be effectively fine-tuned without requiring extensive warmup stages.
This paper makes the following major contributions:
We propose SHAREDLLM, a hierarchical architecture for efficient LLM context window extension. It consists of two models which work collaboratively through shared key-value mechanism with minimal tunable parameters.
We design tree-like structure, called context free, which can express long unstructured context in a coarse-to-fine format. To facilitate this process, we introduces a dynamic con- text tree construction and search algorithm. Given a context and query, it can efficiently transform the context into the hierarchical representation and collect relevant information from that tree.
We conduct comprehensive experimental study to demonstrate the effectiveness of SHAREDLLM. On the settings of both post-pretraining and supervised fine-tuning, SHAREDLLM shows impressive extrapolation property and yields stronger performance than baseline models with superior memory and time efficiency.
METHOD
In this section, we first introduce the overall architecture of our proposed SHAREDLLM in Sec. 2.1 and then elaborate on its two main components, the lower model and upper model in Sec. 2.2 and 2.3
2.1 OVERVIEW
As illustrated in Figure 1 SHAREDLLM adopts hierarchical architecture. The lower model, the "compressor", breaks down the long input context Xe into smaller chunks that are then encoded within limited GPU memory. It then uses the same LLM compress each context chunk into com- pact and structured coarse-to-fing representations in parallel. The upper model, or the "decoder", takes the rear part of the input text (the running context, such as questions) as input. It then integrates the compressed information from the lower model, and finally generates predictions of successive tokens in an auto-regressive manner. The lower and upper models are interconnected via the sharing of key-value (KV) states, which are further integrated at the cross-attention modules in the upper model. To facilitate efficient information gathering and integration, the contextual information pro- cessed by the lower model is organized as a binary tree, referred to the context free, which stores multi- multi--rained information at different levels. Note that the KV compression and transmission occur during the prefilling stage of inference, yet they still improve decoding efficiency because each query
2

---

Published as a conference paper at ICLR 2026
(Rasponse The
Extracted 1x
Context Tree
1'x.d
answer should be...
R1
Encode
Self-attention
N Layers
Down-sample
Self-attantion
Tree Cross Prosss-attention
Layer M
Layer
Layer M
Self-attention
Tree cross-attention
Layer
Layer
Layer
Self-attention
Given the following context...
[Query]]
Chunk Subsequences
Chunk Subsequences
Lower Model (single model with chunk parallelismi
Upper Model
Figure I: Overview of SHAREDLLM. The architecture resembles general encoder-decoder archi- tecture like T5 (Raffel et al. 2020), but the interaction occurs 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).
token attends reduced number of key-value pairs. Specifically, this hierarchical decomposition reduces the theoreical self-attention complexity from quadratic ( to a much more manageable O(n (T/n)² + Tn |S')), where SS') is the highly compressed KV length.
In the following, we elaborate on the lower and upper model. To begin with, we first define some notations enhance clarity and readability. Let X {I1, .2....., ...; .....T)) represent the entire input sequence, where 7 denotes the sequence length. We call the LLM whose context window to be extended as "target LLM". Following previous setting (Yen al. 2024). we split these tokens into two continuous parts: X = concat ([Xc: XDl), where the past context X and the running text XD serve as the inputs to the lower and upper models, respectively. Moreover, the past context X is further divided into n smaller and non-overlapping chunks, denoted by C1, C2,..., C.. namely, where Ci U U U...UU.. C, =xc and C. n,, 0, + j. The chunk size is controlled to fit within the lower model's context window, allowing the lower model to fully utilize its encoding capacity.
2.2 LOWER MODEL
The lower model is a small pretrained LLM, implemented as the first M shallow layers of the target LLM. It independently encodes and compresses each past context chunk C, from the set of chunks (C and constructs context tree that stores multi-grained information across various levels. The encoding process for all chunks {Ci)a= is fully paralleled boost the speed. Below, we detail the context tree structure and its Sflciency-enhanced gupendde dynamic construction, and the tree search process.
Context Tree. The motivation build the context tree intuitive and problem-driven. Given a text chunk C; and a task-specific query, the task-related information is often distributed unevenly across the chunk of text. For instance, to summarize a given passage, one should pay more attention to the topic sentences, collect messages from them and rephrase produce the answer, rather than focuses much on narrative details. Whereas in the task of passkey finding. detailed relations are more important than theme paragraphs. To this end, aim for the contextual representations to capture fine- grained details for the relevant portions of the text, while encoding only coarse- coarse-grained information for the less relevant parts. The tree structure is the best fit simulate this process: the
3

---

Published as a conference paper at ICLR 2026
splitting of nodes resembles splitting larger text chunks into smaller ones, from which we can isolate
fine-grained information.
The root node of context tree contains the entire chunk C, {za, .... 2t), where Ip (s S p S t) denotes a token, s and are the start and end index of that chunk; and each other node consists of a sub-sequence of the chunk Cj. Then we introduce how to build the child nodes from a parent node. Specifically, for any non-leaf node that contains 7 tokens (Ta+1,.... Ty+th, ). at the training phase, we split it into two sub-sequences construct its left child and right child as:
Cpazent = [ Clef= = ( Cright =
(I)
Here we adopt a random splitting by setting 6 13 -c and ~ N2), where is predefined
hyperparameter. This random perturbation serves two critical purposes: first, it acts as a form of structural data augmentation during training. preventing the model from overfitting to fixed splitting positions; second, it mitigates the risk of hard-culting semantic boundaries (such as splitting in the middle of a crucial entity name) by forcing the model to learn robust representations regardless of exact boundary locations, as concluded in (Zhang et al., 2025). At test time, the noise e is fixed to zero for deterministic splitting. One can continue this process until arriving at the limited tree depth. Next, building upon this static tree, construct more efficient guery-dependent dynamic tree.
Query-Dependent Dynamic Tree Construction and Search. A task-specific query is typically highly rele- vant to certain tree nodes while being less relevant to oth- ers. For highly relevant nodes, further expansion is nec- essary to extract fine-gre iine-grained information. In contrast, for less relevant nodes, expansion is unnecessary. Thus, in- stead of constructing an entire static context tree as afore- mentioned, we build a query-dependen dynamic tree that expands only the relevant nodes, shown in Figure 2 significantly saving both GPU memory and time.
Preserved Node
Discarded
SellactedUUneelleteed branch
Selected Node
2
Starting from the root node, we perform depth-first splitting and search process. Each node sequence is first divided into two subsequences according to Eq. [ We then use a non-parametric policy ı to decide the next se- lected node based on the subsequences, Flest and Fright, and a query sequence y.
Figure 2: An running example of our tree (depth=3). Each box indexed by i represents the ith iteration of node split and selection.
(2) Here the policy т determines whether the left or right child of the node will be selected. The unselected sibling node is marked as "preserved" and will not be expanded further. Note, the root node is always selected to ensure expansion. For policy т, it is task-specific. Specifically, for language modeling tasks (where the LLM behaves like the non-SFT model), we keep selecting the right branch to simulate the useful A-shape pattern (Han et al. 2024 Ge et al.) 2024):
((Tiers:Teigne) + left or right,
Friggh = rright.
For instruction-fae -folooiinn tasks (where the L.M serves as the supervised finetuned version), where queries are explicit and available, π selects the node with higher semantical similarity to the query:
arg max (ssm(hs,-1 hh)) ((Ieet:::iihhe).. de (laft,ri)hh
where sim(.,) represents the cosine similarity of two vectors. The hidden vector Te at the last position of a sequence embedded by either the lower or upper model. Specifically, this involves a short forward pass through one self-attention layer in the lower model for Niz, and the upper model for hg. Once the selected node is determined, the process continues with that node, repeating the procedure until reaching leaf nodes. At this point, both the left and right child are marked as "preserved".
For each preserved node, we feed its associated context into the lower model to obtain a collection of key-value (KV) states from all M layers, denoted as S (K, V), where K, V Mxxd
4

---

Published as a conference paper at ICLR 2026
represents the key and value states for all M layers. Here, 1 is the sequence length, and d is the hidden dimension. Next, we perform a uniform downsampling along the length dimension to retain only a portion of the KV states, resulting in S' = (K', V"), where K', V E RR xd and l' are the downsampled length. We implement this uniform downsampling via fractional strided selection to ensure equidistant preservation of context tokens without introducing additional parametric pooling layers. The compression ratio a for the node is defined as a 1/1'. For the context tree, we apply a constant compression ratio Ow for all preserved nodes at level w. but the ratio diminishes progressively from top to bottom, i.e., aw aw+1- In our implementation, we set aw 20 This approach creates coarse-tor to-fiee distribution of semantic information from top to down: nodes higher levels possess longer subsequences and are compressed with higher compression ratio, corresponding to more coarse-grained information, while on the contrary, nodes closer to the bottom store fine-grained information. The overall compression ratio 3 of a tree is defined as the ratio of the chunk length |C| to the total length of the compressed KV states:
Clumww Σ No
|C|
(5)
where The is the number of preserved nodes at level w, and Iy is the compressed length of each preserved node at level w. For the convenience of parallel processing, set B to be the same value for all n context trees. Experimental results in Section demonstrate that this compression ratio can reach high 8, significantly improving efficiency.
2.3 UPPER MODEL
The upper model shares a similar architecture with the full-layer version of the base model, except for the inserted cross-attention layers that interact with the lower model, illustrated in Figure]
Position-aware Cross-attention on the Context Tree. In Section 2.2 we can obtain a sequence of tree-structural representations S' = {S .... S/, ) for chunks {Ci]. F==1+ where S (K) V stands for the representations of chunk C,. Since the sequence of chunk keys IC = [K K is produced from ordered chunks {C1, ..., C, their positional information should be aware the chunk level by the query to maintain the global chronological order of the original text. To this end, assign the following chunk-level positional indices Q and K::
Pq = (n, nn.,....,n
Px = =0,00 (0,0,..., 0,1, - -
(6)
[xD]
|C/1/8
|Ca1/8
|Ca1/3
Here view the upper model's query Q as chunk and endow it with the largest positional index because Q is encoded from Xp which is behind all context chunks X in the raw input sequence X. We then apply rotary positional embedding (RoPE) to and K. according these block indices. By doing so, the cross-attention mechanism naturally respects the relative distances between the query and each compressed chunk.
In the cross-attention layer, we calculate attention results between the query Q and concatenated KVs integrate their carried context information into the running context for more coherent lan- guage modeling. To ensure smooth integration with the upper model's original self-attention states, the cross-attention output is fused residually:
O cross (toss_attn(( Q, con.cat(K,;..................... K,,),, (ooncat(VVii--. ViD)).
Training We use the standard language modeling loss during training, which maximizes the log probability of the ground-truth tokens the target sequences Xtar, conditioned on the context X and all preceding tokens Zet from Xp:
f=
Σ log P(xr((c; Xear
For language modeling data, Xisr = XD. i.e., the target tokens are all tokens in Xp, excluding the first token. For mnstruction-folloving data, Xp includes both the instruction Xist and the annotated response X In this case, we set Xtar = Xress meaning that we optimize only for the response tokens, while the instruction text is masked during loss calculation.
5

---

Published as a conference paper at ICLR 2026
3 EXPERIMENT
3.1 SETUP
We highlight some key experimental settings in this section. For more detailed information, please refer to Section A.I
Dataset For language modeling, we follow Yen al. (2024) prepare the training data by sam- pling a subset of 20B (1%) tokens from RedPajama (Together 2023). Due to the copyright issue, the books3 subset is no longer available and thus excluded from our training set. We will give an analysis towards the impact by this in Section A.4 The sampled texts are truncated to 8, 192 tokens
for training. In SFT, we follow Zhang et al. (0255 to prepare the dataset. More details can be found in the appendix.
Training We initialize the upper model with short-context LLMs, such as LLaMA-2-7B. LLaMA- 3-8B and Mistral-7B. The lower model is initialized with the weights of the first M layers from the same LLM,, where we set M = 4 in language modeling and M = 16 in SFT. We train SHAREDLLM accumulation of 16 steps (global batch size is 128) for language modeling and step (global batch size 8) for SFT. The cross-attention layers remain fully tunable, while we opt to train the upper model's top N - M self-attention layers in language modeling as post-injection aggregation for faster convergence.
8x A800 GPU machine. The batch size set to per GPU with gradient
Baseline Methods. For post-pretraining. we compare with other baselines in the same category which have extrapolation abilities, such as Positional Interpolation (Chen et al. (023), YaRN (Peng et al. 2023) and CEPE (Yen et al.] 2024). For SFT, we additionally compare with training-based methods, like StreamingLL.M (Xiao et al. 2024b), LongAlpaca (Chen et al.) 2024), and Activation Beacon (Zhang et al. 2025). as well the advanced inference time method, SnapKV (Li et al.
2024) and OmniKV (Hao et al. 2025).
3.2
MAIN RESULTS
We first report the results on language modeling at various input lengths, which compares the ex- trapolation (length generalization) capability among methods. All perplexity values reported in Ta- bles and 2 are averaged over 1000 examples, except for the 128K length on which we test only 10 examples due to the data scarcity (Yen et al., [2024; /Zhang et al., 2025). The results unveil our model's s strong extrapolation capability- successfully avoids perplexity explosion even when tested at the 128K-token length, though only having seen up to 8K-token sequences during training. Notably, SHAREDLLM outperforms CEPE in nearly all cases except the run at 128K tokens on ProofPile, showcasing the effectiveness of the introduced self-injection mechanism. Moreover, the improvement over Activation -eanoon more pronounced than over CEPE, as CEPE experiences an additional pretraining stage and warmup stage to align the hidden space between its encoder and decoder. In contrast, SHAREDLLM can directly be finetuned from publicly available off-the-shet checkpoints, which significantly reduces training costs.
Long-context Understanding Benchmarks. We continue to test the supervised fine-tuned ver- sion of SHAREDLLM on tasks from LongBench (Bai et al. 2023) and InfiniBench (Zhang et al. 2024b). The two benchmarks comprise variety of long- context tasks and cover various input lengths, which help to quantify both task and length generalizability in a unified manner.
For LongBench, we report the average scores on all 14 English tasks from 5 categories, including single-document QA (SD-QA), multi-document QA (MD-QA), summarization (Summ.), few- shot learningrreasoning (FS) and code-completion (Code). For InfBench, we report the results on three representative tasks: Mathematical Find (Math.F), English Multi-Choice (EN.MC) and Retrieval of Numbers (Ret.N). SIJAREDLLM outperforms or matches other advanced instruction- tuned long-context baselines across all five categories. In Table SHAREDLLM surpasses advanced baselines both benchmarks, showing superior capabilities in tackling extremely long inputs. We note that middle truncationn,,,on,, a common practice previous works, could reduce the difficulty of some tasks and improve the performance (Zhang al. 2025). especially decoder-only models,

---

Published as a conference paper at ICLR 2026
Table Language modeling results (perplexity) the continual pretraining setting on downsam- pled RedPajama. Best results on context-extender 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.
PG19
ProofPile
Base Model
4K
8K
32K
128K
4K
8K
32K
128K
4K
8K
32K
128K
LLaMA-2-322K (Together
2023
3.58
3.34
2.96
0OM
6.93
6.81
7.04
0OM
2.87
2.58
2.47
0OM
Pl (Chen et ul. 20233
3.49
3.21
2.77
0OM
6.97
6.77
6.89
00M
2.77
2.64
2.51
0OM
YaRN Weng of aral. 2022
3.35
3.09
2.58
0OM
6.85
6.62
6.91
0OM
2.82
2.56
2.47
0OM
CEPE
Yen 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
LLaMA-3.1
3.17
3.26
2.63
3.12
6.77
6.52
6.84
6.03
2.58
2.54
2.52
2.48
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 dash ("-").
PG19
ProofPile
CodeParrot
Base Model
Method
4K
16K
32K
100K
4K
16K
32K
100K
4K
16K
32K
100K
StreamingI.I.M
9.21
9.25
9.24
9.32
3.47
3.51
3.50
3.55
2.55
2.60
2.54
2.56
LLaMA-2
LongAlpacu-
9.96
2.54
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
SHAREDI ...
8.68
8.01
7.96
8.24
3.36
3.24
3.21
3.19
2.33
2.25
2.23
2.36
StreamingLLMM
9.58
9.63
9.52
9.55
4.08
4.19
4.16
4.23
2.99
3.05
3.13
3.02
LongAlpaca- 16K
10.21
10.39
00M
3.26
3.34
00M
3.05
3.21
00M
Mistral-7B
Activation Beacon
9.35
9.41
9,39
9.48
3,82
3.64
3.69
3.72
2.96
2.85
2.74
2.92
SHAREDLLM
8.97
9.02
8.98
9.05
3.58
3.38
3.74
2.71
2.68
2.58
2.76
as the relevant information for many tasks is located the head or rear of the entire context rather than the middle part.
3.3 TIME AND MEMORY EFFICIENCY
SHAREDLLM shows
high computational effi- ciency in terms of both speed and GPU memory utilization. As Figure visualizes, we compare average inference time (ms) and memory con- sumption (GB) produced by SHAREDLLM against other advanced baseline models from the archi- tecture types of stream- ing (Zhang al. 2025). encoder-decoder al., (024) and vanilla with positional encod-
Memory
Time (ms)
cow
Infernnce Time (mt)
100
60
8K
16K
32K
128K
48
84
16K
35K
128K
YRRN
HCH CEPE
Activation Boacon
-ID SnunoLLLaAA.Gan)
Figure 3: Comparison of memory usage (left) and total inference time 100 examples (right) between SHAREDLLM and other training time baseline methods. The data is collected by running tiny ex- periment on 100 examples corresponding lengths. "OOM" means out-of-memory exception triggered during test time.
(Yen
ing (Peng et al., (023) that have shown competitive
performance in prior evaluations. YaRN (Peng et al. 2023). which exploits the same fully attention as vanilla auto-regressive LLaMA, has O(L*) time and space complexity. The squared complexity makes it the only model that triggers the out-of-memory (OOM) exception at 128K length. Activation Beacon (Zhang et al. 2025). which adopts the streaming processing paradigm, maintains a minimum constant memory O(l) under different imput lengths L. where is the sliding window length, predefined constant hyperparameter. However, Activation Beacon is incompatible with
7

---

Published as a conference paper at ICLR 2026
Table 3: Evaluation results of different SFT methods on 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.
LongBench Summ.
InfBench En.MC
Model
Model
SDQA
MDQA
FS
Code
Math.F
Ret.N
22.60
24.70
48.10
2.85
22.79
1.85
StreamingL
21.47
22.22
22.20
50.05
48.00
32.31
5,23
LongAlpaca- 16K
28.70
28.10
27.80
63.70
56.00
6.23
25.74
4.87
LLaMA-2
SnapKV
24.05
22.98
17.25
16.11
58.87
9.95
28.83
OmniKV
23.86
22.77
21.09
35.74
49.37
8.81
26.25
3.66
Activation Beacon
28.27
28.45 44.44
25.15
61.00
57.75
12.14
32.05
80.58
SHAREDLLM
28.83
30.93
25.76
63.50
59.93
13.82
33.65
82.79
Buse
23.10
16,20
23.17
48.20
46.10
3.57
20.65
5.41
Streaming
26.19
16.65
23.48
48.23
45.98
7.26
18.84
9.75
LongAlpacaa166k
27.05
17.33
26.18
51.97
52.28
5.41
21.19
12.48
Mistral-7B
SnapKV
22.87
16.43
16.47
19.74
54.09
4.73
16.18
15.71
OmniKV
22.95
16.87
21.36
42.85
41.90
3.81
19.77
14.98
Activation Beacon
29.89
18.04
25.92
52.36
$1.81
14.72
28.71
62.37
SHAREDLLM
30.75
19.81
27.43
54.92
53.74
16.12
29.80
65.73
Base
5.12
7.95
26.13
68.75
56.04
9.93
24.17
49.85
StreumingL LM
6.73
8,56
26.85
68.32
54.83
11.27
35.81
52.85
LongAlpaca-166
21.41
12.45
27.74
70.72
60.05
12.03
25.28
16.13
3LaMA-3
SnapKV
3.31
6.52
19.96
21.05
66.71
7.82
17.73
43.51
OmniKV
4.54
8,21
20.77
32.19
57.92
8,29
21.16
41.10
Activation Beacon
22.08
13.75
29.06
70.67
61.14
15.56
37.17
95.18
SHAREDLLM
22.62
14.32
28.94
71.45
63.57
17.26
36.99
97.31
also due its specialized attention paradigm, which causes sharp increment in inference time as input size grows. CEPE can process past context chunks in parallel, but these chunks must be passed through all its encoder layers (24-layer RoBERTa CEPE) and layer-wise linear projections to obtain the final hidden states for cross-attention, leading to even slower inference speed than non-parallel Activation Beacon. In contrast, SHAREDLLM avoids such redundancy through shallow-layer compression and injection, which exhibits significant speed-up and limited memory consumption.
FlashAttention [Dao 2023
3.4 ABLATION STUDY
Validation of Design Choices. We conduct more experiments on the following ablative settings to validate the rationale behind the design choices: 1) the choice of context information injection lay- ers; 2) other configurations, including the effect from the contextual information collection policy π (only for followiinnn tasks), the noise in node splitting, and the addition of chunk-level po- sitional indices during cross-attention. Regarding the layers selected to transmit KV cache for cross- attention, our implementation, which adapts the continuous bottom strategy and injects the context information in the bottom M layers, obtains the strongest performance over the other two choices, not to mention its outstanding efficiency from the shortest forwarding and back-pros pook-propagating path. For other settings, shown in the bottom rows, performance significantly drops after removing any of the three items. Among the three items, the query query-aware aware information gathering mechanism plays the most crucial role, as removing causes the largest performance drop query-drived tasks. In addition, the decoder's awareness of the sequential order of chunks essential since the key-value pairs produced by the encoder fed in shuffled manner and must be accurately ordered via po- sitional indices. Finally, introducing noise serves as an effective regularizer during training and also contributes improved overall performance.
Architecture Hyperparameters We further examine SHAREDLLL''s sensitivity to some key hyperparameters, such as tree height and token compression ratio. The performance fluctuation on the same two tasks across these configurable hyperparameters confgurrbbeeyyerparaaeeeet reflects the sensitivity to hyperparameters, indicated by the inconsistent trend when tree hight is less than and compression ratio smaller than 8.
depicted in Figure The figure
8

---

Published as a conference paper at ICLR 2026
Table 4: Ablative Studies on dif- ferent configurations of structural information injection. The best values in each category and set- tings consistent with our defaults are highlighted in bold.
Figure 4: Results on arxiv-32K (perplexity) and MD-QA (av- erage F1) when configuring with different tree heights (left) and compression ratios (right) to SHAREDLLM. The values on the horizontal axis represent these individual variables. The value from the default configuration are highlighted in bold.
Configuration
arxiv
MD-QA
Performance tree height
Performance compression ratio
Andv-3XX
andu-32<
Default
2,46
30.93
2.5
MD4A
510
2.55
MD:0A
31.3
Continuous Top
2.61
28.66
915
Interleaving
2.57
29.15
2 2,00
313
w/o query-aware
29.27
w/o noise
2.51
30.08
25
w/o chunk pid
2.49
29.81
233
The left bar chart reveals the importance of the proper tree height. If the height is excessively small, then the tree is undersplit and the chunk size so large that only coarse- grained information is preserved and received by the upper model, while task-related fine-grained information is lost. Conversely, if the tree is too high, then the tree is oversplit and the leaves carry minor details, which are less useful for tasks demanding a global view of the context, the downstream performance also degrades significantly. A similar trend can be captured with the global compression ratio B. Although the perplexity declines when all KVs (3 = 1) are retained for cross-attention as more semantic information can be utilized, the query-aware information gathering ability deteriorates and thus the MDQA score becomes lower.
Besides the effect task performance, we also conduct more experiments to explore how these configurations impact speed and memory Appendix C
4 RELATED WORK
Building Long-context Language Models. There are two prevalent routines empower LLMs' capability to process extremely long text: directly pretraining on long- long-context corpus (Touvron et al.) 2023a Dubey al. 2024 Jiang et al. 2023 GLM et 2024 Yang al., 2025) or adapting short context- window LLMs to longer context lengths via combined various techniques (Tworkowski al., 2024). The former approach consumes enormous amounts of data and computational re- sources, while the latter makes room for researchers to explore more flexible optimization strate- gies (Fu et al., 2024). Adaptation methods intend to mimic short input scenarios when the input text is actually long. Typical implementations include positional encoding (PE) rescaling (Press et al.) 2021 Chen et al. 2023) Peng al., (023) and positional index rearranging (Xiao et al. 2024bb Ding et al.) 2023 An et al., 2024, He et al. 2024 Both adjust the attention weight distribution to resemble the short-input scenarios. Another line of work compresses past tokens sequentially
into dense representations (Chevalier al., 2023) Zhang et al., 2025 Gao et al.] 2025), , serving
as next-step input or storing them in an external retrievable memory (Wu et al.] 2022]] Xiao et al. 2024a). Yen et al. (2024) utilizes small model (Liu 2019) for context compression to enable higher parallelism and minimize latency. However, this heterogeneous architecture necessitates extra pre- training and warmup stages to stabilize the fine-tuning process. Packer al. (2023) introduces system-level approach that introduces a hierarchical architecture along with a predefined set of I/O operations, enabling LLMs to offload, store, and retrieve long-range contextual information while maintaining bounded active context within the model's window. Nevertheless, the effectiveness of this mechanism is fundamentally constrained by the capability of the underlying backbone LLM. In contrast to these works, our method directly tunes off-the-shelf models to compress context into structural representations for query-aware retrieval. Powered by efficient architecture design and a fast-forwand forwaringg mechanism, the whole procedure can be fully paralleled online without excessive memory usage.
Efficient Techniques for Long-context Modeling. vanilla self-attention, the space and time complexity grows quadratically (O(T.)) with the input sequence length I., which usually causes

---

Published as a conference paper at ICLR 2026
out-of-memory (OOM) issues GPU clusters when inputs are extremely long. A straightforward solution is to add parameter efficient fine-tuning (PEFT) modules (Chen et al. 2024 Zhang et al. 2024a) to shrink the size of gradient tensors during back-propagation. Many works strive to reduce the memory footprint of attention computation to enhance computational efficiency. Long- former (Beltagy al. 2020) introduces hybrid attention pattern capture local and global mantic features concurrently. (Katharopoulos et al. (020) designs linearized attention that merely demands (L) space accomplish attention computation. FlashAttention (Dao et al.) 2022 Dao 2023) and Paged/tteention (Kwon et al., 2023) maximize the memory efficiency from the system s perspective. More recently. (Xiao et al., 2024b) discovers the "attention sink" phenomenon and
2025
constructs pseudo sink to address the issue under
Similar attention patterns have
been identified (Han et al. 2024; Ge al. 2024 Zhang et al. 2025) and leveraged as principle when sparsifying attention maps during long- long-context modeling. Our work basically follows the effi- cient design principle in three aspects: 1) lightweight architecture through lower-layer self-inject;on; 2) compact representations via structural information extraction and compression; 3) efficient con- struction and retrieval algorithm based on context tree data structure.
5 CONCLUSION
In this work, we present SHAREDLLM, a novel and highly efficient framework that leverages self- injection mechanism to adapt off-the-shelf short-context LLMs for long-conted modeling. To over- come the computational bottlenecks of standard transformers, we integrate multi-grained context compression and query-aware information retrieval into dedicated dynamic context tree structure. This allows the model to selectively isolate fine-gre reee-graied details for relevant chunks while maintain- ing coarse-grainn grrriee summaries for others. Extensive experiments demonstrate that SHAREDLLM not only successfully extrapolates to sequences exceeding 128K tokens but also excels across var- ious language modeling and downstream cestructioo-followiin tasks. Crucially, it achieves this while maintaining superior memory and inference time efficiency compared to existing streaming and encoder-decoder architectures. Furthermore, because SHAREDLLM utilizes shared key-value mechanisms and is initialized from the same foundational layers, it can be directly fine-tuned from existing pre-trained checkpoints. This eliminates the need for costly post-pretraining or complex feature alignment stages, offering a highly accessible learning paradigm. Looking ahead, we hope this paradigm can be broadly generalized to other short-context L.MSs. Future work will explore ex- tending this scalable context- window extension approach to more complex architectural paradigms, such as multimodal large language models, ultimately paving the way for efficient, infinite-conten processing across diverse data modalities.
ETHICAL STATEMENT
All datasets in this paper are publicly available and have been widely tested in previous works. We do not leverage any synthetic data during training or evaluation. Components SHAREDLLM are initialized from the checkpoint of released open-sourced L.MSs and its security has been sufficiently validated when input queries are safe. We also serutinized many sampled outputs and found no harmful information was generated.
REPRODUCIBILITY STATEMENT
We provide many materials for reproduction in the appendix, including training and testing config- urations, pseudo code snippets, and an anonymous code repository, etc. More evidence, such as the full code and model checkpoints, will be released in a later time.
ACKNOWLEDGMENT
This research was supported in part by NUS Start-up Grant A-0010106-00-000 This work was also supported by the National Natural Science Foundation of China under Grant No. 02320106007,
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, llge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint w3XXiv:2303.87777 2023.
Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, and Lingpeng Kong. Training-free long-context scaling of large language models. In Forty-first International Confer- ence on Machine Learning, 2024.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
Tom B Brown. Language models are few-shot learners. arXiv preprint (rXiv:2005. 2020.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306. (tti::330..15555, 2023.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Lon- glora: Efficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, 2024.
Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 3829-3846, 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, al. Palm: Scaling language modeling with pathways. arxiv 2022. arXiv preprint arXiv:2204.02311, 10, 2022.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, al. Scaling instruction-finetue lan- guage models. Journal of Machine Learning Research, 25(70):1-53, 2024.
Colin Clement, Matthew Bierbaum, Kevin P O' Keeffe, and Alexander A Alemi. On the use of arxiv dataset. arXiv preprint arXiv: 00075, 2019.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference Learning Representations, 2023.
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Proceedings of the 35th Neural Informa- tion Processing Systems Conference (NeurIPS), 2022.
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02488. 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXhv:2407.21783, 2024.
Wikimedia Foundation. Wikimedia downloads. URL hctps: Ttttpss///umgs..iiinnedd.. org
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannanch Hajishirzi, Yoon Kim, and Hao Peng.
Data engineering for scaling language models to 128k context. arXiv preprint arXiv:2402.1011 10171, 2024.
11

---

Published as a conference paper at ICLR 2026
Jun Gao, Qi Lv. Zili Wang, Tianxiang Wu, Ziqiang Cao, and Wenjie Li. Uniicl: An efficient icl framework unifying compression, selection, and generation. In Anmual Meeting of the Associa- tion for Computational Linguistics, 2025. BBLLhtttss//ppp..eeatttttccttttt org/ CorpusID:28003500
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. In The Twelfth International Conference on Learning Representations, 2024.
Tao Ge, Jing Hu, Xun Wang. Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in a large language model. ArXiv, abs/2307.06644. 2023. URL hntpps://api. semant www.hnticscolllc. ar. org/ Corpus ID: 25 259847422
Team GLM, Aohan Zeng, Bin Xu, Bowen Wang, Chenhui Zhang, Da Yin, Diego Rojas, Guanyu Feng, Hanlin Zhao, Hanyu Lai, Hao Yu, Hongning Wang, Jiadai Sun, Jiajie Zhang, Jiale Cheng, Jiayi Gui, Jie Tang, Jing Zhang, Juanzi Li, Lei Zhao, Lindong Wu, Lucen Zhong, Mingdao Liu, Minlie Huang, Peng Zhang. Qinkai Zheng, Rui Lu, Shuaiqi Duan, Shudan Zhang, Shulin Cao, Shuxun Yang, Weng Lam Tam, Wenyi Zhao, Xiao Liu, Xiao Xia, Xiaohan Zhang, Xiaotao Gu, Xin Lv, Xinghan Liu, Xinyi Liu, Xinyue Yang, Xixuan Song, Xunkai Zhang, Yifan An, Yifan Xu, Yilin Niu, Yuantao Yang, Yueyan Li, Yushi Bai, Yuxiao Dong, Zehan Qi, Zhaoyu Wang, Zhen Yang, Zhengxiao Du Zhenyu Hou, and Zihan Wang. Chatglm: A family of large language models from glm- 130b to glm-4 all tools, 2024.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseck-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501. 150.:258...19448.... 2025.
Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. LM-infinite: Zero-shot extreme length generalization for large language models. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I; Long Papers), pp.
Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10. 18651166////// 101533/112022.aaace maa.244o.aaall1..222.. URL https: hatp:::///laathhooogg. org/ 2024. www.:aaacllllan -ong. 2227
Jitai Hao, Yuke Zhu, Tian Wang, Jun Yu, Xin Xin, Bo Zheng, Zhaochun Ren, and Sheng Guo. Om- nikv: Dynamic context selection for efficient long- long-context 1lms. In The Thirteenth International Conference on Learning Representations, 2025.
Zhenyu He, Guhao Feng, Shengjie Luo, Kai Yang, Liwei Wang, Jingjing Xu, Zhi Zhang, Hongxia Yang, and Di He. Two stones hit one bird: Bilevel positional encoding for better length extrapo- lation. In Forty-first International Conference on Machine Learning, 2024.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's 's the real context size of your long-context language models? arXiv preprint arXiv:2404.00 046544 2024.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Prin- ciples, pp. 611 211-626, 2023.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37:22947-229700 2024.
12

---

Published as a conference paper at ICLR 2026
Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang, Chenchen Zhang, Ge Zhang. Jiebin Zhang, et al. A comprehensive survey on long context lan- guage modeling. arXiv preprint arXiv:2503.17407, 2025.
Yinhan Liu. Roberta: A robustly optimized bert pretraining approach. arXiv preprint
1rXiv:197..1162,, 2019.
Alexandra Sasha Luccioni and Joseph D Viviano. What's in the box? a preliminary analysis of undesirable content in the common crawl corpus. arXiv preprint arXiv:2105.027.. 2021.
Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient Ilm pretraining and inference with unlimited context length. Advances in Neural Information Processing Systems, 37:718312 71854, 2024.
Charles Packer, Vivian Fang. Shishir G Patil, Kevin Lin, Sarah Wooders, and Joseph E Gonzalez. Memgpt: Towards ilms as operating systems. 2023.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0000 2023.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. 2021.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long- range sequence modelling. In International Conference on Learning Representations, 2020.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 21(140):1-67, 2020.
Adam Roberts, Colin Raffel, Katherine Lee, Michael Matena, Noam Shazeer, Peter Liu, Sharan Narang, Wei Li, and Yanqi Zhou. Exploring the limits of transfer learning with unified text-to- text transformer. Google, Tech. Rep., 2019.
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. httpp::/githul. www Satnn/atsu--iae stt/ssaaffee anford. alpaca, 2023.
Together. Redpajama: An open source recipe to reproduce Ilama training dataset, 2023. URL Data
TogetherAI. Llamm---77--3k-innttte and fine-tuning for 1lama-2 models with together api, 2023.
URL http::/wwww.oottteeee..bb/bb 1ttp:11aama2222 7b- 32tb--32kknnntttueee
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava. Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint S.Xiv:23077.9228 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv.2307.00928. 2023b.
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miloš. Focused transformer: Contrastive training for context scaling. Advances in Neural Information Processing Systems, 36, 2024.
Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M Rush, and Tri Dao. M1: Towards scalable test-time compute with mamba reasoning models. arXiv preprint arXiv:2504. 10449, 2025.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu. Brian Lester, Nan Du, An- drew M Dai, and Quoe V Le. Finetuned language models are zero-shot learners. In International Conference Learning Representations, 2021.
13

---

Published as a conference paper at ICLR 2026
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing trans- formers. In International Conference on Learning Representations, 2022.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Inflim: Unveiling the intrinsic capacity of Ilms for under- standing extremely long sequences with training-free memory. arXiv preprint arXiv:2402.044 04617, 2024a.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Rep- resentations, 2024b.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Mar- tin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation mod- els. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 4643-4663, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024..mm wow../ooo.....-oo.......... 260. URL, https: htt.h/olo.aaatooog...oo...o ott.:/oooooo.oooo/22/mmaaaaaaaao 260
An Yang, Anfeng Li, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint M0Xv:22505.093888, 2025.
Howard Yen, Tianyu Gao, and Danqi Chen. Long-context language modeling with parallel context encoding. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 2588 2610, Bangkok, Thailand, August 2024. Association for Computational Linguistics. URL https://aclanthology org 2024 aol.oe.g..1-11.11222
Peitian Zhang. Ninglu Shao, Zheng Liu, Shitao Xiao, Hongjin Qian, Qiwei Ye, and Zhicheng Dou. Extending llama-3's context ten-fold overnight, 2024a.
Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Long context compression with activation beacon. In The Thirteenth International Conference on Learning
Representations, 2025. URL https: htt///ppervvieeeee net
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. Benchh: Extending long context evaluation beyond 100K tokens. Lei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 15262-15277, Bangkok, Thailand, August August 2024b. Association for Computational acl long. .814
URL /t..itpp::///ellaaooolooo//////
14

---

Published as a conference paper at ICLR 2026
USAGE OF LLMS
We only use LLMs for writing suggestions, and revising purposes, including basic spelling, gram- mar, polishing, and ETRX code formatting. The major workloads in this paper, such as ideation, coding, experiments, and paper writing, are fully completed by ourselves, while leaving LLMs as an auxiliary assistant instead of major contributor this paper.
A MORE IMPLEMENTATION DETAILS
A.1 TRAINING CONFIGURATIONS
In accord with the settings in previous
works (Chen et al. 2024 Yen et al.] 2024 Zhang et al.) 2025). for con- tinual pretrainng, we initialize both lower and upper model with the base version (pretrained, non-finetuned) of LLMs. In SFT, we use their corre- sponding instruction-tunee version as the start point for training.
Table 5: Configurations for training on both tasks.
Item
Continual Pretraining
Supervised Fine-tuning
training epoch
warmup ratio
0.01
0.001
1/5
I/10
chunk size
1024
512
1/16, 1/8,
Adam (3,,8g)
0.9, 0.999
Zero Redundancy Optimizer (ZeRO) stage 3 from DeepSpeed without of- fload is enabled in both training and inference to allocate the memory usage evenly among GPUs. The cross-attention layers remain fully tunable, while opt train upper model's top N M self-attention layers in language modeling as post-injection aggregation for faster convergence. No parameter efficient fine-tuning (PEFT) techniques, such as LoRA, are applied during the training time, PEFT seriously slows down model's convergence (Chen et al., 2024), which consequently costs longer tuning time than partial parameter fine-tune one-tuning. We adopt Adam optimizer with the starting learning rate le 15-5 and cosine scheduler during training.
We list more training configurations that are not specified in the main text Table The sequential values of a are level-wise compression ratio, from level to level 3.
A.2 DATASET STATISTICS
We use different compositions of training dataset in continual pretraining and supervised fine-tuning below.
Downsampled Redpajama. We follow (Yen et al., 2024 and (Tovrron et al. 2023b to prepare our training set. The proportions of data regarding seven domains in the resulted training set are listed Table All documents are truncated by 192 tokens to fit in the pretraining mode.
Table 6: Dataset composition in our downsampled Redpajama (10B) tokens.
Domain
Proportion (%)
Arxiv (Clement al. 2019
2.5
Books (w/ S3) (Rae et al 2020
4.5
C4 (Roberts al. 2019)
15.0
CommonCrawi [Luccroni & Viviano 2021
67.0
Github
4.5
StackExchange
2.0
Wikipedia (Foundation)
4.5
During pretraining, 4K tokens are fed to the lower model and upper model respectively. The lan- guage modeling loss is calculated on the upper model's token prediction.
15

---

Published as a conference paper at ICLR 2026
Mixed Dataset in SFT. This dataset is directly picked from (hang al. 2025 which is mix- ture of RedPajama and LongAlpaca (Chen et al. 2024). LongAlpaca is composed of Stanford Alpaca instruction-follont
dataset (Taori et al. 2023) and author-curated long-cor cong-contett tasks such as summarization and long-document question answering. We follow (Zhang et al., 2025) to filter samples and only preserve those whose lengths range from 1K 8K. The distribution of samples in terms of length is specified in Table
Table 7: Proportion of samples within each length interval.
Length
<2K
2~44K
4~~K
6~-KK
Proportion
47%
29%
8%
16%
Since we found there was an absence training data in fine- ane-grained retrieval tasks, we additionally sample a small set (200 samples) of data from -lama----B--26- training corpus and add them the SFT data collections. This tiny proportion data plays decisive roles in ensuring SHAREDLLM's non-decreasing accuracy the input context length grows.
A.3 ONLINE SPLIT-AND- -EARRCH ALGORITHM
We provide the pseudo code for the online split-and-search algorithm introduced in Section 2.2 from the splitting of the root node till collecting all key-value states for all preserved nodes and all M layers. The full implementation 25 lines of code.
intricate, which can be readily accomplished with around
For the full set of the core code, please refer https::///// battps://anoyyoouut //w.//ousnnmmuuu.o.mmm 4cpen. science/x/ shared red1 1m_ anony-04B) for details. The code snippet in the entire model.py file can also be found in this anonymous repository.
Algorithm Pseudo code of dynamic Construction-ani and-Search.
trees;
depth: tree depth; chunk_lds: the entire input ida for chunk in shape (8, 14)
gamma: hyper parametez adjust the variance gaussian sampling
selected selected Input__dd
chunk
aalcatad_lengt. chunk_idd. all_kvs
for in range (depth)
sample longths laft right child
if < depth: 1:
haif_length last_lengtà // sigma half_ length gamma
delta - zandom. randa * sigma
L_eeft Iright helf_length
(deltaa, half_ int (delta)
apli: the node into two chiedrer. jeft_input_ide sigat_inputtids= iaput__isl::__e input_ids [iightt:
query_aware
flag indicating
selected nodes determined quezy
is query_awaze:
short formard (1-layer)
representation vessors for the query and two nodes
t_g ugper_model
(ife ids, 1), lower_model selested = argmox h_leftl, sim(h_g, h_righe)
left, h_right
1)
selected 1 deterministic example, can change to 0 oz zandom selection
acted inppt_irp___d_ [lert_inpor_dds, right__nput__oon [selected]
seiected_ [1_left, [[egighht[[eeeecesee
preserved_input__e [left_input_idd,, right_input_ids) selected)
clsa:
preserved_inaut__
cat [aast___ipptt_loo caunk -1), 0)
cur
lower (preserved peresseveed_input_.. pastd))_aat___y_aallele
cur level kvg = downsampl (cur_ level level_kva) kva) avs append feur__evel_kvv
cat: concatemation; chunk: split into the specified mumher of chunks
16

---

Published as a conference paper at ICLR 2026
A.4 CONSEQUENCE FROM THE ABSENCE OF BOOK-S3
Book-S3 is a large dataset of copyrighted published books composed by professional writers in various domains. Due the copyright infringement allegations, all online entries to access this corpus have been removed. Prior studies (Yen et al. (024) have shown that the absence of Book-S3 subsets RedPajama corpus casts a negative impact on language modeling results. Here we simply show the comparison in terms of perplexity when SHAREDLLM is trained with and without Book- S3. As Table shows, the baselines without Book-S3 as part of their continual pretraining corpus show inferior results, which is consistent with the observation in Yen et al. (024). We hypothesize that the root cause is that Book-S3 contains many well-structured and logically sound articles written by expert-level writers, which show higher quality and lower noise than data from other domains. Therefore, plays a great role improving language modeling.
Table 8: Perplexity increment as a negative effect from the lack books3. i represents the values in corresponding rows are reproduced from open-sourced code.
Arxiv
PG19
ProofPile
Model
4K
8K
32K
128K
4K
8K
32K
128K
4K
8K 32K
128K
a.L.MA2-7B (4K)
2.60
OOM
6.49
OOM
2.28
OOM
Books3 involved in training
YaRN-2-128K
3.13
2.96
2.34
0OM 2.91
6.15
6.02
6.32
OOM 5.99
2.70
2.47
2.41
0OM
CHPE
2.86
2.84
2.34
6.60
6.24
6.66
2.22
2.33
2.26
2.23
Books3 not involved training
YaRN-2-128K 3.46
3.30
2.57
0OM
6.83
6.59
7.14
0OM
2.85
2.68
2.63
0OM
CEPE
3.03
3.02
2.51
2.97
6.69
6.40
6.80
6.10
2.38
2.43
2.45
2.39
SHAREDLIMM
2.99
2.97
2,46
2.91
6.59
6.31
6.72
6.00
2.36
2.37
2.41
2.46
A.5
DETAILS OF TEST BENCHMARKS
For all inference results, we report the average values of five runs under different random seeds.
RedPajama To test the long-context modeling capability, we tiny proportion of corpus which has never been seen by the model during the continual pretraining period as the test set. The sampled passages are ensured match the corresponding test lengths. We hold a constant 4K input for the upper model while the left long context is passed to the lower model, akin to what we did during pretraining.
Long Bench Bai ct al. 2023 is the first bilingual (English and Chinese), multi-task benchmark for long context understanding. comprises 21 datasets (16 English and Chinese) across subcat- egories, which aims for more rigorous evaluation of long context understanding. These categories encompass single document QA, multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion. The average length of documents 6,711 words in English and 13,386 characters in Chinese.
Infinity-Bncch Zhang et al., 2024b) extends context lengths previous long-context bench- marks from 10K to more than 1O0K tokens. The benchmark is composed of synthetic and realistic tasks that span diverse domains and bilingual (Chinese and English), such as retrival (Ret.), summa- rization (sum), question answering (QA), code and math.
B MORE EXPERIMENTAL RESULTS
B. RESUTLS ON PASSKEY RETRIEVAL.
We further assess the retrieval capability of SharedLLM on the passkey retrieval, or needle-in- haystack (NIAH) task. Following the settings in Yen et al. (024), train new version of
17

---

Published as a conference paper at ICLR 2026
LL that can perform accurate passkey retrieval from the haystacks of the surrounded non- sense. We follow the examples in (Chen et al.) 2024) to set up the single key-value pair test cases. The results averaged on 10 randomly generated NIAH test samples be observed that SHAREDLLAN enjoys the minimal accuracy decay as length extends compared to other baselines, although it has only seen context within 8K length.
shown in Figure [ It can
Accuracy
100
80
80
40
32
Contest Length
2
SharedLLMM
Figure 5: Accuracy comparison on passkey retrieval (single key-value pair) task,
B.2 COMPARISON BETWEEN DIFFERENT ATTENTION MAPS
The introduced self-injection algorithm can also be regarded as attention remapping process, where we use 'continuous-right and "query-aware" node selection strategy for language modeling and long-context understanding respectively. Meanwhile, many concurrent works (Han al. 2024 Xiao et al. 2024b Ge et al. 2024) observed the special A-shape attention map and took advantage of this for acceleration. In fact, the policy selection is not only intuitive but also with the fundamental support from pilot experiments. We report the results of all these choices below:
Table 9; Pilot studies of branch- policies.
Setting
Arxiv
MD-QA
Default
2.46 (=00001)
30.93 (=0.16)
Random
2.61 (#00.13)
28.85 (=0.45)
Fixed Left
2.49 (1:0.02)
29.62 (=00115)
Query-aware
2.48 (=0.00)
29.27 (=0.12)
A-shape
2.52 (#00.0)
29.48 (=0.18)
The results manifest that the selected policy can produce the optimal performance on the downstream
tasks.
C
TIME AND MEMORY EFFICIENCY
C.1
QUANLITATIVE ANALYSIS
Apart from strong performance on downstream tasks, SHAREDLLM demonstrates high computa- tional efficiency terms of both inference speed and GPU memory utilization. We compare these metrics produced by SHAREDL L.M against other representative models of streaming (Zhang et al.) 2025), encoder-decoder (Yen et al. 2024) and vanilla (Peng et al. 2023) architectures that have shown competitive performance in prior evaluations. The results are visualized in Figure
YaRN (Peng et al. 2023), which only modifies the encoding policy but still applies the vanilla multi-head attention LLaMA, shows squared (O(L²)) time and space complexity. Consequently, it triggers the out-of-rne of-mmmory exception at an early stage (128K tokens). Activation Beacon (Zhang et al. 2025), which adopts the streaming processing paradigm, maintains minimum constant ory O() under different input lengths L. where is the sliding window length. However, Activa-
tion Beacon is incompatible with FtashAttention DDo 2023) also due its specialized attention
18

---

Published as a conference paper at ICLR 2026
paradigm, which causes sharp increment in inference time as input size grows. CEPE pro- cess past context chunks in parallel, but these chunks must be passed through all its encoder layers (24-layer RoBERTa in CEPE) and layer-wise linear projections obtain the final hidden states for cross-attention, leading to even slower inference speed than non-parallel Activation Beacon. con- trast, SHAREDL.LM alleviates such redundancy through shallow-layer compression and injection, which exhibits significant speed-up and limited memory consumption.
We have explained the outstanding efficiency of our model by comparing the memory usage and inference speed with other competitors. In this section, we give a more comprehensive analysis towards the inherent factors that may impact model's s efficiency, including compression ratio B, tree height h, the number of shared layers M and the retrieval-basee policy which requires an additional short forward pass.
Table 10: Inference time under various M with constant / 3 and / 8. Our default setting is highlighted in bold.
M
1
2
4
8
16
Time (8)
6.78
9.35
11.81
16.81
25.85
Memory (GB)
21.04
21.50
22.39
24.08
27.82
C.2 EFFICIENCY RESULTT
We rerun experiments measure the forward time and memory cost from language modeling on 8K tokens, adjusting one variable at a time while keeping others at their default values. The results are shown in Table I0 II and 12 Among these factors, the number of injection layers, M, has the most significant impact on both speed and memory: both memory and latency grows as M increases. an opposite, compression ratio 3 and tree height produces nuances effect on both metrics. For example, if we decreases 3 from 64 to (preserve all KVs), the inference time increases by 6.7% while memory increases by 3%. similar trend observed on experiments with tree height h. We speculate that the reason behind these outcomes are partly from the internal optimization in FlashAttention, which efficiently computes attention blockwisely. When the configuration meets its requirement for block size and hidden dimension (e.g., length is divisible by 256).
Table 11: Inference time under various with constant = 3 and M 4. Our default setting is highlighted in bold. For 3 E {1, 2}. we are not able to set levelwise compression ratios and thus we set the compression ratio same as the 3 for every level the tree.
64
16
2
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
Table 12: Inference time under various with constant highlighted in bold.
8 and M
4.
Our default setting is
Time (s)
11.16
11.55
11.81
11.86
Memory (GB)
19.72
22.42
22.39
22.41
We further investigate the potential overhead caused by the extra short forward path aware splitting-an andiingg-nn-search algorithm. As shown in Table [3] we observe that it incurs around 15% over- head in both time and space. We believe this type of overhead can be further eliminated with more careful optimization of the implementation details.
19

---

Published as a conference paper at ICLR 2026
Comparison of time and memory consumption when query-based retrieval is incorporat- ed/not incorporated in SHAREDLLM. h, M and |3 are fixed at the default values.
Table 13:
Setting
Time
Memory
w/o query-aware retrieval
11.81
22.39
w query-aware retrieval
13.18
25.44
The sky is blue. The doud is white. The passkey is A182C3D4ESF6. sky The doud white...
Iteration
0.79
The sky is blue. The doud white. The grass is green. The passkey A182C3D4E5F6 The sky is
The sky blue. The doud is white The grass green...
white...
0.85 The sky is blue. The cloud is white. The grass is green. The passkay A182C3D4E5F6...
0.14
The blue. The cloud white. The grass is green...
2
0.91
0.17 The sky blue. The cloud white. The grass is green.
The blue. The cloud white. The grass is green. The passkey A182C3DMESF6..
Figure 6: living example 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.
D VISUALIZATION OF TREE SPLITTING PROCESS
We provide a living example to demonstrate how the tree is constructed and how the key chunk is retrieved when performing the passkey retrieval task. In this example, we assume that the length of input text is 8,192 and the passkey located between token id 15 and 20. The process is depicted in Figure 6 As the figure shows, SharedLLM first split the entire input 4,096 tokens into two chunks of 2,048 tokens. Then, it computes the correlation scores between the query and subchunks, and finds the first chunk more correlated (0.79 > 0.18). Hence, it repeats the procedure by splitting that chunk into two subchunks of 1,024 tokens. The process iterates until the maximum tree depth has been reached (suppose dmaz 3), where the chunk size is 512. At each iteration, the chunks where the passkey resides always selected due their higher correlation scores.
20


Reviews:
[{"review_id": "tMy1Qnx4t0", "reviewer": "Reviewer_AWmy", "summary": "**Summary**\n\nThis work presents SHAREDLLM, a novel framework addressing the challenge of efficient long-context inference. Its key innovation is a two-stage process involving a lower model, which compresses context chunks via a Context Tree structure, and an upper model, which decodes  successive tokens from the compressed KV-cache. This architecture enables an extended context window while enhancing inference efficiency.", "questions": "**Suggestions**\n\n(1) Improve the expression in the Abstract and Introduction to enhance readability.\n\n(2) Include the memory usage or inference latency of all compared methods in Tables 2 and 3 to ensure a fair comparison.\n\n(3) If feasible, conduct experiments on larger models to validate the scalability of the approach.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "**Strengths**\n\n（1）This study presents a novel method for addressing the challenges of long-context inference, with its architectural details thoroughly elaborated. \n\n（2）The efficacy of the proposed framework is demonstrated through experiments, which confirm its capability to extend the context window and improve inference efficiency. Furthermore, ablation studies provide evidence for the effectiveness of the introduced context information injection mechanism.", "weaknesses": "**Weakness**\n\n(1) The research motivation is not sufficiently clear. The point raised in the introduction—that \"specialized attention patterns may cause incompatibility with high-performance attention implementations\"—does not adequately motivate the proposed method. For instance, prompt compression methods (e.g., ICAE and UniICL) can also improve long-context inference efficiency without requiring specialized attention mechanisms. However, the authors neither compare their method with these alternatives nor clarify the uniqueness of the problem their approach aims to solve.\n\n(2) The abstract requires smoother expression. For example, the definition of \"self-injection\" is repeated, resulting in redundant content. Additionally, the mention of \"sender\" and \"receiver\" concepts in the abstract may lead readers to believe they are important modules, yet these terms are only mentioned there, which causes confusion.\n\n(3) If no information has been overlooked, Tables 2 and 3 do not report the memory usage or inference latency of the baseline methods. Without aligning these factors, it is difficult to ensure a fair comparison between different methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oWzUOyc3TQ", "reviewer": "Reviewer_4Mtb", "summary": "This paper proposes SharedLLM, a novel two-stage LLM framework designed to extend the context window of short-context language models. Specifically, it consists of 1) a lower model (i.e., compressor) that segments long input sequences and compresses each into hierarchical “context trees” and 2) an upper model (i.e., decoder) that retrieves relevant information from them. The mechanism is called self-injection. Cross-attention from query to the shared KV between the same base model’s layers efficiently injects the compressed/selected input context information. SharedLLM demonstrates long-context generalization (up to 128K tokens trained only on 8K) capability and outperforms baselines such as CEPE and Activation Beacon. In addition, the paper reports accelerated inference and memory usage savings.", "questions": "* The proposed mechanism seems to be designed for the prefill phase; it may have a limited impact during decoding.\n* When positional indices are added in the cross-attention module, are they implemented as sinusoidal position embeddings directly added to the key–value states, rather than as RoPE-style rotations?\n* During fine-tuning, are only the cross-attention components trained while other parameters are frozen, or is the model fully fine-tuned?\n* (suggestion) The near-half-randomness in context tree construction is interesting, but its benefits are not entirely clear. Given that later tokens within a chunk often implicitly encode earlier information, maybe a deterministic segmentation strategy, such as splitting at points of large neighboring vector similarity differences, would be better.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The idea of using a single model for two complementary purposes in long-context processing is novel. Leveraging the same base model to ensure compatibility is both intuitive and efficient.\n* The context tree construction, which identifies relevant segments within each chunk without requiring an additional similarity computation module, is original and practical.\n* The method avoids complex attention mechanisms, enabling the reuse of existing optimization techniques such as FlashAttention.", "weaknesses": "* As shown in Figure 4, performance improvements with respect to tree height and compression ratio appear somewhat inconsistent, suggesting potential sensitivity to hyperparameters. Furthermore, although the method is generally robust within moderate ranges, its reliance on rule-based policy selection and heuristic coarse-to-fine downsampling may limit task generalization.\n* The lack of inter-chunk dependency modeling is understandable for parallelization and optimization efficiency, but it may introduce limitations in long context integration.\n* The Passkey task might be too well aligned with SharedLLM’s query-aware design, making it an easier benchmark for this method. Additional justification or analysis would strengthen the empirical evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NjMoYaVgvW", "reviewer": "Reviewer_xDgb", "summary": "This paper proposes a two-LLM architecture for long-context inference. One LLM compresses the context, and the other LLM operates based on compressed context, achieving 2x speedup without hurting accuracy.", "questions": "* The proposed method --- using LLM to compress context --- reminds me of MemGPT and generative agent paper, where they also uses the LLM to compress the context. How do you compare with these two approaches? I understand that you train the model end-to-end, but algorithm-wise is there any reason to believe your approach is better?\n* Given the accuracy number of your system is close to token dropping based approach, and SnapKV, though it is a strong baseline, is not state-of-the-art based on NVIDIA's KVPress measurement (https://huggingface.co/spaces/nvidia/kvpress-leaderboard), I would suspect other baselines may have higher accuracy than SnapKV and be better than your approach. Is there any reason to believe that is not the case?\n* Regarding training --- why SFT instead of RL?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* A more end-to-end approach that also includes LLM training stage.", "weaknesses": "* This is a crowded space, and needs to be more embedded into related work to justify the novelty of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2H1XNQQnIf", "reviewer": "Reviewer_CrBs", "summary": "To address the limitation of finite context windows in Large Language Models (LLMs), this paper proposes the SHAREDLLM framework. Grounded in the concepts of multi-grained context compression and query-aware information acquisition, SHAREDLLM consists of two stacked models derived from the same base short-context LLM: a lower \"compressor\" and an upper \"decoder\". It enables \"self-injection\" through shared key-value (KV) states at the lower layers, thus avoiding redundant computations. A core innovation is the context tree—a binary tree structure that dynamically encodes long context into coarse-to-fine representations, expanding only task-relevant nodes. Trained on 8K-length sequences, the model can generalize to 128K+ tokens. It outperforms baselines such as CEPE and Activation Beacon on language modeling (perplexity) and long-context understanding tasks (LongBench/InfiniBench), with an inference speed 2× faster than streaming architectures and 3× faster than encoder-decoder architectures.", "questions": "Can experiments on 256K tokens be supplemented? If not, theoretical justification for the \"arbitrary length generalization\" claim is required.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The dynamic expansion of the context tree adapts to task requirements, balancing information retention and efficiency, representing a relatively novel idea. It provides researchers with a lightweight alternative, lowering the threshold for long-context research; efficiency improvements expand the application of LLMs in long-text scenarios.", "weaknesses": "The model is only tested on up to 128K tokens, yet it claims to \"generalize to arbitrary lengths\"—supplementary experiments on 256K tokens or theoretical analysis are needed to support this claim.Experiments only provide quantitative results without qualitative examples (e.g., cases where the model correctly identifies key information in passkey retrieval tasks), making it difficult to intuitively demonstrate advantages.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)