Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs
Authors: Woomin Song
Seunghyuk Oh
Sangwoo Mo
Jaehyung Kim
Sukmin Yun
Jung-Woo Ha
Jinwoo Shin
Venue: ICLR 2024 poster (2024)
Keywords: Large language models
Long context handling
Token pruning
Primary area: generative models
TL;DR: We propose a computationally efficient method to extend the context limit of large language models.

Abstract: Large language models (LLMs) have shown remarkable performance in various natural language processing tasks. However, a primary constraint they face is the context limit, i.e., the maximum number of tokens they can process. Previous works have explored architectural changes and modifications in positional encoding to relax the constraint, but they often require expensive training or do not address the computational demands of self-attention. In this paper, we present Hierarchical cOntext MERging (HOMER), a new training-free scheme designed to overcome the limitations. HOMER uses a divide-and-conquer algorithm, dividing long inputs into manageable chunks. Each chunk is then processed collectively, employing a hierarchical strategy that merges adjacent chunks at progressive transformer layers. A token reduction technique precedes each merging, ensuring memory usage efficiency. We also propose an optimized computational order reducing the memory requirement to logarithmically scale with respect to input length, making it especially favorable for environments with tight memory restrictions.  Our experiments demonstrate the proposed method's superior performance and memory efficiency, enabling the broader use of LLMs in contexts requiring extended context. Code is available at https://github.com/alinlab/HOMER.

Published as a conference paper at ICLR 2024
HIERARCHICAL CONTEXT MERGING: BETTER LONG CONTEXT UNDERSTANDING FOR PRE-TRAINED LLMs
Woomin Song Seunghyuk Oh Sangwoo Mo Jaehyung Kim
Sukmin Yun 14 Bung-Woo& Ha Jinwoo Shin!
"KAIST Uniiversity of Michigan *Carnegie Mellon University
"Hanyang University ERICA NAVER
ABSTRACT
Large language models (LLMs) have shown remarkable performance in various natural language processing tasks. However, a primary constraint they face is the context limit, i.e., the maximum number of tokens they can process. Previ- ous works have explored architectural changes and modifications in positional encoding to relax the constraint, but they often require expensive training or do not address the computational demands of self-attention. In this paper, we present Hierarchical cOntext MERging (HOMER), new training-free scheme designed to overcome the limitations. HOMER uses a wwvide-and-ooquee al- gorithm, dividing long inputs into manageable chunks. Each chunk is then processed collectively, employing a hierarchical strategy that merges adjacent chunks at progressive transformer layers. A token reduction technique pre- cedes each merging, ensuring memory usage efficiency. We also propose an optimized computational order reducing the memory requirement to logarithmi- cally scale with respect to input length, making it especially favorable for en- vironments with tight memory restrictions. Our experiments demonstrate the proposed method's superior performance and memory efficiency, enabling the broader use of LLMs in contexts requiring extended context. Code is available at ht httpps://githie com/ al aalinlab/iii Dalla///HOE....
INTRODUCTION
In recent years, large language models (LLMs) have performed exceptionally in various natural language processing tasks (OpenAI, 2023; Touvron et al., 2023). Using this capability, multiple emerging applications are using L.MSs as a central component. However, LLMS have fundamental constraint in their context limit, which means the maximum number of input tokens they can process. The ability to handle long contexts is important for real-world applications: chatbots might need to interpret extensive chat histories, while the user could task code comprehension models to process extensive codebases.
A significant challenge in overcoming the context limit is addressing the quadratic computational burden of the self-attention mechanism. Prior works have attempted reduce the computational cost by altering the model architecture, such as introducing sparse attention (Child et al., 2019; Beltagy al., 2020) linearized attention (Kitaev al., 2020; Katharopoulos et al., 2020). Yet, such methods are often not scalable (Tay et al., 2022), and more importantly, they often require extensive model training, making them difficult to use for large-scale models that are prevalent today.
To overcome this issue, recent works have focused on strategies to extend the context limit of pre- trained state-of-the-an LLMs. However, their major focus has been modifying the positional encod- ing (Chen et al., 2023; Peng et al., 2023), which does not address the quadratic computational cost of self-attention, leaving the efficiency concern unaddressed. Reducing the complexity of pre-trained L.LMs remains an important yet underexplored research question.
* "Equal contribution.
I Work done in KAIST.
Work done in Mohamed hin Zayed University of Artificial Intelligence.

---

Published as a conference paper at ICLR 2024
80
Baselines
Baselines
(Ouns)
18
16IN (FR)
HOME (Chrg) YARN (RK)
HOMER (Ours)
HOMER (Ours)
100
memory (GB) (GB)
Average inferencee time (sec.)
16
Ketnoval Annuracy
80
14
50
60-
12
40
40
10:
30
20
20
01
10
4k
Bk 16k
32k
4k
8k
16k
32k
64k
8k
16k
32k
Evaluation langth (Tokens)
Context length
Context langth
(a) Passkey retrieval accuracy.
(b) Peak memory usage.
(c) Average inference time.
Figure 1: (a) Passkey retrieval accuracy on various context lengths, measured with Llama-2-7b-chat
HOMER maintains reasonable performance for context lengths up to 32K tokens. Detailed compar- isons with more baselines are provided in Table (b) The memory requirement for processing long inputs. (c) Average inference time required for generating 100 tokens conditioned on various con- text lengths. All efficiency measurements are done with a single A100 GPU. The baselines include plain Llama, PI, NTK, and YaRN. Peak memory usage of the baselines at 64k is an estimated value, as they do not in a single A100 GPU. Detailed results are provided Table and Appendix
In this paper, we introduce HOMER (Hierarchical cOntext MERging), a novel technique designed extend the context limit while ensuring computational efficiency. HOMER employs divide-and- conquer approach, dividing the long input into manageable chunks. Unlike previous methodologies (Wang et al., 2023; Bertsch et al., 2023), HOMER does not process these chunks independently. Instead, it employs a hierarchical merging strategy. progressively merging adjacent chunks as they are processed along the transformer layers (see Figure for its illustration). To ensure computational efficiency, apply token reduction before each merging stage.
Furthermore, HOMER can be applied to pre-trained LLMs without any further finetuning. This can be beneficial for practical use scenarios where model finetuning is infeasible, such as in envi- ronments with limited computing resources. Also, data preparation may present another challenge for finetuning due to the scarcity of coherent texts with tens of thousands of tokens. For instance, specialized text data should be prepared to finetune an instruction finctuned or hat-ffinttnned model without severely losing its desired properties.
instuuctionfiintuneed
Through extensive evaluation on downstream tasks and perplexity measurements, we demonstrate that HOMER effectively extend pre-trained LLMs to handle long inputs beyond their con- text limits. We first verify the effectiveness of our method on various downstream tasks, includ- ing passkey retrieval and question answering. We further demonstrate the fluency of HOMER by measuring perplexity on long documents. Finally, we highlight the computational efficiency of HOMER as presented in Figure lb and Figure Ic. In all experiments, we illustrate that HOMER can be used with conventional positional encoding scaling techniques (Chen et al., 2023; bloc97, 2023; Peng et al., 2023), and shows improved performance when used on top of these approaches.
In summary, our contributions are as follows:
We present hierarchical context merging: a memory-efficient context limit extension technique, that be used with pre-trained LLMMs without additional training.
We assess the effectiveness of HOMER through experiments on long inputs. In passkey re- trieval experiments, HOMER shows 80.4% retrieval accuracy for 32k inputs, whereas even the best-performing baseline shows only 22.4% accuracy. HOMER also improves the prediction accuracy on question answering by 3% (32.7% - 35.7%), presenting its capability to perform complex reasoning about the content in the extended context length. In language modeling periments, HOMER is the only method showing low perplexity on inputs up to 64k tokens, while the baselines exhibit severe performance degradation for inputs 32k tokens.
2

