Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than 64times the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.

한 줄 요약

HiLS-Attention은 학습 가능한 청크 선택을 통해 긴 문맥에서도 효율적이고 정확한 어텐션을 제공하는 새로운 스파스 어텐션 메커니즘이다.

핵심 기여도

HiLS-Attention은 랜드마크 토큰을 사용한 학습 가능한 청크 요약을 도입하여, 전체 어텐션과 유사한 성능을 달성함.
8K 길이로 학습된 모델이 4M 길이의 문맥에서 90% 이상의 정확도로 NIAH(Needle-in-a-Haystack) 검색 성능을 유지함.
기존 전체 어텐션 모델을 HiLS-Attention으로 전환할 수 있으며, 50B 토큰의 추가 학습만으로도 성능 유지 및 긴 문맥 추론 능력 향상 가능.
HiLS-Attention은 키-벨류 캐시 접근을 스파스화하여 효율성-성능 간의 트레이드오프를 극복함.

핵심 아이디어

HiLS-Attention은 기존 스파스 어텐션의 주요 문제인 **정확하지 않은 청크 선택**을 해결하기 위해, 청크 요약을 **학습 가능한 랜드마크 토큰**을 기반으로 생성하고, 이를 어텐션 계산에 직접 통합함. 이는 전체 어텐션에서 유도된 청크 중요도를 1차 테일러 전개를 통해 근사하는 방식으로, **학습 가능한 청크 중요도 스코어**를 도출함.

기존 방법은 청크 요약을 단순 평균 풀링으로 생성하고, 선택 후에는 요약 정보를 버리는 반면, HiLS는 요약과 선택 스코어를 **LM 손실 함수에 직접 연결**하여 **엔드투엔드 학습**이 가능하도록 함. 이는 **불필요한 청크를 억제하고, 예측에 유용한 청크에 더 높은 스코어를 할당**할 수 있게 해준다.

HiLS-Attention은 **계층적 인수분해**(hierarchical factorization)를 통해 어텐션 질량을 청크 간 분배한 후, 청크 내 토큰 간 분배함으로써, **전체 어텐션 계산 없이도 효율적인 어텐션 할당**을 가능하게 함.

기술적 접근법

**랜드마크 토큰**(landmark token)을 청크 맨 끝에 추가하여, 청크 요약 키를 생성함.
청크 중요도 스코어는 전체 어텐션 유도 질량의 1차 테일러 전개를 근사한 형태로 계산됨.
어텐션 계산 시, 청크 선택 스코어를 **전방 어텐션 가중치에 통합**하여, LM 손실 함수에 직접 연결함.
**계층적 인수분해**(hierarchical factorization)를 통해 청크 간 및 청크 내 토큰 간 어텐션 질량을 분배함.
**스파스 KV 캐시 접근**을 통해 메모리 사용량을 줄이고, **계산 비용을 O(n log n)** 수준으로 유지함.
기존 전체 어텐션 모델은 **50B 토큰의 추가 학습**만으로 HiLS-Attention으로 전환 가능함.

주요 결과

**345M 규모 모델**에서 HiLS-Attention은 8K 학습 길이로 4M 길이의 문맥에서 **90% 이상의 NIAH 정확도**를 달성함.
**256K 학습 길이에서** HiLS-Attention은 변수 추적(VT) 작업에서 전체 어텐션 대비 **최대 50% 성능 향상**을 보임.
**7B 규모 모델**에서 HiLS-Attention은 50B 토큰의 추가 학습만으로도 **LongBench 벤치마크에서 기존 전체 어텐션 기반 모델을 상회**함.
**RULER 벤치마크**에서 HiLS-Attention은 단일 바늘(S-N), 다중 키-쿼리(MK-MQ), 변수 추적(VT) 작업에서 모두 기존 스파스 어텐션 모델 대비 **더 높은 정확도**를 보임.

의의 및 한계

HiLS-Attention은 **긴 문맥 모델링**에서 **성능-효율성 트레이드오프를 극복**하는 첫 번째 스파스 어텐션 메커니즘으로, **긴 문맥 추론 및 추론 정확도를 동시에 향상**시킨다. 특히, **엔드투엔드 학습 가능한 청크 선택**은 기존 스파스 어텐션의 주요 한계를 해결하며, **긴 문맥에서의 추론 정확도를 크게 향상**시킨다.

하지만, HiLS-Attention은 **초대규모 모델**(예: 100B 이상)에서의 성능은 명시되지 않았으며, **복잡한 다중 단계 추론 작업**에서의 효과성도 추가 실험을 통해 검증이 필요하다. 또한, **랜드마크 토큰 기반 요약**이 모든 유형의 문맥에서 최적의 정보를 유지하는지에 대한 분석도 추가 연구 주제로 남아 있다.

실용적 활용

HiLS-Attention은 **대규모 언어 모델의 긴 문맥 추론**을 요구하는 산업 분야(예: 문서 분석, 대화 시스템, 코드 생성)에서 유용하게 활용될 수 있다. 특히, **메모리 효율적인 추론**이 필수적인 모바일 및 임베디드 환경에서 **긴 문맥 처리 능력을 유지하면서도 계산 비용을 줄일 수 있다**. 또한, **기존 전체 어텐션 모델을 스파스화하여 재사용**할 수