긴 문맥에서 필요한 기억을 먼저 찾는 Latent Sparse Attention

배경과 문제

대규모 언어 모델이 단일 질문 응답을 넘어 소프트웨어 공학이나 개방형 과학 연구 같은 장기 작업을 수행하는 에이전트 흐름으로 이동하면서, 모델이 한 번에 처리해야 하는 맥락 길이가 빠르게 늘고 있다. 최근 모델들은 100만 토큰 이상의 맥락 창을 지원하지만, 이 규모에서는 어텐션 연산이 전체 비용을 지배한다. 전체 어텐션은 T개 토큰을 프리필하는 데 O(T²), 토큰 하나를 디코딩하는 데 O(T)의 계산이 필요하기 때문이다.

이를 줄이기 위해 희소 어텐션이 널리 쓰인다. 각 쿼리가 모든 키 대신 일부 키만 보도록 하는 방식인데, 문제는 어떤 키를 볼지 고르는 인덱서 자체가 비싸다는 점이다. 기존 인덱서는 쿼리마다 앞선 모든 키나 블록을 점수 매겨야 하므로 프리필에서 여전히 O(T²), 디코드에서 O(T)의 비용이 든다. 논문의 Figure 2에 따르면 10M 토큰에서 GLM-5.3-Flash의 희소 어텐션 계층은 프리필 시간의 98%를 인덱싱에 쓴다. 즉 초장문 맥락에서 희소 어텐션의 병목은 어텐션 자체가 아니라 '관련 키를 찾는 일'이다.

기존 접근의 한계

희소 어텐션은 인덱서가 점수 매기는 단위에 따라 토큰 희소, 블록 희소, 의미적 희소로 나눌 수 있다. 토큰 희소는 쿼리마다 모든 키를 점수 매기고, 블록 희소는 연속된 키들을 묶어 블록 단위로 점수를 낸다. 의미적 희소는 위치가 아니라 내용 유사성으로 키를 군집화하고 각 군집의 중심점을 점수 매긴다. 이 중 블록 희소 방식이 최신 프런티어 모델에서 많이 쓰이지만, 여전히 모든 블록을 훑어야 한다.

CSA, QSA, GLM-5.3-Flash SA 같은 방법은 네 키를 하나로 압축해 인덱서 비용을 상수 배로 줄였고, HISA나 LongCat LSA는 블록을 먼저 고른 뒤 선택된 블록 안의 키만 다시 점수 매긴다. 그러나 이런 개선은 모두 인덱서의 작업량을 상수 배로 나눌 뿐, 쿼리당 작업량은 여전히 맥락 길이에 선형으로 증가한다. PISA는 트리 탐색으로 이 선형 스캔을 대체했지만 정확도 평가가 16k 토큰에 그쳤다. 의미적 군집화를 쓰는 ClusterKV, RetroInfer, IceCache 등은 대부분 추론 시점에만 동작하고 프리필을 가속하지 못하며, 저자들이 비교한 문헌에서는 1M 토큰을 넘는 초장문 평가가 없었다고 설명한다.

연구방법을 단계별로

LSA의 핵심은 키를 학습된 잠재 공간에서 의미적으로 군집화하고, 쿼리가 모든 키 대신 군집 중심점만 탐색하게 하는 것이다. D-LSA의 절차를 보면 먼저 인덱서용 키를 L2 정규화하고, 연속된 키 네 개를 합쳐 풀링해 군집화할 대상 수를 줄인다. 그다음 Ward 방법으로 상향식 응집 군집화를 수행해 비슷한 키들을 버킷으로 묶는다. 각 버킷은 중심점과 키 위치 테이블을 캐시에 저장해, 나중에 선택된 버킷을 실제 키 인덱스로 확장할 수 있게 한다.

검색 단계에서는 쿼리도 같은 잠재 공간으로 투영하고 정규화한 뒤, 자기보다 앞선 모든 버킷 중심점과 내적해 점수를 매긴다. 상위 p개 버킷을 골라 그 안의 키들을 최대 k개까지 확장하고, 여기에 싱크 키와 최근 윈도우 키를 더해 최종 어텐션 대상 인덱스 목록을 만든다. 이 목록은 모든 어텐션 헤드가 공유한다. 평면 검색은 모든 버킷을 점수 매기므로 프리필이 여전히 이차 비용이지만, 트리 검색은 버킷을 계층적으로 묶고 빔 서치로 내려가면서 쿼리당 O(log T)개의 중심점만 점수 매겨 프리필을 O(T log T)로 만든다.

D-LSA는 사전학습 모델에 적용하기 위해 별도의 작은 키·쿼리 투영 행렬을 추가하고, 원래 모델은 얼린 채 인덱서만 증류한다. 증류 손실은 원래 어텐션의 헤드 평균 분포와 D-LSA 점수 분포 사이의 KL 발산을 최소화하며, 온도 파라미터로 점수 분포를 날카롭게 맞춘다. N-LSA는 별도 투영 없이 어텐션 자체의 키를 군집화하고, 선택된 키에 대해서만 다음 토큰 예측 손실의 그래디언트를 흘려보내 학습한다. 선택 과정에는 그래디언트를 막아 추가 손실이나 교사 모델 없이도 학습된다.

