Two Are Better than One: Context Window Extension with Multi-Grained Self-Injection

## Abstract

Abstract The limited context window of contemporary large language models (LLMs) remains a huge barrier to their broader application across various domains.
While continual pre-training on long-context data is a straightforward and effective solution, it incurs substantial costs in terms of data acquisition and computational resources.
To alleviate this issue, we propose 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 such as LLaMA-2, termed upper model and lower model.
The lower model functions as a compressor while the upper model acts as a decoder.
The upper model receives compressed, multi-grained context information from the lower model and performs context-aware modeling on the running text.
Information transfer between the compressor and decoder occurs only at the lowest layers to refrain from long forward paths in the lower model and redundant cross-attention modules in the upper model.
Based on this architecture, we introduce a specialized tree-style data structure to efficiently encode, store and retrieve multi-grained contextual information for text chunks. This structure, combined with a search algorithm, enables rapid encoding and retrieval of relevant information from various levels of the tree based on the input query.
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 ( 2 × 2\times 2 × over streaming, 3 × 3\times 3 × over encoder-decoder architectures).
The code is available at https://github.com/Clement25/SharedLLM .

## 1 Introduction

Since the release of GPT-3 (Brown, ), the rapid advancement of large language models (LLMs) (Chowdhery et al., ; Achiam et al., ; Touvron et al., ; ; Dubey et al., ) has revolutionized the NLP research community and transformed various workflows. Pretrained on trillions of tokens, LLMs exhibit remarkable abilities, such as completing unfinished text or code and following human instructions to perform designated tasks after minimal supervised fine-tuning (Wei et al., ; Chung et al., ).

Despite their impressive capabilities, several factors limit their broader application. One major constraint is the context window size (Hsieh et al., ), which refers to the maximum number of tokens an LLM can process smoothly in a single input.
The length of context window is typically set during pretraining—for example, LLaMA and LLaMA-2 have context windows of 2,048 and 4,096 tokens, respectively. When input text exceeds this limit, LLMs may exhibit erratic behavior during inference. Unfortunately, due to GPU memory constraints, high training costs, and the scarcity of long-context training data, LLMs are often pretrained with relatively short context windows. This limitation severely restricts their use in many daily tasks of large context lengths, such as long document summarization and information retrieval, where much longer windows are needed.

Many researchers endeavour to extending the context window of LLMs while minimizing the time, memory, and training costs during both training and inference. One approach involves post-pretraining LLMs on long-context corpora using hundreds of GPUs (TogetherAI, ; Xiong et al., ). Another line of work explores position interpolation (Chen et al., ; Peng et al., ), which rescales the RoPE (Rotary Position Embedding) frequency and attention scores.
While this method is canonical, it still requires long-context continual pretraining.
For example, YaRN (Peng et al., ) extends LLaMA’s context length to 128K tokens by continuing pretraining on 64K-token sequences using full attention. The combination of parameter-efficient fine-tuning (PEFT) and sparse attention (Chen et al., ) accelerates tuning but faces challenges with extrapolation.
Other approaches like streaming-style architectures (Xiao et al., ; Zhang et al., ; Yen et al., ), maintain a constant-sized memory that operates as a sliding window. While this design significantly reduces memory usage, its specialized attention pattern causes incompatibility with high-performance attention implementations like FlashAttention (Dao et al., ; Dao, ), potentially leading to slower inference speeds. Context compression techniques are also widely explored (Zhang et al., ; Yen et al., ). Although they offer high parallelism, improving speed, they tend to consume significant memory, greatly limiting their real applications.

To strike a balance between efficiency and performance, we propose SharedLLM in this paper. SharedLLM features a lightweight hierarchical architecture that consists of one upper model and one lower model.
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 in the hidden space between the two submodules, SharedLLM can be trained from scratch without extra stages for hidden-space alignment.
The lower model compresses past context information into multi-grained representations.
Through layer-wise connections between the lower and upper model, the compressed information can be passed to the upper model for context-aware language modeling.

For better organization and utilization of this multi-grained information, we further propose a dedicated data structure, dubbed context tree.
The trees operate in parallel, with each tree handling an independent text chunk split from the longer raw input.
The tree nodes contain token sequences of varying lengths, which are encoded and downsampled by the lower model to obtain a set of meaningful representations.
The sequence corresponding to each node is split from its parent, making it a subsequence of all its ancestors.
Meanwhile, we set the compression ratio proportional to the sequence length.
As a result, the nodes at higher levels have longer sequence and larger compression ratio, which express more coarse-grained information after encoding, as opposed to nodes at lower levels.
We further develop an algorithm for dynamic construction and search on the tree given the task-specific query.
The algorithm accelerates the information gathering process by identifying the most relevant text segments and encoding them as fine-grained representations, while converting less relevant parts into coarse-grained representations in a depth-first manner.
In the information injection stage, since the obtained context trees are ordered, we define a rule to assign chunk-level position ids for these encoded keys from the lower model, as well as the queries from the upper model. With these unique chunk-level ids, queries can perceive the relative position of trees to generate high-quality output.

With these design principles, SharedLLM presents extraordinary performance on down-stream tasks with high efficiency.
Specifically, on language modeling tasks, trained on text of up to 8K tokens, our model demonstrates excellent extrapolation capabilities when tackling sequences up to 128K-token length.
On other long-context instruction-following tasks with input lengths ranging from thousands to millions of tokens, SharedLLM delivers promising results comparable to several strong baselines.
In terms of system indicators, all experiments with a maximum length of over 200K can be conducted on a single A800 80G GPU. SharedLLM delivers several times the speed of all baselines while maintaining relatively low memory consumption.

## 2 Related Work

#### Long-context Language Models.

