🤖 machine learning

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

SpotAttention은 동결된 사전 학습된 트랜스포머에 부착되어 KL 증류를 통해 어텐션 분포를 학습하는 경량화된 플러그인 블록 희소 라우팅 메커니즘으로, 기존 베이스라인 대비 훨씬 빠른 디코딩 속도와 감소된 메모리 사용량을 통해 최대 128K 토큰까지 정확한 롱 컨텍스트 추론을 가능하게 합니다.

원저자: Huzama Ahmad, Se-Young Yun

게시일 2026-06-23

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Huzama Ahmad, Se-Young Yun

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

당신에게 수백만 권의 책을 읽은 매우 똑똑한 사서(AI 모델)가 있다고 상상해 보세요. 당신이 질문을 던지면, 이 사서는 보통 답을 찾기 위해 자신이 읽었던 모든 페이지를 하나하나 다 기억해내려고 노력합니다.

문제는 도서관이 커질수록, 이 "모든 것을 기억하는" 방식이 믿기지 않을 정도로 느려지고 비용이 많이 든다는 점입니다. 이는 마치 계속해서 커지는 건초더미 속에서 특정 바늘 하나를 찾는 것과 같습니다.

SpotAttention은 사서가 전체 도서관을 다시 읽지 않고도 적절한 페이지를 찾을 수 있도록 돕는 가볍고 효율적인 "조수"입니다. 이것이 어떻게 작동하는지 쉬운 개념으로 나누어 설명해 드리겠습니다.

1. 문제점: "무거운 배낭"

AI가 매우 긴 문서(예: 소설 한 권 전체나 거대한 코드베이스)를 바탕으로 질문에 답하려고 할 때, AI는 지금까지 본 중요한 단어들을 담은 "배방"을 계속 유지해야 합니다.

기존 방식: AI가 새로운 단어를 생각할 때마다, 단서를 찾기 위해 배낭 전체를 테이블 위에 쏟아붓습니다. 문서가 길어질수록 배낭은 무거워지고 테이블은 너무 북적거리게 됩니다. 이로 인해 AI는 실행 속도가 느려지고 비용이 많이 듭니다.
"희소성(Sparse)" 아이디어: 다른 연구자들은 "이봐, AI는 정말로 몇 개의 특정 페이지만 보면 되잖아? 그 페이지들만 보자고"라고 제안했습니다. 하지만 어떤 페이지를 볼지 결정하는 것 자체가 이미 느리고 비싼 작업이었습니다. 그것은 마치 어떤 페이지를 읽을지 결정하기 위해 새로운 사람을 고용하는 것과 같았는데, 그 결정하는 사람이 책 전체를 읽는 것만큼이나 느렸습니다.

2. 해결책: "포착자(The Spotter)" (SpotAttention)

저자들은 사서에게 부착될 수 있는 작고 매우 빠른 "포착자(Spotter)"인 SpotAttention을 만들었습니다.

학습 방법: 포착자는 처음부터 새로 배워야 하지 않습니다. 포착자는 "전문가" 사서(사전 훈련된 AI)가 작업하는 모습을 관찰합니다. 이를 통해 전문가는 자연스럽게 어떤 페이지를 쳐다볼지를 모방하며, 어떤 페이지를 주목할지 알아내는 법을 배웁니다.
마법 같은 기술: 모든 단어를 다 읽어서 무엇을 남길지 결정하는 대신, 포착자는 텍스트 덩어리(예: 문단)를 보고 빠르게 점수를 매깁니다. 이는 마치 책장을 슥 훑어보고는 "나는 이 세 권의 책만 꺼내면 돼. 나머지는 그냥 선반에 두면 되겠어"라고 즉시 판단하는 사서와 같습니다.

3. "Dual Top-p" 규칙: 스마트한 예산

이 논문은 Dual Top-p라는 영리한 규칙을 소개합니다. 사서에게는 쳐다볼 수 있는 페이지의 예산이 있다고 상상해 보세요.

기존 방식: 어떤 시스템은 "상위 50%의 페이지만 볼 수 있다"라고 말합니다. 이는 경직되어 있습니다. 때로는 80%가 필요하고, 때로는 20%가 필요합니다.
SpotAttention 방식: 포착자는 페이지의 "중요도"를 살펴보고 이렇게 말합니다. "좋아, 이 특정 질문을 위해서는 처음 몇 페이지(시작 부분), 마지막 몇 페이지(방금 말한 내용), 그리고 가장 중요한 중간 페이지들을 챙겨야 해."
이는 동적으로 예산을 조정합니다. 질문이 단순하면 적은 페이지를 가져오고, 복잡하면 더 많은 페이지를 가져옵니다. 이 과정은 별도의 느린 결정 단계 없이 자동으로 이루어집니다.