---

Published as a conference paper at ICLR 2024
We demonstrate the efficiency of our approach and analyze the source of computational sav- ings. Utilizing an optimized computation order, memory requirement scales logarithmically with respect to the imput sequence length, reducing the memory requirement by over 70%.
We show that our method is compatible with the conventional RoPE-scaling methods in a plug-in manner, and using them together achieves an additional performance gain.
RELATED WORK
Long-range transformers. Classical methods for long-range transformers primarily focus on re- ducing the quadratic computational cost of self-attention, such sparse attention (Dai et al., 2019; Child et al., 2019; Rae et al., 2019; Qiu et al., 2019; Beltagy et al., 2020; Zaheer et al., 2020), or linearized attention (Kitaev al., 2020; Katharopoulos et al., 2020; Wang al., 2020; Choromanski et al., 2021). However, these approaches fundamentally change the underlying architecture, and has not been proven be scalable for large models (Tay et al., 2022).
Extension of LLM context lengths. As the context limit of L.MSs has become a critical problem, line of concurrent works emerged, focusing on efficiently extending the context length of LLMs, with most works focusing on Llama (Touvron et al., 2023). Most works focus on scaling the Rotary Position Embedding (RoPE) (Su al., 2021). Chen al. (2023) and kaiokendev (2023) concur- rently discovered the Position Interpolation method (PI), which involves linearly interpolating the -TK-aware
position ids. bloc97 (2023) suggested an NTK aware scaling method (NTK) which further alters the base of RoPE. Peng et al. (2023) further extended NTK-aware scaling, suggesting another RoPE scaling method, YaRN. Several works additionally alter the attention mechanism by either applying mask (Han et al., 2023) or setting an upper bound on the distance between tokens (Su, 2023).
While all methods are known to work without further training, we consider PI, NTK, and YaRN as main baselines they directly compatible with Flash Attention 2 (Dao, 2023), easily enabling memory-efficient inference on long inputs. We also emphasize that our work is orthogonal to these work, and can be further applied top of these methods to further improve performance.
Divide-and-conquer approaches. Approaches to overcome the quadratic computation problem in long context modeling while using the same quadratic self-attention mechanism are to divide the long input into multiple chunks, and most methods process the chunks independently. Inspired by Fusion-in-Decoder (Izacard & Grave, 2020), SLED (Ivgi al., 2023) independently encodes mul- tiple chunks and feeds all of them to the decoder. Similarly, Unlimiformer (Bertsch et al., 2023) introduces k-NN search on the encoder outputs, reducing the number of visible tokens at infer- ence time.
LLMs including Memorizing transformers (Wu et al., 2022) and LongMem (Wang al., 2023) take a similar approach of individually forwarding each chunk, and retrieve the cached hidden states for further use. Most of these methods, except for Unlimiformer, require method-specific finetuning.
Token reduction. Token reduction methods have been widely studied in the field of efficient vision transformers. The key idea of these methods is to progressively reduce the number of tokens order to reduce computation, resulting in more efficient training and inference. Two main approaches in this direction are either pruning the redundant tokens (Liang et al., 2022) or merging them (Bolya et al., 2022). To the best of our knowledge, this is the first work to apply token reduction to extend the context limit of large language models.
3 HIERARCHICAL CONTEXT MERGING
In this section, we illustrate the detailed procedure of our proposed method, Hierarchical cOntext MERging (HOMER); novel and efficient method for extending the context limit of large language models (LLMs). As visualized in Figure 2, HOMER consists of two steps: (i) hierarchical merging of the intermediate hidden states, which we call context embeddings, and (ii) further refinement of the lower-layer embeddings by propagative refinement to produce a compact, fixed-length embed- ding for each layer, which can be seamlessly integrated as a typical kv-cache (Chen, 2022). We first introduce the key idea of hierarchical merging in Section 3.1. Then, we explain propagative refinement in Section 3.2. Finally in Section 3.3, we introduce an optimized computation order to further reduce the memory requirement scale logarithmically with the input length.
3

---

Published as a conference paper at ICLR 2024
Transformer
Context merging
Context merging
Gather according to top-layer indices
Chunk
Chunk
Chunk 3
Chunk
Refined layer-wise embeddings
Hierarchical merging of context embeddings
Propagative refinement
An overview of the proposed hierarchical context merging. We first divide long context into multiple chunks and independently forward them through the early transformer layers. In the intermediate layers, merge multiple chunks by concatenation, forming new, merged chunk. To keep the chunk length bounded, we apply token reduction on the original chunks to make them shorter, prior to merging. This process is repeated until all chunks chunk. Finally, we further refine the lower-layer embeddings to get a compact fixed-length, layer- wise embedding. The embedding can then be used like standard kv-cache (Chen, 2022).
Figure 2:
merged into single
3.1 HIERARCHICAL MERGING OF CONTEXT EMBEDDINGS
We propose a priide-and-conuue approach to handle the quadratic computation of self-attention more efficiently. We divide the long input into multiple chunks and process the local chunks with the usual self-attention. Although some previous studies have adopted a similar approach (Ivgi et al., 2023; Bertsch al., 2023), they independently handle each chunk, possibly restricting the richness of the intermediate embeddings they only have access to local information. In contrast, we pro- gressively merge adjacent chunks as they move through the transformer layers, enabling the chunks to see each other. However, naïvely concatenating the adjacent chunks lengthens the resulting chunk and adds a significant computational burden. Thus we propose to use a token reduction technique to shorten each chunk before merging.
By hicrarchically reducing and merging the context embeddings, our method bypasses the quadratic computations required by the self-attention mechanism. This approach not only aims at computa- tional efficiency but also preserves the richness of the context. The detailed process of hierarchical context merging is carried out as follows.
Division of long context into multiple chunks. The first step of our method is to divide the long context into uniform chunks. However, simply slicing the input into chunks encounters issues in the network's 's initial layers where each chunk cannot see each other. This approach restricts most tokens from accessing the starting instructions, harming the resulting embeddings' quality. Moreover, the tokens at the end miss the global context, which is essential for generating subsequent tokens. We address this by attaching the initial and concluding parts of the prompt to every segment (i.e., treating them as shared prefixes and suffixes), ensuring each chunk contains the instruction and the ending tokens.
Token reduction individual chunks. To keep the resulting chunk's length short after merging, we adopt token reduction techniques, which have been widely studied the field of efficient vision transformers. For vision transformers (Dosovitskiy et al., 2021), dropping the tokens that receive minimal attention from the [CLS] token (i.e. the classification token) is known to be a simple and effective token pruning method (Haurum al., 2023). Inspired by this, propose to prune the tokens receiving minimal attention from the final token in each chunk. the chunks contain affixes, do not prune the tokens corresponding the affixes.
In practice, we identified a position bias in simple attention-based pruning where tokens near the end often receive higher attention weights. To rectify this, We incorporate calibration technique
4