결과와 숫자의 의미

지식·추론·LongBench v2의 보고된 점수는 기준 모델의 95% 신뢰구간 안에 있었다. 이것이 모든 장문 과제의 정확도가 동일함을 증명하지는 않는다. Qwen-3.8-27B에서 D-LSA(Flat)는 1M 토큰까지 전체 어텐션과 5포인트 이내 차이를 유지했고, 10M 토큰에서는 RULER 63.9 대 59.9, PaperTrail 29.1 대 19.8로 오히려 앞섰다. GLM-5.3-Flash의 Flat 방식은 1M까지 3포인트 이내였지만 2M~5M에서는 최대 8포인트 낮았고, 10M의 RULER에서는 4포인트 이내였다. 다만 MultiKey 2, MultiKey 3, CWE 같은 어려운 검색 과제는 모든 방법에서 10M 토큰에서 5점 이하로 떨어져 전체 평균을 끌어내렸다.

속도 측면에서 한 어텐션 계층의 프리필 속도 향상은 맥락이 길수록 커진다. GLM-5.3-Flash에서 D-LSA는 1M 토큰에서 3.5배, 10M 토큰에서 22.7배 빨라졌고, Qwen-3.8-27B에서는 10M 토큰에서 전체 어텐션 대비 72.4배 빨라졌다. 디코드 속도 향상은 10M 토큰에서 각각 4.6배와 42.5배다. 이는 한 계층만 측정한 값이며, 전체 모델의 TTFT와 TPOT 향상은 이보다 작다. 예를 들어 GLM-5.3-Flash의 10M 토큰 TTFT는 CP8에서 2.2배, TP8에서 5.8배 감소했다. 트리 검색의 프리필은 10M 토큰을 넘어서야 평면 검색을 앞지르며, 100M 토큰에서는 기준 희소 어텐션보다 90배 이상 빠를 것으로 외삽된다.

N-LSA는 약 0.5B 파라미터의 하이브리드 모델에서 128k 토큰 맥락으로 학습되었다. 검증 손실은 전체 어텐션과 같은 1.085에 도달했고, NSA(1.117), PISA(1.100), 슬라이딩 윈도우(1.154)보다 낮았다. RULER 128k 토큰에서 N-LSA는 32.2로 전체 어텐션의 34.7보다 2.5포인트 낮았고, NSA와 PISA는 각각 19.2와 14.6으로 떨어졌다. 한 계층의 순전파+역전파는 128k 토큰에서 전체 어텐션보다 1.6배 빨랐지만, 최대 메모리는 계층당 6.8GiB로 전체 어텐션의 5.8GiB보다 높았다.

구체적인 예시

PaperTrail은 저자들이 만든 벤치마크로, 실제 과학 논문 모음에 가짜 사실을 심어 모델이 문헌 검색처럼 사실을 찾아 연결하는 능력을 시험한다. 예를 들어 한 가짜 논문이 'AlloKV'라는 KV 캐시 제거 방법을 소개하고, 2~3주 뒤 다른 가짜 논문이 그 값을 수정하며 'SuBA'로 이름을 바꾼다. 이후 호스트 논문들은 새 이름을 인용하거나 옛 값을 그대로 인용한다. 질문은 한 논문에서 사실을 회상하는 1단계, 두 논문을 연결하는 2단계, 세 논문 이상의 흔적을 추적하는 3단계로 나뉜다.

10M 토큰 PaperTrail에서 Qwen-3.8-27B의 D-LSA(Flat)는 1단계와 2단계 질문에서 전체 어텐션을 앞섰다. 1단계는 53.3 대 32.2, 2단계는 32.2 대 22.2였다. 반면 3단계 질문은 모든 방법에서 19.0 이하로 떨어졌다. 이는 맥락이 길어질수록 여러 논문에 흩어진 사실을 추적하는 과제가 얼마나 어려운지 보여준다. 별도의 '자료에 답이 없는 질문' 평가는 10M이 아닌 512k 토큰에서 수행했다. 이때 D-LSA의 점수는 각 기준 모델과 3.7포인트 이내였으며, 그 결과만으로 군집화가 오답을 막는 원인이라고 단정할 수는 없다.

성능 예시로, 10M 토큰에서 D-LSA(Flat)의 인덱서는 버킷 중심점 1,310,720개를 저장하는데 246MiB가 든다. 같은 맥락에서 GLM-5.3-Flash의 기존 인덱서는 331MiB, QSA 방식은 642MiB를 쓴다. D-LSA는 버킷당 약 8개 키를 하나의 중심점으로 요약하고 인덱스 차원을 64로 줄여 메모리도 절약한다. 컨텍스트 병렬 처리에서는 모든 랭크가 버킷 캐시를 복제해 통신을 줄이는데, 8개 GPU에 나눈 GLM의 10M 토큰 설정에서 GPU당·계층당 KV 캐시 1.25GiB에 버킷 캐시 246MiB가 더해져 약 19%의 메모리 오버헤드가 생긴다.

