[논문리뷰]DSAS: A Universal Plug-and-Play Framework for Attention Optimization in Multi-Document Question Answering (NeurIPS, 2025)

Jiakai Li, Rongzheng Wang, Yizhuo Ma, Shuang Liang, Guangchun Luo, and Ke Qin.2025. DSAS: A Universal Plug-and-Play Framework for Attention Optimizationin Multi-Document Question Answering (NeurIPS, 2025)

1. Problem Statement

Multi-hop question answering (multi-document question answering, MHQA)은 사용자의 질문에 대해 답을 하기 위해 여러 개의 문서(document) 혹은 증거(evidence)를 필요로 하는 질의 응답 문제를 의미한다. 이러한 증거들은 모두 직접적으로 사용자의 질의와 연관되어 있을 수 있지만, 여러 증거들이 순차적으로 연결이 되어 있거나, 복잡한 형태로 엉켜 llm의 추론을 어렵게 만든다. 예를 들어, When was the city where the creator of the Nobel Prize was born founded?라는 질문을 reasoning chain으로 구성할 경우 세 개의 연속된 질문으로 분해되는데, 여기서 Stockholm이라는 단어는 질문에 직접적으로 드러나지 않아 retrieval가 이와 관련된 문서를 놓칠 수 있기 때문이다. 따라서, 리트리버가 이 문서를 제대로 검색하게 질문의 부족한 정보를 보완하거나 LLM이 학습한 internal knowledge를 활용해서 추론을 진행해야 한다.

DSAS는 특히 MHQA를 해결함에 있어 LLM이 겪는 두 가지 문제를 쟁점으로 해결하고자 한다.

• [Long-range dependency modeling]: 긴 문맥 내에서 핵심 정보를 선택적으로 집중하지 못하고 의미적 연결이 약해진다.
• [Lost-in-the-middle]: 긴 입력의 중간 부분에 위치한 정보에 대해 모델의 어텐션이 약화되어 정답 생성 품질이 저하된다.

이러한 한계를 극복하기 위해 Dual-Stage Adaptive Sharpening (DSAS)라는 새로운 plug-and-play형 어텐션 최적화 프레임워크를 제안한다. 이 방법은 모델 구조 변경이나 추가 학습 없이, 기존 LLM의 어텐션을 개선한다.

2. Limitations of Existing Works

• [긴 문맥에서의 Attention Dilution]: 긴 입력을 단순히 연결(concatenate)하면, 중요한 토큰 간의 상호작용이 불필요한 토큰에 의해 희석되는 attention dilution 현상이 발생함. 이는 중요한 문단 간 의미적 연결이 약화되어, 정답 생성의 품질이 떨어지게 됨.
• [Long-range Dependency Truncation]: StreamingLLM, LM-Infinite 등의 기법은 long-context 문제를 해결하려 하지만, 글로벌 토큰 상호작용을 희생하여 장거리 의미 연결이 깨짐. 그 결과 MHQA 성능이 저하됨.
• [Lost-in-the-middle 현상의 지속]: LLM이 긴 입력의 중간에 위치한 정보를 처리할 때 성능이 급격히 떨어진다고 지적함. LongAlign 등의 기존 방법은 추가 학습 데이터와 파인튜닝이 필요하여 보편적 적용성(universality)이 떨어짐.
• [Task-specific fine-tuning 의존성]: 대부분의 기존 방법은 모델 구조를 수정하거나 별도의 훈련을 필요로 하여, 다양한 LLM에 쉽게 적용되지 않음.

3. Motivations

3.0. Preliminaries - Information Flow

먼저 \(\ell\)번째 레이어의 어텐션 스코어와 어텐션 가중치 행렬을 각각 \(A_{h, \ell}^S, A_{h, \ell}^W\)로 두고, 모든 헤드에 대해 합산(sum)하여 레이어 수준의 행렬 \(A_\ell^S\)와 \(A_\ell^W\)를 정의한다. 이 때 정보 흐름의 방향(information flow)을 \(A_\ell^W(i,j)\)로 정의하고, 이는 \(i\)번째 토큰에서 \(j\)번째 토큰으로의 정보 흐름을 의미한다.

