Paper: Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs (ICLR-2024)

Reviewer comment (Experiments):
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?)

Author response:
**[W2] HOMER+YaRN results are not presented for Q&A (+ YaRN seems to not be working well on this task, why?)**

We conducted additional experiments of HOMER+YaRN for Q&A. HOMER+YaRN achieves an accuracy of 0.360, improving the performance of YaRN (0.310). Since we focused on demonstrating HOMER's applicability in enhancing the best-performing baselines, only HOMER+NTK was explored in the paper.

YaRN introduces multiple hyperparameters, including the temperature t required for scaling the attention weights, and alpha/beta values required for applying NTK-by-parts. We believe that the suggested choice of the hyperparameters can be suboptimal for some tasks, as they were empirically searched by minimizing the perplexity.

Reviewer comment (Theory):
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.

Author response:
**[Q1] Why do we need to reduce the lower layer embeddings, and why do the embeddings have to be the same length for every layer?**

This is because our method focuses on compressing the long context into shorter embeddings, reducing the computational burden of handling long context. Without refinement, the lower layer embeddings would remain large, resulting in large memory requirements and computational overhead in future decoding steps. Therefore, we further reduce the lower layer embeddings with propagative refinement.

We decided to make the embedding dimension the same in every layer, in order to make it more compatible with standard KV-cache (which typically has fixed length across all layers).

Reviewer comment (Novelty):
Similarly section 3.3 is a bit too laconic to be fully understandable: how is this binary-tree inspired computation innovative? What is described in Appendix/figure 3 seems to me a straightforward way to efficiently compute the hierarchical merging process.

Author response:
**[Q2] Section 3.3 requires more explanation. How is binary-tree-inspired computation innovative? It seems like a straightforward way to efficiently compute the hierarchical merging process.**

In typical Transformer models, all tokens at a single layer are computed in parallel. A similar implementation of HOMER would involve processing multiple chunks concurrently. The memory consumption can be large in this implementation, although it will already require less GPU memory than the baselines thanks to token reduction of HOMER.

Section 3.3 introduces an alternative computation ordering for HOMER, where chunks are processed sequentially rather than in parallel. This sequential handling of chunks, made feasible by our divide-and-conquer approach, significantly lowers GPU memory usage. We have refined the manuscript to more clearly convey this aspect.

It is also important to note that while the sequential computation may slightly increase the inference time, HOMER provides speedup (up to 162.6%) even with the memory-efficient computation order, thanks to its substantial computation reduction.

Reviewer comment (Experiments):
In 4.4 (ablation) it is not 100% clear what the uncalibrated (as opposed to calibrated) pruning criteria is

Author response:
**[Q5] What does ‘uncalibrated’ mean in section 4.4?**

It means that the calibration technique described in section 3.1 (second paragraph of ‘Token reduction on individual chunks’) is not applied. In other words, the raw attention logits are used as significance scores for pruning.

Reviewer comment (Evaluation):
For Llama-2-13B, the gain on both perplexity and passkey retrieval is very marginal for all the context lengths below 32K.

Author response:
**[W1] For Llama-2-13B, the gain on both perplexity and passkey retrieval is very marginal for context lengths below 32k.**

Our primary goal with HOMER is to extend the context limit of large language models, focusing on very long contexts that were previously unmanageable. Hence, the performance comparisons under long context lengths (e.g., >= 32k) are more important, and the gain of HOMER is **substantial** for both perplexity and passkey retrieval in this regime.

Furthermore, the gain of HOMER is often significant even in the case of shorter context lengths (e.g., < 32k). For example, at 16k pass key retrieval task, we achieve 97.4% accuracy, while the best-performing baseline, YaRN, does 83.6%. Furthermore, we emphasize that HOMER achieves superior performance with greater efficiency (51.4% memory savings and 42.0% speedup at 16k context), which is its important additional advantage.

Reviewer comment (Evaluation):
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.

Author response:
**[W2] HOMER alone performs worse than the baselines for multiple-choice QA task.**

It is important to note that compatibility with other RoPE scaling methods is one of our main advantages. Combined with these methods, we **always** outperform the baselines for every experiment on extended context.

We also emphasize that applying HOMER consistently gives performance gain, both when applied on the plain model and NTK. This demonstrates HOMER’s ability to successfully leverage the extended context.

Furthermore, we highlight that even HOMER alone achieved superior performance in most of our experiments. Namely, HOMER alone already showed the best performance among all baselines in 15 out of 19 experiments reported in our paper.

Reviewer comment (Novelty):
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.

Author response:
**[W3] The biggest concern is the complexity of the method, compared to the RoPE scaling methods. This could make HOMER very hard for wide adoption in the community.**

HOMER is very easy-to-use, requiring only one line of code. This is made possible by providing a ‘patch’ function (analogous to [1]), which automatically edits an existing model to use HOMER. We believe that HOMER will be easily adopted by the community, thanks to the simple and straightforward implementation.

We also emphasize that our method is training-free, which is a great advantage in terms of its impact. For example, our current implementation of Llama can already be directly applied to all Llama derivatives (e.g. Vicuna, Alpaca, Koala, etc.) in a plug-and-play manner. This is a significant advantage compared to other long-context methods [2][3][4][5][6][7][8], which necessitate fine-tuning.

