Kimi Linear An Expressive, Efficient Attention Architecture
Kimi Team et al., 2025 (arXiv:2510.26692)
트랜스포머와 어텐션 메커니즘은 뛰어난 성능으로 레전드가 되었습니다. 그런 어텐션 메커니즘에도 약점이 많이 있습니다. 특히 길이가 긴 문맥을 다룰 때 시퀀스 길이 \(n\)에 대한 \(O(n^2)\) 시간복잡도는 치명적인 약점입니다. 시간이 흐르고 LLM이 백만 토큰 규모의 입력을 처리하게 되면서 이 단점이 서서히 드러났습니다. 강화학습(RL) 기반의 추론 스케일링이 중요해지면서 이 문제는 더욱 심각해졌죠.
선형 어텐션은 우아한 해결책으로 보였지만 표현력이 부족합니다. 일반 어텐션을 따라잡기 어렵습니다. 이 논문은 '과연 선형 어텐션으로 충분한가?'라는 오래된 질문에 새로운 답을 제시합니다. Kimi Linear는 세밀한 게이팅 메커니즘과 하이브리드 아키텍처를 통해, 처음으로 일반 어텐션을 모든 시나리오에서 능가하는 성과를 달성했습니다.
요약
핵심 아이디어: Kimi Delta Attention(KDA)이라는 향상된 선형 어텐션 모듈을 제안하고, 이를 기존 MLA(Multi-Head Latent Attention)와 3:1 비율로 하이브리드한 구조입니다.
기술 스펙:
- 모델 규모: 3B 활성화 파라미터 / 48B 총 파라미터 (MoE 구조)
- 훈련 데이터: 1.4T 토큰 (본 평가용), 5.7T 토큰 (최종 체크포인트)
- 훈련 배치: 3,200만 토큰, 4,096 토큰 컨텍스트 윈도우
- 평가 매트릭: MMLU, RULER, GSM8K, MATH500, AIME 2025 등 다양한 벤치마크
- 핵심 개선: KV 캐시 75% 감소, 1M 컨텍스트에서 6배 디코딩 가속
논문 상세
1. 문제의식: 선형 어텐션의 딜레마
표준 어텐션은 쿼리와 모든 키-값 쌍을 비교하므로 O(T²) 복잡도를 가집니다. 일반 어텐션으로 100만 토큰을 처리하려면 엄청난 연산 비용이 필요하죠. 반면 선형 어텐션은 O(T) 복잡도를 제공하지만, 표현력이 부족해서 기존 모델들을 따라잡지 못했습니다.
하이브리드 접근법(일부 레이어는 일반 어텐션, 일부는 선형)이 제안되었지만, 대부분 제한된 규모에서만 평가되었거나 공정한 비교가 부족했습니다. 저자들은 이 문제를 체계적으로 해결하기로 결심했습니다.
2. Kimi Delta Attention (KDA): 세밀한 메모리 제어
2.1 델타 규칙의 진화
선형 어텐션의 기본은 행렬 상태 \(S_t \in \mathbb{R}^{d_k \times d_v}\)를 누적하는 것입니다:
\[S_t = S_{t-1} + k_t v_t^\top\]
이는 온라인 학습으로 볼 수 있습니다. 하지만 이 방식은 메모리가 무한정 커져서 오래된 정보가 간섭을 일으킵니다.
DeltaNet은 이를 개선했습니다. 재구성 손실에 대해 경사도 하강을 수행함으로써:
\[S_t = (I - \beta_t k_t k_t^\top) S_{t-1} + \beta_t k_t v_t^\top\]
이것은 고전적인 델타 규칙이며, 메모리를 선택적으로 수정합니다.
**Gated DeltaNet (GDN)**은 스칼라 망각 게이트 \(\alpha_t \in [0, 1]\)을 추가했습니다:
\[S_t = \alpha_t (I - \beta_t k_t k_t^\top) S_{t-1} + \beta_t k_t v_t^\top\]
2.2 KDA의 혁신: 채널별 세밀한 게이팅
KDA는 스칼라 게이트를 채널별 벡터 게이트로 확장합니다:
\[S_t = \left(I - \beta_t k_t k_t^\top\right) \text{Diag}(\alpha_t) S_{t-1} + \beta_t k_t v_t^\top\]
여기서 \(\text{Diag}(\alpha_t) \in \mathbb{R}^{d_k \times d_k}\)는 각 특성 차원이 독립적인 망각률을 가집니다. 이는 RoPE의 차원별 다른 회전 주파수처럼, 각 차원에 다른 위치 인코딩을 효과적으로 부여합니다.
실제 계산에서 이는 다음을 의미합니다:
\[q_t^\top \left(\prod_{j=i+1}^{t} (I - \beta_j k_j k_j^\top) \text{Diag}(\alpha_j)\right) k_i\]
각 차원 \(d\)에서 누적 감쇠:
\[\gamma_i^\text{(d)} = \prod_{j=1}^{i} \alpha_j^{(d)}\]
이렇게 세밀한 제어를 통해, KDA는 중요한 정보는 오래 보존하고 불필요한 정보는 빨리 잊을 수 있습니다.
2.3 하드웨어 효율성: DPLR의 최적화
선형 어텐션의 일반화는 Diagonal-Plus-Low-Rank(DPLR) 구조입니다:
\[S_t = (D - a_t b_t^\top) S_{t-1} + k_t v_t^\top\]
하지만 일반 DPLR은 계산 비용이 높고 수치적으로 불안정합니다 (역수 계산 때문).
KDA의 핵심 통찰: \(a_t = \beta_t k_t\), \(b_t = k_t \odot \alpha_t\)로 제한하면, 이들을 인수분해할 수 있습니다:
\[S_t = \left(\text{Diag}(\alpha_t) - \beta_t k_t k_t^\top\right) S_{t-1} + k_t v_t^\top\]
이 제약으로 인해 이차 청킹이 4개에서 2개로 줄고, 3개의 행렬 곱셈을 제거할 수 있어, DPLR 대비 약 2배의 속도 향상을 달성합니다.
3. 하이브리드 아키텍처 설계
KDA 혼자서도 선형 어텐션의 근본적 한계가 있습니다: 긴 문맥에서 정확한 정보 검색이 어렵습니다. 따라서 저자들은 KDA와 기존 MLA(Multi-Head Latent Attention)를 3:1 비율로 교대로 배치했습니다.
[KDA] → [KDA] → [KDA] → [MLA] → [KDA] → ...
왜 3:1인가? 논문의 절제 연구(Ablation Study)에 따르면:
- 0:1 (순수 MLA): 기준선
- 1:1 (동일 비율): 검증 손실 증가
- 3:1 (최적): 최저 손실, 최고 효율
- 7:1 (더 많은 선형): 훈련 손실은 비슷하지만 검증 손실 악화
MLA 레이어의 No Position Encoding (NoPE): 흥미로운 설계 선택은 MLA에 위치 인코딩을 적용하지 않는 것입니다. 모든 위치 정보 부호화를 KDA에 위임함으로써:
- 장문맥에서 RoPE 기저 주파수 튜닝 불필요
- 컨텍스트 윈도우 확장 시 간편
- 더 안정적인 위치 편향 분포
4. 실험 결과: 모든 척도에서의 우수성
4.1 합성 작업: 기본 능력 검증
복잡한 벤치마크 전에, 세 가지 합성 작업으로 기초를 확인했습니다:
회문(Palindrome): 토큰 수열을 역순으로 재현. 선형 어텐션의 약점인 정확한 복사 능력을 테스트합니다.
다중 쿼리 연관 검색(MQAR): 여러 쿼리에 대해 문맥 내 다양한 위치에서 관련 값 검색. 언어 모델 성능과 상관관계가 높습니다.
스택 상태 추적: 64개의 독립 LIFO 스택을 관리하며 PUSH/POP 연산 추적.
결과: KDA는 모든 작업에서 Gated DeltaNet(GDN)을 상회했고, 수열 길이 증가(256→2,048)에 따른 성능 저하가 가장 완만했습니다.
4.2 사전훈련 성능: 단문맥과 다양성
1.4T 토큰으로 훈련한 결과:
벤치마크 |
MLA |
GDN-H |
Kimi Linear |
|---|---|---|---|
MMLU |
71.6 |
72.2 |
73.8 |
MMLU-Pro |
47.2 |
47.9 |
51.0 |
BBH |
71.6 |
70.6 |
72.9 |
GSM8K |
83.7 |
81.7 |
83.9 |
CEval (중국어) |
79.3 |
79.1 |
79.5 |
4.3 장문맥 성능: 결정적 우위
이것이 바로 하이브리드 구조의 가치를 보여주는 순간입니다. 128k 토큰 컨텍스트:
벤치마크 |
MLA |
GDN-H |
Kimi Linear (RoPE) |
Kimi Linear |
|---|---|---|---|---|
RULER |
81.3 |
80.5 |
78.8 |
84.3 |
MRCR |
22.6 |
23.9 |
22.0 |
29.6 |
RepoQA |
63.0 |
63.0 |
66.5 |
68.5 |
평균 |
52.2 |
51.2 |
51.8 |
54.5 |
NoPE의 효과: Kimi Linear (RoPE)는 Kimi Linear보다 장문맥에서 성능이 떨어집니다. 이는 위치 편향이 KDA를 통해 분산되면 더 유연하고 확장성 있다는 것을 시사합니다.
4.4 강화학습: 추론 확장성
흥미로운 발견은 RL 트레이닝 중입니다. AIME 2025와 MATH500 테스트에서:
Kimi Linear는 MLA보다 더 빠른 수렴과 더 높은 최종 성능을 달성했습니다. 특히 장형 생성이 필요한 추론 작업에서 선형 어텐션의 효율성이 도움이 된 것으로 보입니다.
4.5 효율성: 실제 배포의 게임 체인저
디코딩 속도 (배치 크기 1):
- 4K 토큰: Kimi Linear와 GDN-H 비슷 (MLA 2.2배 빠름)
- 128K 토큰: Kimi Linear 3.98배 빠름
- 1M 토큰: Kimi Linear 6.3배 빠름 (시간당 토큰 11.48ms → 1.84ms)
메모리: KV 캐시 75% 감소로, 더 큰 배치 크기 지원 가능. 실제로 1M 컨텍스트에서 배치 처리량이 6배 향상됩니다.
5. 기술적 심화: 위치 인코딩으로서의 선형 어텐션
논문의 흥미로운 이론적 기여 중 하나입니다. RoPE는 회전 행렬의 누적 곱을 통해 상대적 위치를 인코딩합니다:
\[\text{RoPE: } q_t^\top \left(\prod_{j=i+1}^{t} R_j\right) k_i\]
여기서 \(R_j\)는 블록 대각 회전 행렬입니다.
GDN/KDA도 유사한 구조를 가지지만, 회전 행렬 대신 데이터 의존적이고 학습 가능한 전환 행렬을 사용합니다:
\[\text{GDN: } q_t^\top \left(\prod_{j=i+1}^{t} (I - \beta_j k_j k_j^\top) \alpha_j\right) k_i\]
이는 RoPE의 직교성 제약을 완화하면서, 컨텍스트 길이 외삽 문제를 잠재적으로 해결할 수 있습니다. RoPE는 고정된 주파수를 가져서 훈련 길이에 과적합되기 쉽지만, KDA는 동적으로 조정할 수 있죠.
6. 한계와 향후 방향
현재 한계:
- 순수 선형 어텐션도 여전히 긴 문맥 정보 검색에서 완벽하지 않음 (하이브리드 필요)
- 3:1 비율은 경험적 최적값일 뿐, 보편적 최적 비율은 아님
- 스파스 어텐션(Sparse Attention)과의 비교 부족
향후 연구 방향:
- 선형 어텐션 상태 확장 기법과의 결합
- 스파스 + 선형 어텐션의 하이브리드
- 더 큰 규모(100B+)에서의 확장성 검증
결론
Kimi Linear는 선형 어텐션의 오래된 문제를 새로운 관점에서 해결합니다. 세밀한 채널별 게이팅과 최적화된 하드웨어 구현, 그리고 절제된 하이브리드 설계를 통해 모든 평가 시나리오에서 기존 어텐션을 능가했습니다.
- 일관된 우수성: 단문맥, 장문맥, RL 모든 영역에서 최고 성능
- 실용적 효율성: 1M 토큰에서 6배 디코딩 가속, 메모리 75% 감소
- 공정한 평가: 동일한 훈련 조건에서 체계적인 비교
- 오픈소스: KDA 커널과 vLLM 통합, 사전훈련 체크포인트 공개
참고 자료: