SEAL: Scaling to Emphasize Attention for Long-Context Retrieval

Your reply makes the advantage of SEAL clearer. I would keep my overall score, but improve the score of contribution.

Generalization capability of SEAL

Regarding the generalization of SEAL discussed in Section 8, we conducted additional experiments on the more complex Knowledge-QA task type from LongBench. We kindly ask you to refer to the results, which demonstrate that SEAL consistently improves performance in both Single Doc QA and the more complicated Multi Doc QA tasks.

Specifically, the in-domain results in Table 3 were obtained by tuning the scale with format information from certain sub-tasks of LongBench. These results exhibit significant average score improvements across QA tasks, demonstrating that SEAL performs well even on more complex tasks and within a broader scope.

Regarding your concern about "training data is also the same types of tasks," the out-of-domain results in Table 3 address cases where this is not true. For these experiments, the attention scale trained on the format of the Line retrieval (retrieving numerical data) was applied to the LongBench tasks. Despite LongBench being an English-based QA task, we observed slight average score improvements. This finding highlights that SEAL remains applicable to out-of-domain tasks through enhanced retrieval performance, further showcasing its generalization capability.

Moreover, in addition to the Documentation QA results we reported, we are conducting experiments on more diverse task types as requested by the reviewer. However, due to the time constraints of the rebuttal process, we are submitting the remainder of our responses first. If experiments on additional tasks are completed during the rebuttal period, we will provide an updated response accordingly.

Advantages of SEAL and comparison to PEFT

By comparing SEAL-L, SEAL-H, and SEAL-C, we empirically validated that we accurately identified the core components of retrieval and devised the optimal method to update them, ultimately proposing a cost-effective fine-tuning technique.

Please note that our primary focus, as mentioned in the global response, is not on proposing a more efficient PEFT method but rather on introducing the SEAL pipeline, analyzing its impact on retrieval performance, and identifying key components through this process.

Also, our SEAL pipeline is a solution that is not bound to any specific PEFT method. We hope it is recognized that SEAL-H, C, L (LoRA), and D (DoRA) are all approaches initially proposed by us.

We acknowledge that our contributions might not have been clearly conveyed, which may have led to the misconception that this paper is about proposing an efficient PEFT technique. To address this, we have revised the manuscript's flow and tone to improve the overall writing quality. We hope this is taken into account.

Description about the procedures of SEAL

Thank you for your insightful comment. We have added a brief description of the procedures of SEAL to Line 77 of the introduction section of the revised manuscript, which is written in blue color. We invite you to review this updated content in the revised manuscript.

Improving visibility in Figure 4

We have revised the relevant content in Figure 4 of the updated manuscript, where the modified part is marked with a blue box. Our initial intent was to align the y-axis scales of Figures 4(a) and 4(b) to emphasize that the direct effect of MLPs dominates over Attention Heads. We have modified the caption to preserve this intent while ensuring that the changes are more clearly observable.

In response to the reviewer's request, we conducted additional experiments on the tasks of the latest and more refined benchmark, LV-Eval [1]. We apologize for the delayed response.

We reviewed various long-context benchmarks, such as ZeroSCROLLS [2] and L-Eval [3], and found that many of the datasets in these benchmarks (e.g., Qasper, MuSiQue, NarrativeQA) already overlap with the datasets of LongBench [4], which we had already experimented on. Therefore, we selected some sub-tasks from the more recent LV-Eval benchmark to minimize overlap with those in LongBench.

Importantly, to address the reviewer's two concerns: (1) SEAL is only applicable for classic retrieval tasks like Needle-in-a-Haystack and (2) the training data is also the same types of tasks, we used the SEAL-H scales trained on the data from triviaqa_e sub-task of LongBench, which was used in the experiments in Section 8 of our paper.

The results are as follows:

loogle-SD-mixup and cmrc-mixup tasks are Single-hop QA tasks. loogle-CR-mixup task is Multi-hop QA task.

