💻 computer science

Uncertainty-gated selection for block-sparse attention

이 논문은 모호한 top-k 점수를 가진 쿼리에 대해 선택된 키 블록을 동적으로 확장함으로써, 다양한 모델 아키텍처 전반에서 밀집(dense) 효율성을 거의 유지하면서도 롱 컨텍스트 검색 정확도와 재현율을 크게 향상시키는 블록 희소 어텐션(block-sparse attention)을 위한 불확실성 게이트 라우터를 소개한다.

원저자: Thomas Rossi

게시일 2026-07-10

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Thomas Rossi

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대한 건초더미 속에서 특정 바늘을 찾으려고 노력하는 모습을 상상해 보세요. 하지만 당신에게는 아주 작은 손전등 하나뿐이며, 매우 엄격한 규칙이 있습니다. 한 번에 아주 적은 양의 건초 더미만 비출 수 있다는 규칙이죠. 이것은 현대 AI가 초장문 문서(예: 소설 한 권 전체)를 한꺼번에 읽으려고 할 때 직면하는 문제와 정확히 일치합니다.

문제점: "근시안적인" 손전등

대부분의 AI 모델은 시간을 절약하기 위해 **블록 희소 어텐션(block-sparse attention)**이라는 기술을 사용합니다. 100,000단어짜리 이야기를 모든 단어마다 다 읽는 대신, 이야기를 여러 덩어리(블록)로 나누고 '선택기(selector)'를 사용하여 가장 흥미로운 상위 k개의 덩어리를 골라냅니다.

하지만 여기 함정이 있습니다. 이 선택기는 **근시안적(myopic)**입니다. 예를 들어, 선택기가 두 개의 건초 더미를 보고 있다고 가정해 봅시다. 더미 A의 점수는 9.9이고, 더미 B의 점수는 9.8입니다. 규칙은 "상위 1개를 골라라"라고 말합니다. 선택기는 즉시 더미 A를 선택하고 더미 B를 버려버립니다.

이것은 잘못된 결정이라는 것이 이 논문의 주장입니다. 만약 더미 B가 실제로 질문에 대한 답을 담고 있다면 어떨까요? 두 점수의 미세한 차이가 그저 우연이었을 뿐이라면요? 일단 더미 B가 버려지면, AI는 결코 그것을 다시 가져올 수 없습니다. 이는 마치 탐정이 어떤 단서가 거의 완벽해 보이지 않는다는 이유로 버렸다가, 나중에 그 단서가 사건을 해결하는 데 꼭 필요했다는 것을 깨닫게 되는 것과 같습니다.

해결책: "불확실성 게이트형" 스마트 스위치

저자들은 토마스 로시(Thomas Rossi)가 이끄는 팀과 함께 **불확실성 게이트 선택(Uncertainty-Gated Selection)**이라는 영리한 해결책을 제안합니다. 이것은 손전등에 "신뢰도 측정기"를 추가하는 것과 같습니다.

AI는 최종 결정을 내리기 전에 다음과 같이 자문합니다. "내가 올바른 덩어리를 고르고 있다는 확신이 있는가?"

신뢰도 체크: AI는 상위 덩어리들의 점수를 살펴봅니다. 만약 상위 덩어리가 두 번째로 높은 덩어리보다 훨씬 점수가 높다면(큰 격차), AI는 확신을 가집니다. 그리고 규칙대로 상위 k개의 덩어리만 선택합니다.
"잠깐, 혹시?" 하는 순간: 만약 상위 덩어리와 두 번째로 좋은 덩어리의 점수가 거의 비슷하다면(미세한 격차), AI는 깨닫습니다. "어라, 잘 모르겠는데! 내가 정답을 놓칠 수도 있겠어!"
안전망: AI가 확신이 없을 때, 특별한 규칙을 실행합니다: "예산을 두 배로 늘려라!" 단순히 k개의 덩어리를 고르는 대신, 해당 이야기 부분에 대해서는 2k개의 덩어리를 확보합니다. 더 안전하게 가기 위해 에너지를 조금 더 쓰는 것입니다.

이것은 AI 전체를 바꾸는 마법 같은 주문이 아닙니다. AI가 이미 사용 중인 선택 방식 위에 얹혀 있는 작고 똑똑한 레이어입니다. 마치 조종사가 혼란스러워 보일 때만 운전대를 잡는 부조종사와 같습니다.

논문이 실제로 밝혀낸 것 (증거)

저자들은 단순히 추측한 것이 아니라, 네 가지 다른 AI 모델(Qwen 및 Mistral 포함)과 두 가지 주요 테스트 세트를 통해 검증했습니다. 수치가 말해주는 결과는 다음과 같습니다.

