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

왜 중요한가

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

핵심 기여

학습 가능한 청크-질량 대리변수와 랜드마크 기반 요약의 도입

청크 내부 키들의 가중합으로 정의되는 압축 키와 그에 대응하는 엔트로피 편향을 랜드마크 쿼리로 계산하여 LogSumExp 형태의 청크 질량을 1차 선형 서러게이트로 근사했다. 이 대리변수는 포워드 패스의 청크-레벨 가중치에 직접 포함되어 LM 손실로 역전파될 수 있다. 결과적으로 청크 선택이 종단간 학습되어 관련 청크를 더 정확히 골라낼 수 있다.

계층적 소프트맥스로 포워드 패스 인자분해

어텐션 가중치를 청크 간 분배(inter-chunk)와 청크 내부 분배(intra-chunk)로 분리하여 계산 복잡도를 줄였다. inter-chunk 항에는 학습된 대리질량을 사용하고 intra-chunk 항은 선택된 토큰들에 대한 정규화로 처리한다. 이 구조로 인해 전체 토큰 쌍에 대한 완전한 QK 계산 없이도 정확한 어텐션 분배가 가능하다.

하드웨어 친화적 커널 설계로 인접 쿼리의 청크 통합

인접 쿼리들이 높은 청크 중복을 갖는 점을 이용해 M개의 인접 쿼리 블록에 대해 선택된 청크들의 합집합만 로드하고 병렬로 연산하도록 커널을 설계했다. 이로써 Tensor Core의 타일 활용도를 높이고 메모리 접근을 줄여 실제 추론에서 대규모 문맥에서 수십 배의 속도 이득을 확보했다. 구현은 GQA 제약 없이 M×G 패킹으로 효율성을 확보한다.

짧은 계속학습으로 Full-attention 모델을 HiLS로 전환하는 방법

기존 풀 어텐션 모델에 랜드마크 토큰과 경량 어댑터만 추가한 뒤 약 50B 토큰의 연속 학습으로 HiLS 동작을 확보하는 절차를 제시했다. 이 방식은 대체로 짧은 재학습으로 원래의 단기 문맥 성능을 보존하면서도 초장문 외삽 능력을 획득하게 한다. 경량 모드에서는 랜드마크 임베딩과 저랭크 프로젝션만 학습하여 비용을 극적으로 제한할 수 있다.

이론적 근거로서 LogSumExp의 1차 선형화

LogSumExp(chunk)의 테일러 전개를 통해 쿼리와 학습된 청크 요약 키의 내적과 엔트로피 편향의 합으로 근사할 수 있음을 수학적으로 전개했다. 이 근사는 청크 질량의 평균-최댓값 간의 이질성을 보정하는 편향항을 포함하여 다양한 로그잇 분포에 대해 충분한 표현력을 가진다. 근사 유도는 청크 요약의 표현 설계와 학습 가능성을 이론적으로 뒷받침한다.

핵심 아이디어 이해하기

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

방법론