4. 결과: 속도와 메모리

저자들은 이 방식을 여러 대규모 AI 모델(특히 Qwen 제품군)에 테스트하여 매우 긴 문맥(최대 128,000 단어)을 처리했습니다.

속도: 128,000 단어 길이에서, SpotAttention은 표준 방식(FlashAttention)보다 3.9배 더 빨랐고, 현재 가장 뛰어난 대안인 Twilight보다 1.8배 더 빨랐습니다.
정확도: 대부분의 텍스트를 건너뜀에도 불구하고, AI는 모든 내용을 다 읽었을 때와 마찬가지로 정답을 맞혔습니다. 즉, "지능"을 잃지 않았습니다.
메모리: 포착자의 "노트"(결정을 내리는 데 사용하는 캐시)는 특수 압축을 통해 아주 작은 크기로 줄일 수 있으며, 이 과정에서 정확도를 잃지 않습니다. 이는 컴퓨터 메모리를 크게 절약해 줍니다.

5. 차별점

재학습 불필요: AI를 처음부터 다시 가르칠 필요가 없습니다. 이미 완성되어 작동 중인 AI에 이 작은 포착자를 꽂기만 하면 됩니다.
규칙 기반이 아닌 학습 기반: 기존 방식은 고정된 규칙(예: "항상 중간은 건너뛰어라")을 사용했습니다. 반면 SpotAttention은 무엇이 중요한지에 대한 패턴을 학습하므로 더 똑똑하고 유연합니다.
어디서나 작동: 저자들은 다양한 크기의 AI 모델(40억 파라미터의 작은 모델부터 320억 파라미터의 큰 모델까지)에 대해 테스트를 진행했으며, 모든 모델에서 잘 작동했습니다.

요 요약하자면: SpotAttention은 AI 모델이 긴 문서를 읽을 때 집중해야 할 가장 중요한 부분을 빠르게 식별하도록 돕는 가볍고 학습된 조수입니다. 이는 AI가 지능을 떨어뜨리지 않으면서도 긴 텍스트를 처리할 때 훨씬 더 빠르고 저렴하게 실행될 수 있도록 만듭니다.

기술 요약: SpotAttention

문제 정의

사전 학습된 대규모 언어 모델(LLM)에서 긴 문맥(long-context) 능력은 전체 코드베이스를 처리하거나 방대한 증거에 기반하여 답변을 도출하는 등의 작업을 가능하게 하는 표준이 되었습니다. 그러나 셀프 어텐션(self-attention)의 계산 비용은 시퀀스 길이에 따라 비효로하게 증가합니다. 프리필(prefill) 단계에서는 계산량이 시퀀스 길이의 제곱에 비례하여 증가하며, 디코드(decode) 단계에서는 매 새로운 토큰마다 컨텍스트 크기에 따라 선형적으로 증가하는 키-값(KV) 캐시를 다시 읽어야 합니다. FlashAttention은 전체 어텐션 행렬을 실체화하지 않음으로써 I/O 효율성을 개선했지만, 근본적인 이차(quadratic) 프리필 비용이나 단계별 메모리 읽기의 선형적 증가 문제를 해결하지는 못했습니다.

기존의 희소 어텐션(sparse attention) 방식은 관련 있는 과거 토큰의 부분 집합에만 주의를 기울임으로써 이를 완화하려고 시도합니다. 그러나 이 부분 집합을 선택하는 과정 자체가 종종 비용이 많이 듭니다. Quest는 토큰 블록의 점수를 매기고, Twilight은 top- $p$ 프루닝(pruning)을 추가하지만, 선택기(selector) 자체가 새로운 병목 현상이 됩니다. 또한, 훈련이 필요 없는(training-free) 선택기들은 수작업으로 설계된 휴리스틱에 의존하는 경우가 많아 일반화가 어렵거나 여러 단계의 선택 과정을 거쳐야 하므로 누적 오버헤드가 발생할 수 있습니다.

방법론: SpotAttention