다음으로 세 가지 주요 컴포넌트; 문단(passage, \(p\)), 질문(question, \(q\)),타겟(target, \(t\))의 상호작용을 보기 위해 두 개의 Top-K 기반 정보흐름 지표 \(\mathcal I_{p^m, q}\), \(\mathcal I_{p^m, t}\)를 정의한다. 문단은 \(p^1, \dots, p^C\)로 표기하고, 임의의 \(m\)번째 문단 \(p^m\)의 토큰 시퀀스의 각 인덱스는 \(\{ p_s^m, p_s^m+1, \cdots, p_e^m \}\) 로 표기한다. 이 때, \(s\)와 \(e\)는 각각 시작과 끝 토큰의 position을 의미한다.

$$\mathcal I_{p^m, q} = \frac{1}{Q}\displaystyle\sum \text{Top-K} \Big(\Big\{ \displaystyle\sum_{i \in q}A_\ell^W (i, j) \vert j \in p^m \Big\} \Big),$$ $$\mathcal I_{p^m, t} = \frac{1}{Q}\displaystyle\sum \text{Top-K} \big( \{ A_\ell^W (t,j) \vert j \in p^m \} \big), \quad m \in \{1, \dots, C \}$$

이 지표는 문단 \(p_m\)이 질문(\(q\)) 및 타깃(\(t\))으로 보내는 Top-K 의미 유입도를 의미한다. 여기서 \(Q\)는 질문 길이이고, 두 지표의 값이 클수록 질문 정렬(to \(q\)) 및 정답 형성(to \(t\))에 더 기여함을 의미한다.

3.1. Motivation (1) Paragraph Disparity Level Analysis

대부분의 LLM은 여러 층으로 구성되어 있고, 특히 파라미터가 커질수록 depth가 깊어지는 경우가 대부분이다. 이미 기존의 여러 연구들에 의해 LLM의 디코딩 과정에서 동일한 텍스트여도 각 층이 보는 의미 정보(semantic information) 정보는 다르다는 사실이 알려져 있다. 논문에서는 이에 근거해 각층의 정보 흐름을 측정하여 비교한다. Figure 2는 측정 결과를 보여주며, 이것이 보여주는 것은 결국 얕은 레이어에서는 문단 간 정보 흐름 격차가 작지만, 레이어가 깊어질수록 지원 문단(supporting fact)과 비관련 문단(negative)의 격차가 커지며, 최종적으로 핵심 문단의 정보가 타깃에 집약된다는 사실이다. 참고로 여기서 말하는 타깃이란 answer generation position을 의미하며, 다시 말해 정답을 생성하기 직전의 토큰 위치를 의미한다.

• 각 레이어마다 \(\mathcal I_{p^m, q}, I_{p^m, t}\)를 계산해 문단 수준의 격차(disparity)를 관찰

이 양상은 Hotpotqa, 2WikimultihopQA, MuSiQue 모든 데이터셋에서 동일한 결과를 보여주며, 이는 결국 LLM이 두 단계의 추론 패턴을 가진다는 사실을 말해준다.

1. [질문 정렬]: 초기 레이어에서는 정보가 질문으로 수렴하며, 지원 문단은 비관련 문단보다 더 강한 정보 흐름을 보임. 이는 모델이 의미적으로 더 관련 있는 컨텐츠를 우선적으로 처리할 수 있는 능력을 갖추었음을 보여줌.
2. [정답 형성]: 이후, 레이어가 깊어짐에 딸라 지원 문단으로부터의 정보가 타깃으로 집약되어, LLM이 핵심 문단을 전략적으로 활용해 정답을 생성함.

이러한 관찰은 LLM의 내재적 해석 가능성(inherent interpretability)을 강조하며, 모델이 추론 과정 중에 과제 특화 정보(task-specific information)를 통합할 수 있는 능력을 가지고 있음을 입증한다.