Even though we used the head-wise scale trained specifically for LongBench's triviaqa_e sub-task (which differs entirely from LV-Eval’s sub-tasks), we observed performance improvements at both the 16K and 32K context lengths across many tasks.

This demonstrates two key points:

1. Robust Performance on Challenging Benchmarks: SEAL performs well even on the more challenging LV-Eval benchmark, which consists of single-hop QA tasks and multi-hop QA tasks. This could be the answer to concern (1) above.
2. Generalization Capability: The improved retrieval performance achieved through SEAL extends to other QA tasks, even when using head-wise scales fine-tuned on the task format from different benchmark. This highlights SEAL's strong generalization capability. This could be the answer to concern (2) above.

We hope these experimental results will assist the reviewer and the Area Chair in their evaluation.

References:

[1] LV-Eval: A Balanced Long-Context Benchmark with 5 Length Levels Up to 256K

[2] ZeroSCROLLS: A Zero-Shot Benchmark for Long Text Understanding

[3] L-Eval: Instituting Standardized Evaluation for Long Context Language Models

[4] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding

About SEAL-H, SEAL-C, and SEAL-L

First, thank you for your valuable feedback. Could you kindly clarify which specific aspects you find to be failing to demonstrate the authors' claims? If you provide more details, we will address them thoroughly and promptly.

Our main claim is that SEAL-H and SEAL-C achieve performance comparable to SEAL-L despite using significantly fewer parameters, thereby demonstrating that the core component for improving long-context retrieval performance lies in the scaling of attention components.

Experiments with additional baseline DoRA

In response to the reviewer’s suggestion, we have reported results using [1] DoRA (ICML 2024), a recent variant of LoRA that has gained substantial attention. The results for DoRA (SEAL-D) have been added to Table 1 and Figure 5 of the updated manuscript, which are marked in blue, and we plan to include a detailed explanation in the revised version.

The experimental results show that the performance of SEAL-H and C is comparable to that of LoRA (SEAL-L) and DoRA (SEAL-D). This again highlights that scaling attention heads/channels is a key component in improving long-context retrieval performance.

SEAL's key contributions and the relationship to PEFT

By comparing SEAL-L, SEAL-H, and SEAL-C, we empirically validated that we accurately identified the core components of retrieval and devised the optimal method to update them, ultimately proposing a cost-effective fine-tuning technique.

Please note that our primary focus, as mentioned in the global response, is not on proposing a more efficient PEFT method but rather on introducing the SEAL pipeline, analyzing its impact on retrieval performance, and identifying key components through this process.

Also, our SEAL pipeline is a solution that is not bound to any specific PEFT method. We hope it is recognized that SEAL-H, C, L (LoRA), and D (DoRA) are all approaches initially proposed by us.

We acknowledge that our contributions might not have been clearly conveyed, which may have led to the misconception that this paper is about proposing an efficient PEFT technique. To address this, we have revised the manuscript's flow and tone to improve the overall writing quality. We hope this is taken into account.

References:

[1] DoRA: Weight-Decomposed Low-Rank Adaptation

Explanation of the term “long-context retrieval”

Thank you for your comment. We define the term "long-context retrieval" as performing retrieval tasks under long-context inputs.

Also, in the problem we address, the definition of "long-context" is as follows, as explained in the global response: if the model's context window size is $d$ and the input sequence length is $l$, then long-context refers to a sequence length $l$ where $l<d$ but $l\approx d$.

We agree that a precise definition is necessary, given that this is the first attempt to address this problem. Accordingly, we have added this definition to line 67 of the revised manuscript.

Overall experimental design and use of context extension techniques

In the experiments up to Section 6, no additional context extension techniques were applied. We reported the results of the publicly available long-context models (baseline) and the baseline models enhanced with SEAL. Please note that even though the models we used are designed exclusively for long-context tasks, their performance degrades even when the given context is within the model's limitations ($l < d$). When we apply our technique, SEAL, it shows notable improvement across all context lengths. Therefore, the performance improvements compared to the baseline are solely attributed to SEAL, and isolating experiments for other techniques cannot be conducted.

