선형 어텐션 아키텍처의 메커니즘, 트레이드오프 및 교차층 라우팅

긴 컨텍스트를 다루는 Transformer 계열 모델은 self-attention의 N² 비용 때문에 확장성에 제약이 있다. 선형 어텐션은 쿼리-키 커널을 특징 맵으로 분해해 고정 크기 메모리로 누적함으로써 시퀀스 길이에 대해 선형 시간과 상수 시간 토큰 처리 비용을 가능하게 한다. 본 논문은 이러한 선형 어텐션 계열의 설계 선택이 손실, 처리량, 시퀀스 길이 스케일링, 구현 복잡도 사이에서 어떤 트레이드오프를 만드는지 실험적으로 규명하여 실무자와 연구자가 목적에 맞는 선택을 하도록 돕는다.

Transformer의 self-attention은 각 쿼리가 모든 이전 키와 쌍별 비교를 수행하므로 시퀀스 길이 N에 대해 계산과 메모리 비용이 N²에 비례한다. 이 제약은 긴 문서나 대규모 컨텍스트에서 학습과 추론의 병목이 된다. 선형 어텐션은 쿼리와 키를 특징 맵 ϕ로 변환하여 유사도를 내적 형태로 근사하면 모든 이전 기여를 누적 메모리 W로 합쳐 쿼리당 단일 매트릭스-벡터 곱으로 출력할 수 있다는 관찰에서 출발한다.

DeltaNet 계열은 선형 어텐션의 간섭 문제를 델타 규칙으로 완화한다. 델타 규칙은 현재 키에 대해 메모리가 예측하는 값 v̄을 먼저 계산하고 실제 쓰기값 v와의 차이 r=v−v̄만 메모리에 기록한다. 이렇게 하면 불필요한 중복 쓰기를 줄이고 메모리 간섭을 억제할 수 있으며, 채널별 게이트와 감쇠를 추가하면 망각과 지우기 동작을 더 정밀하게 제어할 수 있다.

교차층 전달 문제는 정보 희석 문제의 확장이다. 낮은 층에서 계산된 유용한 신호가 깊은 층으로 전달되는 동안 표현이 변형되어 회수하기 어려워지는 현상이 관찰되며 기존의 깊이 주의(depth-attention)는 선형 시간 장점을 일부 상실한다. 이 논문은 델타 규칙 내부 신호를 가볍게 공유하는 경로를 설계하여 정렬된 잔류 스트림으로 값을 투사하면 깊이 정렬 문제를 완화하면서 선형 시간 구조를 유지할 수 있음을 보였다.

전체 접근 방식은 동일 데이터·토크나이저·시퀀스 길이와 동일한 토큰 예산에서 여러 선형 어텐션 규칙과 softmax를 매치드 설정으로 비교하는 것이다. 각 아키텍처는 Megatron-LM 기반 구현으로 동일한 학습 파이프라인을 사용했고 하이브리드 스택과 순수 스택 구성을 통해 손실과 처리량의 트레이드오프를 측정했다. 실험은 350M-parameter 클래스에서 15B 토큰 스윕을 중심으로 옵티마이저, 학습률, 시퀀스 길이 확장, 더 큰 DeltaNet 실행과 소수의 다운스트림 평가를 포함한다.

아키텍처 상세는 각 레이어의 키·쿼리·값을 특징 맵 ϕ로 변환하고 고정 크기 메모리 W를 유지하는 공통 순환 메모리 표기법을 사용한다. DeltaNet은 잔차 r만을 쓰는 업데이트 W←W+β r⊗κ를 사용하고 Gated DeltaNet은 스칼라 감쇠 α를 도입해 W←αW+β rα⊗κ로 갱신한다. Kimi Delta는 채널별 감쇠와 Gated DeltaNet-2는 채널별 erase/write 게이트를 추가해 더 세밀한 편집을 가능하게 했다.

교차층 라우팅 실험은 세 경로를 비교했다. CLER은 하위 레이어의 쓰기 잔차를 다음 레이어의 값 타깃에 직접 주입했고 CLER-H는 동일 잔차를 공유된 잔류 스트림으로 투사해 주입했으며 CLVR은 쓰기값 v 자체를 제로 초기화된 투사 P를 통해 잔류 스트림에 더했다. 투사 행렬은 제로로 초기화되어 초기에는 영향이 없고 학습으로만 활성화되도록 설계하여 호스트 아키텍처와의 매칭 비교가 가능하게 했다.

