SAS로 언어 모델링 손실을 활용한 Attention 희소화

왜 중요한가

긴 Context의 Decode에서는 매 생성 단계마다 전체 KV Cache를 읽어 Attention 비용이 Context 길이에 따라 증가합니다. SAS는 사전학습된 Dense LLM의 Backbone을 유지하면서 Selector만 언어 모델링 손실로 학습해, 제한된 Attention 예산에서도 예측에 실제로 기여하는 Context를 고르는 방향으로 희소화를 최적화합니다. Qwen3-4B에서는 512K Context에서 Dense Attention보다 최대 5.6배 낮은 Decode 지연을 기록했습니다.

핵심 기여

언어 모델링 손실로 Context 순위 직접 학습

SAS는 Dense Attention 분포를 층별로 모방하는 대신, Selector 점수를 Attention Logit 내부의 연속 Gate로 연결합니다. 이에 따라 최종 다음 Token 예측 손실이 선택된 Context Block의 상대적 우선순위까지 갱신합니다.

세 가지 핵심 게이팅 설계 확립

Gate를 Softmax 내부에 Log 형태로 넣고, Historical Block 점수를 Softmax로 정규화하며, Hard Binary Mask 대신 연속 점수 차이를 보존합니다. 이 조합이 Current Block과 Historical Context의 균형 및 안정적인 Gradient를 유지합니다.

장문 학습을 위한 Triton Kernel 구현

SAS는 선택된 Block의 Log Gate를 FlashAttention 방식의 Tile 단위 Query-Key 연산에 융합합니다. 전체 Gated Attention 행렬을 물질화하지 않고 선택 Mask, Online Softmax, Block별 Gate Gradient를 처리해 장문 학습의 메모리 부담을 낮춥니다.

제한된 예산에서 여러 과제 성능 개선

Qwen3-4B·8B·14B에서 MATH500, GPQA-Diamond, AIME24·25, LongBench, BFCL, VitaBench를 평가한 결과, SAS는 대부분의 예산과 모델 크기에서 SeerAttention-R보다 높은 점수를 기록했습니다.

핵심 아이디어 이해하기

일반적인 Transformer Attention은 현재 Query가 이전 모든 Key와 Value를 읽으므로, 길이 n인 시퀀스의 Autoregressive Decode 누적 비용이 O(n²)로 증가합니다. Block Sparse Attention은 Context를 여러 Block으로 나눈 뒤 일부만 읽어 한 Query의 비용을 O(|S|)로 줄이지만, 어떤 Block을 남길지는 별도의 Selector가 결정해야 합니다. 문제는 Hard Top-K가 점수 순위만 사용하고 선택 경로를 불연속적으로 만들기 때문에, 최종 언어 모델링 손실이 Selector를 직접 교정하지 못한다는 점입니다.

기존 학습 방식은 Dense Attention이 어느 Block에 확률 질량을 배분했는지를 Selector의 목표로 삼습니다. 그러나 최종 예측은 Attention 확률뿐 아니라 Value 행렬과 여러 Layer의 상호 보완적 선택에도 좌우됩니다. 따라서 한 Layer의 Attention 분포를 잘 모방한 순위가 제한된 예산에서 다음 Token 예측에 가장 유용한 순위와 일치하지 않을 수 있습니다.

SAS는 Selector가 산출한 점수를 먼저 Softmax로 정규화한 뒤 Gate로 바꾸고, 선택된 Block의 Token마다 해당 Gate의 log 값을 Query-Key Logit에 더합니다. 입력은 Query, Context Key, Context Value이며, Selector가 Block 점수를 계산하고 Top-K Block을 고른 다음 softmax(qKᵀ + log g)V를 계산합니다. 이때 Gate가 Softmax 안에서 Attention 질량을 재배치하므로, 언어 모델링 손실의 Gradient가 선택된 Block의 상대적 유용성을 반영해 Selector를 갱신합니다.