SpotAttention은 이미 학습이 완료되어 고정된(frozen) 트랜스포머에 장착할 수 있도록 설계된 플러그인 방식의 블록 희소 라우팅(block-sparse routing) 메커니즘입니다. 이는 모든 풀 어텐션(full-attention) 레이어에 부착되는 경량화된 "선택기(selector 또는 indexer)"로 구성됩니다.

1. 선택기 아키텍처

선택기는 축소된 너비( $d_{idx}$ )와 $H_{idx}$ 개의 헤드를 가진 매우 작은 멀티 헤드 Q-K 스코어러(scorer)입니다. 개별 토큰이 아닌, 쿼리와 키의 블록(예: 블록당 16개 토큰) 단위로 동작합니다.

스코어링 메커니즘: 위치 $q$ 에 있는 쿼리와 위치 $t$ 에 있는 키 블록에 대해, 선택기는 헤드 가중치가 적용된 ReLU 점수를 계산합니다:
$s_s[q, t] = \sum_{h=1}^{H_{idx}} w_h \text{ReLU}(\mathbf{q}^I_h \cdot \mathbf{k}^I_{h,t})$
여기서 $\mathbf{q}^I_h$ 와 $\mathbf{k}^I_{h,t}$ 는 헤드별 프로젝션이며, $w_h$ 는 은닉 상태(hidden state)로부터 유도된 혼합 가중치입니다.
효율성: 블록 단위로 점수를 매김으로써, 선택 과정은 단일 텐서 코어 행렬 곱셈(tensor-core matrix multiplication)으로 실행되어 토큰 단위 연산의 오버헤드를 피할 수 있습니다.

2. KL 증류를 통한 훈련

선택기는 고정된 백본(teacher)의 밀집 어텐션(dense attention) 분포를 대상으로 독립적으로 훈련됩니다.

목적 함수: 훈련은 훈련된 백본(teacher)의 어텐션 분포( $p_\tau$ )와 선택기의 추정 분포( $p_\sigma$ ) 사이의 쿨백-라이블러(Kullback-Leibler, KL) 발산을 최소화합니다.
변형 모델: 논문은 두 가지 KL 변형을 탐구합니다:
- DenseKL: 모든 $T$ 개의 키에 대해 손실을 계산합니다.
- SparseKL: 선택기가 선택한 자체 top- $K$ 집합에 대해서만 분포를 재정규화하여 손실을 계산합니다.
수렴: 선택기는 다양한 백본 크기(4B~32B) 및 아키텍처 전반에 걸쳐 약 200 스텝 이내에 빠르게 수렴하며, 훈련된 분포의 형태를 일치시키는 법을 학습합니다.

3. 이중 top- $p$ 선택 규칙 (Dual Top- $p$ Selection Rule)

추론 시, 선택기는 모든 과거 블록의 점수를 매깁니다. 어떤 블록에 주의를 기울일지 결정하기 위해, SpotAttention은 선택기의 교정된 분포에서 직접 예산을 읽어오는 이중 top- $p$ 규칙을 도입합니다.

예약된 블록 (Reserved Blocks): 핵(nucleus) 선택에 의해 지배되지 않도록 고정된 프리픽스 "싱크(sink)" 블록(예: 처음 128개 토큰)과 고정된 서픽스 "최신성(recency)" 블록(예: 마지막 256개 토큰)은 항상 유지됩니다.
핵 선택 (Nucleus Selection): 잔여 분포(싱크 및 최신성 블록을 제외한 블록)에 대해 top- $p$ (nucleus) 샘플링이 적용됩니다. 선택은 누적 확률 질량이 임계값 $p$ 에 도달할 때까지 진행됩니다.
동적 예산 (Dynamic Budget): 이 접근 방식은 별도의 프루닝 단계 없이도 콘텐츠에 적응하는 쿼리당, 레이어당 가변 예산을 생성합니다. 최소 하한선( $K_{min}$ )은 핵(nucleus)이 너무 작아지는 퇴보적인 경우를 방지합니다.

4. 양자화

선택기의 K-캐시는 INT4 또는 FP4 마이크로스케일로 양자화될 수 있습니다. 선택기의 캐시는 백본의 KV 캐시와 별도로 존재하므로, 백본의 정밀도나 정확도에 영향을 주지 않으면서 선택기의 메모리 사용량을 약 3.5배 줄일 수 있습니다.

주요 기여

