LSA가 초장문 맥락에서 키 탐색 비용을 클러스터 검색으로 줄인다

무엇인가

이 논문은 100만 토큰을 넘는 초장문 맥락에서 희소 어텐션의 병목이 키를 찾는 인덱서에 있다고 본다. 전체 어텐션은 프리필 O(T^2), 디코드 O(T)이고, 키를 일부만 고르는 희소 어텐션도 인덱서가 각 쿼리마다 모든 선행 키나 블록을 점수화하면 프리필 O(T^2), 디코드 O(T)가 남는다. 저자들은 10M 토큰에서 기존 희소 어텐션 한 층의 프리필 시간 중 98%가 인덱싱에 쓰인다고 제시하며, 저장이 아니라 탐색 비용을 줄이는 문제로 설정한다.

어떻게 동작하나

LSA는 인덱스 키를 학습된 잠재 공간에서 내용 기반으로 클러스터링하고, 쿼리가 모든 키 대신 클러스터 중심을 점수화하게 한다. 구체적으로 인덱스 키를 L2 정규화한 뒤 연속 n_pool개를 합산해 풀 키를 만들고, n_cluster개 연속 키로 이루어진 클러스터 블록마다 Ward 방법으로 상향 병합해 버킷을 만든다. 각 버킷은 중심과 키 위치 테이블을 버킷 캐시에 저장한다. 검색은 쿼리를 정규화해 버킷 중심과 내적하고 상위 p개 버킷을 고른 뒤, 선택 버킷을 가장 좋은 키부터 확장해 k개의 장거리 키를 채운다. 여기에 sink 키 n_sink개와 최근 키 윈도 n_w개를 더하며, 윈도를 클러스터 블록 하나로 두어 아직 클러스터링되지 않은 키를 덮는다. 어텐션은 이 합집합에만 수행되고 인덱스 목록은 모든 헤드가 공유한다. Flat은 한 층 버킷을 모두 점수화하고, Tree는 빔 서치로 계층을 내려가 쿼리당 O(log T) 중심만 점수화해 인덱서 프리필을 O(T log T)로 만든다.

무엇과 다른가

D-LSA는 이미 학습된 모델의 어텐션 키·쿼리를 그대로 두고, 별도 인덱스 투영으로 작은 차원 d_idx=64의 키·쿼리를 만든다. 쿼리 헤드는 합산해 토큰당 하나의 인덱스 쿼리를 만들고, Qwen-3.8-27B에서는 RoPE가 걸린 차원을 제외한 부분만 읽는다. 기존 모델은 동결한 채 인덱스 투영과 층별 온도 τ만 증류한다. 손실은 동결 모델의 헤드 평균 어텐션이 허용 키들에 분포한 것을 목표로 하고 D-LSA 점수의 softmax 사이 KL을 최소화하며, 목표는 어텐션 가중치 상위 2048개 키로 잘라 재정규화한다. 학습 없이 키의 주성분 방향을 투영으로 쓰는 D-LSA(Training-Free)도 있다. N-LSA는 인덱스 투영 없이 어텐션 자체의 키를 클러스터링하고, 단일 KV 헤드에서 버킷 중심에 대한 각 쿼리 헤드 최대 유사도로 버킷을 점수화한다. 선택을 통한 그래디언트는 끊고 NTP 손실만으로 학습하므로 추가 파라미터, 교사, 보조 손실이 없다.

어떻게 쓰나

실험은 전체 어텐션 모델 Qwen-3.8-27B와 희소 어텐션 모델 GLM-5.3-Flash에서 Full Attention, 4개 키 블록을 모두 점수화하는 Block Sparse Attention(Flat), 키 블록 위 트리를 쓰는 Block Sparse Attention(Tree), D-LSA(Flat/Tree), D-LSA(Training-Free)를 비교한다. D-LSA는 k=4096 장거리 키를 기본으로 하고 GLM에서는 1M 이전 2048개, 이후 4096개를 쓴다. Qwen-3.8-27B에서 D-LSA(Flat)는 1M 토큰까지 Full Attention과 5점 이내이고, 10M에서 RULER 63.9 대 59.9, PaperTrail 29.1 대 19.8로 오히려 앞선다. GLM-5.3-Flash에서는 1M까지 3점 이내, 2M~5M에서 최대 8점 뒤지고, 10M RULER에서는 4점 이내다. 저자들은 이 비교가 모델과 함께 학습된 기존 희소 어텐션에 유리하다고 밝힌다. D-LSA(Tree)는 Qwen에서는 Flat을 따르지만 GLM 10M RULER에서 55.8 대 66.8로 떨어지고, Block Sparse Attention(Tree)은 512k부터 뒤처져 1M에서 약 63에 그친다.