전체 접근은 청크 요약을 학습 가능한 랜드마크 쿼리와 그 결과로 얻은 압축 키 및 엔트로피 편향으로 구성하는 것이다. 각 청크에 랜드마크 토큰을 추가하여 그 랜드마크의 쿼리로 청크 내부 키들에 대해 자체 어텐션을 수행하면 가중합 키 k'_c와 엔트로피 b'c가 계산된다. 이 두 값이 합쳐져서 청크 질량의 선형 서러게이트 ŝ{i,c}=q_i^T k'_c / sqrt(d) + b'c를 구성하고, 이 값으로 top-K 청크를 선택한다. 수식 원문: ŝ{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)를 내적·합산하여 ŝ{i,c}를 얻고 이 값을 exp로 변환해 상위 K를 선택한 뒤 선택된 청크들에 대해 intra-chunk softmax로 토큰 가중치를 계산한다. 예시: d=64일 때 q_i^T k'c / sqrt(64)=2.0이고 b'c=0.7이면 ŝ{i,c}=2.7이고 exp(2.7)≈14.9로 다른 청크와의 상대적 질량을 비교할 수 있다. 해석: 이 흐름은 청크 수준의 중요도를 상수 비용으로 추정하고, 선택된 청크 내부에서만 토큰 수준 연산을 수행함으로써 전체 비용을 크게 절감한다. 학습 측면에서 ŝ{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 페이징 구현이 필요해 실무 환경으로 이전하려면 엔지니어링 비용이 수반된다.

실무 활용

HiLS-Attention은 긴 문맥을 처리하는 LLM 서비스의 추론 비용과 지연을 줄이면서 검색 정확도를 유지하거나 개선할 수 있는 기술적 기반을 제공한다. 경량 랜드마크 튜닝 모드는 기존 모델에 최소한의 재학습만으로 초장문 외삽 능력을 부여할 수 있어 실제 전환 비용을 제한한다. 구현은 청크 크기, top-K, 로컬 윈도우 등 하이퍼파라미터 조정으로 다양한 운영 제약에 맞출 수 있다.

대규모 문서 아카이브에서 실시간 질의응답을 수행하는 장문 검색 기반 챗봇에서 응답 지연과 메모리 비용을 동시에 절감할 수 있다.
장기간 대화 이력을 유지해야 하는 에이전트 시스템에서 중요 토큰을 정확하게 회수해 장기 추론의 일관성을 높일 수 있다.
로그·코드베이스·연구 문헌 등 초장문 문맥을 요구하는 분석 파이프라인에서 전체 토큰에 대한 풀 어텐션 대신 희소 접근으로 비용을 제한할 수 있다.

코드 공개 여부: 공개

코드 저장소 보기

키워드

HiLS-Attention(계층적 랜드마크 희소 어텐션)chunk-wise sparse attention(청크 단위 희소 어텐션)landmark token(랜드마크 토큰)LogSumExp linearization(LogSumExp 선형 근사)HoPE(부분 RoPE positional encoding)RULER(장문 검색 벤치마크)

용어 해설

청크 단위 희소 어텐션(Chunk-wise Sparse Attention): — 전체 시퀀스를 일정 길이의 청크로 분할하고 각 쿼리가 모든 토큰 대신 일부 청크만 선택해 어텐션을 수행하는 방식이다. 선택된 청크들에만 KV를 로드해 계산과 메모리 비용을 고정 예산으로 유지하며 긴 컨텍스트를 처리할 수 있게 한다. 청크 선택의 정확도가 성능과 extrapolation에 직접적인 영향을 미친다.
랜드마크 토큰(Landmark Token): — 각 청크에 추가되는 특수 토큰으로서 해당 청크의 요약 표현을 생성하기 위해 쿼리 역할을 수행한다. 이 토큰으로부터 계산된 요약 키와 엔트로피 편향이 청크 수준의 중요도 점수를 추정하는 핵심 입력이 된다. 학습 중 LM 손실로 역전파되어 청크 선택을 종단간으로 최적화할 수 있게 한다.
LogSumExp 선형 근사(LogSumExp Linearization): — 토큰 수준의 exp(로그잇) 합으로 정의되는 청크 질량의 로그를 1차 테일러 전개를 통해 쿼리·압축 키 내적과 엔트로피 편향의 합으로 근사하는 수학적 표현이다. 이 표현은 청크 질량을 상수 비용으로 추정할 수 있게 하여 정확한 청크 선택을 가능하게 한다. 근사항은 청크 요약 키와 편향항으로 계산되어 학습 가능한 파라미터로 처리된다.
계층적 소프트맥스(Hierarchical Softmax): — 전체 어텐션을 먼저 청크 단위의 inter-chunk 소프트맥스와 각 청크 내부의 intra-chunk 소프트맥스 두 단계로 인자분해하는 방식이다. inter-chunk 항은 학습된 청크 질량 대리변수로 치환되어 청크 선택이 포워드 패스에 영향을 미치게 설계된다. 이 구조로 인해 LM 손실이 청크 요약과 선택 점수로 직접 전파되어 종단간 학습이 가능해진다.
저계수 쿼리 보정(Low-Rank Query Calibration): — 토큰 수준 쿼리와 압축된 청크 키의 표현 차이를 줄이기 위해 추가되는 저랭크 어댑터이다. 입력 숨김 상태에 대해 up/down 프로젝션을 통해 쿼리를 보정하여 청크 질량 대리변수의 정확도를 개선한다. 모델 크기에 비해 경량이며 길이 외삽 성능과 퍼플렉시티 개선에 기여한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

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