계층형 인덱싱으로 효율적인 미세입자 희소 어텐션
HISA는 기존 토큰‑레벨 희소 어텐션의 인덱싱 비용을 두 단계의 계층적 검색으로 대체한다. 먼저 블록 수준에서 평균 풀링된 대표 벡터를 이용해 상위 m개 블록을 선별하고, 그 블록 내부에서만 원래 인덱서를 적용해 최종 k개의 토큰을 선택한다. 이 과정은 학습 없이 바로 적용 가능하며, 32K 길이에서 2배, 128K 길이에서 4배의 커널 속도 향상을 보인다. 선택 정확도와 최종 모델 성능은 기존 DSA와 거의 동일하고, 토큰 선택 집합의 Io…
저자: Yufei Xu, Fanxu Meng, Fan Jiang
본 논문은 초대형 언어 모델이 4K‑128K 토큰 이상의 긴 컨텍스트를 처리할 때 발생하는 “인덱서 병목”을 해결하기 위해 HISA(Hierarchical Indexed Sparse Attention)라는 새로운 인덱싱 메커니즘을 제안한다. 기존 DeepSeek Sparse Attention(DSA)와 같은 토큰‑레벨 희소 어텐션은 각 쿼리마다 전체 프리픽스 토큰을 스코어링해 상위 k개를 선택한다. 이 과정은 토큰 수 L에 대해 O(L) · L = O(L²) 연산을 요구해, 컨텍스트가 128K 토큰을 넘어가면 인덱서가 전체 연산 비용의 대부분을 차지한다.
HISA는 이 문제를 두 단계의 계층적 검색으로 재구성한다. 첫 단계는 블록‑레벨 코스 필터링이다. 프리픽스를 크기 B인 연속 블록으로 나누고, 각 블록의 인덱싱 키를 평균 풀링해 대표 벡터 \tilde{k}_b를 만든다. 쿼리와 동일한 가중치·쿼리 표현을 재사용해 \tilde{k}_b와의 내적을 계산함으로써 블록 스코어 J_{t,b}를 얻는다. 여기서 상위 m개 블록을 선택하면 전체 토큰 중 m·B개만이 후보 집합 Ω_t에 남는다. 이 단계는 블록 수 ⌈L/B⌉에 비례하므로 O(L/B) 비용이다.
두 번째 단계는 토큰‑레벨 정밀 스코어링이다. 후보 토큰 Ω_t에 대해서만 기존 DSA와 동일한 스코어링 식 I_{t,s}=∑_j w_{t,j}·ReLU(q_{t,j}^⊤k_s) 를 적용한다. 최종적으로 상위 k개의 토큰을 선택해 T^H_t를 만든다. 이 단계는 O(m·B) 비용이며, 전체 레이어에 대해 O(L·m·B)로 축소된다.
복잡도 분석 결과, HISA의 전체 레이어 비용은 O(L²/B + L·m·B)이며, B와 m을 적절히 조정하면 O(L²) 대비 수십 배의 절감이 가능하다. 특히 B≪L, m≪L/B인 경우(초장문 상황) 효율이 크게 향상된다.
실험에서는 세 가지 축을 평가했다. (1) 커널 레벨 지연: TileLang GPU 커널 구현으로 32K 컨텍스트에서 2×, 128K에서 4× 속도 향상을 확인했다. (2) Needle‑in‑a‑Haystack 검색 정확도: HISA는 블록‑스파스 대비 높은 정확도를 유지했으며, 원본 DSA와 거의 동일한 성능을 보였다. (3) LongBench 전반적인 다운스트림 성능: 파인튜닝 없이 바로 적용했을 때 HISA‑버전 DeepSeek‑V3.2는 원본과 차이가 거의 없었다. (4) 선택 일관성: HISA와 DSA가 선택한 토큰 집합의 평균 IoU가 99 % 이상으로, 계층적 필터링이 실질적인 토큰 손실을 일으키지 않음을 입증했다.
또한 HISA는 완전한 드롭‑인 교체 방식이다. 인덱서 모듈만 교체하면 되며, KV 캐시 레이아웃, Sparse MLA 연산, 모델 파라미터 등은 전혀 변경되지 않는다. 따라서 기존 서비스 파이프라인에 바로 적용 가능하고, 추가 학습이나 파인튜닝이 필요 없다는 실용적 장점이 있다.
관련 연구와 비교했을 때, 블록‑스파스 방법(예: MoBA, Native Sparse Attention)은 블록 단위로 전체 토큰을 참석시켜 세밀한 토큰 선택이 불가능하고, 토큰‑레벨 방법(예: QUEST, H2O)도 여전히 전체 프리픽스 스코어링을 수행한다. HISA는 블록을 “검색 가속기”로만 사용하고, 최종 토큰 선택은 기존 DSA와 동일하게 유지함으로써 두 접근법의 장점을 결합한다.
마지막으로, 논문은 향후 연구 방향으로 동적 블록 크기·예산(m) 조정, 하드웨어 특화 블록 풀링, 멀티‑헤드 간 협업 필터링, 그리고 1M 토큰 수준의 초초대형 컨텍스트에 대한 확장성을 제시한다. HISA는 현재 LLM 서비스에서 긴 컨텍스트를 효율적으로 다루기 위한 실용적인 솔루션으로 자리매김할 가능성이 크다.
원본 논문
고화질 논문을 불러오는 중입니다...
댓글 및 학술 토론
Loading comments...
의견 남기기