SEAL: Scaling to Emphasize Attention for Long-Context Retrieval

Authors: Changhun Lee, Jun-gyu Jin, YoungHyun Cho, Eunhyeok Park
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: large language models, long context, retrieval, attention, supervised fine-tuning

Abstract:
In this work, we introduce a novel approach called Scaling to Emphasize Attention for Long-context retrieval (SEAL), which enhances the retrieval performance of large language models (LLMs) over extended contexts. Previous studies have shown that each attention head in LLMs has a unique functionality and collectively contributes to the overall behavior of the model. Similarly, we observe that specific heads are closely tied to long-context retrieval, showing positive or negative correlation with retrieval scores. Built on this insight, we propose a learning-based mechanism using zero-shot generated data to emphasize these heads, improving the model's performance in long-context retrieval tasks. 
By applying SEAL, we can achieve significant improvements in in-domain retrieval performance, including document QA tasks from LongBench, and considerable improvements in out-of-domain cases.
Additionally, when combined with existing training-free context extension techniques, SEAL extends the context limits of LLMs while maintaining highly reliable outputs, opening new avenues for research in this field.

Decision: reject
Reviewer scores: 5, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1729087048466

thread_type:
official_review

<thread_object>
creator: reviewer_xi5c

<content>
summary:
This paper proposes a novel and practical method, SEAL, to improve the long-context retrieval ability of LLMs. 

First, through perturbation experiments, it finds a certain attention head or a certain channel in it can cause a positive or negative effect on long-context retrieval accuracy.

Second, it demonstrates directly scaling the hidden states of these heads or channels can indeed improve the retrieval accuracy of LLMs. 

Third, it adds trainable scale factors into the model and use a small amount of samples of retrieval tasks to fine-tune the model. The results show SEAL can remarkably improve the long-context retrieval ability of LLMs.

strengths:
1. This paper discovers that a certain attention head can cast a remarkable positive or negative effect on long-context retrieval accuracy, even as well as a certain channel. This is interesting and helpful for us to further understand the role of the internal modules of LLMs.

2. The proposed method, SEAL, is very cost-effective, which only needs very few training samples and tuned parameters.

3. There are enough evaluation results of various models to demonstrate the method’s effect.

weaknesses:
1. Narrow scope

The method seems to only be applicable for classic retrieval tasks such as NIAH, and the training data is also the same types of tasks. It will not be surprising that this leads to an improvement, since this task has been too simple, fixed and formulaic, which may represent a narrow application scope for this method. It would be better to train and test on more tasks such as Knowledge-QA.

2. No unique advantages

The author should empirically test whether the time or space required by SEAL is significantly less than that of LoRA. Otherwise it cannot show significant superiority of SEAL compared to LoRA. Because the parameters tuned by LoRA are already very few. Though SEAL can theoretically tune much less parameters, it may not significantly save much time. 

3. There is little detailed description about the procedures of the method in the abstract or introduction. This will make it hard for readers hard to grasp the method quickly. There usually should be a paragraph included in the introduction to describe the specific operation of the method.

4. The curve of data points in Figure 4 (a) may be too small, making it hard to clearly see the changes.

questions:
1. Can you demonstrate the unique advantages of your method compared to LoRA through more experiments?
2. Can you train and test on more various task types to demonstrate the generalization of your method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729422953770

thread_type:
official_review

<thread_object>
creator: reviewer_k91h

<content>
summary:
This paper proposes an approach called Scaling to Emphasize Attention for Long-context retrieval (SEAL), which emphasizes specific heads or channels (attention outputs) particularly related to long-context retrieval by efficiently adjusting the strength of each attention component. The authors claimed that SEAL achieves significant improvements in in-domain retrieval performance and cross-domain document QA tasks, also extends the context limits of LLMs while maintaining highly reliable outputs.

strengths:
1. This paper proposes SEAL to efficiently adjusting the strength of each attention component, and achieves superior performance to various LLM baselines in long-context retrieval.
2. The content, figures, and tables of the paper provide a detailed explanation and analysis of the motivation, methods, and experiments, facilitating the readers' understanding.