There are two prevalent routines to build LLMs that are capable of processing extremely long text: directly pretraining on large corpus of targeted context length from scratch (Touvron et al., ; Dubey et al., ; Jiang et al., ; GLM et al., ) or adapting short context-window LLMs
to longer context lengths via combined various techniques (Tworkowski et al., ).
The former approach consumes tremendous data and computational resources, while the latter allows for more convenience and flexibility for researchers and developers to explore potential optimization to the default settings (Fu et al., ).
The core idea behind these adaptations is to *mimic* short input scenarios (i.e., length within the model’s text window) when the input length exceeds window size.
Attention map manipulation is the most common approach for this goal, which can be realized via positional encoding (PE) rescaling, such as ALiBi (Press et al., ), positional interpolation (PI) (Chen et al., ) and YaRN (Peng et al., ),
or positional index rearranging (Xiao et al., ; Ding et al., ; An et al., ; He et al., ).
Both directly or indirectly adjust attention scores to be similar as the short-input scenarios so that the model can handily deal with.
Another line of works compress past tokens sequentially into dense representations (Chevalier et al., ; Zhang et al., ) as input at the next step or store them in an external retrievable memory (Wu et al., ; Xiao et al., ) to reduce the input lengths. Yen et al. () utilizes small model such as RoBERTa (Liu, ) for context encoding to boost speed and enable higher parallelism. However, this heterogeneous architecture necessitates meticulous task design for the extra pretraining and warmup stages to stabilize the fine-tuning process.
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-forwarding mechanism, the whole procedure can be fully paralleled online without excessive memory usage, which greatly cuts down the latency during inference time.

#### Efficient Methods for Long-context Modeling.

In vanilla self-attention, the space and time complexity grows quadratically ($O(L^{2})$) with the input sequence length $L$, which can cause out-of-memory (OOM) issues on GPU clusters.
A straightforward solution is to add parameter efficient fine-tuning (PEFT) modules (Chen et al., ; Zhang et al., ; ) to shrink the size of gradient tensors during backward propagation.
Many works strive to reduce the memory footprint of attention computation to enhance computational efficiency.
Longformer (Beltagy et al., ) introduces a hybrid attention pattern to capture local and global semantic features concurrently. Katharopoulos et al. () designs linearized attention that merely demands $O(L)$ space to accomplish attention computation. FlashAttention (Dao et al., ; Dao, ) and PagedAttention (Kwon et al., ) maximize the memory efficiency from system’s perspective.
More recently, Xiao et al. () discovers the “attention sink” phenomenon and proposes streaming-llm to address high perplexity issue in generation under window-attention. Our work basically follows the efficient design principle in three aspects: 1) lightweight architecture through lower layer self-injection; 2) compact structural representations via structural information extraction and compression; 3) efficient construction and retrieval algorithm based on the proposed context tree structure.

## 3 Method

Figure: Figure 1: Overview of SharedLLM. The architecture resembles general encoder-decoder architecture like T5 (Raffel et al., ), 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).
Refer to caption: x1.png

In this section, we first introduce the overall architecture of our proposed SharedLLM in Sec. 3.1, and then elaborate on its two main components, lower model and upper model in Sec. 3.2 and 3.3.

### 3.1 Overview

As illustrated in Figure 1, SharedLLM adopts a hierarchical architecture, akin but not identical to classical encoder-decoder models. The lower model, or the “compressor”, breaks down the long input context $X_{C}$ into smaller chunks that can be processed within limited GPU memory. It then uses the same LLM model to compress each context chunk into compact and structured 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, then integrates the compressed information from the lower model, and finally predicts future tokens in an auto-regressive manner.

The lower and upper models are connected via shared key-value (KV) states and cross-attention modules between corresponding layers. To enable efficient and effective information retrieval and integration, the context information processed by the lower model is organized into a binary tree, referred to as the context tree, which stores multi-grained information at different levels.
This structure allows the upper model to leverage its processed running text to efficiently retrieve relevant information from the binary tree based on a depth-first search algorithm. The retrieved information is then integrated with the input through cross-attention, enabling the model to answer questions or perform language modeling.

Figure: Figure 2: An running example of our tree (depth=3). The digits mark the step indices in the split-and-search procedure.
Refer to caption: x2.png

In the following, we elaborate on the lower and upper model. To begin with, we first define some notations to enhance clarity and readability. Let $X=\{x_{1},x_{2},...,x_{T}\}$ represent the entire input sequence, where $T$ denotes the sequence length.
In comply with previous setting (Yen et al., ), we split these tokens into two continuous parts: $X=\texttt{concat}([X_{C};X_{D}])$, where the past context $X_{C}$ and the running text $X_{D}$ serve as inputs to the lower and upper models, respectively.
Moreover, the past context $X_{C}$ is further divided into $n$ smaller and non-overlapping chunks denoted by ${C_{1},C_{2},...,C_{n}}$, namely, where $C_{1}\cup C_{2}\cup...\cup C_{n}=X_{C}$ and $C_{i}\cap C_{j}=\emptyset,\forall i\neq j$. The chunk size is controlled to fit within the lower model’s context window—e.g., 4,096 tokens for LLaMA-2-7B (Touvron et al., )—allowing the lower model to fully utilize its encoding capacity.

### 3.2 Lower Model

The lower model is a small pretrained LLM, implemented as the first $M$ shallow layers of LLaMA-2.
It independently encodes and compresses each past context chunk $C_{i}$ from the set of chunks $\{C_{i}\}_{i=1}^{n}$, and constructs a context tree that stores multi-grained information across various levels.
The encoding for all chunks $\{C_{i}\}_{i=1}^{n}$ is fully paralleled to boost the speed.
Below, we detail the context tree structure and its efficiency-enhanced query-dependent dynamic construction, and the tree search process.

#### Context Tree.

The motivation to build the context tree is intuitive and problem-driven. Given a text chunk $C_{i}$ 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 to 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, we aim for the contextual representations to capture fine-grained details for the relevant portions of the text, while encoding only coarse-grained information for the less relevant parts.
The tree structure is the best fit to simulate this process: the spltting of nodes resembles splitting larger text chunks into smaller ones, from which we can get more fine-grained information.

In the context tree, its root node contains the entire chunk $C_{i}=\{x_{s},...,x_{t}\}$ where $x_{p}\ (s\leq p\leq t)$ denotes a token, $s$ and $t$ are the start and end index of that chunk; and each other node consists of a sub-sequence of the chunk $C_{i}$. Then we introduce how to build the child nodes from a parent node. Specifically, for any non-leaf node that contains $l$ tokens $\{x_{u+1},...,x_{u+l}\}$,
at training phase, we split it into two sub-sequences for constructing its left child and right child as:

$$ $C_{\texttt{parent}}=\{x_{u+k}\}_{k=1}^{l},\quad C_{\texttt{left}}=\{x_{u+k}\}_ {k=1}^{b},\quad C_{\texttt{right}}=\{x_{u+k}\}_{k=b+1}^{l}.$ (1) $$

Here we adopt a random splitting by setting $b=\lfloor\frac{l}{2}-\epsilon\rfloor$ and $\epsilon\sim\mathcal{N}(0,\sigma^{2})$ where $\sigma$ is a predefined hyperparameter, since random lengths can slightly improve the performance as concluded in Zhang et al. ().
At test time, the noise $\epsilon$ is fixed to zero. One can continue this process until arriving at the limited tree depth. Next, building upon this static tree, we construct a more efficient query-dependent dynamic tree.

#### Query-Dependent Dynamic Tree Construction and Search.

A task-specific query is typically highly relevant to certain tree nodes while being less relevant to others. For highly relevant nodes, further expansion is necessary to extract fine-grained information. However, for less relevant nodes, expansion is unnecessary. Thus, instead of constructing an entire static context tree as aforementioned, we build a query-dependent dynamic tree that expands only the relevant nodes, as shown in Figure 2, significantly saving both GPU memory and time.

Starting from the root node, we perform a depth-first splitting and search process. Each node sequence is first split into two subsequences according to Eq. (1). We then use a non-parametric policy $\pi$ to decide the next selected node based on the two subsequences, ${\bm{x}}_{\texttt{left}}$ and ${\bm{x}}_{\texttt{right}}$, and a query sequence ${\bm{y}}$:

$$ $\pi(({\bm{x}}_{\texttt{left}},{\bm{x}}_{\texttt{right}}),{\bm{y}})\to\texttt{ left}\ \text{or}\ \texttt{right},$ $$

Here the policy $\pi$ 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 $\pi$, it is task-specific. Specifically, regarding language modeling task, since there are no explicit queries (i.e., ${\bm{y}}=\emptyset$), we simply set $\pi$ to be deterministic:

$$ $\pi(({\bm{x}}_{\texttt{left}},{\bm{x}}_{\texttt{right}}),{\bm{y}})\equiv \texttt{right}.$ $$

For instruction-following tasks, such as question-answering, where queries like questions are available, $\pi$ selects the node with higher similarity to the query in the hidden space:

$$ $\pi(({\bm{x}}_{\texttt{left}},{\bm{x}}_{\texttt{right}}),{\bm{y}})= \operatorname*{arg\,max}_{\phi\in\{\texttt{left},\texttt{right}\}}(\mathbf{sim }({\bm{h}}_{{\bm{x}}_{\phi}},{\bm{h}}_{{\bm{y}}})),$ $$

where $\mathbf{sim}(\cdot,\cdot)$ represents the cosine similarity between two vectors. The hidden vector ${\bm{h}}$ at the last position of a sequence is 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 ${\bm{h}}_{{\bm{x}}_{\phi}}$ and the upper model for ${\bm{h}}_{\bm{y}}$. 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 $\mathbf{S}=\{\mathbf{K},\mathbf{V}\}$, where $\mathbf{K},\mathbf{V}\in\mathbb{R}^{M\times l\times d}$ represent the key and value states for all $M$ layers. Here, $l$ 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 $\mathbf{S}^{\prime}=\{\mathbf{K}^{\prime},\mathbf{V}^{\prime}\}$, where $\mathbf{K}^{\prime},\mathbf{V}^{\prime}\in\mathbb{R}^{M\times l^{\prime}\times
d}$ and $l^{\prime}$ is the downsampled length. The compression ratio $\alpha$ for the node is defined as $\alpha=l/l^{\prime}$. For the context tree, we apply a constant compression ratio $\alpha_{w}$ for all preserved nodes at level $w$, but the ratio diminishes progressively from top to bottom, i.e., $\alpha_{w}>\alpha_{w+1}$. In our implementation, we set $\alpha_{w}=2\alpha_{w+1}$. Specific value of $\alpha_{w}$ can be found in Appendix A.1.
This approach creates coarse-to-fine distribution of semantic information from top to down: nodes at higher levels possess longer subsequences and are compressed with a higher compression ratio, corresponding to more coarse-grained information, while on the contrary, nodes closer to the bottom stores fine-grained information.

The overall compression ratio $\beta$ of a tree is defined as the ratio of the chunk length $|C|$ to the total length of the compressed KV states:

$$ $\beta=\frac{\sum l_{w}n_{w}}{\sum l^{\prime}_{w}n_{w}}=\frac{|C|}{\sum l^{ \prime}_{w}n_{w}}$ $$

where $n_{w}$ is the number of preserved nodes at level $w$, and $l^{\prime}_{w}$ is the compressed length of each preserved node at level $w$.
For the convenience of parallel processing, we set $\beta$ same for all $n$ context trees.
Experimental results in Section 4 demonstrate that this compression ratio can reach as high as 8, significantly improving efficiency.

### 3.3 Upper Model

The upper model mainly inherits from the LLaMA architecture, which consists of $N$ (32 for LLaMA-2-7B) self-attention layers with slight modifications.
As illustrated in Figure 1, for each one of the $M$ shallow layers, we add a cross-attention module on the top of the vanilla self-attention layer for information fusion.

#### Position-aware Cross-attention on the Context Tree.

In Section 3.2, we can obtain a sequence of tree-structural representations $\mathcal{S^{\prime}}=\{\mathbf{S}^{\prime}_{1},...,\mathbf{S}^{\prime}_{n}\}$ for $n$ chunks $\{C_{i}\}_{i=1}^{n}$, where $\mathbf{S}_{i}^{\prime}=\{\mathbf{K}_{i}^{\prime},\mathbf{V}_{i}^{\prime}\}$ stands for the representations of chunk $C_{i}$.
Since the sequence of chunk keys $\mathcal{K}=\{\mathbf{K}_{1}^{\prime};...,\mathbf{K}_{n}^{\prime}\}$ is produced from ordered chunks $\{C_{1},...,C_{n}\}$, their positional information should be aware at chunk level by the query.
We assign the following chunk-level positional indices to $\mathbf{Q}$ and $\mathcal{K}$:

$$ $\mathbf{P}_{\mathbf{Q}}=\{\underbrace{n,n,...,n}_{|X_{D}|}\},\quad\mathbf{P}_{ \mathcal{K}}=\{\underbrace{0,0,...,0}_{|C_{1}|/\beta},\underbrace{1,1,...,1}_{ |C_{2}|/\beta},\underbrace{n-1,n-1,...,n-1}_{|C_{n}|/\beta}\}.$ (2) $$

Here we view the upper model’s query $\mathbf{Q}$ as one chunk and endow with the largest positional index, because $\mathbf{Q}$ is encoded from $X_{D}$ which is behind all context chunks $X_{C}$ in the raw input sequence $X$.
We will show in Section 4.4 that this setting also facilitates chunk-level extrapolation and answer text production in downstream tasks.

We then conduct cross attention between the query $\mathbf{Q}$ and concatenated KVs to integrate their carried context information into running context for more coherent language modeling:

$$ $O=\mathrm{cross_attn}(\mathbf{Q},\texttt{concat}([\mathbf{K}^{\prime}_{1};... ;\mathbf{K}^{\prime}_{n}]),\texttt{concat}([\mathbf{V}^{\prime}_{1};...; \mathbf{V}^{\prime}_{n}])).$ (3) $$

#### Training

We use the standard language modeling loss during training, which maximizes the log probability of the ground-truth tokens in the target sequences $X_{\mathrm{tar}}$, conditioned on the context $X_{C}$ and all preceding tokens $x_{<t}$ from $X_{D}$:

$$ $\mathcal{L}=-\sum_{x_{t}\in X_{\mathrm{tar}}}\log P(x_{t}|X_{C};x_{<t}).$ $$

For language modeling data, $X_{\mathrm{tar}}=X_{D}$, i.e., the target tokens are all tokens in $X_{D}$, excluding the first token. For instruction-following data, $X_{D}$ includes both the instruction $X_{\textrm{inst}}$ and the annotated response $X_{\mathrm{res}}$. In this case, we set $X_{\mathrm{tar}}=X_{\mathrm{res}}$, meaning that we optimize only for the response tokens, while the instruction text is masked during loss calculation.

## 4 Experiments

### 4.1 Setup

#### Initialization

We initialize the upper model with LLaMA-2-7B in language modeling and LLaMA-2-Chat-7B in supervised fine-tuning (SFT), in consistent with previous works (Chen et al., ; Yen et al., ; Zhang et al., ). The lower model is initialized with the weights of bottom $M$ layers from the same checkpoint as the upper model, where we set $M=4$ in language modeling and $M=16$ in SFT.

#### Dataset

For language modeling, we follow Yen et al. () to prepare the training data by sampling a subset of 20B (1%) tokens from RedPajama’s all 7 domains (Together, ).
Due to the copyright issue, the books3 subset in Books domain (books3 + PG19) is unavailable and thus excluded from our training set, yet we do not renormalize sampling probability across domains.
As a result, the proportion of PG19 increases in the final dataset compared with the default setting in Touvron et al. ().
The sampled texts are truncated to 8,192 tokens for training.
In SFT, we follow Zhang et al. () to use the same mixed dataset composed of downsampled RedPajama and LongAlpaca (Chen et al., ), where the input length is filtered to range between 1200 to 8192 tokens, following Zhang et al. ().

#### Training

We train SharedLLM on an 8$\times$ A800 GPU machine. The batch size is set to 1 per GPU with gradient accumulation of 16 steps (global batch size is 128) for language modeling and 1 step (global batch size is 8) for SFT.
Zero Redundancy Optimizer (ZeRO) stage 3 from DeepSpeed without offload is enabled in both training to distribute the memory allocation among GPUs.
The cross-attention layers remain fully tunable, while we opt to 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 both training, as PEFT seriously slows down model’s convergence (Chen et al., ), which actually requires longer tuning time than partial parameter fine-tuning to reach the optimum.
We adopt AdamW optimizer with the starting learning rate $1e^{-5}$ and cosine scheduler during training.
The chunk size is set to 1,024 for langauge modeling or 512 in SFT, with tree height $h=3$ and compression ratio $\beta=8$.
For other configurations and hyperparameters, please refer to Appendix A.1 for more details.

#### Baselines

We compare our model with many strong baselines. Specifically, in language modeling tasks, we compare SharedLLM with StreamingLLM (Xiao et al., ), AutoCompressor (Chevalier et al., ), LongAlpaca (Chen et al., ), LongLlama (Tworkowski et al., ) and LongChat (Li et al., ), Activation Beacon (Zhang et al., ) and CEPE (Yen et al., ). In understanding tasks, we additionally include LM-Infinite (Han et al., ) and InfLLM (Xiao et al., ) as baselines. For more details, please refer to Appendix

**Table 1: Perplexity of models trained on mixed dataset. “OOM” means out-of-memory exception raised during inference. Excessively large perplexities ($>10^{2}$) are hidden with a dash (“-”).**
| Model | PG19 | ProofPile | CodeParrot |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 4K | 16K | 32K | 100K | 4K | 16K | 32K | 100K | 4K | 16K | 32K | 100K |  |
| 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 |
| AutoCompressor | 11.80 | - | - | OOM | 4.55 | - | - | OOM | 3.47 | - | - | OOM |
| LongAlpaca-16K | 9.96 | 9.83 | - | OOM | 3.82 | 3.37 | - | OOM | 2.81 | 2.54 | - | OOM |
| LongLlama | 9.06 | 8.83 | OOM | OOM | 2.61 | 2.41 | OOM | OOM | 1.95 | 1.90 | OOM | OOM |
| LongChat-32K | 9.47 | 8.85 | 8.81 | OOM | 3.07 | 2.70 | 2.65 | OOM | 2.36 | 2.16 | 2.13 | 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.98 | 8.15 | 7.96 | 8.24 | 3.36 | 3.24 | 3.21 | 3.19 | 2.33 | 2.25 | 2.25 | 2.36 |