350M-parameter, 15B-token 매치드 스윕에서 Kimi Delta Attention 하이브리드 스택을 Muon 옵티마이저로 학습한 설정이 보고된 최저 최종 검증 손실 2.273을 기록했다. Muon은 모든 호스트에서 매치드 설정 대비 최종 손실을 낮추는 경향이었고, 하이브리드 스택은 일반적으로 순수 선형 스택보다 손실을 개선했지만 처리량 손실을 초래했다. 처리량 측면에서는 순수 Gated DeltaNet을 AdamW로 학습한 설정이 가장 빠른 상대 속도를 보였으나 최종 손실은 더 높았다.

학습률 및 옵티마이저 어블레이션에서 옵티마이저와 학습률이 결합된 영향이 컸다. Muon은 일반적으로 낮은 학습률 대역(약 3×10⁻⁴)에 더 안정적이었고 AdamW 기반 선형 어텐션은 더 높은 학습률(약 10⁻³)을 선호하는 경향을 보였다. 단일 학습률 기본값을 사용하면 아키텍처 비교가 왜곡될 수 있어 각 조합별 튜닝이 필요함이 드러났다.

시퀀스 길이 확장 측정에서 소프트맥스는 4k에서 32k로 증가할 때 반복 시간 성장률이 가장 컸고 순수 Gated DeltaNet는 거의 일정한 반복 시간을 유지해 긴 컨텍스트에서의 스케일링 우위를 실증했다. 교차층 라우팅 결과는 초기 CLER이 유의미한 개선을 주지 못했으나 공유된 잔류 스트림으로 쓰기값을 투사하는 CLVR이 소폭의 최종 검증 손실 감소를 일관되게 나타냈다. CLVR의 이득은 모델을 더 오래 학습하거나 더 큰 호스트에서는 축소되는 경향이 관찰되었다.

전체 아키텍처는 각 토큰에서 q, k, v를 산출하고 k, q를 특징 맵 ϕ로 변환해 고정 크기 메모리 W∈ℝ^{d_v×d_ϕ}을 유지하는 구조다. 선형 어텐션은 y=Wϕ(q)로 출력하고 W는 시퀀스 진행에 따라 누적 또는 감쇠·수정되어 갱신된다. DeltaNet은 먼저 메모리 예측 v̄=W_{i-1}κ를 계산하고 실제 v와의 잔차 r=v−v̄만을 β r⊗κ 형태로 기록하는 델타 규칙을 사용해 간섭을 줄였다.

Gated DeltaNet은 스칼라 감쇠 α를 도입하여 W←αW+β r_α⊗κ 형태로 갱신하고 Kimi Delta는 채널별 감쇠 D_α를 도입해 키-특징 차원에서 부분적 망각을 가능하게 했다. Gated DeltaNet-2는 채널별 erase b와 write w를 분리해 erase 방향 e=b⊙κ과 write 대상 z=w⊙v를 정의하고 잔차 r=z−W̃ e를 쓰는 방식으로 더 세밀한 편집 동작을 허용한다.

CLVR 구현은 각 순환 메모리 레이어가 내부 신호 s (쓰기값 또는 잔차)를 모델 차원으로 투사하는 제로 초기화된 행렬 P를 학습한다. 투사된 신호 ε=Ps는 잔류 스트림 h에 더해져 이후 모든 레이어와 출력 헤드가 읽을 수 있게 된다. 제로 초기화는 라우팅이 학습 초기에 없는 기본 호스트 행동을 보존하고 모델이 필요한 경우 경로를 활성화하도록 한다.

실험적 설정은 FineWeb-Edu 데이터셋과 LLaMA2 토크나이저, bf16 정밀도, GH200 기반 자원에서의 Megatron-LM 런처를 사용해 재현 가능하게 고정되었다. 핵심 하이퍼파라미터로는 시퀀스 길이 4096, 글로벌 배치 128, 350M급 모델에서는 층수 20–24, 옵티마이저로 AdamW와 Muon을 비교했다. Muon은 NormMuon 2차 모멘트 처리와 스펙트럴 스케일링, Nesterov 모멘텀 0.95를 포함한 적응형 설정으로 보고되었다.

논문이 명시적으로 언급한 한계는 단일 시드 단일 런 비교에 의존해 표준편차나 시드 평균이 없다는 점이다. 하이퍼파라미터 커버리지가 모든 아키텍처에 균등하지 않아 옵티마이저와 학습률 상호작용을 함께 고려해야 한다는 제약이 있다. 또한 보고된 수치는 학습 반복 시간과 처리량에 초점을 맞추었고 실제 추론 속도나 디코딩 메모리 발자국에 대한 실험은 포함하지 않아 추론 측면의 확장성 평가는 별도 검증이 필요하다.