3.2. Motivation (2) Answer Quality Level Analysis

모델이 생성한 답변이 맞으면 (EM=1) Good, 그리고 관사나 수사같은 단어를 제외하고 생성된 답변에 정답과 오버랩되는 단어가 하나도 없으면(F1=0) Bad 로 정의한다. 이후 질문과 타깃의 정보 흐름 지표를 모든 레이어에 걸쳐 집혜하고, 지원/비관련 문단 기준의 평균값을 Good/Bad 두 그룹으로 비교하였다. 직관적으로 Good 에 해당하는 샘플들은 \(\mathcal I_{p^s, q}\)와 \(\mathcal I_{p^s, t}\)는 크고 반대로 \(\mathcal I_{p^n, q}, \mathcal I_{p^n, t}\)는 매우 작은 경향성을 보여주었다. Bad 그룹에서도 지원 문단에 대한 정보 흐름이 비관련 문단보다 항상 더 크다는 것이다. 즉, 모델은 기본적으로 지원 문단을 구분하지만, 그 유입 강도가 충분히 크지 않아 최적 답변으로 연결되지 못한 경우가 많다는 것을 의미한다. (y축 양 옆의 scale이 달라서 주의!!!!!)

4. Methodology

4.1. Model Overview

이러한 관찰을 바탕으로 본 논문에서는 Dual-Stage Adaptive Sharpening (DSAS)는 두 단계로 어텐션을 최적화하는 방법을 제안한다. DSAS는 두 모듈로 구성되어 있다.

• Contextual Gate Weighting: 질문과 타깃이 핵심 문단을 더 강하게 보도록 증폭
• Reciprocal Attention Suppression: 핵심↔비관련 문단의 상호작용을 억제

4.2. Contextual Gate Weighting

• 입력: 각 문단의 어텐션 스코어 행렬 \(A_\ell^S\)
• 출력: 문단별 contextual gate weight \(w_m\)

[(i) 서브 매트릭스 계산]

$$ M \in \mathbb{R}^{2Q \times p_m} = \underbrace{\left[A_{\ell}^{S}(q, p_m)\right]}_{\text{Question Matrix}} \;\parallel\; \underbrace{\left[A_{\ell}^{S}(t, p_m)^{\uparrow(Q)}\right]}_{\text{Target Matrix}}, $$

Contextual Gate Weighting (CGW)는 레이어 \(\ell\)에서 어텐션 정보로 문단의 가중치를 계산한다. 먼저 \(\ell\)번째 레이어의 어텐션 스코어 행렬 \(A_\ell^S\)에서, 문단 \(p^m\) 이 질문과 타깃에 얼마나 연결되는지를 보기 위한 서브매트릭스 \(M\)을 만든다. 참고로 타깃 토큰이란, answer generation position (i.e., the final token in the input), 다시 말해 정답을 생성하기 직전의 토큰 위치를 의미한다.

• \(Q\): 질문 토큰 길이
• \(p^m\): \(m\)번째 문단의 토큰 길이(문단 토큰 수)
• \(A_\ell^S (q, p^m)\): 질문 토큰들이 문단 토큰을 보는 스코어 서브매트릭스, \(\mathbb R^{Q\times p^m}\)
• \(A_\ell^S (t, p^m)\): 타깃 토큰이 문단 토큰을 보는 스코어 벡터, \(\mathbb R^{1\times p^m}\)
• \(\uparrow(Q)\): 타깃 벡터를 질문 길이 \(Q\)에 맞게 확장; \((1, p^m) \rightarrow (Q, p^m)\) 크기로 복제

• $$

  \(: concatenation. 위에 질문-문단 스코어, 아래에 타깃-문단 스코어가 쌓인\)2Q \times p^m$$ 행렬이 됨.

이제 이 \(M\)을 이용해 문단 \(p^m\)의 결합 정보흐름을 계산한다. 논문은 M의 각 문단토큰(열) 별로 스코어를 합친 뒤, 그 중 Top-K개의 값을 평균내는 형태로 정의한다.

