Title: DSAS: A Universal Plug-and-Play Framework for Attention Optimization in Multi-Document Question Answering
Venue: NeurIPS 2025
Accepted: yes
Keywords: large language model

attention dilution

attention optimization
Abstract: While large language models (LLMs) show considerable promise across various fields, they have notable limitations in handling multi-document question answering (Multi-doc QA) tasks. The first challenge is long-range dependency modeling, where LLMs struggle to focus on key information in long texts, which weakens important semantic connections. Second, most LLMs suffer from the ''lost-in-the-middle'' issue, where they have difficulty processing information in the middle of long inputs. Current solutions either truncate global dependencies or demand costly finetuning, ultimately lacking a universal and simple solution for these challenges. To resolve these limitations, we propose Dual-Stage Adaptive Sharpening (DSAS) containing two modules. (i) The Contextual Gate Weighting (CGW) module alleviates ''lost-in-the-middle'' by assessing paragraph relevance through layer-wise attention tracking and position-aware weighting. (ii) The Reciprocal Attention Suppression (RAS) module enhances focus on critical paragraphs by suppressing information exchange between key and irrelevant texts, thus mitigating the limitations in long-range dependency modeling. Extensive experiments on four benchmarks demonstrate DSAS's efficacy across mainstream LLMs (Llama, Qwen, Mistral, and Deepseek), with an average F1-score improvement of 4.2% in Multi-doc QA tasks on Llama-3.1-8B-Instruct and Qwen2.5-14B-Instruct. Ablation studies confirm the essential contributions of both the CGW and RAS modules. In addition, detailed discussions in the Appendix further validate the robustness and scalability of DSAS.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a plug-and-play method to improve LLM's performance on the Multi-doc QA task. They accumulate document-wise attention weights (they call it information flow) to inspect the attention patterns on document context in a multi-doc QA task. They get two observations: 1) support documents generally receive more attention, and 2) correct answers pay more attention to the support document compared to incorrect answers. Therefore, they propose a plug-and-play method to automatically detect support documents by accumulating attention weights and then reweighting supporting documents and negative documents by the accumulated attention, which upweighs the attention on supporting ones and downweighs the attention on negative ones. Otherwise, they propose a Gaussian-based score to reward middle-position documents to deal with the lost-in-the-middle problem. They also propose a Reciprocal attention suppression module, which penalizes the inter-document attention. They conducted experiments on several multi-doc QA datasets and showed that their methods can beat vanilla LLM without their method.\n\nSTRENGTHS AND WEAKNESSES:\nstrengths:\n\ni. This method is plug-and-play and easy to implement. I believe it will be very attractive and interesting if we can get performance improvement by simply manipulating the attention weights.\n\nii. This method is not only applicable to multi-document question answering, but also has the potential to be applied to RAG and general long-context generation tasks.\n\nweakness:\n\ni. I believe there are some mistakes in equations 7 and 8. E(*) is the cumulative function of the half-normal distribution with mu=0 and sigma=1/sqrt(2) with support R^+. But p_e - mu and p_s - mu do not guarantee to be positive. In Figure 4, it seems that you just want to use the cumulative distribution function of a general standard normal distribution. I do not understand why you write in this way. Additionally, in equation (8), it seems that you want to estimate \\gemma by (F(x_2) - F(x_1))/2\\sqrt(2)(x_2 - x_1) where x_2, x_1 is the normalized p_e and p_s, and F is the cumulative distribution function of N(0,1). This is naive because mu is set to the half-length and the derivative is high when x=0 (i.e., p_s and p_e near the middle). But I am confused about the discounting factor 2\\sqrt(2), won't this make \\gemma a quite small number and make g_m almost 1 everywhere?\n\nii. Ambiguous symbol usage, which is confusing. For example, A^{S}_{h,l} in equation 3, mu and sigma in equation 6 and equations 7&8. \\omega_m in equation 11, A^{S}_{h,l} in 12.\nquality:1 \n\niii. Comparing only to one single baseline is not convincing enough. What are the implementation details of the only baseline method? Is it just vanilla LLM? Why do not compare with other baseline methods? I believe that there are several attention manipulation-based works in long-context generation or RAG, for example [1]. Are there any other baselines that can be implemented, such as masking documents lower than a threshold?\n\niv. In equation 1, you accumulate attention weights of question-length for each j without normalization, while in equation, there is only one attention weight for each j. However, in Figure 2, I_q and I_t show similar magnitudes. Why does this happen?\n\n[1] Found in the Middle: Calibrating Positional Attention Bias Improves Long Context Utilization\n\nQUESTIONS:\nPlease refer to the questions in the \"Strengths And Weaknesses\" section.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to identify key issues in long-range dependency modeling, particularly in multi-document question answering, and presents some experimental techniques to improve LLM performance targeting these deficiencies. The key problems identified are that language models have difficulty processing information located in the middle of long inputs and struggle to focus on important information.\n\nThe authors propose a dual-stage mechanisms to address these issues. First, Contextual Gate Weighting: for each paragraph, the method measures information flow with respect to the question, ranks the paragraphs, and multiplies the attention scores by learned weights that combine contextual relevance and position-aware factors. The second mechanism is Reciprocal Attention Suppression, which aims to mute the attention between key and irrelevant paragraphs by downscaling the cross-attention between them.\n\nThe paper runs several experiments on open-source models and multi-document question answering datasets and reports gains in F1 scores in the majority of the settings.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\nThe paper identifies key problems in the question answering field, and I agree that both of these research directions are important for the community.\n\nThe proposed method requires minimal engineering and no training runs, so it's quite easy to use.\n\nThe language of the paper is clear, and the flow is easy to understand. Attention reweighing has been studied in the community before, but the techniques presented provide an interesting additional reuse case in an experimental setting. The authors have also done a good job citing the relevant previous work in the area.\n\nThe F1 improvements, although in some cases quite marginal, are consistent and certainly improve confidence in the method.\n\n### Weaknesses\n\nThe method needs deep access to the attention stack and major changes to the inference loop. Although it is training-free, it can only be applied to open-source transformer models with the usual infrastructure. The computational overhead is also not trivial. I am also not sure how the latency will be affected.\n\nSince the paper is completely experimental and the techniques are quite arbitrary without any theoretical justification, the paper needs to clear a very high bar in the evaluation process and present a generalizable method. The current manuscript doesn’t clear that bar satisfactorily. The proposed methods seem too specific to the hyperparameters, which are not practical to tune for every specific task. This prevents the use of this method more generally. The hyperparameter tuning process needs a lot of clarification about the data splitting and tuning process.\n\nAs the authors mentioned, for some models like LLaMA (Figure 8), some choices of hyperparameters are worse than the baseline. In other scenarios, a lot of the hyperparameter settings are very close to the baseline. So, it does seem like the precise choice plays a very strong role and is quite model-dependent. “Middle size models” is a bit confusing notion and depends on the model architecture, so it's not clear whether the method will be useful or not on a particular model without having to do some kind of holdout evaluation. \n\nOverall, I agree with the problems identified in the paper, which are both important and well-motivated. However, the proposed solutions come across as somewhat ad hoc and appear to address the symptoms rather than the root causes. While I recognize that tackling the underlying issues is substantially more challenging, the current approach may reflect some degree of research overfit.\n\nQUESTIONS:\nCould the authors demonstrate that the same process, with the current parameter settings, generalizes to other domains—even on a small toy dataset—to help build confidence in the method’s robustness?\n\nCould the authors provide empirical data on the computational overhead or latency introduced by the proposed method compared to standard attention?\n\nHow were the hyperparameters tuned, and what was the data splitting strategy used to avoid potential data leakage or overfitting?\n\nDo the authors have any hypotheses for why the method appears to perform best on medium-sized language models only?\n\nLIMITATIONS:\nPlease check above for broad feedback, I don’t think this paper has a negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Dual-Stage Adaptive Sharpening (DSAS), a training-free framework designed to optimize attention in large language models (LLMs) for multi-document question answering. The DSAS framework consists of two main modules: Contextual Gate Weighting (CGW), which weighs paragraph relevance through attention tracking and positional weighting, and Reciprocal Attention Suppression (RAS), which suppresses unnecessary attention between key and irrelevant paragraphs. The method is applicable as a plug-and-play module requiring no architectural modification or extra parameters. The paper presents systematic analysis of information flow in LLMs and validates DSAS across four public benchmarks and several mainstream LLM architectures.\n\nSTRENGTHS AND WEAKNESSES:\n1. This paper is well-written and well-organized. The idea is simple and easy to follow, and is built upon in-depth analysis of layer-wise information flow.\n2. The ablations (Table 3) and hyperparameter studies (Figure 5) dissect the individual contributions of CGW, RAS, and positional weighting.\n\nQUESTIONS:\n1. The position-aware weighting, aimed at the \"lost-in-the-middle\" problem, may excessively prioritize central paragraphs even when relevant evidence occurs at the edges. While the authors mention mitigation via ranking and capping, additional discussion on failure cases of this bias (especially in non-standard input orderings) is lacking.\n2. DSAS is only evaluated on Multi-doc QA tasks, whereas LLMs are also deployed for other long-context applications like summarization. It is suggested to conduct experiments on other tasks. Moreover, the proposed method is not evaluated over general performance benchmarks, such as how DSAS affects math/code or other reasoning tasks.\n3. Figure 5 reveals non-negligible variance in gains depending on choices of K, n, α, β. These variances indicate the proposed method is unstable and could be hard to generalize to unseen domains.\n\nLIMITATIONS:\nsee above.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents DSAS, a universal plug-and-play framework for attention optimization in multi-document question answering. The goal is to improve the long-context understanding ability of large language models in multi-document question answering tasks by introducing some inductive biases for such tasks. Concretely, Contextual Gate Weighting (CGW) and Reciprocal Attention Suppression (RAS) are proposed and combined to address this problem. Extensive experiments on four benchmarks demonstrate DSAS's universality across mainstream LLMs.\n\nSTRENGTHS AND WEAKNESSES:\nPros:\n- This paper is well-written and easy to understand. The problem is clearly presented, and the motivation is well discussed. The long-context understanding ability of LLM is a very important research topic in the community, and this paper focuses on a big and typical series of subtasks (multi-document question answering tasks) in this topic.\n- The paper is well organized, especially the logic of motivation->method->evaluation and analysis. Section 2, information flow analysis on multi-doc QA, identifies key factors for effective multi-doc QA reasoning in a systematic and quantitative manner, making the proposed methods promising and intuitive.\n- The method is lightweight: it doesn't require further training on the LLMs, but just changes how the attention is calculated, based on the properties of multi-document question answering tasks.\n\nCons:\n- It would be better if the experiment and analysis section can include some case study or quantitative analysis to prove that the inherent discrimination capability is really enhanced.\n- I'd suggest adding some experiments to further analysis DSAS on different subsets. For example, how will the number of supporting/negative/total paragraphs affect the effectiveness of DSAS?\n\nQUESTIONS:\n- How are the performance of other tasks if DSAS is applied to the LLMs? Will the performance significantly decrease, or comparable to the baseline without DSAS?\n- Can DSAS be used for context length extension? For example, suppose that a model is trained on 128K context-length, what is the performance if we apply DSAS on it to process 256K-tokens multi-document question answering tasks?\n\nLIMITATIONS:\nYES", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}