---

Published as a conference paper at ICLR 2024
Algorithm
Memory-efficient computation ordering
procedure HIERAACCHCCLLMMEREEEOO))
node a leaf then
emb_list
LnwardLayerssnndde
Get layer-wise hidden states b Propagative refinement
emb_list + Refine(emb_lin list)
else
1 embs +- CrrrurchicalMerggnnode.....
> Recursively merge children
r r.mbss -- HierarchicalMerge node.right)
emb_men merged + concatenate(L_ene _embs, r. r_embs)
D. Merge chunks Get layer-wise hidden states b Propagative refinement
emb list
emb_list + Refine(emb_lig _ist))
end if
return emb list end procedure
inspired by (Zhao al., 2021). By averaging the attention weights of the last token with respect to the tokens at each position ahead, we derive the bias logits. These bias logits are subtracted from the attention logits during the token reduction to refine token pruning. In summary, the final token pruning is performed by pruning fixed number of tokens according to the significance score Ssig defined as follows:
(i)t((1)
Siig
att
biss
where ssig denotes the significance score of a token at position i, l denotes the token's 's attention
logit, and (diet(1) denotes the bias logit corresponding to the token's 's distance from the final token, dist(i).
Merging chunk embeddings. After shortening, adjacent chunks are concatenated to form a unified chunk. This iterative process of reduction and merging across layers ensures individual chunks converge into single chunk at the final layers. If the chunks include affixes, direct concatenation might lead to redundancy; we address this by simply averaging the duplicates.
Handling position ids. Management of position ids is important design choice for our approach. While dynamically scaling the position ids through conventional methods like PI, NTK, and YaRN viable, these techniques tend to underperform with increased scale factors, being less effective for extended contexts. To circumvent this issue, we reuse the same position ids across different chunks. For affixes, we ensure that corresponding tokens in different chunks assigned the same ids for consistency across the chunks.
3.2 PROPAGATIVE REFINEMENT OF LOWER-LAYER EMBEDDINGS
As depicted in Figure 2, the hierarchical context merging produces embeddings characterized by trapezoidal shape. The higher-layer embeddings concise, while the lower-layer remain extended. To further reduce the computational burden for lower layers, we introduce an additional refinement step after token reduction, called propagative refinement.
The process is straightforward: when token is pruned in the upper layers, the corresponding tokens also pruned in the lower-layer embeddings. Therefore, the pruning decision of the upper layers propagates back to the lower layers. The synchronized pruning across layers results in shorter, uniform embeddings for each layer. For better understanding. we have added a detailed illustration in Appendix D demonstrating the process step-by-step. The rationale behind this is intuition that the upper layers have a better ability to identify the important tokens. Thus, we apply pruning in the upper layers and reflect them in the lower layers. After performing hicrarchical merging and propagative refinement, we end up with standardized, fixed-length embeddings for every layer.
Using the refined embeddings for further generation. Conventional implementation of autore- gressive language models often cache the key and value embeddings order to avoid redundant computation. This technique is commonly known kv-caching (Chen, 2022). As the refined em- beddings have the same length for every layer, they can easily be integrated with the kv-cache implementation by simply replacing it for the generation process.
5

---

Published as a conference paper at ICLR 2024
3.3 COMPUTATION ORDER OPTIMIZATION FOR MEMORY-LIMITED EXVIRONMENTS
In typical Transformer models, all tokens at a given layer are computed in parallel. Following this paradigm, a direct implementation of HOMER would also process multiple chunks concurrently. While such implementation of HOMER inherently requires linear memory with respect to the imput length, we propose a more optimized computation order that allows the memory requirement to scale logarithnically. This efficiency is achieved by strategically reordering the processing steps during the hierarchical merging.
While representing each chunk as node, the hierarchical context merging process can be concep- tualized as a traversal on the binary tree from leaves to the root. By adopting the depth-first search (DFS) algorithm to the computation sequence while executing the propagative refinement, we can achieve a computation cost of a logarithmic scale with respect to the length of the input sequence. For clarity, a pseudo-code representation is provided in Algorithm and Figure 3. A comprehensive proof of the memory requirement can be found in Appendix A. Through this approach, extensive inputs can be processed even in resource-constraines setups.
4 EXPERIMENTS
In this section, we demonstrate the effectiveness of the proposed method, HOMER through exten- sive experiments. Section 4.1 contains the passkey retrieval experiments, originally suggested by Mohtashami & Jaggi (2023). This shows our method's ability to utilize the long context to handle downstream tasks. Section 4.2 contains experiments capability to handle more complex and challenging tasks. Section 4.3 demonstrates that HOMER re- mains fluent, even when conditioned on very long contexts. This is done by measuring perplexity on long documents from PG-19 dataset (Rae et al., 2019). Section 4.4 contains ablation study on the key components that make HOMER effective. Finally Section 4.5, we analyze the memory efficiency of our method. Common setup and baselines. We select Llama-2 as our base model, as it is the most widely used and the strongest open-source large language model. We use the pretrained models for lan- guage modeling experiments, and the chat model for evaluation on downstream tasks, which include passkey retrieval and question answering.
question answering. This shows the model's s
Recent works on positional encoding interpolation have shown their ability to extend Llama's con- text limit without training. We set Position Interpolation (PI) (kaiokendev, 2023), NTK-aware scal- ing (bloc97, 2023), and YaRN (Peng et al., 2023) as our main baselines. As these models scale the positional encoding by a constant factor, we define their context limit as the original context limit (4k tokens for Llama-2) multiplied by the scaling factor. In practice, NTK and YaRN are known to be able to process slightly shorter context than the defined context limit (Peng et al., 2023).
For each task, we report the performance of HOMER applied on plain Llama. To further emphasize that method is orthogonal to the positional encoding scaling methods, and can be applied on top of them, we additionally show the performance of HOMER combined with the best-performing baseline for each task.
4.1 PASSKEY RETRIEVAL
In this section, we investigate if HOMER can effectively leverage the long context to handle down- stream tasks. We evaluate this on the passkey retrieval task, originally proposed by Mohtashami & Jaggi (2023). In this task, the model is asked to retrieve a random number (called pusskey) hidden inside distracting texts. The task is widely used to evaluate the maximum context length that the model can effectively handle.
To evaluate the performance at different input lengths, we evaluate the models with inputs of lengths 4k, 8k. 16k, and 32k tokens. We report the retrieval accuracy in Table 1. The result demonstrates that HOMER successfully maintains high accuracy of around 80% for context length up to 32k tokens which is 8 times longer than the pre-trained context length, while significantly outperforming every baseline. Furthermore, it is also evident that the performance can be further improved by applying HOMER top of YaRN, the nest-performing baseline.

---

Published as a conference paper at ICLR 2024
Table 1: Retrieval accuracy on passkey retrieval. Average accuracy on 500 samples are reported. The best values are in bold, and the second-best values are underlined. Empty values indicate NaN.
Context limit
Llama-2-7b-char
Llama-2-130-aae
Method
4K
8K
16K
32K
4K
8K
16K
32K
Plain
4k
1.000
0.000
1.000
0.000
0.000
0.000
Plain HOMER
None
0.990
0.924
0.890
0.776
1.000
0.944
0.882
0.804
PI
8k
0.432
0.356
0.000
0,600
0,544
0.000
0.000
16k
0.006
0.006
0.006
0.000
0.022
0.028
0.018
0.000
NTK
8k
0.812
0.000
0.000
0.000
0.866
0.000
0.000)
0.000
16k
0.516
0.652
0.000
0.000
0.626
0.692
0.000)
0.000
32k
0.106
0.194
0.162
0.000
0.286
0.570
0.442
0.000
YaRN
8k
0,996
0.002
0.000
1.000
0.464
0.000
0.000
16k
0.844
0.756
0.000
0,000
0,980
0.952
0.214
0.000
32k
0,702
0.654
0.696
0.002
0.926
0.888
0.836
0.026
0.678
0.358
0.148
0.026
0.902
0.826
0.364
0.224
YaRN + HOMER
None
0.996
0.984
0.876
0.802
1.000
1.000
0,974
0.860
4.2 QUESTION ANSWERING
In this section, we push HOMER further and eval- uate its performance on a more challenging task: question answering based on long documents. To
Table 2: Accuracy in question answering, as evaluated on the QuALITY validation set. The best results are highlighted in bold.
this end, we measure the model's performance on the validation set of QuALITY (Pang al., 2021).
For baselines with limited context length, the input documents are clipped to fit in the context limit. For NTK and YaRN, we further clip the documents to be 3/4 of their context limit, as they only ca- pable of handling inputs slightly shorter than the claimed context limit. This observation can also be found in Section 4. as NTK and YaRN mod- els could not handle inputs that are long their context limit. For HOMER experiments, we feed
Method
Context limit
Accuracy
Plain Plain + HOMER
4k 0.327
None
0.358
PI
8k
0.366
NTK
8k
0.379
YaRN
8k
0.310
NTK + HOMER
None
0.388
the full context into the model as HOMER has no hard limit on the maximum context length.
We report the prediction accuracy in Table 2. As evident from the table, HOMER effectively ex- tends the context limit, enjoying over 3% of accuracy gain compared to plain Llama. The perfor- mance is further improved when applied on top of the best-performing positional encoding scaling method (NTK), achieving 38.8% accuracy. This demonstrates that language models extended with HOMER could potentially perform more sophisticated reasoning based on the extended context.
4.3 LANGUAGE MODELING
In this section, we investigate the language modeling fluency of HOMER using the perplexity metric. To this end, we sample 25 long documents from the PG-19 dataset (Rae et al., 2019) and measure the perplexity on documents truncated to specified evaluation lengths.
The core of our methodology is the compression of long context into short embeddings. Aligned with this premise, perplexity measured iteratively: preceding contexts are condensed with HOMER, and the perplexity of the subsequent segment is deduced based on these compressed con- texts. Throughout this procedure, the initial 4k tokens were evaluated using unmodified models, with subsequent tokens assessed in 2k token increments. In every experiment, the last 100 tokens of the input are treated as a suffix.
As illustrated in Table 3, HOMER maintains minimal perplexity values across long documents span- ning up to 64k tokens. A more fine- fine-grained perplexity plot is provided Appendix F. While all other methods either suffer from significant degradation beyond certain thresholds (attributed to
7

---

Published as a conference paper at ICLR 2024
Table 3: Perplexity of 25 long documents from PG-19 truncated to the evaluation length. The best values are in bold, and the second-best values are underlined. Empty values indicate NaN.
Context limit
Llama-2-7b 16K
Llama-2-13b
Method
4K
8K
32K
64K
4K
8K
16K
32K
64K
Plain
4k
6.72
6.14
> 10a
> 10"
103
10
Plain HOMER
None
6.72
7.29
7.78
8.43
9.64
6.13
6.60
6.87
7.13
7,59
PL
7.91
8,19
6.96
7.19
102
10
10
16k
10
10
10
10a
10
NTK
6.97
10a
6.26
9.62
10a
10a
10a
16k
7.59
7.95
10a
10a
6.76
7.05
10
10a
10a
32k
8.42
8.97
9.76
10
10a
7.42
7.90
8.45
10
YaRN
6.79
7.40
6.19
6.59
100
10
10
16k
7.00
7.32
8.98
6.36
6.65
7.83
10a
10
32k
7.50
8,05
8.78
10
6.65
7.05
7.40
8.85
10a
64k
8.49
10
10
10
7.17
8.32
10
YaRN HOMER
None
6.79
7.09
7.52
7.95
8.83
6.19
6.51
6.78
7.02
7.44
lower scaling factors) show heightened perplexity even within shorter contexts, HOMER steadily maintains minimal perplexity across extended contexts. This suggests that HOMER is the only method that maintains reasonable fluency even when conditioned on very long contexts. Moreover, HOMER can be seamlessly integrated with conventional positional encoding scaling techniques to further improve performance. As evident from Table 3, applying HOMER on top YaRN yields lower perplexity.
4.4 ABLATION STUDIES
Table 4: Ablation different components. We report the passkey retrieval accuracy for 500 sam- ples, evaluated on 16k contexts. The best values are highlighted bold.
(a) Token pruming criteria.
(b) Lower-layer embedding refinement.
Method
Accuracy
Method
Accuracy
Random
0.006
No refinement
0.002
Attention-based top-K
0.056
Random
0.116
+ calibration
0.890
Layer-wise top-K 0.040
Propagative refinement
0.890
In this section, we demonstrate the effectiveness of the design choices made for our method. Specif- ically, we focus on (1) the proposed token pruning criteria and (2) the method for refining the lower- layer embeddings after applying hierarchical merging. Following the settings of Section 4.1, compared the retrieval accuracy of each candidate.
Effectiveness of our token pruning method. To reduce redundant tokens in the intermediate Trans- former layers, we define a calibrated significance score based on the attention weights each token receives from the last token the chunk. In the pruning procedure, K tokens with the lowest sig- nificance scores are dropped. We demonstrate the effectiveness of this criteria by comparing it to simple baseline, which randomly selects which token to drop. We additionally report the per- formance with uncalibrated pruning criteria to further emphasize the effectiveness of significance weight calibration. As illustrated in Table 4a, the use of attention ditention-based significance scores and calibration provide an effective proxy for determining the importance of given tokens.
Effectiveness propagative refinement. Another key component of our method is the refinement of the lower-layer embeddings. described as propagative refinement. To evaluate its effectiveness, we compare its performance with three alternative approaches (i) not refining the lower-layer embed- dings, (ii) gathering random tokens, and (iii) gathering tokens according to their significance score at each layer. As illustrated Table 4b, propagative refinement achieves the best performance. We credit this to the ability of upper transformer layers understand high-level information, with their attention weights successfully representing token significance. By selectively providing more significant tokens, propagative pruning reduces computation while improving performance.
8