한계와 오해하기 쉬운 점

가장 흔한 오해는 22.7배나 72.4배 같은 숫자를 전체 모델의 TTFT나 TPOT 향상으로 읽는 것이다. 이는 한 어텐션 계층만 측정한 값이며, 전체 모델에서는 글로벌 어텐션 계층만 빨라지고 다른 계층은 그대로라 향상 폭이 작아진다. 실제로 GLM-5.3-Flash의 10M 토큰 TTFT는 CP8에서 2.2배, TP8에서 5.8배 감소했고, TPOT는 각각 2.1배와 1.9배 감소했다. 또한 H200과 B300, CP8과 TP8은 조건이 다르다. 특히 H200의 GLM 10M TP8 실험은 PP2로 두 노드를 사용한 예외이므로, 서로 다른 설정의 배율을 같은 장비 조건처럼 비교하면 안 된다.

트리 검색이 항상 빠른 것도 아니다. 평면 검색은 프리필에서 모든 버킷을 점수 매기므로 이차 비용이지만, 트리 검색은 계층을 순차적으로 내려가야 해 디코드에서 요청이 하나뿐이면 5M 토큰까지는 평면 검색이 더 빠르다. 논문의 단일 요청 평가에서는 D-LSA(Tree)도 디코드에 평면 검색을 썼다. 또한 100M 토큰 성능은 측정이 아니라 10M 이후를 외삽한 예측이다. 정확도 벤치마크는 부분집합을 사용했고, PaperTrail은 가짜 사실을 실제 논문에 삽입한 합성 과제라는 점도 기억해야 한다.

N-LSA는 약 0.5B 파라미터, 최대 128k 토큰, 평면 검색만으로 학습된 소규모 실험이다. 메모리는 128k 토큰에서 계층당 6.8GiB로 전체 어텐션의 5.8GiB보다 높다. KV 캐시 저장은 여전히 맥락 길이에 선형으로 증가하며, LSA는 키를 찾는 비용만 줄일 뿐 저장 비용은 줄이지 못한다. 트리 검색은 GLM-5.3-Flash에서 1M 토큰을 넘으면 평면 검색보다 정확도가 더 빨리 떨어졌다. 큰 모델, 더 긴 맥락, 트리 검색으로의 네이티브 학습 확장은 미해결 과제로 남아 있다.

실제로 어떻게 활용할까

D-LSA는 사전학습 모델 본체를 고정하고 별도 인덱서만 학습하는 적용 경로다. 계층당 작은 투영 행렬만 352M 토큰으로 증류하면 되는데, 이는 DSA의 밀집 워밍업 2.1B 토큰의 일부에 불과하다. 학습 없이 키의 주성분 방향을 투영으로 쓰는 D-LSA(Training-Free)도 Qwen-3.8-27B에서는 D-LSA(Flat)에 근접한 성능을 보였다. 논문은 공개 코드의 SGLang 통합을 보고하며 텐서 병렬, 컨텍스트 병렬, 파이프라인 병렬과 함께 사용할 수 있다.

운영 환경에서는 컨텍스트 병렬 처리로 KV 캐시를 여러 GPU에 나누고, 버킷 캐시는 모든 랭크에 복제해 통신을 줄인다. 디코드 단계는 CUDA 그래프로 캡처해 커널 실행 오버헤드를 줄였고, 맥락 길이 구간별로 그래프를 여러 개 만들어 짧은 요청이 긴 그래프의 스레드 블록을 불필요하게 실행하지 않도록 했다. 클러스터링은 512개 토큰마다 한 번씩만 수행되어 디코드 비용에 분산 상각된다. 이러한 설계는 10M 토큰 규모의 서빙에서 실제 TTFT와 TPOT 감소로 이어졌다.

N-LSA는 별도 인덱서 투영이나 교사 모델 없이 다음 토큰 예측 손실만으로 학습되므로, 처음부터 효율적인 초장문 모델을 만들고자 할 때 유용하다. 다만 현재는 소규모 검증에 그쳤고, 트리 검색을 네이티브 학습에 통합하거나 더 큰 모델로 확장하는 것은 후속 과제다. LSA는 KV 캐시 압축, 계층 간 인덱스 공유, 다중 토큰 예측 등과 직교하므로 함께 결합할 수 있다. 초장문 맥락에서 어텐션 비용을 실질적으로 낮추고자 하는 연구자나 엔지니어에게 LSA는 공개 구현과 자신의 장문 과제로 성능·메모리·정확도를 함께 검증해 볼 연구 선택지다.