계층적 랜드마크 희소 어텐션(HiLS-Attention): 무한 컨텍스트 모델링을 향한 접근

긴 문맥을 다루는 Transformer의 계산 비용은 시퀀스 길이에 따라 제곱적으로 증가하여 실제 적용에 제약을 초래한다. HiLS-Attention은 청크 수준의 학습 가능한 요약과 계층적 인자분해로 어텐션 연산을 희소화하면서도 청크 선택을 언어모델 손실로 직접 최적화하여 성능 손실 없이 문맥 길이를 수십에서 수백 배로 확장 가능하게 했다. 이로 인해 초장문 대화, 긴 문서 검색, 장기 추적 과제에서 비용과 지연을 동시에 개선할 수 있는 실용적 경로가 마련되었다.

Transformer의 self-attention은 시퀀스 내 모든 토큰 쌍의 상호작용을 계산하므로 길이가 N일 때 연산량이 O(N^2)으로 증가한다. 이 비용 구조는 문맥이 매우 길어질수록 메모리와 시간 병목을 유발하여 실제 추론과 학습의 확장을 제한한다. 따라서 긴 문맥을 실용적으로 다루려면 전체 토큰 쌍을 계산하지 않는 희소화 전략이 필요하다. 청크 단위 희소화는 긴 시퀀스를 고정 크기 청크로 나눈 뒤 쿼리가 일부 청크만 선택해 해당 청크 내부 토큰에만 어텐션을 수행하는 방식으로 연산을 고정 예산으로 유지한다. 기존 방법들은 청크 선택을 위해 평균 풀링 같은 비학습 요약을 사용하거나 간접적(distillation) 신호에 의존하여 선택의 정확도가 떨어지는 문제가 있었다. 이로 인해 중요한 청크를 놓치거나 긴 문맥에서 외삽 성능이 급격히 저하되는 현상이 관찰되었다. 핵심 아이디어는 청크 질량을 직접 추정 가능한 학습 가능한 압축 키와 편향으로 근사하고, 그 대리질량을 포워드 어텐션 가중치에 포함시켜 LM 손실로 역전파하는 것이다. 구체적으로 각 청크 끝에 랜드마크 토큰을 추가하여 해당 랜드마크의 쿼리로 청크 내부 키들에 대한 어텐션을 계산하고, 그 결과로 얻은 가중합을 압축 키로 사용한다. 압축 키와 해당 분포의 엔트로피를 합하면 LogSumExp(chunk)에 대한 1차 선형 근사값이 나오며 이 값으로 청크를 정렬·선택할 수 있다. 이 접근은 두 단계의 계층적 소프트맥스 구조로 작동한다. 먼저 쿼리는 압축 키들과 상호작용해 각 청크에 할당할 총 어텐션 질량을 결정하고(inter-chunk), 다음으로 선택된 청크 내부에서 토큰 수준의 분배(intra-chunk)를 수행해 최종 가중치를 얻는다. 이렇게 하면 전체 토큰 쌍의 QK 계산 없이도 청크 수준에서의 정확한 라우팅이 가능해지며 학습 과정에서 청크 선택 기준이 직접 최적화되어 외삽 능력이 크게 향상된다.

전체 접근은 청크 요약을 학습 가능한 랜드마크 쿼리와 그 결과로 얻은 압축 키 및 엔트로피 편향으로 구성하는 것이다. 각 청크에 랜드마크 토큰을 추가하여 그 랜드마크의 쿼리로 청크 내부 키들에 대해 자체 어텐션을 수행하면 가중합 키 k'_c와 엔트로피 b'c가 계산된다. 이 두 값이 합쳐져서 청크 질량의 선형 서러게이트 ŝ{i,c}=q_i^T k'_c / sqrt(d) + b'c를 구성하고, 이 값으로 top-K 청크를 선택한다. 수식 원문: ŝ{i,c} = q_i^T k'_c / sqrt(d) + b'_c. 변수 의미: q_i는 현재 토큰의 쿼리 벡터(차원 d), k'_c는 청크 c의 압축 키(차원 d), b'_c는 청크 c의 엔트로피 편향(스칼라), sqrt(d)는 스케일링 인자이다. 계산 흐름: 입력 q_i와 모든 청크의 (k'_c,b'c)를 내적·합산하여 ŝ{i,c}를 얻고 이 값을 exp로 변환해 상위 K를 선택한 뒤 선택된 청크들에 대해 intra-chunk softmax로 토큰 가중치를 계산한다. 예시: d=64일 때 q_i^T k'c / sqrt(64)=2.0이고 b'c=0.7이면 ŝ{i,c}=2.7이고 exp(2.7)≈14.9로 다른 청크와의 상대적 질량을 비교할 수 있다. 해석: 이 흐름은 청크 수준의 중요도를 상수 비용으로 추정하고, 선택된 청크 내부에서만 토큰 수준 연산을 수행함으로써 전체 비용을 크게 절감한다. 학습 측면에서 ŝ{i,c}가 포워드 어텐션에 포함되므로 LM 손실의 그래디언트가 랜드마크 임베딩과 Q-Cal 어댑터로 직접 전파된다. Q-Cal은 원래 쿼리 q_i에 대해 저랭크 업/다운 프로젝션을 적용해 q̂_i를 얻는 경량 모듈로 청크 수준 점수 보정에 사용된다. 하드웨어 관점에서는 인접 쿼리의 청크 선택이 높은 중복을 보이는 특성을 이용해 M개의 인접 쿼리를 묶어 선택된 청크들의 합집합만 로드한 뒤 배치 연산으로 처리한다. 이 패킹은 Tensor Core 타일을 효율적으로 채우고 메모리 중복 접근을 줄여 실제 추론 성능을 개선한다.