---

Published as a conference paper at ICLR 2024
Table 5: Peak memory usage for long inputs. All measurements are taken on single A GPU, with Flash Attention 2 (Dao, 2023) applied. The baselines include plain Llama, PI, NTK, and YaRN. We report a single value for all baselines as they share the same memory requirement.
Peak GPU memory (GB)
Setup
4k
16k
32k
64k
Baselines
18.2 16.3 (-10.8%)
23.6
34,4
56.7
80
HOMER
16.5 (-30.1%)
16.7 (-51.4%)
17.6 (-68.9%)
21.3 (at least -73.4%)
Baselines
19.2 (+5.5%
20.1 (-14.6%)
20.3 (-41.1%)
20.8 (-63.4%)
22.5 (at least -71.8%)
4.5
COMPUTATIONAL EFFICIENCY
In this section, we discuss the computational efficiency offered by our methodology, with primary focus on memory efficiency. We first demonstrate the computational efficiency of our method by measuring the peak GPU memory usage while processing long inputs. In the following part of the section, we discuss the four key mechanisms that bring efficiency gains.
The peak memory usage for HOMER and baselines is illustrated in Table 5. Note that report single number for all baselines (Plain Llama, PI, NTK, YaRN) because the baselines only modify the positional encoding, making difference in the peak GPU memory usage. For fair com- parison, all methods are tested with Flash Attention 2 (Dao, 2023) enabled. As shown in the table, HOMER significantly reduces memory requirements, reducing the peak memory usage by over 70% when running inference on 64k inputs.
The first source of our efficiency gains the chunking mechanism. We circumvent the quadratic computation associated with self-attention by processing each chunk separately at the earlier layers. Token reduction is our second source of computation reduction. As our algorithm progressively reduces the number of tokens, fewer tokens have to be processed in the upper layers, reducing the computational overhead. The third source of computation saving is that HOMER outputs con- cise embeddings, optimizing the subsequent self-attention computation during the generation phase. Compared to naive forwarding of the complete input, our compact embeddings significantly mini- mize the size of kv-cache, thus optimizing the computation process. Finally, the memory require- ment is further reduced from linear to logarithmic with respect to the input length, thanks to the optimized computation ordering described in Section 3.3. Additional discussion the inference speed is provided in Appendix E.
5 CONCLUSION
In this paper, we introduced Hierarchical cOntext MERging (HOMER), a novel method that effi- ciently addresses the context limit issue inherent in large language models (LLMs). By employing a strategic divide-and-conquer technique, HOMER prunes redundant tokens, creating compact em- beddings while maintaining the richness of information. This approach, validated by our experi- ments, has proven to be memory-efficient and effective, enabling the handling of extended contexts up to 64k tokens with significantly reduced memory requirements.
Limitations and future work. Although our work focuses on training maining-free extension of context limit, there is no fundamental limit on our method making finetuning impossible. We believe that further improving our method with small-data finetuning can additionally boost performance, and the resulting model would enjoy both the extended context limit and reduced memory requirements.

---

Published as a conference paper at ICLR 2024
ETHICS STATEMENT
While large language models (LLMs) have become new paradigm of Al academia and business industry by showing remarkable contributions, their critical limitations still remain such as halluci- nation, biased content generation, and unintended toxicity. The proposed research on long context windows does not address this limitation directly. The influences of longer context limits on the LLM limitations should be more explored, which is one of significant future research directions.
REPRODUCIBILITY STATEMENT
We outline the implementation details (including detailed algorithm, prompt design, and hyper- parameters) and experiment setups (including tasks, datasets, and metries) Section and Ap- pendix B. We also release the source codes.
ACKNOWLEDGEMENTS AND DISCLOSURE OF FUNDING
This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) telligence Graduate School Program (KAIST); N0.2021---,0068, Artificial Intelligence Innovation Hub; No. 202-000099,,,, Few-shot Learning of Casual Inference in Vision and Language for Deci- sion Making).
Artificial In-
REFERENCES
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- long-document transformer. arXiv preprint arX5v:200.005750, 2020.
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew R Gormley. Unlimiformer: Long-range transformers with unlimited length input. arXiv preprint arXiv:2305.01625, 2023.
bloc97. Ntk aware scaled rope allows Ilama models have extended (8k+) context size with- out any fine-tuning and minimal perplexity degradation. https:://// //www//wwwwww// reddit. com/ r/ r/Local MLL.aM////omeetttt /1//z/1// 11l11/llz///// h2wp///ttaaaaaaremm.. httt:/ttaaareeelccllleeeeee 120scaled_ropee
11 @lama_models__0 have/, 2023.
Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. arXiv preprint arXiv:2210.09461, 2022.
Carol Chen. Transformer inference arithmetic. Transforme--iffeeeeee-- thmetic/. 2022.
https://kigg: hyt.lpppp.///og/////////
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595, 2023.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint 1rXiv:1944.10500 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint 2XXv::0009.1474. 2021.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv:1901.02860. 2019.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691 2023.
10

