Authors
- Yushi Bai¹†, Qian Dong¹†, Ting Jiang², Xin Lv², Zhengxiao Du², Aohan Zeng¹², Jie Tang¹, Juanzi Li¹
¹ Tsinghua University
² Z.ai
Abstract
장문(long-context) 에이전트(agentic) 워크플로우는 대규모 언어 모델(LLM)의 대표적인 활용 사례로 자리 잡았으며, 이에 따라 어텐션(attention)의 효율성은 추론 속도와 서비스 비용 모두에 있어 핵심적인 요소가 되었다. 희소 어텐션(sparse attention)은 이러한 문제를 효과적으로 해결하는 접근법이며, DeepSeek Sparse Attention (DSA) 는 이를 실제 서비스 환경에 적용한 대표적인(production-grade) 솔루션이다. DSA에서는 경량(lightweight) Lightning Indexer가 각 쿼리에 대해 가장 관련성이 높은 상위 $k$개의 토큰을 선택함으로써, 핵심 어텐션 연산의 복잡도를
$$O(L^2)$$
에서
$$O(Lk)$$
로 감소시킨다.
그러나 인덱서(indexer) 자체는 여전히
$$O(L^2)$$
의 계산 복잡도를 가지며, 모든 계층(layer)에서 독립적으로 실행되어야 한다. 이는 연속된 계층들에서 선택되는 상위 $k$개의 토큰이 매우 유사하다는 사실을 고려하면 비효율적이다.
본 논문에서는 이러한 계층 간 중복성(cross-layer redundancy)을 활용하는 IndexCache를 제안한다. IndexCache는 전체 계층을 두 종류로 구분한다.
-
Full layer: 자체 인덱서를 실행하는 계층
-
Shared layer: 가장 가까운 Full layer에서 계산한 상위 $k$ 인덱스를 그대로 재사용하는 계층
이를 통해 대부분의 계층에서 인덱서 계산을 제거할 수 있다.
또한 이러한 계층 구성을 결정하고 최적화하기 위해 두 가지 상호보완적인 방법을 제안한다.
첫 번째는 Training-free IndexCache로, 가중치 업데이트 없이(calibration만 사용) 보정(calibration) 데이터셋에서 언어 모델링 손실(language modeling loss)을 직접 최소화하도록 탐욕적 탐색(greedy search) 알고리즘을 수행하여 어떤 계층에 인덱서를 유지할지를 결정한다.
두 번째는 Training-aware IndexCache이다. 이 방법은 다중 계층 증류 손실(multi-layer distillation loss)을 도입하여, 유지되는 각 인덱서가 자신이 담당하는 모든 계층의 평균 어텐션 분포를 학습하도록 한다. 이를 통해 단순한 교차(interleaved) 배치 방식만으로도 모든 계층에 인덱서를 유지하는 경우와 동등한 정확도를 달성할 수 있다.
30B 규모의 DSA 모델에서 수행한 실험 결과, IndexCache는 인덱서 계산의 75%를 제거하면서도 성능 저하를 거의 발생시키지 않았다. 그 결과 기존 DSA 대비 최대 1.82배의 prefill 속도 향상과 1.48배의 decode 속도 향상을 달성하였다. 또한 이러한 결과는 실제 서비스 규모의 GLM-5 모델을 대상으로 한 예비 실험(Figure 1)에서도 동일하게 확인되었다.
1. 서론
Self-Attention 메커니즘(Vaswani et al., 2017)은 현대 대규모 언어 모델(LLM)의 핵심 구성 요소이다. 그러나 시퀀스 길이에 대해 이차(quadratic) 복잡도를 갖기 때문에, 긴 문맥(long context)에서의 추론에서는 근본적인 병목이 된다. LLM이 장문의 Chain-of-Thought 추론, 다단계 에이전트(agentic) 워크플로우, 그리고 웹 규모의 검색 증강 생성(Retrieval-Augmented Generation, RAG)과 같이 긴 문맥을 요구하는 환경에 점점 더 많이 활용됨에 따라, 모델 성능을 유지하면서 어텐션 비용을 줄이는 것은 매우 중요한 연구 과제가 되었다.
희소 어텐션(sparse attention)은 이에 대한 원리적인 해결책을 제공한다. 모든 이전 토큰에 대해 어텐션을 수행하는 대신, 각 쿼리는 가장 관련성이 높은 일부 토큰만 선택하여 어텐션을 수행한다. 최근 제안된 여러 방법(Yuan et al., 2025; Lu et al., 2025; Zhao et al., 2025; Gao et al., 2026; Team et al., 2026) 가운데 DeepSeek Sparse Attention (DSA) (Liu et al., 2025)는 실제 서비스 환경에서 사용할 수 있는(production-grade) 학습 가능한(trainable) 희소 어텐션 기법으로 주목받고 있다.
희소 토큰 선택을 위해 DSA는 각 계층(layer)에 Lightning Indexer라는 추가 모듈을 도입한다. 이 모듈은 모든 이전 토큰에 대해 점수를 계산하고, 그중 상위 $k$개의 토큰을 선택하여 이후의 핵심 어텐션(core attention)에 전달한다. 이를 통해 계층별 핵심 어텐션의 계산 복잡도는
$$O(L^2)$$
에서
$$O(Lk)$$
로 감소하며, 추가적인 지속 사전학습(continued pre-training)을 통해 모델 성능도 유지할 수 있다.
그러나 인덱서(indexer) 자체는 여전히 모든 계층에서
$$O(L^2)$$
의 계산 복잡도를 가진다. 인덱서는 메인 어텐션보다 FLOP당 계산 비용은 저렴하지만, 총 $N$개의 계층에 대해 수행되므로 전체 비용은
$$O(NL^2)$$
이 된다. 이 비용은 문맥 길이에 대해 여전히 이차적으로 증가하며, 전체 어텐션 계산 비용에서 무시하기 어려운 비중을 차지한다.
논문의 Figure 1에서 30B DSA 모델을 프로파일링한 결과를 보면, 전체 지연(latency)에서 인덱서가 차지하는 비중은 문맥 길이가 길어질수록 급격히 증가하며, 특히 Prefill 단계에서 이러한 현상이 두드러진다. 반면 나머지 계산들은 상대적으로 완만하게 증가한다. 이는 긴 문맥에서 DSA를 더욱 빠르게 만들기 위해서는 인덱서 계산을 줄이는 것이 핵심임을 보여준다.
본 연구의 핵심 아이디어는 인덱서가 선택하는 상위 $k$개의 토큰이 인접한 계층들 사이에서 매우 높은 유사성을 가진다는 관찰에서 출발한다. 이러한 현상은 기존의 완전 어텐션(full attention) 모델에서도 보고된 계층 간 토큰 선택 안정성(cross-layer token selection stability) 과 일치한다(Deshmukh et al., 2025; Gao et al., 2026).
기존 연구들은 이러한 안정성을 활용하기 위해 일부 Full Attention Anchor Layer에서 계산한 인덱스를 다른 계층이 재사용하도록 하였다. 그러나 이러한 방법은 DSA에는 직접 적용될 수 없다. DSA에서는 완전 어텐션이 별도로 존재하지 않고, 경량 Lightning Indexer가 사실상 Full Attention의 역할을 대신하기 때문이다.
저자들은 DSA에서도 동일한 현상이 존재하는지 실험적으로 확인하였다. 모든 계층 쌍에 대해 Top-$k$ 인덱스의 중복률(overlap) 을 계산한 결과(Appendix A), 인접한 계층들은 선택된 토큰의 70~100%를 공유하는 것으로 나타났다. 또한 히트맵에서는 서로 높은 중복률을 보이는 계층 군집(layer clusters)이 명확하게 형성되었으며, 이는 대부분의 인덱서 계산이 중복되고 있음을 시사한다.
이러한 관찰은 다음과 같은 중요한 질문을 제기한다.
DSA에서 대부분의 인덱서를 제거하고, 소수의 인덱서만 유지한 채 나머지 계층들이 그 결과를 재사용해도 모델 성능을 유지할 수 있을까?
본 논문은 IndexCache를 제안하여 이 질문에 그렇다(Yes) 라고 답한다. IndexCache는 계층 간 인덱스 재사용(cross-layer index reuse) 을 통해 DSA에서 최대 75%의 인덱서 계산을 제거한다.
IndexCache는 계층을 두 종류로 나눈다.
-
F (Full) layer: 자체 인덱서를 유지하는 계층
-
S (Shared) layer: 가장 가까운 이전 Full layer에서 계산한 Top-$k$ 인덱스를 그대로 재사용하는 계층
추론 시에는 단 하나의 조건문(branch)만 추가하면 되며(Figure 2), 이를 통해 대부분의 인덱서 계산을 제거할 수 있다.
이를 위해 저자들은 두 가지 상호보완적인 방법을 제안한다.
실험적으로, 30B DSA 모델을 9개의 장문 및 추론 벤치마크에서 평가한 결과,
- Training-free IndexCache(Greedy Pattern Search 사용)와
- Training-aware IndexCache
모두 전체 인덱서의 25%만 유지하면서도 성능 저하는 거의 발생하지 않았다.
그 결과, 200K 문맥 길이에서 기존 DSA 대비
- 최대 1.82배의 Prefill 속도 향상
- 최대 1.48배의 Decode 속도 향상
을 달성하였다.
또한 744B 규모의 GLM-5를 대상으로 한 예비 실험에서도 이러한 접근법의 확장성이 확인되었으며, 긴 문맥 성능 저하 없이 최소 1.3배 이상의 속도 향상을 달성하였다.
2. 예비 지식 (Preliminary)
2.1 DeepSeek Sparse Attention
DeepSeek Sparse Attention (DSA) (Liu et al., 2025)는 각 어텐션 계층을 선택(selection) 과 계산(computation) 의 두 단계로 분리한다.
먼저 경량 Lightning Indexer가 다중 헤드(Multi-head) ReLU-gated 내적(dot product)을 이용하여 현재 쿼리(query)와 모든 이전 토큰(preceding tokens) 사이의 점수를 계산한다. 이후 점수가 가장 높은 상위 $k$개의 위치(top-$k$ positions) 를 선택한다.
그 다음 핵심 어텐션(core attention)은 선택된 희소 부분집합(sparse subset)에 대해서만 수행된다. 이를 통해 계층별 핵심 어텐션의 계산 복잡도는
$$O(L^2)$$
에서
$$O(Lk)$$
로 감소한다.
여기서
$$k = 2048 \ll L$$
이며, $L$은 입력 시퀀스의 길이를 의미한다.
Lightning Indexer는 높은 효율성을 목표로 설계되었다. 이를 위해
- 적은 수의 어텐션 헤드(few heads),
- 저랭크(low-rank) 투영,
- FP8 연산
을 사용한다. 그 결과 FLOP당 계산 비용은 메인 Multi-head Latent Attention (MLA) (Liu et al., 2024a)보다 약 한 자릿수(order of magnitude) 이상 저렴하다.
DSA는 MLA 위에 구현되며, 2단계의 지속 사전학습(two-stage continued pre-training) 을 통해 학습된다.
첫 번째 단계는 Dense Warm-up 단계이다.
이 단계에서는 모든 다른 모델 파라미터를 고정(freeze)한 상태에서 인덱서만 학습한다. 학습 목표는 각 계층의 전체 어텐션(full attention) 분포를 집계한 결과를 교사(teacher)로 사용하는 KL Divergence 기반 Distillation이다.
두 번째 단계는 Sparse Training 단계이다.
이 단계에서는 실제 Top-$k$ 선택을 활성화한 뒤 모델 전체를 공동으로 최적화한다. 이때 인덱서는 detached computational graph 상에서 Distillation gradient를 받아 학습된다.
이러한 계산량 절감에도 불구하고, 인덱서 자체는 여전히
$$O(L^2)$$
의 계산 복잡도를 가진다.
모든 계층에서 인덱서는 자신의 Top-$k$ 집합을 결정하기 위해 모든 이전 토큰에 대해 독립적으로 점수를 계산해야 한다.
따라서 총 $N$개의 계층을 갖는 모델에서는 전체 인덱서 계산 비용이
$$O(NL^2)$$
이 된다.
긴 문맥(long context)에서는 이 비용이 전체 어텐션 연산 예산(attention budget)에서 상당한 비중을 차지하게 된다.
이로부터 자연스럽게 다음과 같은 질문이 제기된다.
정말 모든 $N$개의 계층에서 각각 독립적으로 인덱서를 실행해야 하는가?
계층 간 존재하는 중복성을 활용할 수는 없는가?
2.2 토큰 선택의 계층 간 안정성 (Cross-Layer Stability of Token Selection)
이 질문에 대한 실마리는 보다 일반적인 경험적 관찰에서 얻을 수 있다.
Transformer에서는 중요한 토큰들의 집합이 연속된 계층들 사이에서 놀라울 정도로 안정적(stable) 이라는 사실이 알려져 있다.
Deshmukh et al. (2025)와 Gao et al. (2026)는 모두 인접한 계층들이 대부분의 Top-$k$ 어텐션 질량(attention mass) 을 공유한다는 사실을 관찰하였다.
이를 활용하기 위해 이들 연구는 일부 Anchor Layer만 Full Attention을 수행하고, 그 사이의 중간 계층들은 Anchor Layer에서 계산된 Top-$k$ 인덱스를 그대로 재사용하도록 설계하였다.
그러나 이러한 접근법은 중요한 전제를 가진다.
즉, 중요한 토큰을 식별하기 위한 기준(oracle)으로 Full Attention이 존재해야 한다는 것이다.
반면 DSA에서는 Full Attention이 완전히 제거되었으며, 그 역할을 경량 Lightning Indexer가 대신하고 있다.
따라서 아직까지 충분히 연구되지 않은 새로운 질문이 등장한다.
Lightning Indexer의 출력 역시 계층 간 안정성(cross-layer stability)을 가지는가?
만약 그렇다면,
- Full Attention Oracle 없이도
- 동일한 인덱스 공유(sharing) 원리를 적용하여
- 중복되는 인덱서 계산을 제거할 수 있을 것이다.
더 나아가 다음과 같은 질문도 자연스럽게 제기된다.
- 성능 저하가 발생하기 전까지 최대 어느 정도까지 인덱스를 재사용할 수 있는가?
- 공격적인(Index Reuse 비율이 높은) 인덱스 재사용으로 인해 발생하는 성능 손실을 모델 학습을 통해 보완할 수 있는가?
다음 장에서는 이러한 문제를 해결하기 위한 두 가지 새로운 상호보완적 방법을 제안한다.
3. 방법 (Method)
표기법 (Notation)
Transformer의 전체 계층 수를 $N$, 입력 시퀀스 길이를 $L$, 그리고 각 쿼리마다 선택되는 토큰의 개수를 $k$라고 하자.
계층 $\ell$에서 Lightning Indexer는 쿼리 위치 $t$에 대해 점수 벡터
$$I_t^{(\ell)} \in \mathbb{R}^{L}$$
를 생성한다.
이 점수 벡터로부터 상위 $k$개의 인덱스 집합을 다음과 같이 얻는다.
$$T_t^{(\ell)} = \operatorname{Top\mbox{-}k}(I_t^{(\ell)})$$
본 논문에서는 모든 실험에서
$$k = 2048$$
을 사용한다.
또한,
- 계층 $\ell$에서 위치 $t$의 집계된(aggregated) 어텐션 분포를
$$p_t^{(\ell)}$$
로 표기한다.
이는 여러 헤드의 softmax attention weight를 평균하여 얻는다.
반면 Lightning Indexer의 출력 분포는
$$q_t^{(\ell)}=\operatorname{Softmax}(I_t^{(\ell)})$$
로 정의한다.
개요 (Overview)
IndexCache는 DSA의 전체 $N$개 계층을 두 가지 역할(role)로 나눈다.
이를 길이 $N$의 이진(binary) 패턴 문자열
$$c=c_1c_2\cdots c_N,\qquad c_\ell\in\{F,S\}$$
로 표현한다.
각 기호의 의미는 다음과 같다.
• F (Full)
이 계층은 기존과 동일하게 Lightning Indexer를 유지한다.
모든 이전 토큰에 대해 새로운 Top-$k$ 인덱스
$$T_t^{(\ell)}$$
를 계산한 뒤,
선택된 토큰들에 대해서만 Sparse Core Attention을 수행한다.
즉, 표준 DSA와 동일하게 동작한다.
• S (Shared)
이 계층은 Lightning Indexer를 제거한다.
대신 가장 가까운 이전 Full layer에서 계산한 인덱스를 그대로 재사용한다.
즉,
$$T_t^{(\ell)}
\leftarrow
T_t^{(f(\ell))}$$
이며,
$$f(\ell)=\max\{j<\ell\mid c_j=F\}$$
이다.
이렇게 전달받은 인덱스를 이용하여 바로 Sparse Core Attention을 수행한다.
첫 번째 계층은 항상 F로 지정한다.
이는 최초의 Top-$k$ 인덱스를 생성해야 하기 때문이다.
추론(inference) 시에는 S layer가 단순히
- Indexer의 forward pass를 생략하고
- 가장 가까운 이전 F layer에서 계산한 index tensor를 그대로 사용하면 된다.
Figure 2에서 볼 수 있듯이, 기존 추론 루프에서 필요한 수정은 매우 단순하다.
각 계층마다
- Indexer를 실행할지,
- 아니면 캐시된 인덱스를 복사할지
를 결정하는 조건문 하나만 추가하면 된다.
이 방법의 핵심 설계 문제는 패턴 $c$를 어떻게 결정하느냐이다.
만약 대부분의 계층이 안전하게 인덱스를 공유할 수 있다면,
전체 Indexer 계산 비용인
$$O(NL^2)$$
의 상당 부분을 제거할 수 있으며,
Sparse Core Attention의 계산량인
$$O(NLk)$$
은 그대로 유지된다.
이를 위해 저자들은 두 가지 방법을 제안한다.
-
Training-free 방법
- 이미 학습된 DSA 모델에서 Greedy Search를 수행하여 패턴 $c$를 결정한다.
- (Section 3.1)
-
Training-aware 방법
- Multi-layer Distillation Loss를 이용하여 Cross-layer Sharing에 적합하도록 Indexer를 함께 학습한다.
- (Section 3.2)
3.1 Training-Free IndexCache
이미 학습된(pretrained) DSA 모델이 주어졌다고 하자.
목표는 가능한 많은 계층을 Shared layer(S) 로 바꾸면서도 모델 성능 저하를 최소화하는 패턴
$$c$$
를 찾는 것이다.
먼저 가장 직관적인 접근이 왜 실패하는지 설명한 뒤,
Greedy Search 알고리즘을 제안한다.
3.1.1 Uniform Interleaving이 최적이 아닌 이유
가장 단순한 방법은 Uniform Interleaving이다.
예를 들어
처럼
매 $r$번째 계층만 Indexer를 유지하는 방식이다.
예를 들어
$$r=4$$
이면 4개 계층마다 하나만 F가 된다.
그러나 이 방식에는 중요한 문제가 있다.
각 계층의 Indexer 중요도는 동일하지 않다.
실험적으로,
- 특히 네트워크의 초기 계층,
- 그리고 구조가 전환되는(transitional) 구간
에서는 Indexer를 제거했을 때 성능이 크게 감소하는 반면,
다른 계층에서는 거의 영향이 없다.
따라서 Uniform Interleaving은
- 중요한 Indexer를 제거하고,
- 반대로 거의 필요 없는 Indexer를 남기는
비효율적인 선택을 할 수 있다.
그 결과 모델 성능이 눈에 띄게 감소한다
(정량적 비교는 Section 4.3).
이러한 관찰은
모델 스스로 어떤 Indexer가 중요한지를 판단하게 하자
라는 아이디어로 이어지며,
이를 위해 Greedy Layer Selection 알고리즘을 제안한다.
3.1.2 Layer Selection Algorithm
저자들은 Greedy Search를 이용하여
F layer를 하나씩 S layer로 바꾸어 나간다.
각 선택의 기준은
작은 Calibration Set에서 측정한 Language Modeling Loss
이다.
LM Loss를 다운스트림 성능의 대리 지표(proxy)로 사용한다.
Calibration Set
학습 데이터에서
$$B$$
개의 Mini-batch를 미리 저장(cache)해 둔다.
모든 후보 패턴은 항상 동일한 Mini-batch에서 평가된다.
따라서 Loss 차이는
데이터 차이가 아니라
패턴 변경 자체의 영향만을 반영한다.
패턴
$$c$$
에 대한 손실은
모델
$$M$$
과 전체 Batch
$$D$$
에 대해 Forward Pass를 수행하여 얻는다.
논문에서는 이를
$$\operatorname{EVALLOSS}(M,D,c)$$
로 표기한다.
Greedy Search 과정
초기 상태는
$$c_\ell = F
\qquad
(\forall \ell)$$
즉,
모든 계층이 Full layer인 기존 DSA이다.
이후
$$K$$
단계 동안 반복한다.
여기서
$$K$$
는 제거하려는 Indexer 개수이다.
예를 들어 전체 Indexer의 1/4만 유지하려면
$$K=\frac{3N}{4}$$
이다.
매 단계에서
- 첫 번째 계층을 제외한 모든 F layer를 순회한다.
- 각 계층을 임시로 S로 변경한다.
- LM Loss를 계산한다.
- Loss 증가가 가장 작은 계층을 선택한다.
- 그 변경을 확정(commit)한다.
Algorithm 1에서 전체 절차를 제시한다.
계산 복잡도
모든 계층이 F인 상태에서
모든 계층이 S가 될 때까지 탐색하면
Forward Pass는 총
$$\frac{N(N-1)}{2}$$
번 수행된다.
이는 상당히 큰 비용이다.
논문에서는 이를 줄이기 위해
Pipeline Parallelism을 활용한다.
모델을
$$P$$
개의 Block으로 나눈 뒤,
각 Block의 첫 번째 계층은 항상 F로 고정한다.
그리고 매 단계마다
각 Block에서 가장 좋은 후보를 하나씩 선택하여 확정한다.
이렇게 하면
한 단계에서 최대
$$P$$
개의 계층을 동시에 S로 바꿀 수 있으며,
전체 Forward Pass 수는 약
$$P$$
배 감소한다.
Greedy Search의 특성
Greedy Search는 전역 최적해(global optimum)를 보장하지는 않는다.
그럼에도 불구하고 저자들은 세 가지 흥미로운 특성을 반복적으로 관찰하였다.
-
Uniform Interleaving보다 항상 성능이 우수하다.
동일한 Indexer 유지 비율에서도 Greedy Search가 더 좋은 성능을 보인다(Table 2).
-
계층 중요도가 자연스럽게 드러난다.
오른쪽 그림(유지 비율 1/2, 1/4, 1/8 표시)을 보면,
Greedy Search 과정에서의 LM Validation Loss는
- 처음 약 20단계에서는 거의 증가하지 않는 "쉬운(easy)" 계층,
- 약 35단계 이후 급격히 증가하는 "중요한(critical)" 계층
으로 명확히 구분된다.
이는 Indexer의 중요도에 자연스러운 순서(ordering)가 존재함을 시사한다.
-
Calibration Set이 달라도 결과가 매우 안정적이다.
즉,
Greedy Search가 찾아내는 중요도 순위는 특정 데이터셋에 의한 우연한 결과가 아니라,
모델 자체가 가진 고유한 특성(intrinsic property) 임을 의미한다.
또한 LM Loss는 실제 다운스트림 태스크 성능과 양의 상관관계를 가지므로,
Layer Selection을 위한 유효한 Proxy Metric으로 사용할 수 있다.
3.2 Multi-Layer Distillation을 이용한 Training-Aware IndexCache
Training-free IndexCache는 추가적인 가중치 업데이트가 필요 없다는 장점이 있지만, 근본적인 한계를 가진다.
기존 DSA에서는 각 Lightning Indexer가 오직 자신의 계층만을 위해 학습되었다. 따라서 하나의 Indexer를 여러 계층이 공유하도록 만들면, 원래의 학습 목표와 맞지 않아 성능이 감소할 수 있다.
반면 DSA를 처음부터 학습하거나(Training from Scratch), 지속 사전학습(Continued Pre-training)을 수행하는 경우에는 더 나은 방법이 가능하다.
즉,
유지되는 하나의 Indexer가 여러 계층을 동시에 담당하도록 직접 학습시키는 것
이다.
Single-layer Distillation에서 Multi-layer Distillation으로
기존 DSA 학습(Section 2.1)에서는
계층 $\ell$의 Lightning Indexer가 같은 계층의 평균 어텐션 분포
$$p_t^{(\ell)}$$
를 모방하도록 KL Divergence Distillation을 수행한다.
손실 함수는
$$L_I
=
\sum_t
D_{\mathrm{KL}}
\left(
p_t^{(\ell)}
\;\|\;
q_t^{(\ell)}
\right)$$
이다.
여기서
- (p_t^{(\ell)}): Teacher(Full Attention)의 평균 어텐션 분포
- (q_t^{(\ell)}): Lightning Indexer의 출력 분포
이다.
논문에서는 이를 여러 계층으로 일반화한다.
계층 $\ell$이 유지되는 Full layer(F) 라고 하자.
그리고
$$\ell+1,\;
\ell+2,\;
\cdots,\;
\ell+m$$
이 모두 Shared layer라고 하자.
이들은 모두
$$T_t^{(\ell)}$$
를 그대로 재사용한다.
따라서 이 하나의 Indexer는
자신의 계층뿐 아니라
모든 Shared layer에서도 좋은 Top-$k$를 생성해야 한다.
이를 위해 다음과 같은 Multi-layer Distillation Loss를 정의한다.
$$L_I^{\mathrm{multi}}
=
\frac{1}{m+1}
\sum_{j=0}^{m}
\sum_t
D_{\mathrm{KL}}
\left(
p_t^{(\ell+j)}
\;\|\;
q_t^{(\ell)}
\right)$$
직관적으로,
기존 방식은
"Layer $\ell$ 하나만 잘 맞추도록"
학습하였다.
반면 Multi-layer Distillation은
"Layer $\ell$, $\ell+1$, ..., $\ell+m$ 모두에서 동시에 좋은 Top-$k$를 생성하도록"
학습한다.
즉,
특정 계층 하나에 과적합(overfitting)되는 대신,
여러 계층이 공통적으로 중요하게 생각하는 토큰을 선택하도록 만든다.
평균 분포에 대한 Distillation과 Gradient가 동일함
자연스러운 의문이 생긴다.
여러 개의 KL Loss를 더하면 서로 간섭(interaction)이 발생하지 않을까?
논문은 그렇지 않음을 증명한다.
먼저 여러 Teacher Distribution의 평균을 정의한다.
$$\bar p_t
=
\frac{1}{m+1}
\sum_{j=0}^{m}
p_t^{(\ell+j)}$$
그리고 이 평균 분포 하나만 Teacher로 사용하는 손실을 정의한다.
$$L_I^{\mathrm{avg}}
=
\sum_t
D_{\mathrm{KL}}
\left(
\bar p_t
\;\|\;
q_t^{(\ell)}
\right)$$
Proposition 1
다음이 성립한다.
$$\nabla_\theta
L_I^{\mathrm{multi}}
=
\nabla_\theta
L_I^{\mathrm{avg}}$$
즉,
두 손실 함수는 Gradient가 완전히 동일하다.
증명
KL Divergence는
$$D_{\mathrm{KL}}(p\|q)
=
\sum_s
p(s)
\log
\frac{p(s)}{q(s)}$$
로 정의된다.
여기서
Teacher 분포
$$p$$
는 고정되어 있으므로,
미분 시 Entropy 항은 사라진다.
즉,
$$\nabla_\theta
D_{\mathrm{KL}}
(p\|q)
=
-
\nabla_\theta
\sum_s
p(s)
\log q(s)$$
이다.
이를 Multi-layer Loss에 적용하면
$$\nabla_\theta
L_I^{\mathrm{multi}}
=
-
\frac{1}{m+1}
\sum_{j=0}^{m}
\sum_t
\nabla_\theta
\sum_s
p_t^{(\ell+j)}(s)
\log
q_t^{(\ell)}(s)$$
이 된다.
평균을 앞으로 모으면
$$\nabla_\theta
L_I^{\mathrm{multi}}
=
-
\sum_t
\nabla_\theta
\sum_s
\left(
\frac{1}{m+1}
\sum_{j=0}^{m}
p_t^{(\ell+j)}(s)
\right)
\log
q_t^{(\ell)}(s)$$
괄호 안은 바로 평균 분포
$$\bar p_t$$
이다.
따라서
$$\nabla_\theta
L_I^{\mathrm{multi}}
=
\nabla_\theta
L_I^{\mathrm{avg}}$$
가 된다.
증명 완료.
□
해석 (Interpretation)
이 결과는 매우 중요한 의미를 가진다.
Multi-layer Distillation은
단순히
여러 개의 KL Loss를 평균낸 휴리스틱(heuristic)
이 아니다.
오히려
여러 계층의 Attention Distribution 중심(centroid)을 Teacher로 사용하는 Distillation과 정확히 동일하다.
즉,
Lightning Indexer는
각 계층을 따로따로 맞추는 것이 아니라,
여러 계층이 공통적으로 중요하다고 생각하는 토큰들을 선택하도록 학습된다.
다시 말해,
Shared layer 전체에서 사용할 수 있는 **Consensus Top-$k$**를 예측하도록 학습되는 것이다.
왜 실제 구현에서는 (L_I^{\mathrm{multi}})를 사용하는가?
두 손실 함수는 Gradient가 동일하므로
이론적으로는 어느 쪽을 사용해도 된다.
하지만 실제 구현에서는
$$L_I^{\mathrm{multi}}$$
를 사용한다.
이유는 메모리 효율 때문이다.
Shared layer에서는
현재 Full layer의 예측값
$$q^{(\ell)}$$
만 전달하면 된다.
반면
$$L_I^{\mathrm{avg}}$$
를 사용하려면
평균 Distribution을 계산하기 위해
$$p^{(\ell)}$$
$$q^{(\ell)}$$
를 모두 전달해야 한다.
이는 불필요한 메모리 사용량과 추가적인 계산 비용을 발생시킨다.
따라서 실제 구현에서는
Gradient는 동일하면서 구현이 더 단순한
$$L_I^{\mathrm{multi}}$$
를 채택하였다.
학습 과정 (Training)
전체 학습 절차는 기존 DSA와 동일한 2단계 학습을 따른다.
1. Warm-up 단계
- Full layer(F)의 Lightning Indexer만 학습한다.
- 나머지 모든 모델 파라미터는 고정(freeze)한다.
- 손실 함수는
$$L_I^{\mathrm{multi}}$$
를 사용한다.
2. Sparse Training 단계
실제 Top-$k$ 선택을 활성화한 뒤,
Indexer는 계속해서
$$L_I^{\mathrm{multi}}$$
로 학습한다.
이때 KL Divergence는 선택된 Top-$k$ 토큰에 대해서만 계산한다.
동시에,
나머지 모델 파라미터들은 일반적인 Language Modeling Loss(LM Loss) 를 이용하여 함께 최적화한다.
4. 실험 (Experiments)
4.1 실험 설정 (Setup)
모델 (Model)
실험에 사용한 DSA 모델은 GLM-4.7-Flash*의 기본(base) 모델로부터 시작하여 2단계 학습(two-stage training) 을 거쳐 구축되었다.
GLM-4.7-Flash는 다음과 같은 특성을 가진다.
- 30B-A3B MoE 모델
- Multi-head Latent Attention (MLA) 사용
- 47개의 Transformer Layer
평가 성능은 원래의 GLM-4.7-Flash와 거의 동일한 수준이다(Zeng et al., 2026의 Table 6 참고).
Training-free IndexCache
Greedy Pattern Search는
- SFT(Supervised Fine-Tuning) 데이터
- Batch Size = 768
- Context Length = 200K
에서 계산한 토큰당 Validation Loss를 기준으로 수행하였다.
Training-aware IndexCache
처음부터 DSA 전체 학습 파이프라인을 수행하는 것은 매우 많은 계산 자원이 필요하다.
따라서 저자들은
GLM-4.7-Flash 기본 모델에서 직접 시작하여
SFT 데이터(200K Context Length)를 이용해 DSA 모델로 학습하였다.
학습은 다음 두 단계로 구성된다.
-
Dense Warm-up
-
Sparse Training
이 단축된 학습 파이프라인(shortened pipeline)은 전체 DSA 학습과 거의 동일한 성능을 달성하면서도, Training-aware IndexCache를 평가하기에는 충분하였다.
평가 벤치마크 (Evaluation)
Long-context 벤치마크 (5개)
- MRCR v2
- GraphWalks
- LongBench v2
- RULER
- AA-LCR
일반 및 추론 벤치마크 (4개)
- AIME 2025
- GPQA-Diamond
- LiveCodeBench v6
- IFBench
평가 세부 설정은 Appendix D에 제시되어 있다.
4.2 End-to-End 추론 속도 향상
추론 성능은
- 30B DSA 모델
- SGLang
- dp attention 활성화 (dp size = 8)
- NVIDIA H100 노드
에서 측정하였다.
비교 대상은
- 기존 DSA
- IndexCache (Indexer 유지율 1/2)
- IndexCache (Indexer 유지율 1/4)
이다.
다음 네 가지 Context Length에서 측정하였다.
평가 지표는 세 가지이다.
-
Prefill Latency
- 첫 번째 토큰이 생성될 때까지 걸리는 시간(Time-To-First-Token)
-
Single-request Decode Throughput
- GPU 하나당 하나의 요청만 처리할 때의 Decode 속도
-
Total Decode Throughput
- KV Cache를 완전히 활용(약 GPU당 800K 토큰)했을 때의 전체 Decode 처리량
결과는 Table 1과 Figure 3에 제시되어 있다.
Prefill 성능
IndexCache는 문맥 길이가 길어질수록 더욱 큰 Prefill 가속 효과를 보였다.
200K Context에서는
Prefill Latency가
- DSA : 19.5초
- IndexCache (1/4) : 10.7초
로 감소하였다.
즉,
1.82배의 속도 향상을 달성하였다.
이는 Prefill 단계의 병목이던 Indexer 계산의 75%를 제거했기 때문이다.
심지어 Context Length가 10K처럼 짧은 경우에도
1.27배의 속도 향상을 얻었다.
저자들은 200K보다 더 긴 문맥에서는 속도 향상이 더욱 커질 것으로 예상한다.
Decode 성능
긴 문맥에서 Decode 속도 향상도 매우 크다.
200K Context에서
- DSA : 58 token/s
- IndexCache (1/4) : 86 token/s
로
1.48배의 속도 향상을 달성하였다.
그 이유는 DSA의 Decode 과정에서도
매 토큰마다 전체 문맥에 대해 Indexer를 실행해야 하기 때문이다.
문맥이 길어질수록 이 Indexer가 병목이 되며,
IndexCache는 바로 이 병목을 제거한다.
KV Cache가 가득 찬 상태에서는
IndexCache (1/4)가
모든 Context Length에서
22~51%의 Decode Throughput 증가를 보였다.
특히 200K에서는
- DSA : 197 token/s
- IndexCache : 297 token/s
으로
1.51배 증가하였다.
744B 규모의 GLM-5에서도 비슷한 결과가 나타났다.
100K 이상의 Context에서
IndexCache (1/4)는
모두 최소 1.3배 이상의 속도 향상을 달성하였다.
전체적으로,
IndexCache는 긴 문맥 추론(long-context inference) 에서 특히 큰 효과를 보였다.
4.3 Training-Free IndexCache 결과
Table 2에서는 Training-Free IndexCache를 평가하였다.
비교한 Indexer 유지율은
이며,
각각에 대해
- Uniform Interleave
- Greedy Search Pattern
을 비교하였다.
(Appendix B)
Greedy Search는 Long-context 성능을 거의 완전히 회복한다
Uniform Interleave는 Indexer를 많이 제거할수록 Long-context 성능이 크게 감소한다.
예를 들어
Long Avg는
- 원래 DSA : 50.2
- Uniform 1/2 : 47.4
- Uniform 1/4 : 43.0
으로 감소하였다.
반면 Greedy Search Pattern은 이러한 성능 저하를 거의 제거하였다.
Long Avg는
- 1/4 유지 : 49.9
- 1/2 유지 : 50.3
으로,
원래 DSA와 거의 동일하였다.
즉,
Indexer를 몇 개 남기느냐보다 어떤 계층의 Indexer를 남기느냐가 훨씬 중요하다.
그러나 Indexer를 1/8만 유지하면
성능 저하가 상당히 커졌다.
Long Avg는
- Uniform : 35.3
- Greedy : 46.1
이었다.
Greedy Search가 손실을 크게 줄여주기는 하지만,
이 정도의 극단적인 Sparsity에서는 성능 저하를 완전히 막을 수는 없었다.
긴 Chain-of-Thought 추론 능력은 유지된다
Uniform 1/8을 제외하면
모든 설정에서
G&R Avg는
기존 DSA(74.6)와 1점 이내의 차이만 보였다.
특히
Greedy 1/4는 오히려
으로 DSA보다 높은 점수를 기록하였다.
이는 불필요한 Indexer 계산을 제거하는 것이
추론 과정에서 약한 형태의 Regularization 효과를 낼 수도 있음을 시사한다.
즉,
IndexCache는
Long-context 효율성을 얻기 위해
일반적인 추론 능력을 희생하지 않는다.
4.4 Training-Aware IndexCache 결과
Training-Aware IndexCache는
Section 3.2의 Multi-layer Distillation을 이용하여 학습하였다.
평가한 유지율은
이며,
모두 Uniform Pattern을 사용하였다.
또한
1/2 유지율에서 두 가지 Ablation을 수행하였다.
-
Uniform 대신 Greedy Pattern 사용
-
Cross-layer Distillation 제거
(각 Indexer를 자기 Layer만 Distillation)
결과는 Table 3에 제시되어 있다.
주의할 점은
여기서의 DSA Baseline 역시
단축 학습 파이프라인을 사용했기 때문에,
Table 2의 DSA와는 약간의 성능 차이가 있다.
Training-aware IndexCache는 DSA와 동일한 성능을 달성한다
Uniform Pattern만 사용해도
1/2 유지율에서
Long Avg는
- DSA : 51.0
- IndexCache : 51.6
으로 오히려 더 높았다.
G&R Avg도
로 거의 동일하였다.
1/4 유지율에서도
두 평균 성능 모두
Baseline과 0.4% 이내의 차이만 보였다.
즉,
모델은
Index Sharing 구조에 적응하도록 학습될 수 있다.
재학습하면 Layer Pattern의 중요성이 사라진다
Training-Free에서는
Greedy Search가 반드시 필요하였다.
그러나 Training-aware에서는
Uniform Pattern만으로도
Greedy보다 같거나 오히려 더 좋은 성능을 얻었다.
예를 들어
1/2 유지율에서
Long Avg는
- Uniform : 51.6
- Greedy : 50.6
이었다.
그 이유는 다음과 같다.
Training-Free에서는
각 Layer가 자신의 Indexer 출력에 강하게 의존한다.
따라서 이전 Layer의 Index를 사용하면
분포 변화(distribution shift)가 발생하여 성능이 감소한다.
Greedy Search는
이러한 민감한 Layer를 피하는 방식으로 동작한다.
반면 Training-aware에서는
- Shared Layer는 상속받은 Index에 적응하도록 학습되고,
- Full Layer의 Indexer는 여러 Layer에서 사용할 수 있는 Index를 생성하도록 학습된다.
이러한 공동 적응(joint adaptation)을 통해
Layer마다 존재하던 민감성이 사라지며,
단순한 Uniform Pattern만으로도
기존 DSA와 동일한 성능을 얻을 수 있게 된다.
Cross-layer Distillation은 실제로 효과가 있다
Cross-layer Distillation을 제거하면
Long Avg는
로 감소하였다.
특히
AA-LCR은
으로 크게 감소하였다.
이는 Multi-layer Distillation이 실제로 매우 중요한 역할을 함을 보여준다.
각 Indexer가
여러 Layer의 평균 Attention Distribution(centroid)을 학습함으로써,
특정 Layer에만 특화되지 않고
여러 Layer에서 공통적으로 사용할 수 있는 **Consensus Top-$k$**를 생성하게 된다.
4.5 확장성 실험 (Scaling Experiment)
Training-Free IndexCache를
744B(Active 40B) 규모의 GLM-5에도 적용하였다.
GLM-5는 기본적으로 DSA를 사용한다.
Table 4는
5개의 Long-context 벤치마크 결과를 보여준다.
전체적인 경향은 30B 모델과 동일하였다.
- Uniform Interleave는 공격적인 Index 제거 시 성능이 감소한다.
- Greedy Search는 성능을 대부분 회복한다.
흥미롭게도
1/2 유지율에서는
Uniform Pattern도
Long Avg를 거의 유지하였다.
이는 단순히
균등 패턴이 우연히 중요한 Indexer Layer를 건너뛰지 않았기 때문이라고 저자들은 설명한다.
반면 Greedy Pattern은
더 안정적인 결과를 보였다.
- 1/2 유지 : 78.7
- Baseline : 78.4
로 약간 더 좋았으며,
1/4 유지에서도
로 거의 동일하였다.
또한 Artificial Analysis Index의 모든 평가에서
Indexer를 절반만 유지한 IndexCache를 평가한 결과,
원래 GLM-5와 거의 동일한 성능을 보였다(Figure 1).
저자들은 가까운 미래에
이 Production-scale 모델에도
Training-Aware IndexCache를 적용할 계획이다.
이미 Training-Free 버전만으로도 Baseline과 동일한 품질을 달성했기 때문에,
Training-aware 학습을 적용하면
이러한 성능을 더욱 안정적으로 유지하면서,
실제 서비스 환경에서의 추론 효율성도 더욱 향상될 것으로 기대하고 있다.
5. 관련 연구 (Related Work)
5.1 효율적인 어텐션 (Efficient Attention)
Self-Attention의 이차(quadratic) 계산 비용을 줄이는 것은, 특히 긴 시간 범위(long-horizon)의 에이전트(agent) 활용이 증가하는 현재 LLM 연구에서 핵심적인 주제이다.
Training-free Sparse Attention
Training-free 희소 어텐션 방법들은 추론 단계에서만 희소성을 도입한다.
이를 위해 다음과 같은 기법들이 사용된다.
- 고정된 희소 패턴(fixed sparse patterns)
- 휴리스틱 기반 토큰 제거(eviction) 전략
- 경량 중요도 추정(lightweight importance estimation)
(Zhang et al., 2023; Xiao et al., 2024; Ribar et al., 2024; Tang et al., 2024; Jiang et al., 2024; Xiao et al., 2025; Xu et al., 2025; Lai et al., 2025; Fu et al., 2025; Zhu et al., 2025; Yang et al., 2025c; Zhang et al., 2025)
그러나 이러한 방법들은 학습(training)과 추론(inference)의 불일치(training-inference mismatch) 를 초래하며, 긴 문맥에서는 오류가 누적(error accumulation)될 수 있다(Hu et al., 2026).
Trainable Sparse Attention
반면 Trainable Sparse Attention은 학습 단계에서부터 희소성을 모델에 내재화한다.
대표적인 방법으로는
- 학습 가능한 게이팅(gating) 메커니즘(Gao et al., 2024; 2025)
- End-to-End Sparse Pre-training(Yuan et al., 2025)
- 블록 단위 Mixture Routing(Lu et al., 2025; Zhao et al., 2025; Team et al., 2025b)
- Full Attention에서 Sparse Attention으로의 Distillation(Liu et al., 2025; Shen et al., 2025)
등이 있다.
본 논문의 기반이 되는 DeepSeek Sparse Attention (DSA) (Liu et al., 2025)는
Full Attention으로부터 경량 Lightning Indexer를 Distillation하여,
각 Query에 대해 Top-$k$ 토큰을 선택하도록 한다.
이를 통해 핵심 어텐션(Core Attention)의 계산 복잡도를
$$O(Lk)$$
로 줄인다.
희소화(sparsity) 이외에도,
Hybrid Architecture는 계산량이 큰 Quadratic Attention Layer의 개수를 줄이는 접근을 사용한다.
예를 들어 Quadratic Layer와 다음과 같은 구조를 번갈아(interleave) 배치한다.
5.2 계층 간 공유 (Cross-Layer Sharing)
최근 연구들은 인접한 Transformer Layer의 표현(representation)이 매우 높은 일관성(consistency)을 가진다는 사실을 보여주었다.
이러한 구조적 특성은 계산 중복을 줄이고 추론을 가속하기 위해 활용되고 있다.
대표적인 연구로는
- TidalDecode (Yang et al., 2025a)
- LessIsMore (Yang et al., 2025b)
- OmniKV (Hao et al., 2025)
- DELTA (Zarch et al., 2025)
등이 있다.
이들은 주기적으로 배치된 Anchor Layer에서 계산한 Top-$k$ 인덱스를 이후 Sparse Decoding Layer들이 재사용하도록 한다.
Kascade (Deshmukh et al., 2025)는
계층 간 유사도(cross-layer similarity) 행렬에 대해 동적 계획법(dynamic programming) 을 적용하여,
최적의 Anchor Layer를 선택하는 방법을 제안하였다.
또한 정확도를 유지하기 위해서는 Head-aware Remapping이 매우 중요함을 보였다.
그러나 이러한 기존 방법들은 모두
Anchor Layer에서 Full Attention을 수행하여 정확한 Top-$k$를 계산한다
는 전제를 갖는다.
한편, 또 다른 연구 흐름은
Cross-layer KV Cache Sharing이다.
이 방법들은 여러 Layer가 동일한 Key-Value Cache를 공유하도록 하여 메모리 사용량을 줄인다.
대표적인 연구는
- Sun et al., 2024
- Brandon et al., 2024
- Liu et al., 2024b
- Qiao et al., 2025
- Zuhri et al., 2025
- Wu et al., 2025
등이다.
HySparse (Gao et al., 2026)는
위 두 가지 접근을 통합하였다.
즉,
- Full Attention Layer
- Sparse Layer
를 번갈아 배치하고,
Sparse Layer는
- Top-$k$ Block Index
- KV Cache
를 모두 상속받는다.
그러나 지금까지의 모든 방법은
Full Attention Layer를 Oracle로 사용해야 한다.
반면 DSA에서는 Full Attention 자체가 완전히 제거되어 있다.
IndexCache와 기존 연구의 차이점
IndexCache는 기존 연구와 두 가지 측면에서 다르다.
첫째,
Oracle 자체가 훨씬 저렴하다.
기존 연구는
$$O(L^2)$$
의 Full Attention Score를 공유하지만,
IndexCache는
DSA의 경량 Lightning Indexer 출력만 공유한다.
따라서 계산 비용이 훨씬 낮다.
둘째,
Index Sharing 구성을 체계적으로 최적화하는 방법을 함께 제안한다.
구체적으로
비록 본 논문은 DSA를 대상으로 구현하였지만,
IndexCache의 핵심 아이디어는
고정된 Sparse Pattern이 아니라 동적인 Token Selection을 사용하는 모든 Sparse Attention으로 확장될 수 있다.
예를 들어
- MoBA (Lu et al., 2025)의 Block-level Selection
- NSA (Yuan et al., 2025)의 Dynamic Token Selection
에서도 동일한 방식으로 계층 간 Index 재사용(Cross-layer Reuse) 을 적용할 수 있을 것으로 기대된다.
6. 결론 (Conclusion)
본 논문에서는 IndexCache를 제안하였다.
IndexCache는 토큰 선택을 담당하는 Lightning Indexer의 계층 간 중복성(cross-layer redundancy) 을 활용하여 Sparse Attention을 가속하는 방법이다.
IndexCache는 전체 Layer를
-
소수의 F(Full) Layer
-
다수의 S(Shared) Layer
로 나눈다.
이를 통해
전체 Indexer 계산 비용
$$O(NL^2)$$
의 최대 75%를 제거할 수 있으며,
추론 과정에서는 단 하나의 조건문(branch) 만 추가하면 된다.
또한 이러한 계산량 감소에도 불구하고 모델 성능 저하는 발생하지 않았다.
보다 넓은 관점에서 보면,
본 연구는 지금까지 Full Attention을 Oracle로 사용하는 경우에만 적용되던 계층 간 공유(cross-layer sharing) 의 원리가,
Sparse Attention에도 자연스럽게 확장될 수 있음을 보여준다.
향후 DeepSeek-V3.2, GLM-5와 같이 Sparse Attention이 차세대(frontier) LLM의 기본 구조(default) 가 되어감에 따라,
저자들은 Cross-layer Index Reuse가 효율적인 추론 파이프라인의 표준 구성 요소로 자리잡을 것으로 기대하고 있다.
A. 계층 간 Top-(k) 인덱스 중복도 (Cross-Layer Top-(k) Index Overlap)
DSA에서 Top-(k) 인덱스 선택이 계층 간 얼마나 중복되는지를 실험적으로 검증하기 위해, 각 계층의 Lightning Indexer가 선택한 Top-(k) 인덱스 사이의 쌍별(pairwise) 중복률(overlap ratio) 을 계산하였다.
구체적으로, 모든 계층 쌍 ((i, j))에 대해 다음 값을 측정하였다.
$$\frac{|T^{(i)} \cap T^{(j)}|}{k}$$
여기서
- (k = 2048)
- (T^{(i)}) : 계층 (i)가 선택한 Top-(k) 토큰 집합
- (T^{(j)}) : 계층 (j)가 선택한 Top-(k) 토큰 집합
이다.
이 값은 길이 200K인 입력 768개로 구성된 Calibration Set에 대해 평균을 계산하였다.
Figure 4. 30B DSA 모델의 모든 계층 쌍에 대한 Top-(k) 인덱스 중복률. Greedy Search로 찾은 1/4 유지율(IndexCache) 의 공유 블록(shared blocks)은 빨간색 박스로 표시되어 있다.
Figure 4는 47개 계층을 갖는 30B DSA 모델의 중복률을 히트맵(heatmap)으로 시각화한 것이다. 또한 Greedy Search로 찾은 1/4 IndexCache 패턴에서의 공유 블록을 빨간색 박스로 표시하였다.
이 그림에서는 다음과 같은 몇 가지 중요한 특징을 확인할 수 있다.
• 대각선(diagonal) 근처에서 높은 중복률
인접한 계층들은 0.7~1.0의 매우 높은 중복률을 보인다.
이는 연속된 계층들이 거의 동일한 토큰 집합을 선택한다는 것을 의미하며, 계층 간 Top-(k) 선택의 높은 안정성을 확인해 준다.
• 블록 구조(Block Structure)
히트맵에서는 서로 높은 중복률을 보이는 계층 군집(layer clusters) 이 뚜렷하게 나타난다.
대표적인 예는 다음과 같다.
- Layer 3–5
- Layer 6–8
- Layer 17–30
- Layer 31–36
등이다.
이는 모델이 여러 기능적 블록(functional blocks) 으로 구성되어 있으며, 각 블록 내부에서는 토큰 선택이 매우 일관적이라는 것을 시사한다.
• 불균일한 감소(Uneven Decay)
중복률은 블록 내부에서는 천천히 감소하지만,
블록 경계(block boundary) 를 넘으면 훨씬 빠르게 감소한다.
이는 일부 전이(transition) 계층이 어텐션의 초점을 크게 바꾸는 역할을 수행한다는 것을 의미한다.
• 초기 계층과 후반 계층의 차이(Early–Late Distinction)
히트맵의 좌하단(bottom-left) 과 우상단(top-right) 영역은 매우 어둡게 나타난다.
즉,
$$\text{Overlap} \le 0.4$$
이다.
이는 초기 계층과 후반 계층이 근본적으로 서로 다른 토큰 집합에 주목(attend) 한다는 것을 보여준다.
다시 말해,
모델의 앞부분과 뒷부분은 서로 다른 정보를 처리하며, 동일한 Top-(k) 인덱스를 공유하기 어렵다는 점을 시사한다.
Authors
¹ Tsinghua University
² Z.ai
Abstract
장문(long-context) 에이전트(agentic) 워크플로우는 대규모 언어 모델(LLM)의 대표적인 활용 사례로 자리 잡았으며, 이에 따라 어텐션(attention)의 효율성은 추론 속도와 서비스 비용 모두에 있어 핵심적인 요소가 되었다. 희소 어텐션(sparse attention)은 이러한 문제를 효과적으로 해결하는 접근법이며, DeepSeek Sparse Attention (DSA) 는 이를 실제 서비스 환경에 적용한 대표적인(production-grade) 솔루션이다. DSA에서는 경량(lightweight) Lightning Indexer가 각 쿼리에 대해 가장 관련성이 높은 상위$k$ 개의 토큰을 선택함으로써, 핵심 어텐션 연산의 복잡도를
에서
로 감소시킨다.
그러나 인덱서(indexer) 자체는 여전히
의 계산 복잡도를 가지며, 모든 계층(layer)에서 독립적으로 실행되어야 한다. 이는 연속된 계층들에서 선택되는 상위$k$ 개의 토큰이 매우 유사하다는 사실을 고려하면 비효율적이다.
본 논문에서는 이러한 계층 간 중복성(cross-layer redundancy)을 활용하는 IndexCache를 제안한다. IndexCache는 전체 계층을 두 종류로 구분한다.
이를 통해 대부분의 계층에서 인덱서 계산을 제거할 수 있다.
또한 이러한 계층 구성을 결정하고 최적화하기 위해 두 가지 상호보완적인 방법을 제안한다.
첫 번째는 Training-free IndexCache로, 가중치 업데이트 없이(calibration만 사용) 보정(calibration) 데이터셋에서 언어 모델링 손실(language modeling loss)을 직접 최소화하도록 탐욕적 탐색(greedy search) 알고리즘을 수행하여 어떤 계층에 인덱서를 유지할지를 결정한다.
두 번째는 Training-aware IndexCache이다. 이 방법은 다중 계층 증류 손실(multi-layer distillation loss)을 도입하여, 유지되는 각 인덱서가 자신이 담당하는 모든 계층의 평균 어텐션 분포를 학습하도록 한다. 이를 통해 단순한 교차(interleaved) 배치 방식만으로도 모든 계층에 인덱서를 유지하는 경우와 동등한 정확도를 달성할 수 있다.
30B 규모의 DSA 모델에서 수행한 실험 결과, IndexCache는 인덱서 계산의 75%를 제거하면서도 성능 저하를 거의 발생시키지 않았다. 그 결과 기존 DSA 대비 최대 1.82배의 prefill 속도 향상과 1.48배의 decode 속도 향상을 달성하였다. 또한 이러한 결과는 실제 서비스 규모의 GLM-5 모델을 대상으로 한 예비 실험(Figure 1)에서도 동일하게 확인되었다.
1. 서론
Self-Attention 메커니즘(Vaswani et al., 2017)은 현대 대규모 언어 모델(LLM)의 핵심 구성 요소이다. 그러나 시퀀스 길이에 대해 이차(quadratic) 복잡도를 갖기 때문에, 긴 문맥(long context)에서의 추론에서는 근본적인 병목이 된다. LLM이 장문의 Chain-of-Thought 추론, 다단계 에이전트(agentic) 워크플로우, 그리고 웹 규모의 검색 증강 생성(Retrieval-Augmented Generation, RAG)과 같이 긴 문맥을 요구하는 환경에 점점 더 많이 활용됨에 따라, 모델 성능을 유지하면서 어텐션 비용을 줄이는 것은 매우 중요한 연구 과제가 되었다.
희소 어텐션(sparse attention)은 이에 대한 원리적인 해결책을 제공한다. 모든 이전 토큰에 대해 어텐션을 수행하는 대신, 각 쿼리는 가장 관련성이 높은 일부 토큰만 선택하여 어텐션을 수행한다. 최근 제안된 여러 방법(Yuan et al., 2025; Lu et al., 2025; Zhao et al., 2025; Gao et al., 2026; Team et al., 2026) 가운데 DeepSeek Sparse Attention (DSA) (Liu et al., 2025)는 실제 서비스 환경에서 사용할 수 있는(production-grade) 학습 가능한(trainable) 희소 어텐션 기법으로 주목받고 있다.
희소 토큰 선택을 위해 DSA는 각 계층(layer)에 Lightning Indexer라는 추가 모듈을 도입한다. 이 모듈은 모든 이전 토큰에 대해 점수를 계산하고, 그중 상위$k$ 개의 토큰을 선택하여 이후의 핵심 어텐션(core attention)에 전달한다. 이를 통해 계층별 핵심 어텐션의 계산 복잡도는
에서
로 감소하며, 추가적인 지속 사전학습(continued pre-training)을 통해 모델 성능도 유지할 수 있다.
그러나 인덱서(indexer) 자체는 여전히 모든 계층에서
의 계산 복잡도를 가진다. 인덱서는 메인 어텐션보다 FLOP당 계산 비용은 저렴하지만, 총$N$ 개의 계층에 대해 수행되므로 전체 비용은
이 된다. 이 비용은 문맥 길이에 대해 여전히 이차적으로 증가하며, 전체 어텐션 계산 비용에서 무시하기 어려운 비중을 차지한다.
논문의 Figure 1에서 30B DSA 모델을 프로파일링한 결과를 보면, 전체 지연(latency)에서 인덱서가 차지하는 비중은 문맥 길이가 길어질수록 급격히 증가하며, 특히 Prefill 단계에서 이러한 현상이 두드러진다. 반면 나머지 계산들은 상대적으로 완만하게 증가한다. 이는 긴 문맥에서 DSA를 더욱 빠르게 만들기 위해서는 인덱서 계산을 줄이는 것이 핵심임을 보여준다.
본 연구의 핵심 아이디어는 인덱서가 선택하는 상위$k$ 개의 토큰이 인접한 계층들 사이에서 매우 높은 유사성을 가진다는 관찰에서 출발한다. 이러한 현상은 기존의 완전 어텐션(full attention) 모델에서도 보고된 계층 간 토큰 선택 안정성(cross-layer token selection stability) 과 일치한다(Deshmukh et al., 2025; Gao et al., 2026).
기존 연구들은 이러한 안정성을 활용하기 위해 일부 Full Attention Anchor Layer에서 계산한 인덱스를 다른 계층이 재사용하도록 하였다. 그러나 이러한 방법은 DSA에는 직접 적용될 수 없다. DSA에서는 완전 어텐션이 별도로 존재하지 않고, 경량 Lightning Indexer가 사실상 Full Attention의 역할을 대신하기 때문이다.
저자들은 DSA에서도 동일한 현상이 존재하는지 실험적으로 확인하였다. 모든 계층 쌍에 대해 Top-$k$ 인덱스의 중복률(overlap) 을 계산한 결과(Appendix A), 인접한 계층들은 선택된 토큰의 70~100%를 공유하는 것으로 나타났다. 또한 히트맵에서는 서로 높은 중복률을 보이는 계층 군집(layer clusters)이 명확하게 형성되었으며, 이는 대부분의 인덱서 계산이 중복되고 있음을 시사한다.
이러한 관찰은 다음과 같은 중요한 질문을 제기한다.
본 논문은 IndexCache를 제안하여 이 질문에 그렇다(Yes) 라고 답한다. IndexCache는 계층 간 인덱스 재사용(cross-layer index reuse) 을 통해 DSA에서 최대 75%의 인덱서 계산을 제거한다.
IndexCache는 계층을 두 종류로 나눈다.
추론 시에는 단 하나의 조건문(branch)만 추가하면 되며(Figure 2), 이를 통해 대부분의 인덱서 계산을 제거할 수 있다.
이를 위해 저자들은 두 가지 상호보완적인 방법을 제안한다.
Training-free IndexCache
Training-aware IndexCache
실험적으로, 30B DSA 모델을 9개의 장문 및 추론 벤치마크에서 평가한 결과,
모두 전체 인덱서의 25%만 유지하면서도 성능 저하는 거의 발생하지 않았다.
그 결과, 200K 문맥 길이에서 기존 DSA 대비
을 달성하였다.
또한 744B 규모의 GLM-5를 대상으로 한 예비 실험에서도 이러한 접근법의 확장성이 확인되었으며, 긴 문맥 성능 저하 없이 최소 1.3배 이상의 속도 향상을 달성하였다.
2. 예비 지식 (Preliminary)
2.1 DeepSeek Sparse Attention
DeepSeek Sparse Attention (DSA) (Liu et al., 2025)는 각 어텐션 계층을 선택(selection) 과 계산(computation) 의 두 단계로 분리한다.
먼저 경량 Lightning Indexer가 다중 헤드(Multi-head) ReLU-gated 내적(dot product)을 이용하여 현재 쿼리(query)와 모든 이전 토큰(preceding tokens) 사이의 점수를 계산한다. 이후 점수가 가장 높은 상위$k$ 개의 위치(top-$k$ positions) 를 선택한다.
그 다음 핵심 어텐션(core attention)은 선택된 희소 부분집합(sparse subset)에 대해서만 수행된다. 이를 통해 계층별 핵심 어텐션의 계산 복잡도는
에서
로 감소한다.
여기서
이며,$L$ 은 입력 시퀀스의 길이를 의미한다.
Lightning Indexer는 높은 효율성을 목표로 설계되었다. 이를 위해
을 사용한다. 그 결과 FLOP당 계산 비용은 메인 Multi-head Latent Attention (MLA) (Liu et al., 2024a)보다 약 한 자릿수(order of magnitude) 이상 저렴하다.
DSA는 MLA 위에 구현되며, 2단계의 지속 사전학습(two-stage continued pre-training) 을 통해 학습된다.
첫 번째 단계는 Dense Warm-up 단계이다.
이 단계에서는 모든 다른 모델 파라미터를 고정(freeze)한 상태에서 인덱서만 학습한다. 학습 목표는 각 계층의 전체 어텐션(full attention) 분포를 집계한 결과를 교사(teacher)로 사용하는 KL Divergence 기반 Distillation이다.
두 번째 단계는 Sparse Training 단계이다.
이 단계에서는 실제 Top-$k$ 선택을 활성화한 뒤 모델 전체를 공동으로 최적화한다. 이때 인덱서는 detached computational graph 상에서 Distillation gradient를 받아 학습된다.
이러한 계산량 절감에도 불구하고, 인덱서 자체는 여전히
의 계산 복잡도를 가진다.
모든 계층에서 인덱서는 자신의 Top-$k$ 집합을 결정하기 위해 모든 이전 토큰에 대해 독립적으로 점수를 계산해야 한다.
따라서 총$N$ 개의 계층을 갖는 모델에서는 전체 인덱서 계산 비용이
이 된다.
긴 문맥(long context)에서는 이 비용이 전체 어텐션 연산 예산(attention budget)에서 상당한 비중을 차지하게 된다.
이로부터 자연스럽게 다음과 같은 질문이 제기된다.
2.2 토큰 선택의 계층 간 안정성 (Cross-Layer Stability of Token Selection)
이 질문에 대한 실마리는 보다 일반적인 경험적 관찰에서 얻을 수 있다.
Transformer에서는 중요한 토큰들의 집합이 연속된 계층들 사이에서 놀라울 정도로 안정적(stable) 이라는 사실이 알려져 있다.
Deshmukh et al. (2025)와 Gao et al. (2026)는 모두 인접한 계층들이 대부분의 Top-$k$ 어텐션 질량(attention mass) 을 공유한다는 사실을 관찰하였다.
이를 활용하기 위해 이들 연구는 일부 Anchor Layer만 Full Attention을 수행하고, 그 사이의 중간 계층들은 Anchor Layer에서 계산된 Top-$k$ 인덱스를 그대로 재사용하도록 설계하였다.
그러나 이러한 접근법은 중요한 전제를 가진다.
즉, 중요한 토큰을 식별하기 위한 기준(oracle)으로 Full Attention이 존재해야 한다는 것이다.
반면 DSA에서는 Full Attention이 완전히 제거되었으며, 그 역할을 경량 Lightning Indexer가 대신하고 있다.
따라서 아직까지 충분히 연구되지 않은 새로운 질문이 등장한다.
만약 그렇다면,
더 나아가 다음과 같은 질문도 자연스럽게 제기된다.
다음 장에서는 이러한 문제를 해결하기 위한 두 가지 새로운 상호보완적 방법을 제안한다.
3. 방법 (Method)
표기법 (Notation)
Transformer의 전체 계층 수를$N$ , 입력 시퀀스 길이를 $L$ , 그리고 각 쿼리마다 선택되는 토큰의 개수를 $k$ 라고 하자.
계층$\ell$ 에서 Lightning Indexer는 쿼리 위치 $t$ 에 대해 점수 벡터
를 생성한다.
이 점수 벡터로부터 상위$k$ 개의 인덱스 집합을 다음과 같이 얻는다.
본 논문에서는 모든 실험에서
을 사용한다.
또한,
로 표기한다.
이는 여러 헤드의 softmax attention weight를 평균하여 얻는다.
반면 Lightning Indexer의 출력 분포는
로 정의한다.
개요 (Overview)
IndexCache는 DSA의 전체$N$ 개 계층을 두 가지 역할(role)로 나눈다.
이를 길이$N$ 의 이진(binary) 패턴 문자열
로 표현한다.
각 기호의 의미는 다음과 같다.
• F (Full)
이 계층은 기존과 동일하게 Lightning Indexer를 유지한다.
모든 이전 토큰에 대해 새로운 Top-$k$ 인덱스
를 계산한 뒤,
선택된 토큰들에 대해서만 Sparse Core Attention을 수행한다.
즉, 표준 DSA와 동일하게 동작한다.
• S (Shared)
이 계층은 Lightning Indexer를 제거한다.
대신 가장 가까운 이전 Full layer에서 계산한 인덱스를 그대로 재사용한다.
즉,
이며,
이다.
이렇게 전달받은 인덱스를 이용하여 바로 Sparse Core Attention을 수행한다.
첫 번째 계층은 항상 F로 지정한다.
이는 최초의 Top-$k$ 인덱스를 생성해야 하기 때문이다.
추론(inference) 시에는 S layer가 단순히
Figure 2에서 볼 수 있듯이, 기존 추론 루프에서 필요한 수정은 매우 단순하다.
이 방법의 핵심 설계 문제는 패턴$c$ 를 어떻게 결정하느냐이다.
만약 대부분의 계층이 안전하게 인덱스를 공유할 수 있다면,
전체 Indexer 계산 비용인
의 상당 부분을 제거할 수 있으며,
Sparse Core Attention의 계산량인
은 그대로 유지된다.
이를 위해 저자들은 두 가지 방법을 제안한다.
Training-free 방법
Training-aware 방법
3.1 Training-Free IndexCache
이미 학습된(pretrained) DSA 모델이 주어졌다고 하자.
목표는 가능한 많은 계층을 Shared layer(S) 로 바꾸면서도 모델 성능 저하를 최소화하는 패턴
를 찾는 것이다.
먼저 가장 직관적인 접근이 왜 실패하는지 설명한 뒤,
Greedy Search 알고리즘을 제안한다.
3.1.1 Uniform Interleaving이 최적이 아닌 이유
가장 단순한 방법은 Uniform Interleaving이다.
예를 들어
처럼
매$r$ 번째 계층만 Indexer를 유지하는 방식이다.
예를 들어
이면 4개 계층마다 하나만 F가 된다.
그러나 이 방식에는 중요한 문제가 있다.
각 계층의 Indexer 중요도는 동일하지 않다.
실험적으로,
에서는 Indexer를 제거했을 때 성능이 크게 감소하는 반면,
다른 계층에서는 거의 영향이 없다.
따라서 Uniform Interleaving은
비효율적인 선택을 할 수 있다.
그 결과 모델 성능이 눈에 띄게 감소한다
(정량적 비교는 Section 4.3).
이러한 관찰은
라는 아이디어로 이어지며,
이를 위해 Greedy Layer Selection 알고리즘을 제안한다.
3.1.2 Layer Selection Algorithm
저자들은 Greedy Search를 이용하여
F layer를 하나씩 S layer로 바꾸어 나간다.
각 선택의 기준은
작은 Calibration Set에서 측정한 Language Modeling Loss
이다.
LM Loss를 다운스트림 성능의 대리 지표(proxy)로 사용한다.
Calibration Set
학습 데이터에서
개의 Mini-batch를 미리 저장(cache)해 둔다.
모든 후보 패턴은 항상 동일한 Mini-batch에서 평가된다.
따라서 Loss 차이는
데이터 차이가 아니라
패턴 변경 자체의 영향만을 반영한다.
패턴
에 대한 손실은
모델
과 전체 Batch
에 대해 Forward Pass를 수행하여 얻는다.
논문에서는 이를
로 표기한다.
Greedy Search 과정
초기 상태는
즉,
모든 계층이 Full layer인 기존 DSA이다.
이후
단계 동안 반복한다.
여기서
는 제거하려는 Indexer 개수이다.
예를 들어 전체 Indexer의 1/4만 유지하려면
이다.
매 단계에서
Algorithm 1에서 전체 절차를 제시한다.
계산 복잡도
모든 계층이 F인 상태에서
모든 계층이 S가 될 때까지 탐색하면
Forward Pass는 총
번 수행된다.
이는 상당히 큰 비용이다.
논문에서는 이를 줄이기 위해
Pipeline Parallelism을 활용한다.
모델을
개의 Block으로 나눈 뒤,
각 Block의 첫 번째 계층은 항상 F로 고정한다.
그리고 매 단계마다
각 Block에서 가장 좋은 후보를 하나씩 선택하여 확정한다.
이렇게 하면
한 단계에서 최대
개의 계층을 동시에 S로 바꿀 수 있으며,
전체 Forward Pass 수는 약
배 감소한다.
Greedy Search의 특성
Greedy Search는 전역 최적해(global optimum)를 보장하지는 않는다.
그럼에도 불구하고 저자들은 세 가지 흥미로운 특성을 반복적으로 관찰하였다.
Uniform Interleaving보다 항상 성능이 우수하다.
동일한 Indexer 유지 비율에서도 Greedy Search가 더 좋은 성능을 보인다(Table 2).
계층 중요도가 자연스럽게 드러난다.
오른쪽 그림(유지 비율 1/2, 1/4, 1/8 표시)을 보면,
Greedy Search 과정에서의 LM Validation Loss는
으로 명확히 구분된다.
이는 Indexer의 중요도에 자연스러운 순서(ordering)가 존재함을 시사한다.
Calibration Set이 달라도 결과가 매우 안정적이다.
즉,
Greedy Search가 찾아내는 중요도 순위는 특정 데이터셋에 의한 우연한 결과가 아니라,
모델 자체가 가진 고유한 특성(intrinsic property) 임을 의미한다.
또한 LM Loss는 실제 다운스트림 태스크 성능과 양의 상관관계를 가지므로,
Layer Selection을 위한 유효한 Proxy Metric으로 사용할 수 있다.
3.2 Multi-Layer Distillation을 이용한 Training-Aware IndexCache
Training-free IndexCache는 추가적인 가중치 업데이트가 필요 없다는 장점이 있지만, 근본적인 한계를 가진다.
기존 DSA에서는 각 Lightning Indexer가 오직 자신의 계층만을 위해 학습되었다. 따라서 하나의 Indexer를 여러 계층이 공유하도록 만들면, 원래의 학습 목표와 맞지 않아 성능이 감소할 수 있다.
반면 DSA를 처음부터 학습하거나(Training from Scratch), 지속 사전학습(Continued Pre-training)을 수행하는 경우에는 더 나은 방법이 가능하다.
즉,
이다.
Single-layer Distillation에서 Multi-layer Distillation으로
기존 DSA 학습(Section 2.1)에서는
계층$\ell$ 의 Lightning Indexer가 같은 계층의 평균 어텐션 분포
를 모방하도록 KL Divergence Distillation을 수행한다.
손실 함수는
이다.
여기서
이다.
논문에서는 이를 여러 계층으로 일반화한다.
계층$\ell$ 이 유지되는 Full layer(F) 라고 하자.
그리고
이 모두 Shared layer라고 하자.
이들은 모두
를 그대로 재사용한다.
따라서 이 하나의 Indexer는
자신의 계층뿐 아니라
모든 Shared layer에서도 좋은 Top-$k$를 생성해야 한다.
이를 위해 다음과 같은 Multi-layer Distillation Loss를 정의한다.
직관적으로,
기존 방식은
학습하였다.
반면 Multi-layer Distillation은
학습한다.
즉,
특정 계층 하나에 과적합(overfitting)되는 대신,
여러 계층이 공통적으로 중요하게 생각하는 토큰을 선택하도록 만든다.
평균 분포에 대한 Distillation과 Gradient가 동일함
자연스러운 의문이 생긴다.
논문은 그렇지 않음을 증명한다.
먼저 여러 Teacher Distribution의 평균을 정의한다.
그리고 이 평균 분포 하나만 Teacher로 사용하는 손실을 정의한다.
Proposition 1
다음이 성립한다.
즉,
두 손실 함수는 Gradient가 완전히 동일하다.
증명
KL Divergence는
로 정의된다.
여기서
Teacher 분포
는 고정되어 있으므로,
미분 시 Entropy 항은 사라진다.
즉,
이다.
이를 Multi-layer Loss에 적용하면
이 된다.
평균을 앞으로 모으면
괄호 안은 바로 평균 분포
이다.
따라서
가 된다.
증명 완료.
□
해석 (Interpretation)
이 결과는 매우 중요한 의미를 가진다.
Multi-layer Distillation은
단순히
이 아니다.
오히려
즉,
Lightning Indexer는
각 계층을 따로따로 맞추는 것이 아니라,
여러 계층이 공통적으로 중요하다고 생각하는 토큰들을 선택하도록 학습된다.
다시 말해,
Shared layer 전체에서 사용할 수 있는 **Consensus Top-$k$**를 예측하도록 학습되는 것이다.
왜 실제 구현에서는 (L_I^{\mathrm{multi}})를 사용하는가?
두 손실 함수는 Gradient가 동일하므로
이론적으로는 어느 쪽을 사용해도 된다.
하지만 실제 구현에서는
를 사용한다.
이유는 메모리 효율 때문이다.
Shared layer에서는
현재 Full layer의 예측값
만 전달하면 된다.
반면
를 사용하려면
평균 Distribution을 계산하기 위해
를 모두 전달해야 한다.
이는 불필요한 메모리 사용량과 추가적인 계산 비용을 발생시킨다.
따라서 실제 구현에서는
Gradient는 동일하면서 구현이 더 단순한
를 채택하였다.
학습 과정 (Training)
전체 학습 절차는 기존 DSA와 동일한 2단계 학습을 따른다.
1. Warm-up 단계
를 사용한다.
2. Sparse Training 단계
실제 Top-$k$ 선택을 활성화한 뒤,
Indexer는 계속해서
로 학습한다.
이때 KL Divergence는 선택된 Top-$k$ 토큰에 대해서만 계산한다.
동시에,
나머지 모델 파라미터들은 일반적인 Language Modeling Loss(LM Loss) 를 이용하여 함께 최적화한다.
4. 실험 (Experiments)
4.1 실험 설정 (Setup)
모델 (Model)
실험에 사용한 DSA 모델은 GLM-4.7-Flash*의 기본(base) 모델로부터 시작하여 2단계 학습(two-stage training) 을 거쳐 구축되었다.
GLM-4.7-Flash는 다음과 같은 특성을 가진다.
평가 성능은 원래의 GLM-4.7-Flash와 거의 동일한 수준이다(Zeng et al., 2026의 Table 6 참고).
Training-free IndexCache
Greedy Pattern Search는
에서 계산한 토큰당 Validation Loss를 기준으로 수행하였다.
Training-aware IndexCache
처음부터 DSA 전체 학습 파이프라인을 수행하는 것은 매우 많은 계산 자원이 필요하다.
따라서 저자들은
GLM-4.7-Flash 기본 모델에서 직접 시작하여
SFT 데이터(200K Context Length)를 이용해 DSA 모델로 학습하였다.
학습은 다음 두 단계로 구성된다.
Dense Warm-up
Sparse Training
이 단축된 학습 파이프라인(shortened pipeline)은 전체 DSA 학습과 거의 동일한 성능을 달성하면서도, Training-aware IndexCache를 평가하기에는 충분하였다.
평가 벤치마크 (Evaluation)
Long-context 벤치마크 (5개)
일반 및 추론 벤치마크 (4개)
평가 세부 설정은 Appendix D에 제시되어 있다.
4.2 End-to-End 추론 속도 향상
추론 성능은
에서 측정하였다.
비교 대상은
이다.
다음 네 가지 Context Length에서 측정하였다.
평가 지표는 세 가지이다.
Prefill Latency
Single-request Decode Throughput
Total Decode Throughput
결과는 Table 1과 Figure 3에 제시되어 있다.
Prefill 성능
IndexCache는 문맥 길이가 길어질수록 더욱 큰 Prefill 가속 효과를 보였다.
200K Context에서는
Prefill Latency가
로 감소하였다.
즉,
1.82배의 속도 향상을 달성하였다.
이는 Prefill 단계의 병목이던 Indexer 계산의 75%를 제거했기 때문이다.
심지어 Context Length가 10K처럼 짧은 경우에도
1.27배의 속도 향상을 얻었다.
저자들은 200K보다 더 긴 문맥에서는 속도 향상이 더욱 커질 것으로 예상한다.
Decode 성능
긴 문맥에서 Decode 속도 향상도 매우 크다.
200K Context에서
로
1.48배의 속도 향상을 달성하였다.
그 이유는 DSA의 Decode 과정에서도
매 토큰마다 전체 문맥에 대해 Indexer를 실행해야 하기 때문이다.
문맥이 길어질수록 이 Indexer가 병목이 되며,
IndexCache는 바로 이 병목을 제거한다.
KV Cache가 가득 찬 상태에서는
IndexCache (1/4)가
모든 Context Length에서
22~51%의 Decode Throughput 증가를 보였다.
특히 200K에서는
으로
1.51배 증가하였다.
744B 규모의 GLM-5에서도 비슷한 결과가 나타났다.
100K 이상의 Context에서
IndexCache (1/4)는
모두 최소 1.3배 이상의 속도 향상을 달성하였다.
전체적으로,
IndexCache는 긴 문맥 추론(long-context inference) 에서 특히 큰 효과를 보였다.
4.3 Training-Free IndexCache 결과
Table 2에서는 Training-Free IndexCache를 평가하였다.
비교한 Indexer 유지율은
이며,
각각에 대해
을 비교하였다.
(Appendix B)
Greedy Search는 Long-context 성능을 거의 완전히 회복한다
Uniform Interleave는 Indexer를 많이 제거할수록 Long-context 성능이 크게 감소한다.
예를 들어
Long Avg는
으로 감소하였다.
반면 Greedy Search Pattern은 이러한 성능 저하를 거의 제거하였다.
Long Avg는
으로,
원래 DSA와 거의 동일하였다.
즉,
그러나 Indexer를 1/8만 유지하면
성능 저하가 상당히 커졌다.
Long Avg는
이었다.
Greedy Search가 손실을 크게 줄여주기는 하지만,
이 정도의 극단적인 Sparsity에서는 성능 저하를 완전히 막을 수는 없었다.
긴 Chain-of-Thought 추론 능력은 유지된다
Uniform 1/8을 제외하면
모든 설정에서
G&R Avg는
기존 DSA(74.6)와 1점 이내의 차이만 보였다.
특히
Greedy 1/4는 오히려
AIME 2025
GPQA-Diamond
으로 DSA보다 높은 점수를 기록하였다.
이는 불필요한 Indexer 계산을 제거하는 것이
추론 과정에서 약한 형태의 Regularization 효과를 낼 수도 있음을 시사한다.
즉,
IndexCache는
Long-context 효율성을 얻기 위해
일반적인 추론 능력을 희생하지 않는다.
4.4 Training-Aware IndexCache 결과
Training-Aware IndexCache는
Section 3.2의 Multi-layer Distillation을 이용하여 학습하였다.
평가한 유지율은
이며,
모두 Uniform Pattern을 사용하였다.
또한
1/2 유지율에서 두 가지 Ablation을 수행하였다.
Uniform 대신 Greedy Pattern 사용
Cross-layer Distillation 제거
(각 Indexer를 자기 Layer만 Distillation)
결과는 Table 3에 제시되어 있다.
주의할 점은
여기서의 DSA Baseline 역시
단축 학습 파이프라인을 사용했기 때문에,
Table 2의 DSA와는 약간의 성능 차이가 있다.
Training-aware IndexCache는 DSA와 동일한 성능을 달성한다
Uniform Pattern만 사용해도
1/2 유지율에서
Long Avg는
으로 오히려 더 높았다.
G&R Avg도
로 거의 동일하였다.
1/4 유지율에서도
두 평균 성능 모두
Baseline과 0.4% 이내의 차이만 보였다.
즉,
모델은
Index Sharing 구조에 적응하도록 학습될 수 있다.
재학습하면 Layer Pattern의 중요성이 사라진다
Training-Free에서는
Greedy Search가 반드시 필요하였다.
그러나 Training-aware에서는
Uniform Pattern만으로도
Greedy보다 같거나 오히려 더 좋은 성능을 얻었다.
예를 들어
1/2 유지율에서
Long Avg는
이었다.
그 이유는 다음과 같다.
Training-Free에서는
각 Layer가 자신의 Indexer 출력에 강하게 의존한다.
따라서 이전 Layer의 Index를 사용하면
분포 변화(distribution shift)가 발생하여 성능이 감소한다.
Greedy Search는
이러한 민감한 Layer를 피하는 방식으로 동작한다.
반면 Training-aware에서는
이러한 공동 적응(joint adaptation)을 통해
Layer마다 존재하던 민감성이 사라지며,
단순한 Uniform Pattern만으로도
기존 DSA와 동일한 성능을 얻을 수 있게 된다.
Cross-layer Distillation은 실제로 효과가 있다
Cross-layer Distillation을 제거하면
Long Avg는
로 감소하였다.
특히
AA-LCR은
으로 크게 감소하였다.
이는 Multi-layer Distillation이 실제로 매우 중요한 역할을 함을 보여준다.
각 Indexer가
여러 Layer의 평균 Attention Distribution(centroid)을 학습함으로써,
특정 Layer에만 특화되지 않고
여러 Layer에서 공통적으로 사용할 수 있는 **Consensus Top-$k$**를 생성하게 된다.
4.5 확장성 실험 (Scaling Experiment)
Training-Free IndexCache를
744B(Active 40B) 규모의 GLM-5에도 적용하였다.
GLM-5는 기본적으로 DSA를 사용한다.
Table 4는
5개의 Long-context 벤치마크 결과를 보여준다.
전체적인 경향은 30B 모델과 동일하였다.
흥미롭게도
1/2 유지율에서는
Uniform Pattern도
Long Avg를 거의 유지하였다.
이는 단순히
균등 패턴이 우연히 중요한 Indexer Layer를 건너뛰지 않았기 때문이라고 저자들은 설명한다.
반면 Greedy Pattern은
더 안정적인 결과를 보였다.
로 약간 더 좋았으며,
1/4 유지에서도
로 거의 동일하였다.
또한 Artificial Analysis Index의 모든 평가에서
Indexer를 절반만 유지한 IndexCache를 평가한 결과,
원래 GLM-5와 거의 동일한 성능을 보였다(Figure 1).
저자들은 가까운 미래에
이 Production-scale 모델에도
Training-Aware IndexCache를 적용할 계획이다.
이미 Training-Free 버전만으로도 Baseline과 동일한 품질을 달성했기 때문에,
Training-aware 학습을 적용하면
이러한 성능을 더욱 안정적으로 유지하면서,
실제 서비스 환경에서의 추론 효율성도 더욱 향상될 것으로 기대하고 있다.
5. 관련 연구 (Related Work)
5.1 효율적인 어텐션 (Efficient Attention)
Self-Attention의 이차(quadratic) 계산 비용을 줄이는 것은, 특히 긴 시간 범위(long-horizon)의 에이전트(agent) 활용이 증가하는 현재 LLM 연구에서 핵심적인 주제이다.
Training-free Sparse Attention
Training-free 희소 어텐션 방법들은 추론 단계에서만 희소성을 도입한다.
이를 위해 다음과 같은 기법들이 사용된다.
(Zhang et al., 2023; Xiao et al., 2024; Ribar et al., 2024; Tang et al., 2024; Jiang et al., 2024; Xiao et al., 2025; Xu et al., 2025; Lai et al., 2025; Fu et al., 2025; Zhu et al., 2025; Yang et al., 2025c; Zhang et al., 2025)
그러나 이러한 방법들은 학습(training)과 추론(inference)의 불일치(training-inference mismatch) 를 초래하며, 긴 문맥에서는 오류가 누적(error accumulation)될 수 있다(Hu et al., 2026).
Trainable Sparse Attention
반면 Trainable Sparse Attention은 학습 단계에서부터 희소성을 모델에 내재화한다.
대표적인 방법으로는
등이 있다.
본 논문의 기반이 되는 DeepSeek Sparse Attention (DSA) (Liu et al., 2025)는
Full Attention으로부터 경량 Lightning Indexer를 Distillation하여,
각 Query에 대해 Top-$k$ 토큰을 선택하도록 한다.
이를 통해 핵심 어텐션(Core Attention)의 계산 복잡도를
로 줄인다.
희소화(sparsity) 이외에도,
Hybrid Architecture는 계산량이 큰 Quadratic Attention Layer의 개수를 줄이는 접근을 사용한다.
예를 들어 Quadratic Layer와 다음과 같은 구조를 번갈아(interleave) 배치한다.
Sliding Window Attention (sink 포함 또는 미포함)
Linear Attention
State Space Layer
5.2 계층 간 공유 (Cross-Layer Sharing)
최근 연구들은 인접한 Transformer Layer의 표현(representation)이 매우 높은 일관성(consistency)을 가진다는 사실을 보여주었다.
이러한 구조적 특성은 계산 중복을 줄이고 추론을 가속하기 위해 활용되고 있다.
대표적인 연구로는
등이 있다.
이들은 주기적으로 배치된 Anchor Layer에서 계산한 Top-$k$ 인덱스를 이후 Sparse Decoding Layer들이 재사용하도록 한다.
Kascade (Deshmukh et al., 2025)는
계층 간 유사도(cross-layer similarity) 행렬에 대해 동적 계획법(dynamic programming) 을 적용하여,
최적의 Anchor Layer를 선택하는 방법을 제안하였다.
또한 정확도를 유지하기 위해서는 Head-aware Remapping이 매우 중요함을 보였다.
그러나 이러한 기존 방법들은 모두
는 전제를 갖는다.
한편, 또 다른 연구 흐름은
Cross-layer KV Cache Sharing이다.
이 방법들은 여러 Layer가 동일한 Key-Value Cache를 공유하도록 하여 메모리 사용량을 줄인다.
대표적인 연구는
등이다.
HySparse (Gao et al., 2026)는
위 두 가지 접근을 통합하였다.
즉,
를 번갈아 배치하고,
Sparse Layer는
를 모두 상속받는다.
그러나 지금까지의 모든 방법은
Full Attention Layer를 Oracle로 사용해야 한다.
반면 DSA에서는 Full Attention 자체가 완전히 제거되어 있다.
IndexCache와 기존 연구의 차이점
IndexCache는 기존 연구와 두 가지 측면에서 다르다.
첫째,
Oracle 자체가 훨씬 저렴하다.
기존 연구는
의 Full Attention Score를 공유하지만,
IndexCache는
DSA의 경량 Lightning Indexer 출력만 공유한다.
따라서 계산 비용이 훨씬 낮다.
둘째,
Index Sharing 구성을 체계적으로 최적화하는 방법을 함께 제안한다.
구체적으로
Training-free Greedy Search
Training-aware Multi-layer Distillation
비록 본 논문은 DSA를 대상으로 구현하였지만,
IndexCache의 핵심 아이디어는
고정된 Sparse Pattern이 아니라 동적인 Token Selection을 사용하는 모든 Sparse Attention으로 확장될 수 있다.
예를 들어
에서도 동일한 방식으로 계층 간 Index 재사용(Cross-layer Reuse) 을 적용할 수 있을 것으로 기대된다.
6. 결론 (Conclusion)
본 논문에서는 IndexCache를 제안하였다.
IndexCache는 토큰 선택을 담당하는 Lightning Indexer의 계층 간 중복성(cross-layer redundancy) 을 활용하여 Sparse Attention을 가속하는 방법이다.
IndexCache는 전체 Layer를
소수의 F(Full) Layer
다수의 S(Shared) Layer
로 나눈다.
이를 통해
전체 Indexer 계산 비용
의 최대 75%를 제거할 수 있으며,
추론 과정에서는 단 하나의 조건문(branch) 만 추가하면 된다.
또한 이러한 계산량 감소에도 불구하고 모델 성능 저하는 발생하지 않았다.
보다 넓은 관점에서 보면,
본 연구는 지금까지 Full Attention을 Oracle로 사용하는 경우에만 적용되던 계층 간 공유(cross-layer sharing) 의 원리가,
Sparse Attention에도 자연스럽게 확장될 수 있음을 보여준다.
향후 DeepSeek-V3.2, GLM-5와 같이 Sparse Attention이 차세대(frontier) LLM의 기본 구조(default) 가 되어감에 따라,
저자들은 Cross-layer Index Reuse가 효율적인 추론 파이프라인의 표준 구성 요소로 자리잡을 것으로 기대하고 있다.
A. 계층 간 Top-(k) 인덱스 중복도 (Cross-Layer Top-(k) Index Overlap)
DSA에서 Top-(k) 인덱스 선택이 계층 간 얼마나 중복되는지를 실험적으로 검증하기 위해, 각 계층의 Lightning Indexer가 선택한 Top-(k) 인덱스 사이의 쌍별(pairwise) 중복률(overlap ratio) 을 계산하였다.
구체적으로, 모든 계층 쌍 ((i, j))에 대해 다음 값을 측정하였다.
여기서
이다.
이 값은 길이 200K인 입력 768개로 구성된 Calibration Set에 대해 평균을 계산하였다.
Figure 4. 30B DSA 모델의 모든 계층 쌍에 대한 Top-(k) 인덱스 중복률. Greedy Search로 찾은 1/4 유지율(IndexCache) 의 공유 블록(shared blocks)은 빨간색 박스로 표시되어 있다.
Figure 4는 47개 계층을 갖는 30B DSA 모델의 중복률을 히트맵(heatmap)으로 시각화한 것이다. 또한 Greedy Search로 찾은 1/4 IndexCache 패턴에서의 공유 블록을 빨간색 박스로 표시하였다.
이 그림에서는 다음과 같은 몇 가지 중요한 특징을 확인할 수 있다.
• 대각선(diagonal) 근처에서 높은 중복률
인접한 계층들은 0.7~1.0의 매우 높은 중복률을 보인다.
이는 연속된 계층들이 거의 동일한 토큰 집합을 선택한다는 것을 의미하며, 계층 간 Top-(k) 선택의 높은 안정성을 확인해 준다.
• 블록 구조(Block Structure)
히트맵에서는 서로 높은 중복률을 보이는 계층 군집(layer clusters) 이 뚜렷하게 나타난다.
대표적인 예는 다음과 같다.
등이다.
이는 모델이 여러 기능적 블록(functional blocks) 으로 구성되어 있으며, 각 블록 내부에서는 토큰 선택이 매우 일관적이라는 것을 시사한다.
• 불균일한 감소(Uneven Decay)
중복률은 블록 내부에서는 천천히 감소하지만,
블록 경계(block boundary) 를 넘으면 훨씬 빠르게 감소한다.
이는 일부 전이(transition) 계층이 어텐션의 초점을 크게 바꾸는 역할을 수행한다는 것을 의미한다.
• 초기 계층과 후반 계층의 차이(Early–Late Distinction)
히트맵의 좌하단(bottom-left) 과 우상단(top-right) 영역은 매우 어둡게 나타난다.
즉,
이다.
이는 초기 계층과 후반 계층이 근본적으로 서로 다른 토큰 집합에 주목(attend) 한다는 것을 보여준다.
다시 말해,
모델의 앞부분과 뒷부분은 서로 다른 정보를 처리하며, 동일한 Top-(k) 인덱스를 공유하기 어렵다는 점을 시사한다.