추론 시에는 학습된 연속 순위를 Hard Top-K Block Index로 바꾸고 선택된 KV Block만 읽습니다. 따라서 학습 중에는 연속 Gate로 순위 정보를 보존하고, 추론 중에는 고정된 Attention 예산으로 계산량을 제한하는 구조입니다. 예를 들어 Qwen3-4B의 512K Context에서 Batch 1 Decode 지연이 Dense Attention 대비 5.6배 낮아졌습니다.

방법론

Context Token을 Block 크기 64로 나누고, 항상 유지하는 Current Block B₀와 Selector가 순위를 매기는 Historical Block으로 구분합니다. Selector는 Query와 각 Block의 Key를 입력으로 받아 점수 s를 산출하고, g = softmax(s)로 Historical Block Gate를 정규화합니다. Top-K Gate에 해당하는 Block과 Current Block의 Token만 선택 집합 S로 구성합니다.

선택된 Token에 Broadcast된 Gate를 Attention Logit에 더해 o_SAS = softmax(qK_Sᵀ + log g_S)V_S를 계산합니다. Current Block의 Gate는 1로 고정되어 Log Bias가 0이며, Historical Block은 정규화된 Log Gate를 받습니다. Softmax 내부에 Gate를 넣으면 원래 Attention 확률이 고정된 채 Value만 축소되는 것이 아니라, 전체 Attention 질량의 배분 자체가 바뀝니다.

Gate Activation은 Historical Block 간 Softmax를 사용합니다. 점수에 같은 상수 c를 더해도 상대적 순위가 유지되며, log g_m = s_m - LSE(s)가 Current Block의 무편향 Gate와 비교 가능한 값을 만듭니다. 반면 Sigmoid Gate는 1로 포화되고 Raw Logit 주입은 0 주변으로 수렴해 Historical Block과 Current Block 사이의 구별 신호가 약해졌습니다.

Ranking Preservation 실험에서는 연속 Soft Gate와 Straight-Through Estimator 기반 Hard Gate를 비교했습니다. 연속 Gate는 선택되지 않은 후보까지 하나의 Softmax 정규화 안에서 제한된 확률 질량을 가지므로 Gradient가 안정적으로 유지됩니다. Hard Gate는 초기 순위가 부정확할 때 선택되지 않은 높은 점수 Block의 기여가 제한된 정규화에서 과도하게 커져 Gradient Norm이 크게 튀는 문제가 생겼습니다.

Training Scope는 전체 Context를 대상으로 하는 방식과 Top-K 선택 Block만 대상으로 하는 Sparse Scope를 비교했습니다. Sparse Scope는 초기 수렴이 느리지만 계산량이 낮고, Table 1의 Qwen3-4B GPQA-Diamond 실험에서 1 epoch 후 Full Scope 54.4와 비교해 54.8의 점수를 기록했습니다. 이후 실험은 비용이 낮은 Sparse Scope를 주로 사용했습니다.

구현에서는 Triton 기반 FlashAttention 스타일 Kernel이 Tile 단위로 Query-Key Logit을 계산하면서 선택된 Historical Block의 Log Gate를 더하고, 비선택 Block을 Mask 처리하며, Current Block은 그대로 둡니다. Backward 단계에서는 선택된 Block 내부 Token의 Logit Gradient를 합산해 Selector로 전달하므로 전체 Attention 행렬을 저장하지 않고도 학습이 가능합니다.

주요 결과

Qwen3-4B·8B·14B의 Reasoning 평가에서 SAS는 낮은 Attention 예산일수록 SeerAttention-R보다 큰 차이를 냈습니다. 예를 들어 예산 1024에서 MATH500은 Qwen3-4B 90.65, 8B 91.27, 14B 92.93이었고, GPQA-Diamond은 각각 50.41, 53.17, 61.14였습니다. 예산 2048의 Qwen3-4B AIME24에서는 SAS 68.85가 SeerAttention-R 55.83보다 13.02포인트 높았습니다.