**Table 2: Perplexity of models trained on downsampled RedPajama. LLaMA-2-32K and YaRN-2-128K have seen sequence as long as up to 32K and 64K tokens respectively at training time, while CEPE and SharedLLM are trained on 8K-token sequences. $\dagger$: results run on the reproduced model following original paper and the released code. Notations share the same meanings with the last table.**
| 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 |  |  |  |  |  |  |  |  |  |  |  |  |
| LLaMA-2-32K | 2.60 | 2.51 | 2.32 | OOM | 6.61 | 6.50 | 6.97 | OOM | 2.46 | 2.22 | 2.27 | OOM |
| YaRN-2-128K | 3.13 | 2.96 | 2.34 | OOM | 6.15 | 6.02 | 6.32 | OOM | 2.70 | 2.47 | 2.41 | OOM |
| CEPE | 2.86 | 2.84 | 2.34 | 2.91 | 6.60 | 6.24 | 6.66 | 5.99 | 2.22 | 2.33 | 2.26 | 2.23 |
| Books3 not involved in training |  |  |  |  |  |  |  |  |  |  |  |  |
| 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 |
| SharedLLM | 2.99 | 2.97 | 2.46 | 2.91 | 6.59 | 6.31 | 6.72 | 6.00 | 2.36 | 2.37 | 2.41 | 2.46 |

### 4.2 Main Results

#### Language Modeling.

We first evaluate our models on the language modeling task with sequence lengths ranging from 4K to 128K using a single A800 80GB GPU.
The evaluation covers four datasets: ArXiv, PG19 (Rae et al., ), ProofPile (Azerbayev et al., ), and CodeParrot Tunstall et al. () under two settings that utilize different training datasets. Under each setting we test on three out of the four datasets, respectively.
The results are posted on Table 1 and 2.
All perplexity values in these tables are averaged over 100 examples except for the 128K length, on which we test only 10 examples due to the data scarcity (Yen et al., ; Zhang et al., ).
For the experiments on encoder-decoder and hierarchical models at 4K length, the input is divided by half (2K/2K) and fed separately into their two submodules.
The results show that our model owns strong extrapolation capability—it avoids perplexity explosion even tested on 128K-token length although it only has seen up to 8K-token sequences during training.
In Table 1, SharedLLM outperforms other baselines trained on mixed dataset 3-10%.
In Table 2, for models trained on RedPajama, the perplexity without books3 (bottom) gets a bit worse than those including books3 in the training set, showing the major contribution to language modeling by books3, consistency with the discovery claimed in Yen et al. ().
Notably, SharedLLM outperforms CEPE in nearly all cases except 128K context length on ProofPile, showcasing the effectiveness of structural self-injection mechanism.
Between the two settings, the improvement over Activation-Beacon is more pronounced than over CEPE, because CEPE experiences an additional pretraining stage to adapt the RoBERTa encoder to the RedPajama corpus and a warmup stage to align the hidden space between encoder and decoder.
In contrast, SharedLLM can directly be finetuned from publicly available off-the-shelf checkpoints, which saves a great amount of training efforts.

#### Long-context Understanding Benchmarks.

We continue to test the supervised fine-tuned verion of SharedLLM on many downstream tasks from InfiniBench (Zhang et al., ) and LongBench (Bai et al., ).
Both benchmarks consist of a variety of long-context tasks established from raw and synthetic datasets.

**Table 3: Evaluation of different methods on Math.Find and En.MC from InfiniBench.**
| Method | Math.Find | En.MC |
| --- | --- | --- |
| LM-Infinite | 5.71 | 30.57 |
| StreamingLLM | 6.00 | 32.31 |
| InfLLM | 11.14 | 31.44 |
| SharedLLM | 13.58 | 33.65 |

On InfiniBench, we are interested in the following two tasks: Math.Find asks a model to retrieve a special value specified in the prompt (e.g., minimum, maximum, medium, etc.), which examines both the precise retrieval and query understanding abilities of the model.
En.MC instructs a model to collect key information from a extremely long passage and choose the correct answer from many candidate options.
We compare SharedLLM with advanced baselines capable of extremely long inputs, as shown in Table 3. SharedLLM surpasses these strong baselines on both tasks (2.44 points or 21.9% on Math.Find, 1.34 points or 4.1% on En.MC over state-of-the-arts), showing excellent capabilities in tackling extremely long input.

For LongBench, We report the categorical scores on all 14 English tasks in 5 categories, including single-document QA (SD-QA), multi-document QA (MD-QA), summarization (Summ.), few-shot learning (FS) and code-completion (Code), as shown in Table 4.
SharedLLM outperforms or matches other advanced instruction-tuned long-context baselines across all five categories.
Particularly, we notice that MD-QA.
We note that truncation from the middle could reduce the difficulty of some tasks and improve the performance (Zhang et al., ), especially on decoder-only models, as relevant information for many tasks is located at the beginning or end of the entire text rather than the middle part.

**Table 4: Evaluation of different methods on LongBench. Text samples are truncated to the test length from middle before generation in some models. We particularly highlight these values in “Test Length” column, as well as model’s training length (“Train Length”). Models in the upper rows follow the conventional “pretrain+finetuned” paradigm while models in the bottom rows are directly trained on mixed dataset without continual pretraining to extend the context window in advance.**
| Method | Train Length | Test Length | SDQA | MD-QA | Summ. | FS | Code |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-2-7B-Chat | 4K | 4K | 24.90 | 22.60 | 24.70 | 60.00 | 48.10 |
| StreamingLLM | 4K | 4K | 21.47 | 22.22 | 22.20 | 50.05 | 48.00 |
| LongAlpaca-16K | 16K | 16K | 28.70 | 28.10 | 27.80 | 63.70 | 56.00 |
| YaRN-128K | 64K | 32K | 24.03 | 24.11 | 19.82 | 60.00 | 62.73 |
| Activation Beacon | 8K | 16K | 28.27 | 28.44 | 25.15 | 61.00 | 57.75 |
| SharedLLM | 8K | 32K | 28.15 | 30.93 | 24.28 | 63.50 | 57.95 |

### 4.3 Time and Memory Efficiency