플러그인 블록 희소 어텐션: 고정된 트랜스포머에 부착되어 KL 증류를 통해 훈련되는 경량 선택기입니다. 이는 단일 텐서 코어 행렬 곱셈으로 실행되어, 기존 연구에서 나타나는 무거운 선택기 오버헤드를 방지합니다.
이중 top- $p$ 규칙: 학습된 분포로부터 직접 쿼리별, 레이어별 예산을 도출하는 콘텐츠 적응형 선택 메커니즘입니다. 안정성을 보장하기 위해 싱크 및 최신성 블록을 예약하며, 이를 통해 정적 top- $K$ 의 정확도를 유지하면서도 더 작은 평균 예산을 달로합니다.
훈련 및 추론 동작: 논문은 선택기가 백본의 레이어별 예산 비대칭성(초기 레이어는 후기 레이어보다 더 많은 컨텍스트에 주의를 기울임)을 학습함을 보여주며, 이는 정적 $K$ 가 복제할 수 없는 패턴입니다. 이 방법은 16K 훈련 길이를 128K 평가 길이까지 정확도 저하 없이 효과적으로 일반화합니다.

결과

이 방법은 Qwen3 (4B–32B), Qwen3.5 (4B–9B)의 5개 백본을 대상으로 RULER, BABILong, InfiniteBench, LongBench-v2의 4가지 벤치마크에서 최대 128K 컨텍스트 길이에 대해 평가되었습니다.

정확도: SpotAttention은 컨텍스트 길이가 128K(훈련 길이의 8배)에 도달할 때까지 밀집 백본의 정확도와 일치하는 성능을 보입니다. 예산 $K=0.5L$ 에서 모든 모델에 대해 밀집 베이스라인의 한 번의 부트스트랩 표준 오차 이내의 성능을 유지합니다.
지연 시간 및 처리량:
- Qwen3-8B에서 $L=128K$ 일 때, SpotAttention은 FlashAttention보다 3.9배, 가장 강력한 훈련 불필요 베이스라인인 Twilight보다 1.8배 더 빠릅니다.
- 디코드 처리량은 128K에서 70 tokens/s, 512K에서 36 tokens/s를 유지하는 반면, FlashAttention과 Twilight는 각각 256K와 512K에서 메모리 부족(OOM)이 발생합니다.
- 선택기 자체는 Twilight의 선택기보다 훨씬 빠릅니다 (128K에서 1.55 ms vs 4.59 ms).
양자화: 선택기의 K-캐시를 INT4 또는 FP4로 양자화하면 측정 가능한 정확도 손실 없이 메모리 사용량을 3.5배 줄일 수 있습니다. 이는 단일 스트림 디코드 지연 시간(백본 KV 대역폭에 의해 결정됨)을 개선하지는 않지만, 메모리 사용량과 프리필 비용을 크게 줄여줍니다.

의의 및 주장

논문은 SpotAttention을 희소 어텐션의 "선택기 병목 현상"을 해결하기 위한 솔루션으로 제시합니다. KL 증류를 통해 교정된 분포를 학습함으로써, 선택기는 각 쿼리와 레이어의 특정 요구 사항에 적응하는 동적 핵 기반 예산을 가능하게 합니다.

저자들의 주장은 다음과 같습니다:

효율성: 긴 컨텍스트(최대 512K)에서 비교 대상 중 가장 높은 처리량을 달성하여, 밀집 모델 및 다른 희소 베이스라인을 멈추게 하는 메모리 제한을 극복합니다.
일반화: 16K 컨텍스트 길이로 훈련된 선택기가 재훈련 없이 128K 이상으로 효과적으로 일반화됩니다.
모듈성: 백본을 재학습할 필요가 없는 "플러그인" 솔루션으로, 기존에 사전 학습된 모델에 적용 가능합니다.
한계: 저자들은 평가된 가장 큰 백본이 32B라는 점과 결과가 현재 영어 벤치마크의 Qwen 제품군에 국한되어 있음을 언급했습니다. 또한, 희소 프리필 지연 시간은 벤치마킹되지 않았으나 아키텍처가 이를 지원한다는 점을 명시했습니다. 또한, 이중 top- $p$ 규칙은 선택기가 잘 교정되어야 한다는 점에 의존하며(SparseKL 훈련 시에는 성립하나 DenseKL 시에는 무너짐), 이를 인정했습니다.