---

Published as a conference paper at ICLR 2024
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: // //operreview. net. detoommmddyy cbFdNTTy.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308.16137, 2023.
Joakim Bruslund Haurum, Sergio Escalera, Graham W Taylor, and Thomas B Moeslund. Which tokens to use? investigating token reduction in vision transformers. arXiv preprint 2023.
arXiv:2308. 046567730 04657,
Maor Ivgi, Uri Shaham, and Jonathan Berant. Efficient long-text understanding with short-text models. Transactions of the Association for Computational Linguistics, 11:284-299, 2023.
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering. arXiv preprint arXiv:2007.01282, 2020.
kaiokendev. Things i'm learning while training superhot. httts::/kaaickenev.. github. 2023.
wttttiifexteddngg ng- www--ooottttttooooo -o
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference on Ma- chine Learning, 2020.
Nikita Kitaev, tukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0445// 2020.
Youwei Liang. Chongjian Ge, Zhan Tong, Yibing Song, Jue Wang, and Pengtao Xie. Not all patches are what you need: Expediting vision transformers via token reorganizations. arXiv preprint arXiv:2202.07880, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture 2016.
models. arXiv preprini
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305. 2023.
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303.08771 2023.
Richard Yuanzhe Pang, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma, Jana Thompson, He He, et Quality: Question answering with long input texts, yes! arXiv preprint 1 2021.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint umXiv:2309.00077 2023.
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise 2019.
self-attention for long document understanding. arXiv preprint
Jack W Rac, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05500 2019.
Uri Shaham, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy. Zeroscrolls: A zero-shot benchmark for long text understanding. arXiv preprint arXiv:2305 /4/96, 2023.
Jianlin Su. Rectified rotary position embeddings. 2023.
Jianlin Su. Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. arXiv preprint a0X 004009,44, 2021.
11