Figure 6 referred to by the reviewer corresponds to the results in Section 7 (SEAL with training-free context length extension), where NTK and Self-Extend were adopted as training-free context extension methods. These methods were applied exclusively in Section 7. This section demonstrates that SEAL can orthogonally contribute to performance improvement when combined with training-free context extension methods for the case of short-context window models + training-free context extension methods + SEAL.

Therefore, the middle subfigure in Figure 6 and the second row of results in Table 2 (labeled NTK and Self-Extend) isolate and evaluate the performance of context extension techniques (NTK and Self-Extend). Once again, no additional context extension methods were applied to the baseline to any of the figures or tables apart from Figure 6 and Table 2.

Logical flow in writing

We have revised not only the parts you pointed out but also improved the overall claim to ensure a more natural logical flow throughout the manuscript. We kindly request you to review the updated version of the manuscript.

“in-domain” and “out-of-domain” experiments in Section 8

When considering tasks A and B, along with their respective format datasets a and b, we define in-domain as the scenario where the scaling factor tuned on dataset a (or b) is applied to task A (or task B). Conversely, out-of-domain refers to situations where the scaling factor tuned on dataset a is applied to a new task, B, or vice versa.

In our evaluation of the generalization ability of SEAL (Section 8), we designated Line-retrieval and LongBench as tasks A and B, respectively. For the out-of-domain experiments in Table 3, we tuned the scaling factor using the line-retrieval format dataset (a) and tested it on LongBench (task B). This represents an out-of-domain scenario. Even under such conditions, SEAL improves retrieval performance, demonstrating slight improvements in the LongBench average score, despite it being a completely different English QA task.

On the other hand, the in-domain experiments in Table 3 involved tuning the scaling factor using the format data from a subset of LongBench subtasks (b) and testing it on LongBench (task B). This in-domain scenario shows significant accuracy improvements.

Thus, in-domain and out-of-domain experiments represent distinct experimental settings with different training datasets, while only task B is commonly used for verification. Including these results in the same table may have caused some confusion, and we will revise the format to reduce potential misunderstandings.

Using different LLMs in Figures 5 and 6

Most importantly, the objectives and experimental setups in Section 6 (Figure 5) and Section 7 (Figure 6) are fundamentally different, which is why different LLM families were used.

In Figure 5, the goal was to evaluate the effectiveness of SEAL by applying it to pre-existing long-context models. Therefore, we used open-source models that support relatively large context windows (16k/32k), such as Mistral, LongChat, and Vicuna.

In contrast, Figure 6 aimed to demonstrate that methods for extending the context window and SEAL can be applied orthogonally. For this purpose, we started with Llama-2, a representative foundation model with a short-context window (4k), and applied other context extension methods (NTK, Self-Extend) alongside SEAL. Consequently, the model families evaluated in Figures 5 and 6 have to be different.

Since the evaluation categories in Figures 5 and 6 are distinct, direct comparison of results is not feasible. However, conducting experiments across such diverse configurations and scenarios demonstrates that SEAL performs well regardless of how long-context models are created (e.g., pre-training for Mistral, fine-tuning for LongChat and Vicuna, and training-free extension for Llama + NTK or Self-Extend) or the model family (e.g., Mistral and Llama). This highlights SEAL’s superiority and its strong generalization capabilities.

Isolating the impact of context extension methods

In the experiments for Figure 5 (Section 6), we did not employ any additional context length extension techniques, including Self-Extend. The experiments up to Section 6 focused on applying SEAL to publicly available long-context models (with $d$=16k or 32k) to isolate and evaluate the pure impact of SEAL compared to the baseline. For further details, we kindly ask you to refer to the response provided above (Overall experimental design and use of context extension techniques).

The dataset setup and significance of the experiments in Table 3

The reason LongBench tasks were used in both the in-domain and out-of-domain experiments is as follows:

Out-of-domain experiments: LongBench, being an English QA task, significantly differs in characteristics from line retrieval, which involves simple retrieval of numerical data. Therefore, LongBench was chosen as the task B case for the out-of-domain experiments.

In-domain experiments: Since LongBench and line-retrieval belong to entirely different domains, the performance improvement in the out-of-domain experiments was relatively small. We were concerned that readers might misunderstand this result and assume that SEAL does not effectively enhance retrieval performance for more complex tasks, such as QA. To address this, we added in-domain experiments using LongBench itself.

As explained in previous responses, the in-domain and out-of-domain experiments are distinct, with only the task for the verification (LongBench) being the same. The important point is, the datasets used to tune the scaling factors are different, line-retrieval format data (a) for the out-of-domain case and LongBench format data (b) for the in-domain case. Therefore, we believe there is no issue with the fairness of the experimental setup.

We are writing to raise a serious concern about one of the reviews we received for our submission. We are addressing this issue not only as an author but also as a supporter of this conference, as we believe the review reflects a significant problem with the process.

First, we would like to thank all reviewers for their time and effort. Reviewing papers is very laborious, and we greatly appreciate the constructive feedback from the other reviewers, which will help us improve our work. However, the review in question is both unhelpful and unfair. It provides little more than superficial comments, fails to engage with the core ideas of our research, and ultimately assigns an unreasonably low score that effectively blocks our paper from acceptance.

From the feedback, it is clear that the reviewer did not make an effort to understand the paper or its contributions. If the paper had been properly read, Weaknesses 2 or Question 3 could not have been raised. Instead of providing meaningful critique, the reviewer has cited minor or surface-level issues (e.g., requesting definitions for relatively intuitive terms, and pointing out typographical errors) as the main reasons for assigning a reject rating. We believe that such issues traditionally belong in the suggestion part of a review and are matters that can be adequately addressed during the rebuttal process. However, during the review discussion period, we did not receive any additional questions or engagement from this reviewer. This lack of interaction has denied us the opportunity to clarify misunderstandings or address their concerns.

To be clear, we are not claiming that our paper is without flaws. We fully acknowledge the valid criticisms raised by other reviewers and are working to address them. However, receiving such a dismissive and low-effort review has made us question the fairness of the review process.

We strongly request that this review be reconsidered in the final decision process. We hope the committee will address this issue and uphold the standards of this esteemed conference.

Thank you for your attention to this matter.

We have carefully considered the reviewer's follow-up response and greatly value this opportunity to enhance our research through constructive discussion. We also understand the importance of writing as emphasized by the reviewer, and we will make every effort to correct flaws up until the deadline of the rebuttal.

Regarding the concern about the “lost-in-the-middle” issue raised by the reviewer in a subsequent response, we would like to emphasize that SEAL is not limited to addressing this specific issue or tasks such as Needle-in-a-Haystack.

First, from a practical perspective, one of SEAL’s major advantages, as reconfirmed by the reviewer, is "enhancing the long-context retrieval capabilities of existing LLMs”. Many users rely heavily on long-context LLMs they already use due to various reasons (e.g., license, fine-tuned on specific data, or compatibility with inference frameworks).

For instance, if a user has invested significant resources in fine-tuning a Mistral-7B-Instruct model with their custom documents, re-fine-tuning this data on a new 128k model would be prohibitively costly. In contrast, SEAL can improve retrieval performance for various existing long-context models at a minimal cost.

Additionally, the majority of LLMs are short-context models (<8k). SEAL offers an additional unique advantage by enabling these existing or the latest short-context models to be transformed into long-context LLMs with very low training costs, as demonstrated in Section 7.

Second, as models become increasingly sophisticated, benchmarks for testing long-context capabilities are also advancing. In Section 8 (In-domain results of Table 3), we demonstrated that SEAL significantly improves performance on a much more complex task: Documentation QA, which involves reading lengthy documentation and performing Question-Answering. This result highlights that SEAL is not limited to the Needle-in-a-Haystack task but is a scalable technique capable of enhancing retrieval performance even in scenarios where specific long-context models struggle with more challenging benchmarks.