전제와 한계

한 층 성능에서 D-LSA(Flat)는 10M 토큰의 GLM-5.3-Flash에서 기존 희소 어텐션 대비 프리필 22.7배, 디코드 4.6배 빠르고, Qwen-3.8-27B에서는 Full Attention 대비 프리필 72.4배, 디코드 42.5배 빠르다. 1M에서는 3.5배, 512k에서는 2.1배로 이득이 줄어든다. B300에서는 각각 12.5배·42.9배(프리필), 5.2배·18.4배(디코드)다. 인덱서 메모리는 10M에서 D-LSA(Flat)가 층당 246MiB로 GLM-5.3-Flash의 331MiB, QSA의 642MiB보다 작다. SGLang에 텐서·전문가·파이프라인 병렬과 프리필·디코드 문맥 병렬을 지원해 통합했고, 문맥 병렬에서는 버킷 캐시를 모든 랭크에 복제해 통신 대신 메모리를 쓰는 설계를 택했다. 10M·CP8에서 각 랭크는 10GiB KV 캐시의 1.25GiB 몫과 246MiB 버킷 캐시 전체를 갖는다. 전체 모델에서는 GLM CP8 10M에서 TTFT 2.2배, TPOT 2.1배, TP8에서 5.8배·1.9배 개선이고, Qwen에서는 TTFT 9.2배, TPOT 2.9배다.

N-LSA는 약 0.5B 하이브리드 모델(24층 중 4층마다 어텐션, 나머지는 Gated DeltaNet)에서 검증했다. 100B 토큰 FineWeb-Edu를 4k로 예비학습한 뒤 12B 토큰 LongDocs를 128k로 학습했고, 쿼리 8헤드가 단일 128차원 KV 헤드를 공유하며 위치 인코딩을 쓰지 않는다. 128k 검증 손실은 N-LSA 1.085로 Full Attention 1.085와 같고, PISA 1.100, NSA 1.117, SWA 1.154보다 낮다. RULER 128k에서는 32.2 대 Full Attention 34.7로 3.5점 이내이며 NSA 19.2, PISA 14.6, SWA 3.0을 크게 앞선다. 한 어텐션 층의 순전파·역전파는 64k에서 1.1배, 128k에서 1.6배 빠르고, 전체 학습은 같은 검증 손실까지 1.24배 빨리 끝난다. 대신 128k에서 층당 최대 메모리는 6.8GiB로 Full Attention의 5.8GiB보다 크고, 짧은 문맥에서는 키 선택 비용이 절감보다 크다.

개발자 관점에서 LSA는 1M~10M 토큰급 RAG·에이전트 서빙에서 어텐션 인덱서가 병목일 때 검토할 만하다. 이미 학습된 모델에는 D-LSA 증류로 작은 인덱스 투영만 붙일 수 있고, 처음부터 학습한다면 N-LSA가 추가 파라미터 없이 NTP 손실만으로 동작한다. 다만 LSA는 키를 찾는 비용만 줄일 뿐 KV 캐시 저장량은 선형으로 늘린다. 저자들은 KV 압축, 층 간 인덱스 공유, 다중 토큰 예측과 직교해 결합 가능하다고 본다. 한계도 명시된다. GLM-5.3-Flash에서 Tree 검색은 1M 이후 정확도가 Flat보다 빨리 떨어졌고, 단일 요청 디코드에서는 약 5M부터, 프리필에서는 10M을 넘어야 Flat을 추월할 것으로 예측된다. N-LSA는 0.5B, 128k, Flat 검색까지만 검증됐고 더 큰 모델·더 긴 문맥·Tree 검색으로의 확장은 열린 문제다. D-LSA도 전체 모델을 희소 어텐션으로 추가 학습하는 단계는 비용 때문에 생략했으며, 증류 이후의 정확도 격차를 희소 학습이 메울지는 검증되지 않았다.