---

Published as a conference paper at ICLR 2024
Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Won Chung. William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q Tran, Dani Yogatama, and Donald Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling? arXiv preprint arXiv:2207.1055// 2022.
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:23077009288 2023.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
Weizhi Wang. Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Aug- menting language models with long-term memory. arXiv preprint arXiv:3306007774, 2023.
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. arXiv preprint arXiv:2203.088: 08911,,,, 2022.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang. et al. Big bird: Transformers for longer sequences. Neural Information Processing Systems, 2020.
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning, pp- 12697-12706. PMLR, 2021.
12

---

Published as a conference paper at ICLR 2024
APPENDIX
Memory-efficient computation order
14
Implementation details
15
C Prompts for downstream tasks
16
D Illustration of propagative refinement
17
Inference speed analysis
17
F Perplexity plot for language modeling experiment
18
G Perplexity evaluation on downstream tasks
19
13

---

Published as a conference paper at ICLR 2024
A MEMORY-EFFICIENT COMPUTATION ORDER
In this section, we outline the proof for the logarithmic memory requirement of the proposed memory-efficient computation ordering suggested in Section 3.3.
Transformer Layers
La layers
L1 layers
4
Lo layers
C tokens
C tokens
C tokens
C tokens
Figure 3:
Hierarchical context merging process conceptualized as binary tree. The top-left num- bers of each node denote the memory-e -efficeenfiieent computation order. Note that propagative refinement must be applied after processing each node to enjoy the optimized memory usage.
A.1 PRELIMINARIES
Problem setup. We conceptualize the hierarchical context merging process as a binary tree. For example, Figure 3 illustrates merging process with 4 input chunks.
Constants. L, refers to the number of layers used for processing a chunk at binary tree height i. I==E
La is the total number of network layers. C is the maximum chunk size. M is the memory required for storing a key-value pair for a single token in a single layer.
Remarks. As the chunk size is bounded, the memory required for forwarding single chunk through a single layer can be treated as constant. Therefore, it suffices to consider the memory required for storing the key-value pairs each layer.
Let Fina VinalMem(h h) be the memory occupied after processing binary tree of height h. As prop- agative refinement reduces the intermediate hidden states be C/2 tokens long, FinalMem(A) is bounded as follows.
C La M
1 LCM 2
PinalMent
2
A.2 PROOF
Proposition. Let PeakMem(h) be the peak memory usage for processing a binary tree of height h. Then,
PeakMem(h) < h +1) LCM.
We prove the proposition using induction. First, consider the leaf node where h = 0. As C tokens passed through L,, layers, the peak memory usage given follows, proving the base case.
PeakMem(O) = L.CM < LCM
Now consider a non-leaf node with h > Processing of non-leaf node consists of three steps: (1) sequentially processing two child nodes, (2) obtaining merged chunk and forwarding it through T.,, layers, (3) applying propagative refinement on the >h L, lower-layer hidden states. As step
14

---

Published as a conference paper at ICLR 2024
(3) is a memory reduction step, PeakMem(h) is the maximum of peak memories of steps (1) and (2).
In step (1) the two child nodes sequentially processed, resulting in the peak memory of F FinalMem(h - 1) + PeakMem(h- - 1).
In step (2) hidden states of length must be held for Σ L, layers, the peak memory is
Cx L x M.
1=0
By applying the induction hypothesis, we get
PeakMem(h) = max
PinalMern iialMMemh - 1) + PeakMem(h 1), x
Σ
Li M
1=0
LCM + (2 + LCM,LCM)
max
((++1) LCM, LCM)
max
2 h LCM.
As the peak memory grows linearly with the tree height, it grows logarithmic with the input sequence length.
B IMPLEMENTATION DETAILS
Context merging schedule. Following the formulation in Figure 3, we detail how many layers assigned to each level of the binary tree. The basic principle is to assign an equal number of layers each node. In practice, we noticed that additionally assigning more layers the leaf nodes helps improve the overall performance. We assign 12 additional layers for 7b models and 20 layers for 13b models.
Calibration. For all models (HOMER and HOMERR+baslinnee calibration is performed using 100 text corpora segments from the validation set and the test set WikiText- (Merity et al., 2016).
Maximum chunk length. In all experiments involving HOMER, the maximum chunk length was set to be half of the context limit.
15

---

Published as a conference paper at ICLR 2024
C PROMPTS FOR DOWNSTREAM TASKS
The detailed prompt format for each downstream task are provided Table and Table
Table 6: Prompt for passkey retrieval task. Slight modifications are made from the original prompt turn it into chat prompt (Touvron al., 2023).
Prefix
[NST]] 

There is an important info hidden inside a lot of irrelevant text.
Find it and memorize them. I will quiz you about the important information there.


<<
Context
The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. (repeat x times)
The pass key is 12323. Remember it. 12323 is the pass key. The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. (repeat y times)
Suffix
What the pass key? The pass key is [/NST]
Table 7: Prompt for question answering task. Basic prompt format follows Shaham et al. (2023). Slight modifications are made to turn it into a chat prompt (Touvron et al., 2023).
Prefix

 

You are provided story and multiple-choice question with 4 possible answers (marked by A, B, C. D). Choose the best answer by writing its corresponding
letter (either A, B, C, or D). Do not provide any explanation.


