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

왜 중요한가

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

핵심 기여

통일된 순환 메모리 표기법으로 다양한 메커니즘 정형화

Softmax attention과 DeltaNet, Gated DeltaNet, Kimi Delta Attention, Gated DeltaNet-2를 한 가지 순환 메모리 관점에서 수식화했다. 이 표기법은 메모리 W, 델타 규칙 잔차 r, 그리고 감쇠·지우기·쓰기 게이트의 역할을 명확히 분리한다. 동일한 표기 하에서 각 설계가 선택성, 망각, 채널 단위 제어, 메모리 세분성에 어떻게 차이를 만드는지 드러났다.

대규모 비교 실험과 하이퍼파라미터 민감도 측정

350M 매개변수 모델을 15B 토큰까지 학습하는 매치드 스윕에서 아키텍처별 검증 손실과 학습 처리량, 학습률 및 옵티마이저 민감도를 비교했다. Muon과 AdamW 간의 학습률 의존성이 아키텍처 성능을 크게 좌우함을 학습률 어블레이션으로 규명했다. 하이브리드 스택과 순수 스택 간의 손실-처리량 트레이드오프를 수치적으로 제시했다.

교차층 라우팅 설계와 실험적 평가

델타 규칙에서 파생한 내부 신호를 층간으로 경량 전달하는 두 축의 실험을 수행했다. 초기 CLER(오류 잔차를 다음 레이어 값 타깃으로 주입)는 유의미한 이득을 주지 못했으며, 동일 신호를 정렬된 잔류 스트림으로 제로 초기화된 투사로 주입하고 쓰기값을 전달하는 CLVR이 소폭의 검증 손실 개선을 일관되게 보여 동작 신호로서의 쓰기값이 더 유용함을 입증했다. 이 경로는 선형 시간 구조를 보존하면서 호스트 아키텍처에 작은 성능 이득을 더했다.

Megatron 기반 구현과 재현 가능한 런 레코드 공개

논문에서 사용한 구현은 tommasocerruti/linear-attention-architectures 레포지토리로 공개되어 Megatron-LM 포크 형태로 재현 가능한 런처와 데이터 파이프라인을 포함한다. 하드웨어, 정밀도, 배치, 토큰 예산, 체크포인트 의미를 런처에 고정하여 비교의 명확성을 확보했다. 재현을 위한 실행 가능한 런처가 제공되어 결과 검증과 추가 실험 확장이 용이하다.

핵심 아이디어 이해하기

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를 포함한 적응형 설정으로 보고되었다.

한계점

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

실무 활용

공개된 Megatron 기반 구현은 연구자가 논문에서 비교한 선형 어텐션 변형을 동일한 데이터·하이퍼파라미터 레시피로 재현할 수 있게 한다. 특히 시퀀스 길이가 긴 작업에서 순수 선형 스택이 반복 시간 면에서 이득을 제공하므로 긴 문맥을 중시하는 실무 환경에서 장점이 있다. 학습률과 옵티마이저 조합에 민감하므로 실무 적용 전 해당 조합의 튜닝이 권장된다.

긴 문서 요약이나 컨텍스트 창이 긴 대화형 시스템에서 학습 및 추론 비용을 절감하는 모델 설계
메모리-제약 환경에서 높은 학습 처리량을 우선시하는 모델 스택 선택
DeltaNet 계열의 경량 레이어를 활용해 깊이 간 정보 공유가 필요한 어플리케이션에서 CLVR로 소폭 성능 개선 시도

코드 공개 여부: 공개

코드 저장소 보기

키워드

Linear attention(선형 어텐션)DeltaNet(DeltaNet)Gated DeltaNet(Gated DeltaNet)Kimi Delta Attention(Kimi Delta Attention)Cross-Layer Value Routing(CLVR)Muon optimizer(Muon 옵티마이저)

용어 해설

선형 어텐션(Linear Attention): — 선형 어텐션은 softmax 커널 대신 쿼리와 키를 특징 맵으로 분해하여 모든 이전 토큰을 누적 메모리로 요약한다. 이 방식은 시퀀스 길이에 대해 선형 시간 복잡도를 제공하므로 긴 컨텍스트에서 학습 및 추론 비용을 낮춘다. 단점으로는 누적된 연관성 간섭과 명시적 정규화 부재로 인한 기억오염 문제가 있다.
델타 규칙(Delta Rule): — 델타 규칙은 메모리가 현재 키에 대해 예측하는 값과 실제 쓰기값의 차이만을 메모리에 기록하는 방식이다. 기존의 단순 누적 쓰기보다 메모리 간섭을 줄이고 오류 보정 성격의 업데이트를 제공한다. 이 규칙은 DeltaNet 계열 아키텍처의 핵심으로 작동하며 선택적 수정과 낮은 간섭을 가능하게 한다.
순환 메모리(Recurrent Memory): — 순환 메모리는 고정 크기 행렬로 과거의 (키, 값) 기여를 누적하여 쿼리와 곱해 출력값을 계산하는 상태다. 한 토큰당 연산이 메모리 업데이트와 매트릭스-벡터 곱으로 축소되어 컨텍스트 길이에 대해 선형 또는 상수 비용을 달성한다. 다양한 게이트와 채널별 조작으로 잔존 정보의 망각과 덮어쓰기 동작을 제어한다.
특징 맵 ϕ(Feature Map ϕ): — 특징 맵 ϕ는 키와 쿼리를 저차원 또는 변환된 표현으로 매핑하여 내적으로 유사도를 근사하도록 설계된다. 이 매핑을 통해 쿼리-키의 쌍별 비교를 누적 메모리와의 단일 곱으로 치환할 수 있다. 적절한 ϕ 선택은 선형 어텐션의 검색 품질과 간섭 수준을 좌우한다.
교차층 값 라우팅(Cross-Layer Value Routing): — 교차층 값 라우팅은 하위 순환 메모리 레이어가 계산한 내부 쓰기값을 공유된 잔류 스트림으로 투사하여 상위 레이어들이 바로 읽을 수 있게 하는 경량 경로다. 투사 행렬은 제로 초기화되어 학습 초기에 영향이 없고 필요한 경우 학습으로 활성화된다. 이 방식은 레이어 간 표현 정렬 문제를 완화하여 값 신호의 깊이 확산을 가능하게 한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

Remember‑R1으로 시각적 어텐션 지속화

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

Remember‑R1으로 시각적 어텐션 지속화