Apart from strong performance on downstream tasks, SharedLLM demonstrates high computational efficiency in terms of both speed and GPU memory utilization.
We compare these metrics produced by SharedLLM against other representative models from the model classes of streaming (Zhang et al., ), encoder-decoder (Yen et al., ) and vanilla Peng et al. () architectures that have shown competitive performance in prior evaluations.
The results are visualized in Figure 3.
YaRN (Peng et al., ), which exploits the same fully attention as vanilla auto-regressive LLaMA, has $O(L^{2})$ time and space complexity. The squared complexity makes it the only model that triggers out-of-memory exception at 128K length.
Activation Beacon (Zhang et al., ), which adopts the streaming processing paradigm, maintains a minimum constant memory $O(l)$ under different input lengths $L$, where $l$ is the sliding window length.
However, Activation Beacon is incompatible with FlashAttention (Dao, ) also due to its specialized attention paradigm, which causes a 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 in 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.

Figure: Figure 3: Comparison of memory usage (left) and total inference time on 100 examples (right) between SharedLLM and other recent baselines. 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.
Refer to caption: extracted/5953228/Figures/mem_and_time.png

### 4.4 Ablation Studies

We consider the following ablative settings to verify the rationale of the design considerations in SharedLLM: 1) tree depth; 2) compression ratio; 3) the collection of context information injection layers; 4) other configurations, including the effect from retrieval policy $\pi$ (only for instruction-following task), the noise in node splitting, and the addition of chunk-level positional indices during cross-attention.

The results are displayed in Table 5, from which we find that both tree depth and compression ratio should be set appropriately to achieve near-optimal performance.
For example, SharedLLM performs best when the tree height is 3.
If the height is too small, i.e., the tree is undersplit and the chunk size is excessively large so that only coarse-grained context information is retained while task-related fine-grained information is not explicit, or too large, i.e., the tree is oversplit and the leaves carry fragmented information which can hardly provide valuable clues for task solving, performance degrades accordingly.
A similar trend can be viewed on global compression ratio $\beta$. While abandoning downsampling KV ($\beta=1$) may bring decline in perplexity, its query-aware information retrieval ability deteriorates.
In terms of injection layer selection, our implementation, which is refer to as *continuous bottom*, injects the context information in the bottom $M$ layers.
In contrary, *Continuous top* injects context information at the topmost $M$ layers (from layer $N-M+1$ to layer $N$).
Interleaving applies cross-attention at regular intervals, such as layer 4, 8, 12, 16…
Among these configurations, SharedLLM wins over continuous top and interleaving on both tasks, indicating the correctness of injection layer selection in SharedLLM.

For other settings, as shown in the bottom rows, removing either of them causes performance drop compared to the default setting, which reveals the contributions of the three design considerations to model’s performance. Among these items, the query-aware information retrieval is the core component for the context-tree so that the performance on MD-QA drops mostly after removing it from the network. The sequential order is similarly important and should be perceived during cross-attention to organize the answer accordingly.

Besides the effect on task performance, we also perform more experiments to explore how these configurations impact speed and memory in Appendix B.

**Table 5: Ablative Studies on different configurations of structural information injection. The best values in each category and settings consistent with our defaults are highlighted in bold.**
| Item | Configuration | Arxiv (32K) | MD-QA |
| --- | --- | --- | --- |
| Tree Height | 2 | 2.51 | 30.15 |
| 3 | 2.46 | 30.93 |  |
| 4 | 2.57 | 29.47 |  |
| Compression Ratio $\beta$ | 1 | 2.43 | 30.55 |
| 4 | 2.48 | 30.28 |  |
| 8 | 2.46 | 30.93 |  |
| 16 | 2.52 | 29.81 |  |
| Injection Layers | Continuous Bottom | 2.46 | 30.93 |
| Continuous Top | 2.61 | 28.66 |  |
| Interleaving | 2.57 | 29.15 |  |
| Other Settings | Default | 2.46 | 30.93 |
| w/o retrieval | - | 29.27 |  |
| w/o noise | 2.51 | 30.08 |  |
| w/o chunk-level pid | 2.49 | 29.81 |  |

## 5 Conclusion

In this work, we present SharedLLM, which leverages a self-injection mechanism to adapt a pair of short-context LLMs for efficient long-context modeling.
By integrating the operations of context compression and key information retrieval into a dedicated binary-tree structure,
SharedLLM excels in language modeling and various downstream instruction-following tasks, while maintaining excellent memory and time efficiency.
Besides, SharedLLM is directly trained from off-the-shelf LLMs, eliminating the need for additional feature alignment steps and making implementation easier.
We hope this learning paradigm can be generalized to other short-context LLMs, offering a scalable approach for a context-window extension to an arbitrary length.

Limitations.
While SharedLLM demonstrates superior performance on both language modeling and long-context benchmarks, as well as high efficiency in terms of time and memory, there are still some limitations. First, although this work strikes a relatively good balance between efficiency and performance at the model architecture level, further improvements could be achieved by optimizing at the system and hardware levels. Second, while a simple and effective retrieval mechanism is implemented in this work, more advanced retrieval techniques, such as BM25 (Robertson et al., ) and Graph-RAG (Edge et al., ), were not explored and may further enhance performance. We aim to pursue these improvements in future research.

## Appendix A Implementation Details

### A.1 Training Configurations

We list more training configurations that are not specified in the main text in Table 6.
The sequential values of $\alpha$ are level-wise compression ratios, from level 1 to level 3.

**Table 6: Configurations for training on both tasks.**
| Item | Language Modeling | Supervised Fine-tuning |
| --- | --- | --- |
| training epoch | 1 | 2 |
| warmup ratio | 0.01 | 0.001 |
| $\sigma$ | l/5 | l/10 |
| chunk size | 1024 | 512 |
| $\alpha$ | $1/16,1/8,1/4$ |  |
| AdamW ($\beta_{1},\beta_{2})$ | 0.9, 0.999 |  |

### A.2 Online Split-and-Search Algorithm

We provide the pseudo code for the online split-and-search algorithm introduced in Section 3.2, from the splitting of root node till collecting all key-value states for all preserved nodes and all $M$ layers.
The code snippet in the entire model.py file can be found in the supplementary material.

Figure: Algorithm 1 Pseudocode of dynamic Construction-and-Search.

### A.3 Dataset Statistics

#### Downsampled Redpajama.

We follow Yen et al. () and Touvron et al. () to prepare our training set. The proportions of data regarding seven domains in the resulted training set are listed in Table 7. Note that for books domain we have excluded S3 due to copyright issues, as we highlighted in Section 4.