큰 승리: 어려운 테스트인 LongBench-v2에서, 기존 방식(단순히 상위 k개 선택)은 0.47의 "페어드 리콜(paired recall)" 점수를 기록했습니다. 이는 올바른 단서를 절반도 안 되는 확률로 찾아냈음을 의미합니다. 새로운 "불확실성 게이트형" 방식은 이 점수를 0.75로 끌어올렸습니다. 이는 28 퍼센트 포인트라는 엄청난 상승입니다.
속도: 불확실성을 확인하는 과정이 속도를 늦출 것이라고 생각할 수도 있습니다. 놀랍게도 그렇지 않습니다. 매우 긴 길이(128K 토큰)에서, 새로운 방식은 전체를 다 읽는 느린 "밀집(dense)" 방식 시간의 0.62배로 실행되었습니다. 기존의 지름길 방식들보다 훨씬 빠르면서도 훨씬 똑똑했습니다.
"건초더미 속의 바늘" 테스트: AI가 숨겨진 특정 사실을 찾아내야 하는 합성 테스트인 RULER NIAH에서, 이 새로운 방식은 완벽하지만 느린 방식이 찾아낸 정답의 0.81에서 0.89 사이를 찾아냈으며, 여전히 훨씬 빠르게 작동했습니다.

논문이 제외한 범위 (금지 구역)

이 방법이 하지 못하는 일을 아는 것도 중요합니다. 저자들도 이 점을 명확히 했습니다.

짧은 이야기를 위한 마법의 해결책은 아닙니다: 저자들은 이야기가 짧아서 AI가 모든 것을 쉽게 볼 수 있었던 LongBench-v1에서도 테스트를 진행했습니다. 그런 경우, 새로운 방식은 도움이 되지 않았습니다. "성능 향상"은 이야기가 너무 길어서 AI가 반드시 까다롭게 골라야만 하는 상황에서만 발생합니다. 여유 공간이 충분하다면 추가적인 확인 작업은 불필요합니다.
"점수 산정" 시스템의 대체물이 아닙니다: 논문은 덩어리를 점수 매기는 두 가지 방식(K-mean과 Quest)을 테스트했습니다. 새로운 방식은 이 두 가지 모두에서 작동했습니다. 어떤 점수 산정 시스템을 사용하든 상관없습니다. "불확실성 체크"가 여러분이 가진 방식이 무엇이든 더 좋게 만들어 줍니다.
모든 것에 대한 완벽한 해결책은 아닙니다: 저자들은 매우 까다로운 추론 작업(예: 3단계 홉(hop)이 필요한 변수 추적)에서는 최고의 모델들도 고전했으며, 새로운 방식이 이를 완전히 해결할 수는 없었다고 인정했습니다. 이는 모델 자체가 더 똑똑해져야 하는 문제이지, 선택기의 문제가 아니라는 점을 시사합니다.

핵심 요약

이 논문은 AI의 의사결정 과정에 간단한 "신뢰도 체크"를 추가함으로써, 점수가 비슷하다는 이유로 중요한 단서를 버리는 일을 막을 수 있다고 제안합니다.

결과는 이 접근 방식이 AI 모델이 긴 텍스트를 읽는 능력을 속도를 늦추지 않으면서도 측정 가능한 수준으로 개선함을 보여줍니다. 이것은 "맹목적인 추측"을 가장 필요할 때 수행하는 "신중한 재확인"으로 바꿔놓았습니다. 저자들은 이 방식이 다양한 유형의 AI 모델과 텍스트 길이에 걸쳐 작동한다는 것을 발견했으며, 이는 때때로 가장 빠른 방법은 언제 속도를 늦춰야 할지 똑똑하게 판단하는 것임을 증명합니다.

기술 요약: 블록 희소 어텐션을 위한 불확실성 게이트 선택 (Uncertainty-Gated Selection for Block-Sparse Attention)

문제 정의

블록 희소 어텐션 메커니즘(예: SSA, Quest, H2O, SnapKV)은 $O(N^2)$ 소프트맥스를 각 쿼리에 대한 키 블록의 top- $k$ 선택으로 대체하여 긴 컨텍스트 언어 모델을 확장합니다. 효율적이기는 하지만, 이 선택 과정은 **근시적(myopic)**입니다. $k$ 번째와 $(k+1)$ 번째 블록의 점수가 거의 비슷할 때, 선택기는 추가적인 예산 없이 결정을 확정해 버립니다. 만약 탈락된 블록에 결정적인 증거(예: 멀티 홉 추론 또는 '바늘 찾기' 검색 시)가 포함되어 있다면, 정보는 돌이킬 수 없이 손실됩니다. 하위 레이어는 이를 복구할 수 없기 때문입니다. 이는 단순한 노이즈 문제가 아니라 구조적인 실패이며, 특히 관련성이 표면적인 쿼리-키 매칭보다는 잠재적 컨텍스트에 의존하는 시나리오에서 성능을 저하시킵니다.

