DSAS: A Universal Plug-and-Play Framework for Attention Optimization in Multi-Document Question Answering

TL;DR

The paper tackles the challenges of long-range dependency modeling and the lost-in-the-middle problem in multi-document QA by introducing DSAS, a training-free plug-and-play framework with two modules: Contextual Gate Weighting CGW and Reciprocal Attention Suppression RAS. Through information flow analysis, it shows how LLMs manage cross-document content and identifies a two-stage reasoning pattern that DSAS exploits to enhance focus on key paragraphs while suppressing irrelevant interactions. DSAS demonstrates universal gains across multiple models and benchmarks, achieving up to 4.2% average F1 improvement without architectural changes or fine-tuning, and exhibits robustness to input order and varying paragraph counts. The work highlights the potential of attention optimization for long-context reasoning and suggests avenues for future integration with retrieval-based approaches and scalable long-context strategies.

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. Notably, DSAS functions as a plug-and-play solution requiring no architectural modifications or extra training parameters. 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.

Figures (11)

• Figure 1: In Multi-doc QA tasks, directly processing long inputs comprising numerous paragraphs with LLMs presents two major challenges: long-range dependency modeling and "lost-in-the-middle", resulting in degraded answer quality. DSAS acts as a plug-in that enhances LLMs through a dual-stage process: (i) reinforcing information flow between supporting passages and both the question and target, and (ii) suppressing interactions between supporting and negative paragraphs.
• Figure 2: Layer-wise information flows of HotpotQA, 2WikiMultiHopQA and MuSiQue tested on Llama-3.1-8B-Instruct. $p^s$ and $p^n$ denote supporting paragraphs and negative paragraphs, respectively. The results of Qwen2.5-7B-Instruct are shown in Appendix \ref{['app:para_ana']}.
• Figure 3: Comparison between mean values of the good and bad instances from the 1000 samples of Llama-3.1-8B. The results of Qwen2.5-7B are shown in Appendix \ref{['app:ans_ana']}.
• Figure 4: Illustration of Dual-Stage Adaptive Sharpening (DSAS), including Contextual Gate Weighting (CGW) and Reciprocal Attention Suppression (RAS).
• Figure 5: Llama-3.1-8B's hyperparameter study of $K$, $n$, $\alpha$, $\beta$ on HotpotQA, 2WikiMultiHopQA, MuSiQue, and LongBench benchmarks. The hyperparameter studies of Llama-3.2-3B, Qwen2.5-7B, and Qwen2.5-14B are shown in Appendix \ref{['app:hyper_full']}.