$$\mathcal I_m^{\text{Comb}} = \frac{1}{K} \displaystyle\sum_{\text{Top-K}} \Big\{ \displaystyle\sum_{i=1}^{2Q} M_{i,j} \Big\}^{p^m}_{j=1}, \quad m \in \{ 1, \dots, C\}$$

\(\mathcal I_m^{\text{Comb}}\)는 문단의 \(j\)에 대해 (질문+타깃)이 그 토큰을 얼마나 강하게 보는지를 정량화 한 것이다. 즉, 이 정보 흐름량이 클수록 질문과 타깃이 이 문단의 일부 토큰들에 강하게 주의를 주고 있다는 뜻이 된다.

[(ii) 내용 점수 계산, Z-normalization]

$$v_m =0.5 \cdot \sigma\Big( \frac{\mathcal I_m^{\text{Comb}} -\mu_{I}}{\sigma_I}\Big)+ 0.5, \quad m \in \{1, \dots, C \}$$

이후 문단마다 스케일이 달라질 수 있으므로 \(\mathcal I^{\text{Comb}}\)를 Z-normalization 후 sigmoid로 0~1 범위로 눌러 \(v_m\)을 만든다. \(\mu_I, \sigma_I\)는 \(\mathcal I^{\text{Comb}}\)의 평균과 표준편차이다. 정규화를 함으로써 \(v_m\)은 내용적 관련성 기반 점수로 전환된다.

Z-normalization (Z-Score normalization): Z-score 정규화는 데이터의 각 값을 평균과 표준편차를 기준으로 변환하여, 평균이 0이고 표준편차가 1인 분포로 만드는 방법입니다. 즉, 어떤 값이 평균에서 얼마나 떨어져 있는지를 표준편차 단위로 표현하는 것

[(iii) Lost-in-the-middle 완하는 위한 위치 점수 \(\gamma_m\) 계산]
CGW는 내용 관련성만으로는 중간 위치 증거가 소실될 수 있다고 보고, U-shaped attention bias를 보정하기 위해 가우시안 분포 기반의 위치 가중을 도입한다. 먼저 확률밀도함수(probability density function,(PDF) \(f(x)\))와 누적분포함수 (cumulative distribution function(CDF), \(F(x)\))를 다음처럼 정의한다.

$$f(x) = \frac{1}{\sigma_p \sqrt{2\pi}} \cdot \exp \Big( - \frac{(x - \mu_p)^2}{2\sigma_p^2} \Big), \quad F(x) = \int_{-\infty}^x f(t) dt$$

여기서 \(\mu_p, \sigma_p\)는 입력 토큰 인덱스 \(\{0, 1, \dots, L-1 \}\)의 평균 및 표준편차로 \(\mu_p = 0.5(L-1)\)과 \(\sigma_p = \sqrt{\frac{L^2-1}{12}}\) 정의된다. 이 때, 문단 \(p^m\)의 토큰 인덱스를 \(\{ p_s^m, p_s^m+1, \dots, p_e^m \}\)이라고 할 때, 위치 점수(positional value) \(\gamma_m\)은 다음과 같다.

$$z_1 = \frac{p_s^m - \mu_p}{\sigma_p}, \; z_2 = \frac{p_e^m - \mu_p}{\sigma_p} \quad \rightarrow \quad \gamma_m = \frac{F(z_2) - F(z_1)}{z_2 - z_1}$$

위치 점수는 직관적으로 문단 구간이 가우시안 분포의 중심(중간)에 가까울수록 커지는 형태이며, 이는 중간 증거를 상대적으로 띄워주는 역할을 한다.

[(iv) 순위 기반 position-aware weight \(g_m\)]

$$ g_m = \begin{cases} \left(\dfrac{0.5C + 1}{\mathrm{rank}_m}\right)^{\gamma_m}, & \text{if } \mathrm{rank}_m \le 0.5C, \\ 1, & \text{otherwise.} \end{cases} $$