기존 방법들은 컷오프(cutoff)를 고정된 결정 규칙으로 취급합니다. 이들은 해당 결정의 불확실성이나, 컷오프가 모호할 때 더 많은 정보를 확보함으로써 얻을 수 있는 **정보의 가치(Value of Information, VoI)**를 고려하지 않습니다.

방법론

본 논문은 추가적인 학습이나 파라미터 없이 기존의 블록 스코어링 백본(SSA 방식의 mean-pooling 또는 Quest의 K-min/K-max 상한값 등) 위에 작동하는 **불확실성 게이트 라우터(uncertainty-gated router)**를 제안합니다. 이 방법은 다음 5단계로 작동합니다:

블록 스코어링 (백본 불가지론적): 키(Keys)를 블록 단위로 그룹화합니다. 기존 백본(예: 평균 풀링 내적 또는 좌표별 상한값)을 사용하여 점수를 계산합니다.
타일별 선택 (Per-Tile Selection): 계산 오버헤드를 줄이기 위해, 선택은 행(row) 단위가 아닌 쿼리 타일(연속된 쿼리 행들의 집합) 단위로 수행됩니다. 타일 점수는 타일 내 행들 중 최대 점수입니다. top- $k$ 선택을 통해 초기 블록 세트( $kv\_idx$ )를 결정합니다.
불확실성 측정 (VoI 신호): 각 타일과 헤드에 대해 **정규화된 컷오프 마진( $\sigma$ $σ$ )**을 계산합니다:
$\sigma = \frac{s_{(k-1)} - s_{(k)}}{s_{(0)} - s_{(k)}}$
여기서 $s_{(i)}$ $s_{(i)}$ 는 정렬된 블록 점수입니다.
- $\sigma \to 1$ : 컷오프가 결정적임 (유지된 세트가 꼬리 부분과 격차가 큼).
- $\sigma \to 0$ : 컷오프가 모호함 ( $k$ 번째와 $(k+1)$ 번째 블록이 동률임).
  이 $\sigma$ 는 top- $k$ 결정으로 인한 예상 손실의 대리 지표 역할을 합니다. 낮은 $\sigma$ 는 더 많은 정보를 획득할 가치가 높음을 나타냅니다.
집계 및 트리거링: 헤드별 $\sigma$ 값들을 평균하여 타일당 $\bar{\sigma}$ 신호를 생성합니다. $\bar{\sigma}$ 가 레이어별 경험적 분위수 임계값( $q$ ) 미만으로 떨어지면 트리거가 활성화됩니다. 이는 결정이 모호했던 "위험한" 타일을 식별합니다.
제어된 확장: 트리거된 타일은 어텐션 세트를 $\rho$ 배수(예: $\rho=2$ )만큼 확장하여, top- $k$ 대신 top- $\rho k$ 블록을 참조합니다. 트리거되지 않은 타일은 기본 예산을 유지합니다. 구현에는 비-트리거 타일을 센티널(sentinel)로 패딩하여 균일한 디스패치 형상을 유지하는 퓨즈드 커널(fused kernel)을 사용하여, 확장이 발생하더라도 비용이 서브 리니어하게 스케일링되도록 합니다.

주요 기여

VoI 공식화: 선택기의 컷오프를 불확실성 하의 결정으로 보는 새로운 백론-불가지론적(backbone-agnostic) 공식을 제안합니다. 이는 타일당 단 하나의 스칼라( $\sigma$ )와 레이어당 하나의 분위수 임계값만을 추가하며, 재학습이나 추가 파라미터를 요구하지 않습니다.
결합 가능성 (Composability): 제안된 라우터가 서로 다른 스코어링 백본(SSA 방식의 K-mean 및 Quest의 K-max)과 결합 가능하다는 것을 실증적으로 보여줍니다. 라우터는 특정 모델에 더 우수한 백본의 성능을 일관되게 끌어올리며, 이는 "더 나은 스코어링"과 "더 나은 예산 할당"이 서로 독립적인 개선 요소임을 입증합니다.
퓨즈드 구현: 선택과 확장을 단일 코드 경로에서 처리하고 형상 균일 디스패치를 유지하는 커스텀 Triton 커널을 사용합니다. 이를 통해 시스템은 128K 컨텍스트에서 밀집(dense) 어텐션 시간의 $0.62\times$ 에서 $0.80\times$ 수준으로 실행될 수 있습니다.
진단적 벤치마킹: 선택기의 품질을 모델의 능력과 분리하기 위해 Pointer-Chase Haystack (PCH) 벤치마크를 도입하였으며, LongBench-v1을 네거티브 컨트롤(negative control)로 사용하여 라우터의 이점이 선택기 예산이 증거 밀도에 비해 타이트할 때 조건부로 발생함을 입증했습니다.