Finally, beyond the performance of LLMs, the interpretability of LLMs is critically important. Through detailed analysis, we have investigated how attention components influence retrieval performance, which could contribute to future research. For example, pre-training long-context model often requires extensive datasets with long sequences and substantial training resources. Our analysis of core components and mechanisms might help make this process more efficient.

We hope this addresses the reviewer’s latest concerns and provides further clarity regarding SEAL’s contribution and impact on a broader scope.

Explanation of "cost-efficiency"

We apologize for the unclear explanation of cost-efficiency. We define cost-efficient as the absence or minimal occurrence of various overhead costs during training and inference. Specifically, cost-efficiency in our paper includes 1) A very small number of learnable parameters, 2) The ability to fine-tune with as few as 50 samples, and 3) Achieving tuning within an hour using a single GPU for the training phase. Also, no additional overhead is incurred during the inference phase.

As noted in the dataset generation explanation in Section 6, tuning with 50 samples on a 7B model takes approximately 40 minutes for a dataset length of 31k and about 10 minutes for a dataset length of 16k. This demonstrates that our method is highly efficient and practical.

Regarding SEAL, the key component identification and fine-tuning processes are not separate. As described in Section 4.2, we design SEAL based on the observation that each attention component needs to be emphasized or suppressed to adjust the influence of attention components. Based on this observation, we proposed variations of SEAL that exclusively focus on the head-related learnable parameters. Thus, no additional time or process is required for component identification.

Please note that our main contribution lies in demonstrating that LLMs inherently possess retrieval capabilities, which can be effectively restored with minimal fine-tuning, and in identifying that attention heads are the key components in this process. We would like to clarify that cost-efficiency is a natural byproduct of our main contributions, rather than cost-efficiency itself is a primary objective.

Broader range of advanced baselines

To the best of our knowledge, our work is the first to improve retrieval performance for long-context that do not exceed the model’s context window. Consequently, we hope you can understand the difficulty of establishing appropriate or fair baselines for comparison. For instance, previous context length extension methods have entirely different objectives and typically require significant additional training costs.

From this perspective, we include SEAL-L as a main baseline to validate the efficiency of our observation in enhancing LLM's performance with a minimal fine-tuning and generated dataset and to explain the differences between SEAL-L and SEAL-H/C, which reveal the effectiveness of head-wise/channel-wise scaling in the context of retrieval enhancement. Please refer to the global response for more details.

In response to the reviewer’s suggestion to include state-of-the-art techniques, we have reported results using [1] DoRA (ICML 2024), a recent variant of LoRA that has gained substantial attention. The results for DoRA (SEAL-D) have been added to Table 1 and Figure 5 of the updated manuscript, and we plan to include a detailed explanation in the revised version.

The experimental results show that the performance of SEAL-H and C is comparable to that of LoRA (SEAL-L) and DoRA (SEAL-D). This highlights that scaling attention heads/channels is a key component in improving long-context retrieval performance, further demonstrating the effectiveness of SEAL.

Once again, one of our main contributions is identifying attention heads and channels as the key components in enhancing the long-context retrieval performance inherent in LLMs through minimal fine-tuning. SEAL can be combined with the superset of key components, PEFT methods, as demonstrated by the LoRA and DoRA extension experiments. These experiments illustrate SEAL’s adaptability and potential for broader application across various PEFT techniques.

Experiments with larger model sizes

Thank you for the valuable feedback regarding the verification of versatility. In response, we have conducted additional experiments on the 13B models of the LongChat and Vicuna model families (longchat-13b-16k, vicuna-13b-v1.5-16k). The results of the Line retrieval task and Needle-in-a-Haystack task have been added to Table 1 and Figure 5 of the revised manuscript, which we encourage you to review. The proposed SEAL method also demonstrates significant improvements with the 13B models, further substantiating its versatility and scalability to larger models.

References:

[1] DoRA: Weight-Decomposed Low-Rank Adaptation