IndexCache: 계층 간 인덱스 재사용을 통한 희소 어텐션 가속화

긴 문맥을 처리할 때 병목이 되는 희소 어텐션의 인덱서 연산을 획기적으로 줄이는 방법을 제시한다. 계층 간의 유사성을 활용해 성능 저하 없이 추론 속도를 높이고 비용을 절감할 수 있어 실무적 가치가 크다.

Transformer의 Self-Attention은 모든 토큰 간의 관계를 계산하므로 시퀀스 길이의 제곱에 비례하는 연산량이 필요하다. 이를 해결하기 위해 DeepSeek Sparse Attention(DSA) 같은 기법은 Indexer라는 가벼운 모듈을 사용해 중요한 상위 k개 토큰만 골라내어 연산량을 줄이지만, 이 Indexer 자체도 모든 레이어에서 독립적으로 실행되며 제곱 단위의 연산량을 소모하는 한계가 있다.

연구진은 인접한 레이어들이 선택하는 중요한 토큰의 집합이 70~100% 일치한다는 계층 간 중복성에 주목했다. IndexCache는 모든 레이어에서 인덱서를 돌리는 대신, 일부 레이어(Full)에서만 인덱스를 계산하고 나머지 레이어(Shared)는 이를 그대로 재사용함으로써 중복된 연산을 제거한다.

단순히 일정 간격으로 공유하는 것이 아니라, 언어 모델의 손실값(LM Loss)을 최소화하는 최적의 레이어 조합을 찾아내거나, 인덱서가 여러 레이어의 요구사항을 동시에 만족하도록 학습시키는 방식을 통해 정확도 손실을 거의 없앴다. 결과적으로 긴 문맥 처리 시 가장 큰 병목이었던 인덱서 연산을 75%까지 줄여 전체 추론 속도를 대폭 개선했다.

IndexCache는 전체 N개의 레이어를 Full(F)과 Shared(S) 역할로 구분한다. F 레이어는 자체 인덱서를 실행해 새로운 top-k 인덱스를 생성하고 이를 캐시에 저장하며, S 레이어는 가장 가까운 이전 F 레이어의 캐시된 인덱스를 그대로 가져와 어텐션 연산을 수행한다.

최적의 레이어 패턴을 찾기 위해 Greedy Search 알고리즘을 사용한다. 모든 레이어를 F로 시작하여, 하나씩 S로 바꿨을 때 캘리브레이션 데이터셋에서의 LM Loss 증가가 가장 적은 레이어를 선택해 나가는 방식이다. [패턴 후보 입력 → 순방향 전파 연산 → LM Loss 출력 → 최소 손실 패턴 확정] 과정을 통해 가중치 업데이트 없이도 효율적인 구조를 결정한다.

학습 기반 방식에서는 Multi-layer Distillation Loss를 도입한다. 특정 F 레이어의 인덱서가 자신이 서빙하는 모든 레이어의 평균적인 어텐션 분포를 예측하도록 학습시킨다. [각 레이어의 어텐션 분포 입력 → 평균값 계산 → KL Divergence 연산 → 인덱서 가중치 갱신] 과정을 통해 인덱서가 여러 레이어에 범용적인 토큰을 선택할 수 있게 한다.

30B DSA 모델 실험 결과, 인덱서의 75%(1/4 유지)를 제거했을 때 200K 문맥 길이에서 프리필(Prefill) 속도는 1.82배, 디코딩(Decode) 속도는 1.48배 향상되었다. 이는 인덱서 연산이 긴 문맥에서 전체 지연 시간의 상당 부분을 차지하기 때문이다.

성능 유지 측면에서, Greedy Search로 찾은 패턴은 단순 균등 간격(Uniform) 방식보다 월등히 우수했다. 1/4 유지 패턴에서도 원래 DSA 모델과 거의 동일한 벤치마크 점수를 기록했으며, 특히 AIME 2025와 GPQA-Diamond 같은 추론 태스크에서는 오히려 성능이 소폭 향상되는 정규화 효과도 관찰되었다.

744B 규모의 GLM-5 모델에 적용했을 때도 유사한 경향이 확인되었다. 1/2 인덱서 유지 시 성능 저하 없이 전체 추론 속도가 약 1.2배 향상되었으며, 이는 대규모 프로덕션 모델에서도 IndexCache의 확장성과 실용성이 높음을 입증한다.

IndexCache는 DSA(DeepSeek Sparse Attention)의 핵심 구성 요소인 Lightning Indexer의 중복성을 해결한다. DSA는 O(L²)의 인덱서와 O(Lk)의 코어 어텐션으로 구성되는데, 문맥이 길어질수록 인덱서의 비중이 커지는 문제를 계층 간 인덱스 공유로 해결한다.

수학적으로 Multi-layer Distillation Loss는 각 레이어의 KL Divergence 합으로 정의되지만, 이는 각 레이어 어텐션 분포의 평균(Centroid)을 타겟으로 하는 단일 증류 손실과 그래디언트가 동일함을 증명했다. 이를 통해 구현 효율성을 높이면서도 여러 레이어를 아우르는 합의된(consensus) top-k를 학습할 수 있다.

레이어 선택의 민감도는 네트워크의 초기 및 전이 영역에서 높게 나타났다. Greedy Search는 이러한 임계 레이어(critical layers)를 자동으로 식별하여 인덱서를 유지함으로써, 단순한 균등 간격 공유가 유발하는 누적 오차와 성능 저하를 방지한다.

구현 측면에서 Shared 레이어는 추가적인 GPU 메모리를 거의 사용하지 않는다. F 레이어에서 계산된 인덱스 텐서를 임시 버퍼에 저장하고 S 레이어에서 이를 참조하기만 하면 되므로, 기존 DSA 할당량 외의 오버헤드가 거의 없다.

1/8 이하의 극단적인 인덱서 유지 비율에서는 Greedy Search를 사용하더라도 긴 문맥 성능의 유의미한 하락이 관찰되었다. 또한, 본 연구는 DSA 아키텍처를 중심으로 검증되었으므로 다른 형태의 희소 어텐션이나 동적 토큰 선택 기법에 대한 일반화 가능성은 추가 연구가 필요하다.