**Table 7: Dataset composition in our downsampled Redpajama (20B) tokens.**
| Domain | Proportion (%) |
| --- | --- |
| Arxiv | 2.5 |
| Books (w/o S3) | 4.5 |
| C4 | 15.0 |
| CommonCrawl | 67.0 |
| Github | 4.5 |
| StackExchange | 2.0 |
| Wikipedia | 4.5 |

#### Mixed Dataset in SFT.

This dataset is directly loaded based on the code from Zhang et al. (), which is a mixed version of RedPajama and LongAlpaca (Chen et al., ).
We follow Zhang et al. () to only filter samples whose lengths range from 1K to 8K. The distribution of sample lengths is below.

**Table 8: Proportion of samples within each length interval.**
| Length | $<$2K | 2$\sim$4K | 4$\sim$6K | 6$\sim$8K |
| --- | --- | --- | --- | --- |
| Proportion | 47% | 29% | 8% | 16% |

### A.4 Baseline Methods

We briefly introduce the baseline methods for comparison in our experiments below.

#### StreamingLLM (Xiao et al., 2024b )

observes the “attention sink” phenomenon during inference time when inputs are extremely long. The authors then devise a training strategy by prepending a special token as dedicated attention sink during pretraining. The model can inference in a streaming manner by holding a constant KV cache, which greatly cuts down the time and memory cost.

#### AutoCompressor (Chevalier et al., 2023a )

compresses past text segments into smaller and compact sequences. The compressed sequences are then added to the next segments as prefix to continue the forward pass.

#### LongAlpaca (Chen et al., 2024 )

introduces a new sparse attention pattern ($S^{2}$-Attn) to approximate the full-attention performance. They further design a lora-based finetuning strategy to further reduce the memory cost during training time.

#### LongLlama

is built upon the foundation of OpenLLaMA and fine-tuned using the Focused Transformer (FoT) method (Tworkowski et al., ), which contrasts the context representations from the same document to those from others. Trained from the CodeLLaMA checkpoint, LongLlama successfully extends the context length from 2K to 256K.

## Appendix B Overhead Analysis

In section 4.3, 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 efficiency, including compression ratio $\beta$, tree height $h$, the number of shared layers $M$ and the retrieval-based policy which requires an additional short forward pass.

**Table 9: Inference time under various $M$ with constant $h=3$ and $\beta=8$. Our default setting is highlighted in bold.**
| $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 |

We rerun our experiments to 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 9, 10 and 11.
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.
As an opposite, compression ratio $\beta$ and tree height $h$ produces nuances effect on both metrics. For example, if we decreases $\beta$ from 64 to 1 (preserve all KVs), the inference time increases by 6.7% while memory increases by 3%.
A similar trend is 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 10: Inference time under various $\beta$ with constant $h=3$ and $M=4$. Our default setting is highlighted in bold. For $\beta\in\{1,2\}$, we are not able to set levelwise compression ratios and thus we set the compression ratio same as the $\beta$ for every level of the tree.**
| $\beta$ | 64 | 32 | 16 | 8 | 4 | 2 | 1 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Time (s) | 11.68 | 11.73 | 11.78 | 11.81 | 11.87 | 12.04 | 12.47 |
| Memory (GB) | 22.20 | 22.20 | 22.20 | 22.39 | 22.40 | 22.35 | 22.97 |

**Table 11: Inference time under various $h$ with constant $\beta=8$ and $M=4$. Our default setting is highlighted in bold.**
| $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 |

We further investigate the potential overhead caused by the extra short forward path in query-aware splitting-and-search algorithm. As shown in Table 12, we observe that the extra forwarding incurs around 15% overhead in both time and space.
We believe this type of overhead can be further eliminated with more careful optimization to the implementation details.

**Table 12: Comparison of time and memory consumption when query-based retrieval is incorporated/not incorporated in SharedLLM. $h$, $M$ and $\beta$ are fixed at the default values.**
| Setting | Time (s) | Memory (GB) |
| --- | --- | --- |
| w/o query-aware retrieval | 11.81 | 22.39 |
| w query-aware retrieval | 13.18 | 25.44 |

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. What criteria guided the choice of retrieval policy in the context tree, and how sensitive is SharedLLM’s performance to different retrieval policy settings?
2. How does SharedLLM perform with different compression ratios, tree depths, and injection layer settings?
3. Have you considered testing SharedLLM with alternative context extension methods, such as position interpolation or memory-augmented architectures?
soundness: 3
strengths: 1. Innovative Multi-Grained Context Extension: Introduces a unique approach to extend context windows using a compressor-decoder model architecture, which efficiently handles large context data.
2. Strong Experimental Results: Demonstrates superior performance on several long-context benchmarks, providing evidence of the model's robustness.
3. High Efficiency: Outperforms other methods in terms of speed and memory usage, making SharedLLM viable for large-scale applications.
summary: The paper introduces SharedLLM, a novel approach for extending the context window of large language models (LLMs) by using a hierarchical architecture that pairs two short-context LLMs. In SharedLLM, one model, called the lower model, acts as a compressor that processes past context into compact, multi-grained representations. The other, the upper model, serves as a decoder that integrates this compressed context with current text to predict future tokens in a context-aware manner. Information is passed from the compressor to the decoder through self-injection layers at specific levels, allowing efficient integration without extensive cross-attention.
weaknesses: 1. The paper provides little information on how the model’s performance changes with different tree depths, compression ratios, and injection layers beyond the default settings. Since these parameters are key to achieving a balance between efficiency and accuracy, a sensitivity analysis would be beneficial.
2. While the paper introduces a query-aware retrieval policy in the context tree for efficient information extraction, it lacks a detailed analysis of how different retrieval policies affect SharedLLM’s performance. An ablation study comparing retrieval policies (e.g., different similarity metrics or selection thresholds) would enhance understanding and offer actionable tuning guidance for practitioners.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: Q1. How do you setup the values of hyperparameters of context tree? for example, their depth? are they sensitive to the inference tasks?

Q2. How do you take care of variable length during training and inference? The dynamic NTK and Yarn used the inference-time ratio on this, but I am not so sure about this in your method, thank you in advance as I am out of curiosity. 

Q3. Robustness of different M values. M is an important model architecture base value, but I am not sure the robustness and meaning of setting different M values, and what is the best practice for this?
soundness: 3
strengths: - the basic empirical setup for evaluation in experimental studies is clean on standard benchmarks, from ppl to instruction-following tasks

