HISA: 세밀한 희소 어텐션을 위한 효율적인 계층적 인덱싱

긴 문맥을 처리할 때 어텐션 연산 자체보다 '어떤 토큰이 중요한지'를 찾아내는 인덱싱 과정이 새로운 병목 현상이 되고 있다. HISA는 이 과정을 계층화하여 연산량을 대폭 줄이면서도 기존 모델의 정확도를 그대로 유지하는 실용적인 해결책을 제시한다.

Transformer의 Self-Attention은 모든 토큰 간의 관계를 계산하므로 시퀀스 길이의 제곱에 비례하는 비용이 발생한다. 이를 해결하기 위해 중요한 토큰만 골라 계산하는 Sparse Attention이 등장했지만, '어떤 토큰이 중요한지'를 판단하기 위해 여전히 모든 토큰의 점수를 매겨야 하는 인덱서(Indexer)가 긴 문맥에서 새로운 병목이 된다.

HISA는 도서관에서 책을 찾을 때 모든 책장을 일일이 확인하는 대신, 먼저 '분야별 서가'를 확인하고 관련 있는 서가 안에서만 책을 찾는 방식을 취한다. 먼저 토큰들을 일정 크기의 블록으로 묶어 대표 벡터를 만들고(Mean Pooling), 쿼리와 블록 대표 벡터 간의 유사도를 측정해 관련 없는 블록을 통째로 제거한다.

이후 살아남은 블록 내부에서만 기존의 정밀한 인덱싱을 수행하여 최종 토큰을 선택한다. 이 방식은 전체 토큰을 모두 훑을 필요가 없으므로 연산량이 시퀀스 길이에 따라 선형적으로 증가하는 수준으로 억제되며, 특히 128K 이상의 초장문 문맥에서 그 효과가 극대화된다.

전체 구조는 블록 분할 및 풀링, 1단계 블록 필터링, 2단계 토큰 정밀화로 구성된다. 입력 시퀀스 L을 크기 B의 연속적인 블록 M개로 나눈다. 각 블록 내의 인덱싱 키들에 대해 평균 풀링(Mean Pooling)을 수행하여 블록 대표 벡터를 생성한다. [블록 내 모든 키 벡터를 더하고 개수로 나누어] → [블록의 전반적인 의미를 담은 하나의 벡터를 얻고] → [이 값이 블록 전체를 대변하는 지표가 된다].

1단계에서는 쿼리와 블록 대표 벡터 간의 내적 점수를 계산한다. [쿼리 벡터와 블록 대표 벡터를 곱해] → [유사도 점수를 산출하고] → [상위 m개의 블록만 남겨 후보군을 좁힌다]. 이때 인과적 마스크(Causal Mask)를 준수하며 첫 블록과 마지막 두 블록은 항상 포함시켜 전역적 문맥과 최신 정보를 보존한다.

2단계에서는 선택된 m개 블록 내의 토큰들에 대해서만 원래의 DSA 스코어링을 적용한다. [제한된 후보 토큰들과 쿼리 간의 정밀 연산을 수행해] → [최종 k개의 토큰 인덱스를 추출하고] → [이를 하위 Sparse MLA 연산자에 전달한다]. 이 과정은 전체 토큰 수 L이 아닌 mB(m << M)에 비례하는 연산만 수행한다.

커널 수준 벤치마크에서 HISA는 32K 컨텍스트에서 2배, 128K 컨텍스트에서 4배의 속도 향상을 달성했다. 이는 시퀀스 길이가 길어질수록 O(L²) 비용을 갖는 기존 방식 대비 HISA의 계층적 구조가 연산 효율성 면에서 압도적임을 보여준다.

Needle-in-a-Haystack 테스트 결과, DeepSeek-V3.2의 원래 인덱서와 거의 동일한 수준의 정보 검색 능력을 유지했다. 특히 블록 단위로만 선택하는 Block-Sparse 방식이 긴 문맥에서 정확도가 급락하는 것과 대조적으로, HISA는 토큰 단위 정밀화를 통해 성능 저하를 방지했다.

LongBench 평가에서도 DSA와 대등한 점수를 기록했으며, 원래 DSA가 선택한 토큰 집합과의 교집합(IoU) 분석에서 99% 이상의 일치도를 보였다. 이는 HISA가 연산량을 줄이면서도 원래 모델이 의도한 어텐션 패턴을 거의 완벽하게 재현함을 의미한다.

HISA는 DeepSeek-V3.2 및 GLM-5에서 채택된 Token-level Sparse Attention의 인덱싱 경로를 최적화하기 위해 설계되었다. 기존의 플랫한 스캔 구조를 O(L/B + mB)의 복잡도를 갖는 계층적 구조로 재설계하여, 시퀀스 길이 L에 대한 의존성을 대폭 낮췄다.

블록 대표 벡터는 KV 캐시 업데이트 시점에 점진적으로 유지 관리될 수 있어 추가적인 오버헤드가 거의 없다. 또한 하위 Sparse MLA 연산자와의 호환성을 위해 출력 데이터 구조를 동일하게 유지하여 모델 아키텍처 변경 없이 적용 가능하다는 것이 큰 장점이다.

1단계 필터링에서 첫 블록과 최근 블록들을 강제로 포함하는 설계는 시스템 프롬프트나 최근 문맥의 중요성을 보존하며, 가변 길이 시퀀스 패킹 시의 경계 처리 문제를 완화한다. 이는 실무적인 서빙 환경에서의 안정성을 고려한 설계이다.

TileLang을 사용하여 GPU 커널을 최적화했으며, 블록 스코어링, 필터링, 토큰 정밀화의 3단계 파이프라인으로 구성된다. 실험 결과 mB >= k 조건만 만족하면 성능 손실 없이 연산 이득을 취할 수 있음이 확인되었다.

블록 내 토큰들의 의미적 일관성이 낮을 경우(예: 문단 경계), 평균 풀링된 벡터가 중요한 토큰을 제대로 대변하지 못해 정보 손실이 발생할 수 있다. 또한 현재는 학습 단계가 아닌 추론 시 교체 방식으로 작동하므로, 블록 스코어링 단계를 함께 학습시킨다면 더 높은 정확도를 기대할 수 있을 것이다.