Finally, we highlight that our computational complexity is significantly lower than the baselines, providing up to 73.4% memory saving while being 162.6% faster. We believe that the significant efficiency gain will also encourage the wide adoption of HOMER.

---

[1] https://github.com/jquesnelle/yarn/blob/master/scaled_rope/patch.py

[2] Landmark Attention: Random-Access Infinite Context Length for Transformers, NeurIPS 2023

[3] Augmenting Language Models with Long-Term Memory, NeurIPS 2023

[4] Recurrent Memory Transformer, NeurIPS 2022

[5] Memorizing Transformers, ICLR 2022

[6] In-context Autoencoder for Context Compression in a Large Language Model, under review for ICLR 2024 (scores: 8/8/5/5)

[7] CLEX: Continuous Length Extrapolation for Large Language Models, under review for ICLR 2024 (scores: 8/6/6/5)

[8] PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training, under review for ICLR 2024 (scores: 6/6/6/5)

Reviewer comment (Evaluation):
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.

Author response:
**[Q1] The inference time of HOMER was not discussed at all in the paper.**

We appreciate your feedback that inference time is very important for LLM applications. In fact, **the inference speed of HOMER is significantly faster than the baselines**, thanks to the extensive computation reduction.

The following 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. We report a single value for all baselines following the setup in Section 4.5. The average inference time, calculated over 25 runs, is expressed in seconds. We also report the percentage of the speedup.

| Tokens | Method      | 8k                    | 16k                   | 32k                    |
|--------|-----------|-----------------------|-----------------------|------------------------|
| 20     | Baselines | 1.879                 | 3.224                 | 6.546                  |
|      | HOMER     | 1.673 (12.3% speedup) | 2.270 (42.0% speedup)    | 3.513 (86.3% speedup)  |
| 50     | Baselines | 3.842                 | 6.028                 | 11.143                 |
|      | HOMER     | 3.026 (27.0% speedup)   | 3.639 (65.6% speedup) | 4.873 (128.7% speedup) |
| 100    | Baselines | 7.149                 | 10.733                | 18.828                 |
|     | HOMER     | 5.355 (33.5% speedup) | 5.930 (81.0% speedup)    | 7.169 (162.6% speedup) |

As evident from the results, HOMER provides a significant speedup (up to 162.6%) compared to the baseline methods. It's important to note that our method is even more effective while generating longer outputs conditioned on longer inputs, underscoring its effectiveness in handling long contexts.

The main source of performance gain in HOMER is the computation reduction, elaborated in Section 4.5. 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 lowers the computational demand during the decoding stage.

We also emphasize that the computational overhead introduced by HOMER is minimal. The hierarchical context merging process involves relatively cheap operations, including matrix subtraction (calibration), gathering by index (calibration, token pruning and propagative refinement), top-k selection (token pruning), and tensor concatenation (merging). Conversely, it reduces more costly operations such as matrix multiplication for computing self-attention.

We also added the above analysis in Appendix E.

Reviewer comment (Writing):
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).

Author response:
**[Q2-1] The explanation of propagative refinement is confusing. How are the refined embeddings gathered at the low-level layers?**

For clarification, let us provide an alternative explanation of the process. In fact, propagative refinement is an additional refinement step taken after token reduction. 

The process is straightforward: when a token is pruned in the upper layers, the corresponding tokens (i.e. tokens that originated from the same input token) are also pruned in the lower-layer embeddings. This ‘propagates’ the pruning decision of the upper layers to the lower layers embeddings, hence the name. The synchronized pruning across layers results in shorter, uniform embeddings for each layer.

For example, consider the situation where the tokens at positions 2, 3, and 5 are pruned in the N-th layer. By applying propagative refinement, we simply drop the tokens at positions 2, 3, and 5 from the lower-layer embedding, too.

We also revised Section 3.2 with the above clarification, and added a step-by-step illustration of the process at Appendix D.

---
**[Q2-2] How does the refined embeddings impact the computation results shown in the ablation Table 4(b)?**

Although our motivation for propagative refinement was for computational efficiency, Table 4(b) shows that it is also beneficial for performance. This aligns with the recent findings in encoder-decoder models, where providing a large number of encoder outputs (that exceed the original context length) to the decoder often degrades performance, while selecting a few, relevant tokens is beneficial (see comparison of Unlimiformer and SLED in [9]). We believe that the superior performance of propagative refinement suggests that a similar phenomenon also happens for the decoder-only models.

[9] Unlimiformer: Long-Range Transformers with Unlimited Length Input, NeurIPS 2023

Reviewer comment (Reproducibility):
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).

Author response:
**[Q3] How is K selected for top-K token pruning?**

K was simply set to be the half of maximum chunk length, so that the resulting chunk after merging has the maximum length (in other words, we selected the largest K possible). For example in case of plain Llama + HOMER, K is 1024 because the maximum chunk length is 2048 (as mentioned in Appendix B). More extensive search on K could provide an even better balance between speed and performance.