Context
(The actual document)
Suffix
Question and Possible Answers: (question)
(a) {choice 1
(b) {choice 2
(c) (choice 3
(d) {choice 4)
Answer:
[/NST]
16

---

Published as a conference paper at ICLR 2024
D ILLUSTRATION OF PROPAGATIVE REFINEMENT
In this section, we provide a comprehensive explanation of propagative refinement suggested in Section 3.2. Figure illustrates the process, where 3 out of 6 tokens are pruned at layer N.
Layer N
2
4
6
3
5
Layer N-1
2
4
6
1
3
5
4
6
Layer 1
1
2
4
6
1
3
5 6
4 6
(a)
(b)
(c)
Figure 4: Illustration of the propagative refinement process.
Initially, shown in part (a), the three least significant tokens (2, 3, and 5) are marked for pruning layer N. Subsequently, in part (b), the corresponding tokens in the lower-layer embeddings are also marked for pruning. Finally, part (c) demonstrates the outcome after pruning, where all marked tokens are eliminated across every layer, resulting uniform, compressed embedding structure composed of just three tokens.
E INFERENCE SPEED ANALYSIS
In this section, we discuss the inference speed of HOMER. Besides reducing memory requirements, HOMER also provides significant speedup due the extensive reduction in computation. Table illustrates the average inference time for HOMER and other baselines. Specifically, we compare the time required to generate 20, 50, and 100 tokens, conditioned on 8k, 16k, and 32k contexts.
Table 8: Inference time for long inputs. All measurements are taken on single A100 GPU, with Flash Attention 2 (Dao, 2023) applied. We also report the percentage of the speedup. We report single value for all baselines' following the setup in Section 4.5. The baselines include plain Llama, PI, NTK, and YaRN. The inference time is averaged over 25 runs.
Average rum time (scconds)
Setup
16k
32k
20 tokens
Baselines"
1.879
3,224
6.546
HOMER
1.673 (12.3% speedup)
2.270 (42.0% speedup)
3.513 (86.3% speedup)
50 tokens
Baselines*
3.842
6.028
11.143
HOMER
3.026 (27.0% speedup)
3.639 (65.6% speedup)
4.873 (128.7% speedup)
100 tokens
Baselines"
7.149
10.733
18.828
HOMER 5.355 (33.5% speedup)
5.930 (81.0% speedup) 7.169 (162.6% speedup)
The main source of performance gain in HOMER, as described in Section 4.5, is the computation reduction. The following points highlight these improvements:
The divide-and-conquer approach circumvents the quadratic computation associated with self- attention.
Token pruning significantly reduces the number of tokens to process, especially in the upper layers.
HOMER compresses long contexts into short embeddings, substantially reducing the size of the kv-cache. This step lowers the computational demand during the decoding stage.
17

---

Published as a conference paper at ICLR 2024
As evident from the results, HOMER provides significant speedup (up to 162.6%) compared the baseline methods. It's important to note that our method is even more beneficial when generating longer outputs conditioned on longer imputs, underscoring its effectiveness in handling long contexts.
We also emphasize that the additional computation introduced by HOMER is minimal. The hier- archical context merging process involves relatively cheap operations, including matrix subtraction (calibration), gathering by index (calibration, token pruning, and propagative refinement), top-k se- lection (token pruning), and tensor concatenation (merging). Conversely, it reduces more costly operations such as matrix multiplication for computing self-attention.
F PERPLEXITY PLOT FOR LANGUAGE MODELING EXPERIMENT
In this section, we provide a fine-grained perplexity plot for the fluency experiment in Section 4.3.
20
Plain
YaRN (8K)
YARN (64K)
PI (8K)
YaRN (16K)
HOMER (Ours)
NTK (9K)
YaRN (32K)
HOMER (Ours) YaRN (8K)
18
16
Perplexity
14
12
8
6
4k
8k
16k
32k
64k
Evaluation length (Tokens)
Figure 5: Perplexity plot on 25 long documents from PG-19 dataset (Rae et al., 2019), measured with Llama-2-7b. HOMER consistently achieves low perplexity across long documents up 64K tokens, demonstrating its ability to remain fluent while conditioned on very long inputs. Detailed comparison with more baselines are provided in Table 3.
18

---

Published as a conference paper at ICLR 2024
G
PERPLEXITY EVALUATION ON DOWNSTREAM TASKS
Plain
YaRN (16K)
YaRN (64K)
HOMER (Ours) YaRN (8K)
YaRN (8K)
YaRN (32K)
HOMER (Ours)
10
Perplexity
10
10"
10
2k4k
8k
16k
32k Evaluation length (Tokens)
48k
64k
Figure 6: Perplexity plot on 100 long samples from passkey retrieval, measured with Llama-2- 7b-chat. HOMER achieves lower perplexity when conditioned on longer inputs, demonstrating its ability to effectively handle the long imputs. Perplexity values on landmark lengths provided in Table
Table 9: Perplexity values for passkey retrieval. Empty values indicate NaN.
Context Limit
2k
4k
8k
16k
32k
64k
Plain
23.12
13.013
10
Plain + HOMER
None
23.22
29.882
23.932
20.897
7.119
4.085
YaRN
23.414
17.338
23.394
> 104
16k
28.183
19.07
15,366
> 102
>> 10
32k
19.248
20.491
28.034
17.187
31.017
10a
22.031
16.963
15.481
36.760
38.786
102
YaRN + HOMER
None
23.414
17.232
22.263
11.317
7.618
2.412
In this section, provide additional perplexity experiments a challenging benchmark where accessing previous long contexts is essential. To achieve this, we reformulated the passkey retrieval task in Section 4. and measured the perplexity of ground -ruth answer phrases (e.g.. 'The passkey is 12321."). The results are demonstrated in Figure 6 and Table 9.
As the results show, HOMER exhibits lower perplexity when conditioned longer contexts, achiev- ing its best performance with 64k inputs. Furthermore, HOMER outperforms the long-conter com- petitors with context lengths of 16k and beyond. These experiments emphasize the efficacy of HOMER in utilizing long contexts, particularly in scenarios where accessing such context is neces- sary.
19


Reviews:
[{"review_id": "Jll1XeauXH", "reviewer": "Reviewer_uJAL", "summary": "This work presents a new approach to handle long context by LLMs without any training, by compressing and merging context chunks in a divide-and-conquer manner. Specifically, the long context is first split into chunks, and a LLM then encodes and compresses the chunk by pruning its tokens according to attention importance, finally concatenating adjacent chunks at each transformer layers. The proposed approach is shown achieving good performance on retrieval and QA tasks. Additionally, it can be viewed orthogonal to other techniques to extend LLM context length via scaling Rotary embeddings, and is shown state-of-the-art performance on QA when combing with other existing techniques together.", "questions": "What if we only look at the attention importance at the final layer directly, and use the top-k positions as the final pruning decisions for every layer? How is this compared to the proposed approach?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed approach brings a new perspective to handle long context by LLMs through context compression and merging, complementary to other popular techniques focusing on Rotary position embeddings.\n  \n- The approach is shown to surpass all baselines on Passkey Retrieval, extending to 8 times longer than the pretrained context length. For QA, it is shown to achieve the best performance combining with another long-context technique NTK.\n  \n- The proposed approach has the advantage of low GPU memory usage, compared with other baselines, due to the compression of context, leading to a reduced size of KV cache.", "weaknesses": "Though effective on the relatively easy Passkey Retrieval, the heavy compression and merging of context might hurt the reasoning capability of LLMs, as shown by Table 3 that the proposed approach alone underperforms other long-context techniques on QA. It is unknown whether this technique would remain useful on other tasks that also require reasoning. More experiments on other tasks could be conducted to demonstrate the advantages as well as the disadvantages of the proposed approach.", "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": "pxpQh4UxYS", "reviewer": "Reviewer_X3So", "summary": "The paper introduces a new method called Hierarchical Context Merging to extend the context limit of LLMs all the way to 64K tokens. The method is training-free, and uses a divide-and-conquer methodology to segment inputs into smaller units. At each following Transformer layer, the segments are fused together through a token reduction technique to reduce the memory footprint. The authors also propose to optimize the computation order to further reduce the memory requirement by 70%. \n\nThe proposed method is evaluated on three sets of tasks: Perplexity on 19 long documents from PG-19, passkey retrieval and question answering. The paper shows superior performance on all three benchmarks against RoPE-based baselines such as PI, NTK and YaRN.", "questions": "I do have a few questions that I need clarifications from the authors:\n\n1. The inference time of HOMER was not discussed at all in the paper. Instead, memory is the major focus. While memory footprint is an important dimension, inference time is very important for LLM applications. Given the extensive manipulations of HOMER on context embeddings at each Transformer layer, I expect a significant overhead occurring due to such forward operations compared to RoPE methods.\n2. The explanation of propagative refinement is confusing: token pruning is done through a forward pass, and the top-K token embeddings are tracked and obtained only at the top-layer after the entire forward pass. It is unclear to me how the refined embeddings are gathered at the low-level layers, and impact the computation results shown in the ablation in Table 4(b).\n3. Minor point: the token pruning section didn’t explain how it is done – it is unclear how the top-K hyperparameter is selected, and the tradeoff between speed and K is made (e.g. keeping a large K at each layer will result in more steps of pruning).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The work shows several strengths:\n1. The proposed method HOMER is the only method which has a reasonable Perplexity on 64K long documents. The performance gain on passkey retrieval at 32K tokens is impressive: 22.4% -> 80.4%.\n2.  The memory reduction up to 70% is also quite impressive.", "weaknesses": "However, I do have several weaknesses regarding the work:\n1. For Llama-2-13B, the gain on both perplexity and passkey retrieval is very marginal for all the context lengths below 32K.\n2. The Accuracy on QA benchmark is very concerning (Table 3): HOMER alone is worse than the NTK method, not showing any performance advantage. 1&2 together raises the question of the generalization of the HOMER method, and the presented results are not comprehensive enough to convince the readers.\n3. The biggest concern is the complexity of the method: compared to RoPE based methods, HOMER involves merging chunks, token reduction, embedding refinement as well as computational order reduction. All these techniques together make the proposed method very hard for wide adoption in the community, hence limiting the impact of the 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": "JePYY4n4X4", "reviewer": "Reviewer_ae6y", "summary": "This paper proposes an approach to extend the context limit of large language models (LLMs). The key idea is to hierarchical context merging, progressively conducting token pruning from bottom to top layers. The token pruning method is simple and training-free, which prunes the tokens with minimal attention from the final token in each chunk, and conducts adjustments (e.g., attention and position bias) to make sure the model to work properly. The results show that it can work on 64k context length without perplexity explosion.", "questions": "See the weakness section", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- A training-free method and can be easily adopted\n- Some empirical results showing the feasibility of the token pruning method.", "weaknesses": "While this paper presents a seemingly promising approach to long context modeling, its major issue lies in the lack of clarity regarding its use case. According to the paper's description, HOMER first breaks down a long context (e.g., 64k) into N chunks, prunes tokens from each chunk, and finally merges them into a compressed and short context. If this is the case, this method is not suitable for streaming decoding scenarios (where the key-value cache of previously decoded tokens is not recalculated, only new decode tokens are computed).\n\nThe key question arises: if this work is not for streaming decoding, then there is no need for such a process. For a 64k length context, wouldn't it be sufficient to ignore the first 62k context and only use the last 2k context? According to the results in Figure 1 of the paper, using only 2k context, the perplexity (PPL) is less than 6 (around 5 point something), whereas after all the effort of merging with HOMER, the PPL is above 8.\n\nGiven this, I question the necessity of using HOMER. Using the last 2K context directly seems to yield better results, casting a large doubt on the value of the method proposed in this paper. If the setting of this paper is not for streaming decoding, then the performance of HOMER should at least be better than the PPL of the 2K context. However, this method requires recalculating the key-value pairs for the previous context, which doesn't seem to align with a streaming decoding setting, making it the most confusing aspect for me.", "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": "nHgpMXpukn", "reviewer": "Reviewer_wcXN", "summary": "Paper addresses long contexts (>4k tokens of a pre-trained model such as LLaMA) in LLMs which are challenging because of the nature of Transformer self-attention.\nA technique called HOMER is introduced: contrarily to previous approaches that approximate the computation of self attention, or that use position interpolation at inference, authors propose a ‘compression’ method which decomposes the long input into manageable chunks; then a hierarchical merging strategy ensures the ‘communication’  between those chunks. One advantage of the approach is that it does not necessarily require fine-tuning. Perplexity experiments are done to assess the efficiency of the method as well as experiments on passkey retrieval and Q&A tasks. Finally it is shown that the approach proposed can be used successfully in addition to position interpolation techniques.", "questions": "-It is unclear to me why PROPAGATIVE REFINEMENT OF LOWER-LAYER EMBEDDINGS is needed. During inference the final layer representation of the long input is computed so why do you actually need to refine the lower layer embeddings afterwards ? In other words, why do you need to have fixed-length embeddings for every layers afterwards ? I could not fully understand this part.\n\n-Similarly section 3.3 is a bit to laconic to be fully understandable: how this binary-tree inspired computation is innovative ? What is described in Appendix/figure3 seems to me a straightforward way to efficiently compute the hierarchical merging process.\n\n-for Q&A it is a pity that HOMER+YaRN results are not presented (+ YaRN seems to not be working well on this task, why?)\n\n-for Q&A i’m wondering if the comparison with YaRN and NTK is really fair as for them, context documents are clipped whereas for HOMER, which is a compression technique, the full context document can be used.\n\n-In 4.4 (ablation) it not 100% clear what the uncalibrated (as opposed to calibrated) pruning criteria is", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "-a new approach to handling long contexts in LLMs that can be used additionally to position interpolation techniques proposed recently\n\n-extensive experiments done with LLaMA2 (ppl and long-context tasks) and comparison with recent (2023) baselines such as YaRN\n\n-convincing results on both ppl and tasks evaluations\n\n-method is orthogonal to position interpolation methods", "weaknesses": "-the parts related to propagative refinement (3.2) and efficiency improvement (3.3) should be clarified/improved\n\n-for Q&A it is a pity that HOMER+YaRN results are not presented (+ YaRN seems to not be working well on this task, why?)", "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)