모든 문단에 무조건 위치 보정을 걸면 중간의 비관련 문단이 존재할 경우 잘못된 정보의 비중도 과하게 키워질 수 있다. 이를 보정하기 위해 상위 50% 문단에만 위치 보정을 적용하며 나머지는 1로 둔다. 문단을 내용 점수 \(v_m\)으로 내림차순 정렬했을 때 순위를 \(rank_m\)이라 할 때, 상위 50% 문단들 중에서도 중간에 있고(높은 \(\gamma_m\)) 순위가 높을수록(낮은 \(rank_m\)) \(g_m\)이 커져 더 강화된다.

[Contextual gate weight \(w_m\)]

$$w_m^{'} = v_m \cdot g_m^{\alpha}$$ $$w_m = (1-\beta) \frac{w^{'}_m - \min(w_i^{'})}{\max(w_i^{'}) - \min(w_i^{'})} + \beta, \quad i \in \{ 1, \dots, C \}$$

최종적으로 내용 점수와 위치 점수의 곱을 가중치로 만들고, min-max 정구화 후 바닥값을 \(\beta\)로 두어 최종 contextual gate weight(\(w_m\))을 얻는다. \(\alpha\)는 위치 가중의 영향도를 조절하는 하이퍼파라미터로, 클수록 \(g_m\)의 영향이 커진다. CGW는 결론적으로 질문/타겟과 실제로 강하게 연결되는 문단을 기준으로 중간 위치 증거를 보정해 lost-in-the-middle 문제를 완화한 뒤, 최종적으로 가중치 \(w_m\)으로 질문/타겟으로 부터 문단 방향으로 스코어를 증폭하여, 정답 생성에 중요한 문단이 질문과 타겟에 더 강하게 상호작용하도록 유도한다.

4.3. Reciprocal Attention Suppression (RAS)

RAS는 비관련 문단과 핵심 문단 간의 어텐션 교차를 억제하여 long-range dependency 혼란을 줄인다.

• 입력: CGW를 통해 계산된 문단별 가중치 \(w_m\)
• 출력: 수정된 어텐션 스코어 행렬 \(A_{h, \ell}^S\)

$$A_{h, \ell}^S (i, j) = \min (w_{m1}, w_{m2}) \cdot A_{h, \ell}^S (i, j)$$

단, 핵심 문단(key paragraph)의 집합을 \(P_{key, \ell}\)라고 하고, 관련 없는 문서 집합을 \(P_{irr, \ell}\)이고, \(m_1 \in P_{key, \ell}\) & \(m_2 \in P_{irr, \ell}\)라고 할 때, RAS의 수식은 위와 같다(만약 \(j < i\)라면, \(\{ p_s^{m_1}, p_s^{m_1} + 1, \dots, p_e^{m_1} \}\)이고, \(j \geq i\)이면 \(\{ p_s^{m_2}, p_s^{m_2} + 1, \dots, p_e^{m_2} \}\)이다.) 이를 통해 불필요한 cross-paragraph interference를 줄이고, 정답 생성을 위한 정보 흐름을 강화한다.

5. Experiments

5.1. Main Results

Table 1은 MHQA데이터셋에 대한 추론 성능을 비교한 것이다. 어텐션 패턴을 기반으로 입력 문단의 위치를 재할당하는 PINE과 비교했을 때, DSAS는 전반적으로 vanilla LLM대비 더 큰 폭으로 성능이 상승하는 것을 볼 수 있다. 또한 3B(Llama-3.2-3B)에서도 F1 score가 평균 +0.9 상승이 확인되고, 32B(Qwen2.5-32B)에서도 평균 +2.5 상승이 되는 것을 볼 수 있다 . 즉, 스케일 전반에서 플러그인 방식으로 성능을 끌어올리는 경향이 보인다.

5.2. Ablation Study

• w/o CGW: CGW(Contextual Gate Weighting) 제거,
• w/o RAS: RAS(Reciprocal Attention Suppression) 제거
• w/o p-a: position-aware weighting 제거

CGW를 제거(w/o CGW)하는 것은 어텐션 재가중 과정에서 질문과 타깃 토큰이 어떤 문단을 핵심 문단으로 판단하는지에 따라 가중치를 부여하는 gate 함수 자체를 적용하지 않고, RAS와 position-aware weighting만 적용하는 셋팅이다. 이 셋팅에서는 질문–핵심 문단–타깃 사이의 정보 증폭 경로가 차단되므로, 모델은 핵심 문단을 명확히 부각하지 못한 상태에서 추론을 수행하게 된다.

RAS를 제거(w/o RAS)하는 것은 핵심 문단(supporting paragraph)와 비관련 문단(negative paragraph) 사이의 불필요한 cross-attention을 억제하는 항을 사용하지 않고, CGW와 position-aware weighting만 적용하는 셋팅이다. 이 경우, 핵심 문단을 강조하는 효과는 유지되지만, 비관련 문단에서 발생하는 noise가 여전히 핵심 문단 표현에 섞여 들어오게 되어, 핵심 문단의 표현 순도가 떨어진 상태로 추론이 진행된다.

마지막으로 position-aware weighting을 제거하는 것은 \(g_m =1\)로 고정하여 문단 위치에 따른 가중 조절을 전혀 사용하지 않고 추론을 진행하는 셋팅이다. 이 상태에서는 DSAS가 중간에 위치한 문단을 상대적으로 더 강조하거나 양끝 문단의 영향을 줄이는 기능이 사라지게된다.

실험 결과 CGW와 RAS가 거의 비슷하게 성능 게인이 가장 크며, CGW가 약간 더 큰 것을 볼 수 있다.

5.3. Performance based on Paragraph Count

DSAS가 핵심적으로 해결하려는 문제 중 하나는 바로 lost-in-the-middle 문제이다. 이를 증명하기 위해 입력으로 들어가는 문단 수를 늘려가면서 성능 변화 실험을 진행하였다. Table 4에서 보이는 것과 같이, 실제로 vanilla LLM대비 문단 수 증가에 따른 성능 감소 폭이 DSAS에서 확연이 작은 것을 볼 수 있다. 이는 DSAS가 negative paragraph에서 오는 정보 흐름을 억제하기 때문에 성능 감소가 더 완만하게 되는 것이다.

5.4. Attention Score Visualization

Figure 11은 HotpotQA에서 여러 핵심 문서(sp1, sp2)와 비관련 문서(np1 ~ np8) 에 대해 타깃과 질문 토큰 사이의 어텐션 스코어를 비교하는 것이다. Figure 11의 (a)는 vanilla LLM의 어텐션 히트맵이고, (b)는 DSAS의 어텐션 히트맵이다. 결론적으로 DSAS가 (i) 질문/타겟이 supporting paragraph를 더 강하게 보게 만들고, (ii) supporting paragraph끼리의 상호작용도 강화하며, 동시에 (iii) supporting↔negative 간 불필요한 상호작용은 억제한다.

Conclusion

Limitations

• [일반화 부족] MHQA는 비교적 짧은 텍스트를 생성하는 short-form generation에 속하므로, long-form generation이나 summarization등의 일반적인 long-context generation task에서도 효과가 유지되는지 검증이 필요함.
• [Latency analysis의 부재] 어텐션 수정으로 인한 memory/time cost에 대한 latency analysis가 부재함.

Contributions

• 추가 학습 없이 LLM의 추론 성능을 향상시킬 수 있는 어텐션 수정 알고리즘을 제안함.
• 어떤 LLM에 상관없이 model-agnostic하게 사용 가능한 알고리즘
• lost-in-the-middle을 position-aware weighting을 통해 효과적으로 억제함
• 추가적인 모듈 없이 LLM만으로도 reranking이 가능하도록 만듦

NR 카테고리 내 다른 글 보러가기