결과

이 방법은 32K에서 128K 컨텍스트에 걸쳐 네 가지 모델(Qwen2.5-14B, Qwen2.5-7B-1M, Mistral-Nemo, Qwen3.6-35B)에 대해 평가되었습니다.

성능 향상:
- LongBench-v2 medium (32K 컨텍스트)에서, Quest 기반 라우터 설정은 top-k 베이스라인의 0.47 대비 **0.75의 페어드 리콜(paired recall)**을 달이트 (28 퍼센트 포인트 상승), McNemar $p < 0.01$ .
- RULER NIAH multi-key에서, Quest 기반 라우터는 Qwen2.5-14B에서 0.98 정확도(top-k의 0.93 대비), Qwen2.5-7B-1M에서 1.00에 도달했습니다.
- 라우터는 각 모델에서 승리한 백본의 성능을 일관되게 향상시켰습니다. 특히, 최적의 백본은 아키텍처에 따라 바뀌었습니다: Qwen2.5/Mistral에서는 Quest의 K-max가 승리한 반면, Qwen3.6에서는 SSA 방식의 K-mean이 승리했습니다(QK-Norm이 크기 분산을 규제하기 때문). 라우터는 두 경우 모두 성능을 높였습니다.
효율성:
- 128K 컨텍스트에서, 라우터는 Qwen2.5-7B-1M에서 밀집 정확도의 0.81, Qwen3.6에서 0.89를 보존하면서, 각각 밀집 어텐션 시간의 $0.62\times$ 및 $0.80\times$ 로 실행되었습니다.
- 더 긴 컨텍스트(64K–128K)에서 전체 연산량 중 어텐션 점유율이 증가함에 따라 "파레토 승리"(유사한 품질에서의 더 빠른 속도)가 나타나며, 이는 암달의 법칙(Amdahl's law)과 일치합니다.
예산 vs 선택성: 절제 연구(Ablation studies)는 라우터의 성능 향상이 예산 증가와 선택적 할당 모두에서 온다는 것을 보여줍니다. 성능 한계에 근접한 모델(예: Qwen-14B)에서는 향상의 대부분이 예산 기반인 반면, 여유가 더 많은 모델(예: Mistral-Nemo)에서는 선택성(모호한 타일에만 예산을 투입하는 것)이 이득의 대부분을 제공합니다.

의의 및 주장

본 논문은 불확실성 게이트 라우터가 근시적인 top- $k$ 선택의 구조적 실패 모드를 완화하는 범용적이고 플러그 앤 플레이(plug-and-play)가 가능한 레이어라고 주장합니다. 그 의의는 다음과 같습니다:

강건성: 수동 튜닝 없이도 다양한 모델 아키텍처의 특정 스코어링 특성(예: Qwen3.6의 QK-Norm 효과 처리)에 적응합니다.
효율성: 128K 컨텍스트에서 밀집 정확도에 근접하면서도 실행 시간을 크게 단축하여, 긴 컨텍스트 추론을 더 실용적으로 만듭니다.
이론적 근거: 선택 문제를 정보의 가치(VoI) 관점에서 프레임화하며, 가우시안 노이즈 가정 하에서 유한 샘플 회귀 상한(finite-sample regret bound)을 제공합니다 (단, 논문은 현재 방법이 작동하는 노이즈 지배 영역에서는 이 상한이 약하다고 언급함).

저자들은 라우터의 성능 향상이 조건부임을 명시적으로 밝힙니다. 즉, 답변 관련 증거의 위치에 비해 쿼리당 예산이 적을 때만 활성화됩니다. 예산이 넉넉한 시나리오(예: 짧은 프롬프트를 가진 LongBench-v1)에서는 라우터가 미미한 이점만을 제공하며, 이는 해당 방법의 특이성을 검증하는 네거티브 컨트롤 역할을 합니다.