weaknesses:
1. The experimental results in Table 1 show that SEAL-H and SEAL-C require fewer parameters than Baseline and SEAL-L. However, their performance does not consistently surpass SEAL-L in long-context scenarios, failing to demonstrate the authors' claims.
2. The experiments only select SEAL-L as the baseline, it should include other PEFT methods for comparison.

questions:
1. Except fewer parameters, what other advantages does SEAL have over LoRA or other PEFT methods? Since the parameters of SEAL-L is also small compared to LLMs, what are the unique application scenarios for SEAL?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729590370606

thread_type:
official_review

<thread_object>
creator: reviewer_xqus

<content>
summary:
The paper introduces SEAL (Scaling to Emphasize Attention for Long-context retrieval), a novel attention scaling approach that improves retrieval performance for long-context tasks in Large Language Models (LLMs). It addresses the challenge of performance degradation over extended contexts, particularly in retrieval tasks. SEAL fine-tunes specific attention heads or channels using a minimal amount of training data, leading to significant improvements in long-context retrieval across various benchmarks. The paper focuses on cost-efficient enhancement of long-context capabilities without altering the model’s learned behavior.

strengths:
1. SEAL presents an innovative approach by leveraging attention head/channel scaling to enhance long-context retrieval.
2. The method uses very few trainable parameters and requires minimal training data, making it highly efficient.

weaknesses:
1. The term “long-context retrieval” is ambiguous. It would be clearer to refer to “retrieval tasks that have long contexts,” which directly emphasizes tasks like passage retrieval or number retrieval.
2. The paper lacks explicit detail about which context extension techniques are used. For example, Figure 6 mentions the use of Self-Extend, but no experiments isolating its performance are provided.
3. Logical Flow in Writing: Certain parts of the paper are difficult to follow due to writing issues such as ambiguous expressions, inconsistent time tense, and occasional typographical errors (e.g., “biases” instead of “bias”).
4. The distinction between “in-domain” and “out-of-domain” in the experiments is confusing. Specifically, if “in-domain” refers to training on retrieval tasks, why are the same datasets used for both “in-domain” and “out-of-domain” experiments?

questions:
1. What specifically constitutes “long-context retrieval”? Could the authors clarify this definition and provide more precise terminology?
2. Why are different LLMs used in Figures 5 and 6? Is there a specific reason for the model changes, and how do these variations impact the comparability of the results?
3. Can the authors provide experiments isolating the effect of Self-Extend in Figure 5 to verify its individual impact on performance?
4. What is the rationale behind using the same datasets for “in-domain” and “out-of-domain” experiments in Table 3? How is “out-of-domain” defined in this context, and what criteria differentiate the two?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730301590175

thread_type:
official_review

<thread_object>
creator: reviewer_fzzr

<content>
summary:
This work focuses on scaling to emphasize attention to long-context retrieval, designed to enhance the retrieval performance of LLMs in handling extended contexts. A cost-effective, learning-based mechanism is proposed to improve the model's performance in long-context retrieval tasks, which emphasizes specific attention heads tailored to retrieval tasks. Experimental results demonstrate superior performance over the compared baselines.

strengths:
1. This paper is well-organized and easy to read. 
2. The proposed method presents a reasonable approach for long-context retrieval by identifying the key components of Transformer architecture to boost retrieval performance. 
3. The approach is practical and has the potential for broad application in various RAG settings.

weaknesses:
1. The term "cost-efficient" is not clearly defined, resulting in ambiguity when assessing the cost-effectiveness of the approach. The strategy of identifying key components initially and subsequently fine-tuning these components may prove to be computationally intensive. It would be beneficial to provide details regarding the computational time involved in this process.
2. A more thorough evaluation would benefit from comparisons with a broader range of advanced baseline models. Currently, the proposed method is compared against only one simple. Including more sophisticated long-context modeling methods and state-of-the-art techniques would better validate the effectiveness of the proposed method.
3. To confirm the versatility of the proposed method, it would be beneficial to conduct experiments on different LLMs of varying sizes.