작은 스케일(345M)에서 HiLS-Attention은 학습 길이(8K)에서 풀 어텐션과 동등한 퍼플렉시티를 달성하면서 RULER 기반 NIAH 작업에서 완전한 회수 성능(100%)을 보였다. 초장문 외삽에서는 8K 학습만으로도 4M 길이까지 90% 이상의 검색 정확도를 유지하는 결과가 관찰되었고, Naive-BSA와 평균풀링 요약을 사용하는 기존 희소법 대비 명확한 우위를 보였다. 이 결과는 Tab.1과 Tab.2의 수치로 구체화되어 있으며 논문은 1M~4M 범위에서도 높은 NIAH 성능을 보고했다. 효율성 측면에서 HiLS는 추론에서 길이가 증가할 때 선형 또는 사실상 상수에 가까운 디코드 지연을 유지하여 긴 문맥에서 큰 속도 이점을 보였다. 실험에서 512K 문맥 기준으로 prefill에서 약 13.5×, 토큰당 디코드에서 약 15.7× 속도 향상이 관찰되었다(동일 하드웨어, H800 기준). 교차점은 약 16K로, 이 길이 이상에서 HiLS가 실질적으로 더 빠른 것으로 나타났다. 대규모 실험에서는 1.4B 모델을 원천적으로 HiLS로 학습한 결과와 7B급 Olmo3 기반의 계속학습 전환 실험이 모두 보고되었다. 원천 학습에서는 단기 문맥 성능이 풀 어텐션과 거의 동일했으며 계속학습 방식으로 50B 토큰 정도를 사용하면 기존 풀-어텐션 모델을 HiLS로 전환하여 단기 성능을 유지하면서 장문 벤치마크에서 우위를 확보할 수 있었다.

전체 아키텍처는 표준 디코더형 Transformer의 어텐션 모듈을 HiLS로 대체하는 방식이다. 각 청크의 끝에 랜드마크 토큰을 추가해 해당 랜드마크의 쿼리로 청크 내부 키들에 대한 자체 어텐션을 수행하고 그 결과로 k'_c와 b'_c를 계산한다. 청크 크기 S=64, top-K=32, 로컬 슬라이딩 윈도우 512와 같은 하이퍼파라미터 설정이 핵심 실험 구성으로 사용되었다. 수학적 기반은 LogSumExp의 1차 테일러 전개로, log ∑j exp(s{j})를 q^T k'_c / sqrt(d) + b'_c로 근사하는 정식화가 중심이다. 여기서 b'_c는 청크 내부 어텐션 분포의 엔트로피로 동작하여 평균-최댓값 간의 보정을 제공한다. 이 근사는 각 청크에 대해 O(S) 연산으로 계산 가능하므로 전체 시퀀스에 대해 O(N) 비용으로 청크 요약을 얻을 수 있다. 하드웨어 최적화는 인접 쿼리 블록을 M 단위로 묶어 각 블록에 대해 선택된 청크의 합집합만 로드하는 전략에 기반한다. 이로써 텐서 코어 연산은 (M×G,d)×(d,S) 형태로 수행되어 작은 GQA 그룹에서도 높은 활용률을 달성할 수 있다. 구현상에는 GQA-aware 그룹화, 메모리 페이징(청크 단위 KV 페이징), 그리고 CUDA/Triton 기반의 루프 패킹이 필요하다. 학습 전략으로는 완전 파라미터 튜닝과 랜드마크 토큰만 학습하는 경량 모드 둘 다 보고되었고, 경량 모드는 전체 파라미터의 1% 미만만 업데이트해도 실용적 성능을 확보했다. 또한 위치 인코딩으로는 RoPE의 일부 차원을 고정하고 나머지는 NoPE로 대체한 HoPE가 HiLS와 결합될 때 장문 외삽 성능을 향상시키는 것으로 관찰되었다. 실험에는 345M, 1.4B, 7B 스케일 모델과 다양한 계속학습 설정이 포함되었다.

랜드마크 토큰을 제거하고 공유 쿼리만 사용할 경우 단기 성능은 유지되나 장문 외삽 성능이 크게 저하되는 경향이 관찰되었다. HiLS의 견고한 외삽 성능은 랜드마크 기반 요약과 HoPE 같은 위치 인코딩의 조합에 의존하며, 이들 구성요소의 대체는 성능 저하를 유발할 수 있다. 또한 커널 수준의 최적화와 청크 기반 KV 페이징 구현이 필요해 실무 환경으로 이전하려면 엔지니어링 비용이 수반된다.