LongBench에서는 수학 데이터로만 학습한 Selector가 분포가 다른 장문 입력에도 전이됐습니다. 예산 2048의 Qwen3-14B 8K+ 구간에서 SAS는 53.9로 SeerAttention-R 51.5보다 2.4포인트 높았고, 예산 4096에서는 평균 56.2로 Dense Attention 56.6에 근접했습니다. BFCL Multi-Turn에서도 예산 2048의 Qwen3-4B에서 SAS 32.50이 SeerAttention-R 29.00보다 3.5포인트 높았습니다.

계속 사전학습한 OLMo3-7B에서는 SAS-RoPE의 평균 점수가 43.28로 OLMo3-512SWA 43.24보다 높고 Dense OLMo3-Base 43.88에 근접했습니다. LongBench 평균은 30.0으로 Dense Base 29.0과 Sliding-Window 28.0을 웃돌았으며, 8K 초과 입력에서 이점이 집중됐습니다.

SGLang의 단일 GPU Decode 측정에서 SAS는 Batch 1 기준 64K, 256K, 512K Context에서 Dense Attention보다 각각 2.4배, 4.6배, 5.6배 빨랐습니다. Batch 8에서는 64K Context에서 약 13배의 속도 향상이 기록됐지만, 512K에서는 Top-K 선택이 전체 Decode 단계의 90%를 차지해 장기적으로 Selector Scoring과 순위 계산이 주요 병목으로 남았습니다.

기술 상세

표준 Attention은 Query q, Key K, Value V를 입력으로 받아 o = softmax(qKᵀ)V를 계산합니다. q는 현재 Token의 Query Vector이고 K와 V는 n개 이전 Token의 Key·Value Matrix입니다. 전체 이전 Token을 읽으므로 한 Decode 단계 비용은 O(n)이고, 길이 n 전체의 누적 비용은 O(n²)입니다. Block Sparse Attention은 선택 Token 집합 S만 읽어 o = softmax(qK_Sᵀ)V_S를 계산하므로 한 단계 비용이 O(|S|), 누적 비용이 O(n|S|)가 됩니다.

SAS에서 Selector는 s = Rθ(q, {K_Bm})으로 Historical Block별 점수를 만들고, g = softmax(s)로 Gate를 정규화합니다. Top-K Gate의 Block Index를 I로 두면 선택 Token 집합은 Current Block B₀와 선택 Historical Block의 합집합입니다. 학습 출력은 o_SAS = softmax(qK_Sᵀ + log g_S)V_S이며, 입력된 Log Gate가 Attention Softmax 내부의 상대 Logit을 바꿔 출력 Value의 기여도를 재배분합니다.

Gate를 Softmax 밖에서 softmax(qKᵀ)(g ⊙ V)처럼 적용하면 기존 Attention 확률 p가 고정되고 Value만 크기 조정됩니다. 반대로 Softmax 안의 Log Gate는 Gate가 적용된 확률 p̃와 출력 o의 차이를 이용해 Block 간 질량을 재배치하는 Gradient를 만듭니다. Table 1의 GPQA-Diamond 실험에서 Inner Softmax는 1 epoch 54.4, Outer Softmax는 41.6으로 측정됐습니다.

Softmax Gate는 log g_m = s_m - LSE(s)로 표현되어 모든 Historical Block의 점수를 정규화하고 Current Block의 Gate 1과 비교합니다. 예를 들어 모든 Selector 점수에 같은 상수 c를 더해도 Softmax의 상대 비율은 바뀌지 않으므로 Global Shift에 불변입니다. Sigmoid는 Gate가 1에 가까워지고 Raw Logit 주입은 점수 분산이 줄어들어 두 방식 모두 구별 가능한 순위 신호가 약해졌습니다.

FlashAttention 스타일 Triton Kernel은 전체 Attention Matrix를 만들지 않고 KV Tile을 순차적으로 읽습니다. 선택된 Historical Tile에는 정규화된 Log Gate를 더하고 비선택 Tile은 Mask하며, Current Tile은 Gate 없이 Online Softmax에 넣습니다. Backward에서는 선택된 Block 내부의 Attention-Logit Gradient를 합산해 Selector 점수로 전달합니다. 추론 단계에서는 Selector가 모든 Cached Block의 순위를 계산하고 Top-K KV Block만 FlashInfer Sparse Decode Kernel로 전달합니다.