questions:
Please refer to the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732354270002

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
## Global Response

First and foremost, we extend our gratitude to each reviewer for their thorough examination of our work. In addition to this global response, we have provided detailed and appropriate answers to each reviewer and have made the necessary revisions to the manuscript (PDF). We would greatly appreciate it if the reviewers could also review the updated document.

In this study, our key contribution is the meaningful discovery that pre-trained large language models (LLMs) already possess sufficient information to handle long-context retrieval tasks, and their retrieval performance can be greatly enhanced with minimal adjustments.

As discussed in the introduction, LLMs face two primary challenges in handling long-context scenarios. Defining the model's context window size as $d$ and the input sequence length as $l$, these challenges are as follows:

1. When $l>d$, the model's performance quickly breaks down for inputs exceeding the context window size.
2. When $l<d$ but $l\approx d$, the retrieval performance degradation issue arise in handling long input contexts.

While the former issue has been addressed by techniques like context extension, the latter topic has remained largely unexplored.

In this work, **we are the first to propose methods specifically for improving retrieval performance in situations where $l<d$**. Our contributions can be summarized as follows:

1. Through circuit analysis, we demonstrate that LLMs inherently contain sufficient information to solve long-context tasks and show that simply scaling the strength of attention heads can significantly restore performance in long-context retrieval tasks.

2. We designed a pipeline to significantly improve retrieval performance at minimal cost by generating only around 50 synthetic data samples using format information from retrieval tasks and fine-tuning using those data.

3. 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.

As this is the first work in this research direction, there are no existing studies or appropriate baselines for comparison. We hope this point is considered. Our primary focus, as mentioned earlier, 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.**
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.

Through this research, we aim to not only significantly enhance the retrieval performance of existing models with minimal effort but also, as demonstrated in Section 7, to enable models to maintain superior performance on even longer sequences by combining our approach with context extension techniques. We believe this study contributes toward realizing such advancements effectively.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733208853430

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Reminder about the end of the comments period

comment:
First and foremost, we sincerely thank each reviewer for their effort in reviewing our work.

We would like to kindly remind you that the deadline for reviewers to leave comments is the end of December 2nd (AoE), which is approximately five hours from now.

We have provided appropriate responses to each reviewer's most recent feedback, and we hope they will be taken into account. If you have any additional feedback or concerns after reviewing our responses, we would greatly appreciate it if you could leave a comment before the deadline.

We will also do our best to respond promptly by the authors' message deadline.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes an efficient method, SEAL, which tunes the scales of attention heads and channels for effective long context retrieval. Due to the limitations in the pretraining phase, LLMs are known for deficiencies in handling very long context window. By learning to weigh the strengths of different attention heads and channels using a small set of generated examples, the scaled model is capable of emphasizing the model components that are beneficial for the retrieval task and de-emphasizing those components negatively affecting the retrieval task. Experimental results demonstrate the advantage of SEAL on line retrieval task and LongBench QA task.

Strengths:
- The proposed scaling method is novel and effective, backed by an interesting observation on a significant performance change when pruning different attention heads or channels.
- Given the tunable parameters only involve scaler weights for attention heads and channels, the method is efficient and could be easily adopted in different models and real applications.
- Experimental results on two different datasets demonstrate the advantage of SEAL over the baseline.

Weaknesses:
- The clarity of the paper writing needs to be improved, such as the description of SEAL-L and SEAL-D, how the method is used in the context extension setting.
- The implication of SEAL is not clearly revealed from the empirical results. Comparing with normal PEFT methods such as LoRA and DoRA, the performance improvement is not very consistent. If the provided LoRA (SEAL-L) and DoRA (SEAL-D) are also proposed by the authors, additional baselines with normal PEFT training should be given to clearly reveal the difference and advantage of the proposed method.
- More datasets/tasks could be incorporated to showcase the method's generalizability.