- ablation studies demonstrate the effectiveness of the approach
summary: The paper presents SharedLLM, which uses two short-context LLMs (derived from the same model family like LLaMA-2) in a hierarchical structure to handle long contexts. My understanding is that this paper has made the following contribution:

- Integrating the concept of "Context Tree": A binary tree structure storing text at different granularities, with higher compression ratios at higher levels

- Architecture: A lower model compresses context into multi-grained representations, while an upper model performs language modeling using this compressed information

- Query-Aware Retrieval: For instruction-following tasks, uses similarity scoring to selectively expand relevant tree nodes

- Layer-wise Connection: Information transfer occurs only at lower layers to reduce computational overhead

Empirical results on language modeling tasks up to 128K tokens and various instruction-following benchmarks.

However, I am a bit concerned with the presentation and also the technical depth explanation.
weaknesses: 1. In my personal opinon, this paper could benefit from the improving the presentation, specifically,

- In figure 1, it is hard for me to capture whether this is a encoder-decoder model or decoder-only model. The role of cross attention is not very clear. I would recommend authors to haven an overview of the model, and then dive into details, and use another figure to explain the context tree to avoid readers getting distracted. At least, you can mark which part is an encoder, and which part is a decoder.

- I really find it hard to understand what is "Tree Cross-attention" in Figure 1.

- The paper emphasizes on "self-injection" but this concept is nowhere in this central figure. 

- there are notations of compression ratio but I am not so sure that they are well explained.

- Authors were arguing "information preservation" but I am not so sure what does this mean?could you elaborate on this and connect this with related work?

2. Experimental setup and results

- There are many hyperparameter setup of the trees constructed, but the intuitive why those hyperparameters were used/set are not well discussed.  See Q1

- What do you handle the drift or variable length during training and inference? See Q2

- Different values of M were set, what's the intuition and what's the best practice? See Q3

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. I'm curious about the result in Table 1 and Table 2. CEPE in line 349 is a re-produced result. Does it mean that the other results are all from the public checkpoint of the according paper? If so, I think the perplexity comparison is meaningless due to different experiment setting. If not so, the perplexity in short context (4k) is highly different, which is also weird.

2. On Infinibench evaluation, why you are only interested in two subtasks? There is not a rationale for that. After all, if you use a bench for evaluation, you usually use the whole subtasks.
soundness: 2
strengths: 1. The paper focuses on an important research area, where LLMs are struggling to achieve effective and efficient long-context training and inference. SharedLLM integrates the compression and retrieval idea, which is a promising research direction in this field.

2. The paper is well written. Although the proposed method includes a long procedure, the method section is organized well to help readers get the core idea. The experiment section includes all the necessary part, including performance, efficiency, and ablation studies.
summary: This paper proposes ShardLLM to reduce heavy training and inference cost in long-context LLMs. Specifically, this method insight is on the multi-grained context compression and query-aware information retrieval. SharedLLM uses  two short-context LLMs, where  the lower model functions as a compressor while the upper model acts as a decoder. The lower model divides the sequence into non-overlapping chunks and makes split-and-search procedure to choose relevant chunks. The selected chunks are then downsampled to reduce KV cache cost. The upper model encodes these keys and values via cross attention with chunk level positional embedding. 

The paper makes experiments on language modeling and supervised fine-tuning to verify the effectiveness of SharedLLM. In language modeling, SharedLLM outperforms other efficient long-context baselines. In supervised fine-tuning, SharedLLM also achieves strong performance in LongBench and InfiniteBench. The efficient experiments show that SharedLLM is cheap and easy to deploy. Finally, the ablation studies discuss the design choice.
weaknesses: 1. The method is very complicated. I admit that "complicated" itself is not sometimes a weakness. However, simple and elegant design is usually working well in model architecture research. Complicated design brings engineering difficulty and optimization problem. SharedLLM uses the last output as the embedding of a chunk which is used for retrieval by the queries. It is highly questionable whether optimization is valid where sparse computation usually has difficulty on gradient estimation.

2. SharedLLM divides context into different chunks and adds downsampling module to compress KV cache. This may hurt performance on some long-context tasks. The experiment part is not strong enough to support SharedLLM's long-context capability. I will give some existed benchmarks and easy examples accordingly:
- Needle-in-a-haystack is now a compulsory evaluation to show long-context's retrieval ability, which is also included in InfiniteBench. However, the experiment only includes two subtasks. The other part is also essential to show model's long-context ability.
- If I query the model to repeat all the previous context, the sparse-divided and compress context may hurt the context information. Moreover, If my first query does not include some context information, and my second query needs that again. Then, SharedLLM will drop it in the first response and can not answer my second question.

3. The experiment also lacks important baselines of KV pruning methods. For example, StreamingLLM is an early baseline which directly drops all the global information while only maintaining the local KV cache. There are many sparse KV cache works, including H2O, SnapKV, FastGen. These works are efficient and memory-friendly, while maintaining parts of long-context capability.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: See Weakness section
soundness: 3
strengths: 1. The SharedLLM introduces an original dual-model architecture with a "context tree" data structure, enhancing the efficiency of context compression and retrieval for large language models.
2. Despite the complexity of the concepts introduced, the paper communicates the workings of the SharedLLM and its underlying mechanisms
summary: The paper introduces SharedLLM, an innovative approach designed to extend the context window of LLMs without incurring the substantial costs associated with continual pre-training on long-context data. The authors claim that SharedLLM achieves comparable or superior results on long-context tasks while significantly reducing memory consumption and increasing processing speed compared to existing methods.
weaknesses: 1. The evaluation only uses the LLaMA-2 model, with no justification for not including more recent or varied models like LLaMA-3.
2. The method proposed in this paper does not seem to outperform other models in Longbench.

Meta-review: This paper proposes a method for longer context length using two LMs, which play role of a compressor and a decoder each other. 

Longer context length without heavy computational costs and data for post-training are critical and important topic. 
However, main concerns are too complicated method, experimental results not solid and insufficent in-depth analysis for supporting the efficacy of the proposed method, which are raised by most reviewers.

AC also agree to reviewers concerns, so recommends rejecting this paper.

Decision: Reject