학습된 순위는 각 Layer의 단독 Attention 질량을 최대화하는 방향과 달랐습니다. SAS는 Layer별로 Dense Attention 질량을 덜 덮었지만, Layer 전체에서 선택한 Block의 합집합이 Full Attention Oracle과 겹치는 Recall은 SeerAttention-R보다 높았습니다. 이는 언어 모델링 손실이 Layer 간 상호 보완적인 Context 선택을 함께 최적화한 결과로 해석됩니다.

실무 활용

SAS는 Dense LLM의 Backbone을 고정한 채 Block Selector만 추가 학습해 긴 Context Decode의 KV 읽기 범위를 제한하는 방식입니다. 저장소와 SGLang 기반 추론 구현이 제공되며, Prefill은 Dense Attention으로 유지하고 Decode 단계에서만 Block Sparse Attention을 적용합니다.

긴 문서 질의응답과 장문 Context 생성에서 KV Cache 읽기 비용 절감
Reasoning 모델의 제한된 Attention 예산 운용
BFCL·VitaBench와 같은 다중 턴 도구 사용 Agent의 장문 상태 처리
기존 Dense LLM을 재설계하지 않는 Post-training 희소화
계속 사전학습 과정에서 Backbone과 Selector를 함께 학습하는 Sparse Attention 구성

코드 공개 여부: 공개

코드 저장소 보기

키워드

SAS(Simple Attention Sparsification)Attention 희소화Top-K 선택Context RankingGated Sparse AttentionLog-space GateFlashAttentionTriton KernelLong-context Understanding

용어 해설

Attention 희소화(Attention Sparsification): — 각 Query가 전체 Context Token을 읽지 않고 일부 Token 또는 Block만 선택해 Attention을 계산하는 기법입니다. 시퀀스 길이에 따른 누적 Attention 비용을 줄이지만, 제한된 예산 안에서 어떤 Context를 남길지 결정하는 선택 품질이 성능을 좌우합니다.
Top-K 선택(Top-K Selection): — Selector가 계산한 중요도 점수에서 상위 K개의 Context Block만 고르는 방식입니다. 선택 집합이 점수 변화에 따라 불연속적으로 바뀌므로, 일반적인 Hard Top-K 경로에서는 언어 모델링 손실의 Gradient가 Selector까지 직접 전달되지 않습니다.
로그 공간 게이트(Log-space Gate): — Selector 점수에서 만든 Gate를 Attention Logit에 log 형태로 더하는 방식입니다. Attention Softmax 내부에서 Block별 확률 질량을 직접 조절하므로, Value를 사후에 단순히 축소하는 방식보다 Context 간 상대적 우선순위를 학습하기에 적합합니다.
FlashAttention: — Attention 행렬 전체를 메모리에 저장하지 않고 Tile 단위의 온라인 Softmax로 Query-Key와 Value 연산을 처리하는 메모리 효율적 Attention 구현입니다. SAS는 여기에 Block별 Log Gate와 선택 Mask를 결합한 Kernel을 추가합니다.
KV Cache: — Autoregressive 생성에서 이전 Token의 Key와 Value를 저장해 다음 Query가 재사용하는 메모리 구조입니다. 긴 Context에서는 매 Decode 단계마다 전체 Cache를 읽는 비용이 커지므로, SAS는 선택된 KV Block만 읽어 비용을 줄입니다.

SAS로 언어 모델링 손실을 활용한 Attention 희소화

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

Declarative Attention으로 KV 캐시 읽기 줄이기

무작위 선택으로 KV 캐시 축출 비용 줄이기

KV cache를 검색 가능한 벡터 공간으로 라우팅

Random Attention으로 KV 캐시 점수화를 제거한 추론 압축

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 기사

Declarative Attention으로 KV 캐시 읽기 줄이기

무작위 선택으로 KV 캐시 축출 비용 줄이기

KV cache를 검색 가능한 벡터 공간으로 라우팅

Random Attention으로 KV 캐시 점수화를 제거한 추론 압축