Attention Residuals: 깊이 방향 선택적 어텐션으로 잔차 연결 재설계

Digest: 현대 LLM의 표준인 **잔차 연결(residual connection)**은 모든 레이어 출력을 동일한 단위 가중치로 누적하여, 깊어질수록 은닉 상태가 O(L)로 성장하면서 개별 레이어의 기여가 점진적으로 희석되는 PreNorm dilution 문제를 야기한다. Kimi Team은 이 고정된 누적을 학습된 softmax 어텐션으로 대체하는 **Attention Residuals(AttnRes)**를 제안한다. 핵심 통찰은 RNN이 시퀀스 차원에서 고정 반복을 어텐션으로 대체하여 성공한 것처럼, 깊이 차원에서도 동일한 전환이 가능하다는 것이다. 각 레이어는 단일 학습 가능한 pseudo-query 벡터 w_l을 통해 모든 이전 레이어 출력에 대해 content-dependent 가중치를 계산하며, 이를 통해 선택적 정보 검색이 가능해진다. 대규모 학습의 메모리 부담을 줄이기 위한 Block AttnRes(레이어를 N개 블록으로 묶어 블록 수준에서 어텐션 수행)는 메모리를 O(Ld)에서 O(Nd)로 줄이면서도 Full AttnRes 성능의 대부분을 보존한다. Scaling law 실험에서 Block AttnRes는 베이스라인 대비 1.25배 컴퓨트 이점(Figure 4)을 보였으며, 48B Kimi Linear 모델에서 GPQA-Diamond +7.5, Math +3.6, HumanEval +3.1 등 모든 다운스트림 태스크에서 일관된 성능 향상(Table 3)을 달성했다.

섹션별 요약

Introduction

표준 잔차 연결은 h_l = h_{l-1} + f_{l-1}(h_{l-1})로, 전개하면 모든 이전 레이어 출력의 균일 가중합이 됨
PreNorm 환경에서 은닉 상태 크기가 O(L)로 성장 → 초기 레이어 정보가 매몰되고, 실제로 상당 비율의 레이어를 pruning해도 손실이 미미함 [11]
핵심 관찰: 깊이 방향 누적과 시퀀스 방향 반복 사이의 형식적 이중성(duality) → 시퀀스에서 RNN→Attention 전환이 성공했듯, 깊이에서도 동일 전환 가능
기여: (1) AttnRes 방법론 + Block AttnRes, (2) 대규모 인프라 최적화, (3) 포괄적 평가

Methods

Full Attention Residuals:

어텐션 가중치: α_{i→l} = φ(q_l, k_i) / Σ_j φ(q_l, k_j), 여기서 φ(q,k) = exp(q^T RMSNorm(k))
각 레이어 l의 query는 학습 가능한 pseudo-query 벡터 q_l = w_l ∈ R
key/value: k_i = v_i = h_1 (임베딩, i=0) 또는 f_i(h_i) (레이어 출력, i≥1)
레이어 입력: h_l = Σ_{i=0}^{l-1} α_{i→l} · v_i
복잡도: O(L²d) 연산, O(Ld) 메모리

Block Attention Residuals:

L개 레이어를 N개 블록(각 S=L/N 레이어)으로 분할
블록 내부: 표준 합산으로 누적 → 블록 표현 b_n = Σ_{j∈B_n} f_j(h_j)
블록 간: softmax 어텐션으로 N개 블록 표현 + 부분합에 대해 선택적 접근
메모리: O(Ld) → O(Nd), 연산: O(L²) → O(N²)
실험적으로 N≈8이면 Full AttnRes 성능의 대부분 회복 (Figure 6)

Two-Phase Computation (Algorithm 1):

Phase 1: 블록 내 모든 S개 레이어의 inter-block 어텐션을 배치로 병렬 처리
Phase 2: intra-block 어텐션을 순차적으로 계산 + online softmax merge
추론 지연 오버헤드 < 2%

Results

Scaling Law (Table 2, Figure 4):

활성 파라미터	토큰 수	Baseline	Block AttnRes	Full AttnRes
194M	38.7B	1.931	1.909	1.899
241M	45.4B	1.895	1.875	1.869
296M	62.1B	1.829	1.809	1.804
436M	87.9B	1.766	1.746	1.737
528M	119.0B	1.719	1.693	1.692

Scaling curve: Baseline L=1.891×C^{-0.057}, Block AttnRes L=1.870×C^{-0.058}, Full AttnRes L=1.865×C^{-0.057}
Block AttnRes가 1.25× 더 많은 컴퓨트로 학습한 Baseline 손실과 동등

48B Kimi Linear 다운스트림 성능 (Table 3):

카테고리	벤치마크	Baseline	AttnRes	차이
General	MMLU	73.5	74.6	+1.1
General	GPQA-Diamond	36.9	44.4	+7.5
General	BBH	76.3	78.0	+1.7
Math&Code	Math	53.5	57.1	+3.6
Math&Code	HumanEval	59.1	62.2	+3.1
Math&Code	MBPP	72.0	73.9	+1.9
Chinese	C-Eval	79.6	82.5	+2.9

특히 다단계 추론(GPQA-Diamond, Math)과 코드 생성(HumanEval)에서 개선이 두드러짐
이는 깊이 방향 선택적 정보 흐름이 합성적(compositional) 태스크에 유리하다는 가설과 일치

Training Dynamics (Figure 5):

출력 크기: Baseline은 깊이에 따라 단조 증가 → AttnRes는 블록 경계에서 리셋되는 유계 주기 패턴
그래디언트 크기: Baseline은 초기 레이어에 불균형적으로 집중 → AttnRes는 훨씬 균일한 분포

Ablation (Table 4, 16-layer 모델):

변형	손실
Baseline (PreNorm)	1.766
DenseFormer	1.767
mHC	1.747
Full AttnRes	1.737
w/ input-dependent query	1.731
w/ input-independent mixing	1.749
w/ sigmoid (softmax 대신)	1.741
w/o RMSNorm	1.743
SWA (W=1+8)	1.764
Block (S=4)	1.746

Discussion

시퀀스-깊이 이중성: 잔차 연결 ↔ RNN 반복의 형식적 대응 관계를 통해 기존 방법들을 통합 이해
- 표준 잔차 = 깊이 방향 linear attention (all-ones 하삼각 행렬)
- Highway = softmax-free stick-breaking attention
- mHC = 행렬 값 상태를 가진 깊이 방향 linear attention
- AttnRes = 깊이 방향 softmax attention (가장 표현력이 높은 변형)
구조화 행렬 관점: depth mixing matrix M ∈ R^{L×L}로 모든 잔차 변형을 통합
- Standard: 고정, rank-L (all-ones)
- Highway: 동적, 1-semiseparable
- mHC: 동적, m-semiseparable
- Full AttnRes: 동적, dense rank-L
- Block AttnRes: 동적, rank N~N+S
최적 아키텍처 변화: AttnRes는 더 깊고 좁은 모델을 선호 (d_model/L_b ≈ 45 vs 베이스라인 ≈ 60, Figure 7)

Insights

주목할 점: “시퀀스에서 RNN→Attention 전환”이라는 역사적 성공을 “깊이 차원”에 정확히 대응시킨 우아한 관점. 이 이중성은 단순 비유가 아니라 구조화 행렬로 형식화됨
연결 고리: Test-Time Training(TTT), Fast Weight Programmers, Hyper-Connections 등 최근 연구들이 모두 이 이중성의 다른 측면을 탐구하고 있음
시사점: Block AttnRes의 N≈8이면 충분하다는 결과는, 깊이 방향 정보 흐름이 극도로 세밀한 레이어 단위 선택보다는 “블록 수준의 기능 그룹” 선택에 가깝다는 것을 시사
비판적 코멘트: 추론 오버헤드 <2%는 인상적이나, 학습 오버헤드 <4%의 정확한 벽시계 시간 비용은 48B 규모에서만 측정됨. 또한 pseudo-query를 0으로 초기화하는 것이 중요하다고 했지만, 이 초기화 민감도에 대한 체계적 분석이 부족

Discussion Points

논쟁점: DenseFormer가 베이스라인과 동등한 성능(1.767 vs 1.766)을 보인 반면 AttnRes는 명확한 개선을 보였다는 것은, cross-layer access 자체보다 input-dependent weighting이 핵심이라는 강력한 증거
검증 필요 가정: “AttnRes가 더 깊은 모델을 선호한다”는 관찰(Figure 7)이 실제 배포 시나리오에서도 유효한지 — 깊은 모델은 순차 추론 지연이 증가하므로 실용적 최적점은 다를 수 있음
후속 연구: (1) linear attention 커널(φ를 분해 가능한 형태로)로 O(L²) → O(L) 감소 가능성, (2) AttnRes의 학습된 패턴이 모델 해석가능성 도구로 활용될 수 있는지, (3) LoRA/QLoRA 같은 파인튜닝 시 AttnRes 가중치의 적응 행태

메타데이터

항목	내용
제목	Attention Residuals
저자	Kimi Team: Guangyu Chen, Yu Zhang, Jianlin Su* 외 34명
소속	Moonshot AI (Kimi)
연도	2026
발표	arXiv:2603.15031
링크	arXiv, GitHub
키워드	Attention Residuals, residual connections, depth-wise attention, PreNorm dilution, Block AttnRes, Kimi Linear

BibTeX:

@article{chen2026attnres,
  title={Attention Residuals},
  author={Guangyu Chen and Yu Zhang and Jianlin Su and others},
  journal={arXiv preprint arXiv:2603.15031},
  year={2026},
  url={https://arxiv.org/abs/2603.15031}
}

왜 이 연구를 하는가?

핵심 질문

깊이 방향의 고정된 균일 누적(잔차 연결)을 학습된 선택적 어텐션으로 대체하면, PreNorm dilution을 해결하고 깊은 Transformer의 정보 흐름을 개선할 수 있는가?

기존 접근법의 한계

한계	설명
선택적 접근 불가	표준 잔차는 모든 레이어 유형(어텐션, MLP)이 동일하게 가중된 누적 상태를 받음 — 유형별 차별화 불가
비가역적 정보 손실	누적으로 합쳐진 정보는 이후 레이어에서 선택적으로 복구 불가능
출력 성장	PreNorm에서 ‖h_l‖이 O(L)로 성장 → 후반 레이어가 점점 더 큰 출력을 학습해야 영향력 유지 → 학습 불안정
DenseFormer 한계	Cross-layer access를 제공하지만 고정된, 입력 비의존적 가중치 → 베이스라인과 동등 성능 (1.767 vs 1.766)
mHC 한계	입력 의존적이나 m개 병렬 스트림 유지 필요 → 메모리 I/O 34d vs AttnRes의 5.5d (Table 1)

핵심 통찰

시퀀스-깊이 이중성: RNN이 시퀀스에서 고정 반복으로 정보를 압축하는 것과, 잔차 연결이 깊이에서 고정 누적으로 정보를 압축하는 것은 구조적으로 동일한 문제. Transformer가 RNN의 고정 반복을 어텐션으로 대체하여 성공했듯, 깊이에서도 동일한 전환이 가능
최소 비용 선택적 접근: 레이어당 단 하나의 d차원 pseudo-query 벡터만으로 모든 이전 레이어에 대한 content-dependent 선택이 가능 — 전체 파라미터 대비 무시할 수 있는 추가 비용

방법 (Method)

프레임워크 개요

graph TB
    subgraph Standard["(a) 표준 잔차 연결"]
        E1[Embedding h₁] --> L1[Layer 1]
        L1 -->|"⊕ (고정 +1)"| L2[Layer 2]
        L2 -->|"⊕ (고정 +1)"| L3[Layer 3]
        L3 -->|"⊕ (고정 +1)"| OUT1[Output]
    end

    subgraph Full["(b) Full AttnRes"]
        E2[Embedding h₁] --> FL1[Layer 1]
        FL1 --> FL2[Layer 2]
        FL2 --> FL3[Layer 3]
        E2 -.->|"α₀→l"| FL2
        E2 -.->|"α₀→l"| FL3
        FL1 -.->|"α₁→l"| FL3
        FL3 --> OUT2[Output]
    end

    subgraph Block["(c) Block AttnRes"]
        E3[Embedding h₁] --> BL1[Block 1<br/>S layers → b₁]
        BL1 --> BL2[Block 2<br/>S layers → b₂]
        BL2 --> BL3[Block N<br/>S layers → bₙ]
        E3 -.->|"softmax attn"| BL2
        E3 -.->|"softmax attn"| BL3
        BL1 -.->|"softmax attn"| BL3
        BL3 --> OUT3[Output]
    end

    style Standard fill:#f9f9f9
    style Full fill:#e8f4f8
    style Block fill:#e8f8e8

핵심 구성요소

1. Softmax Depth-wise Attention

표준 잔차가 모든 이전 출력을 동일 가중치(=1)로 합산하는 것을, softmax 정규화된 어텐션 가중치로 대체한다:

$h_{l} = α_{0 \to l} \cdot h_{1} + \sum_{i = 1}^{l - 1} α_{i \to l} \cdot f_{i} (h_{i})$

여기서 α_{i→l} = exp(w_l^T · RMSNorm(k_i)) / Σ_j exp(w_l^T · RMSNorm(k_j))이며, w_l은 레이어별 학습 가능한 pseudo-query 벡터이다. RMSNorm은 크기가 큰 레이어 출력이 어텐션 가중치를 지배하는 것을 방지한다.

2. Block Partitioning

L개 레이어를 N개 블록으로 분할하여 메모리를 O(Ld) → O(Nd)로 감소:

블록 내: 표준 합산 b_n = Σ_{j∈B_n} f_j(h_j)
블록 간: N개 블록 표현에 대해 softmax 어텐션
실험적으로 N≈8이면 Full AttnRes 대비 0.001 이내 손실 차이 (528M 모델, Table 2)

3. Two-Phase Inference

Phase 1: 블록 내 모든 레이어의 inter-block 어텐션을 단일 배치 행렬곱으로 병렬 처리 → 메모리 접근을 S회에서 1회로 상각
Phase 2: intra-block 어텐션을 순차 처리 + online softmax로 Phase 1 결과와 병합
총 추론 오버헤드 < 2%

4. Cross-Stage Caching (학습)

파이프라인 병렬 시 블록 표현을 로컬 캐시에 저장하여 중복 전송 제거
Peak per-transition 비용: O(C) → O(P), V배 개선
총 학습 오버헤드 < 4%

5. 초기화

모든 pseudo-query w_l을 0으로 초기화 → 초기 α가 균일 → 학습 시작 시 표준 잔차와 동등
이 초기화가 학습 안정성에 중요함을 실험적으로 검증

발견 (Findings)

주요 결과

모델	Scaling Loss (528M)	GPQA-Diamond (48B)	Math (48B)	HumanEval (48B)
Baseline	1.719	36.9	53.5	59.1
Block AttnRes	1.693	—	—	—
Full AttnRes	1.692	—	—	—
AttnRes (48B)	—	44.4 (+7.5)	57.1 (+3.6)	62.2 (+3.1)

핵심 발견

1. 일관된 스케일링 이점: 194M~528M 모든 모델 크기에서 AttnRes가 베이스라인을 outperform하며, Block AttnRes는 1.25× 컴퓨트 이점에 해당하는 손실 감소를 달성한다 (Figure 4). 스케일이 커질수록 Full과 Block의 차이가 줄어듦 (528M에서 0.001).

2. PreNorm dilution 해결: Training dynamics 분석(Figure 5)에서 AttnRes는 (a) 깊이에 따른 출력 크기를 유계로 유지하고, (b) 그래디언트 분포를 균일화함으로써, PreNorm의 두 가지 핵심 문제를 직접 해결한다.

3. 합성적 태스크에서 큰 개선: GPQA-Diamond(+7.5), Math(+3.6), HumanEval(+3.1) 등 다단계 추론이 필요한 태스크에서 개선이 두드러짐. 이는 후반 레이어가 초기 레이어의 정보를 선택적으로 검색하여 활용할 수 있게 되었기 때문.

4. Input-dependent가 핵심: DenseFormer(고정 가중치 cross-layer)는 베이스라인과 동등(1.767 vs 1.766)이지만, AttnRes(입력 의존적 가중치)는 명확한 개선(1.737). 이는 “cross-layer access” 자체가 아니라 content-dependent selection이 핵심임을 보여줌 (Table 4).

5. 학습된 패턴의 해석가능성: Figure 8의 어텐션 가중치 히트맵에서 (a) locality 보존(대각선 우세), (b) 임베딩에 대한 지속적 어텐션, (c) 선택적 skip connection(레이어 4→초기, 15-16→초기) 등 해석 가능한 패턴이 자발적으로 학습됨.

이론적 의의

시퀀스-깊이 이중성의 형식화

이 논문의 가장 중요한 이론적 기여는 depth mixing matrix M ∈ R^{L×L}를 통해 기존의 모든 잔차 변형을 통합된 프레임워크로 정리한 것이다:

방법	가중치 유형	M의 구조	깊이 어텐션 대응
Standard Residual	고정 (=1)	All-ones 하삼각	Linear attention
Highway	동적 (게이트)	1-semiseparable	Stick-breaking attention
mHC	동적 (m 스트림)	m-semiseparable	Matrix-valued linear attention
Full AttnRes	동적 (softmax)	Dense, rank-L	Softmax attention
Block AttnRes	동적 (softmax)	Rank N~N+S	Block-wise softmax attention

이 관점에서 AttnRes는 “시퀀스에서 linear attention → softmax attention” 전환의 깊이 버전으로, 이론적으로 가장 표현력이 높은 변형이다.

최적 아키텍처 재배치

Figure 7의 아키텍처 스윕에서 AttnRes는 최적점을 d_model/L_b ≈ 60에서 ≈ 45로 이동시킨다. 이는 AttnRes가 깊이를 더 효과적으로 활용하여, 같은 파라미터 예산에서 더 깊고 좁은 모델이 유리해진다는 것을 의미한다. 이 발견은 향후 아키텍처 설계에 중요한 시사점을 제공한다.

핵심 용어 정리

용어	정의
Residual Connection (잔차 연결)	h_l = h_{l-1} + f_{l-1}(h_{l-1}) 형태의 skip connection. 그래디언트 소실을 방지하는 심층 네트워크의 핵심 구성요소
PreNorm	레이어 함수 f_l 내부에서 정규화를 적용하는 방식. 학습 안정성을 제공하지만 은닉 상태 크기가 O(L)로 성장하는 dilution 문제 유발
PreNorm Dilution	PreNorm 잔차 연결에서 깊이가 깊어질수록 개별 레이어의 상대적 기여가 줄어드는 현상. ‖h_l‖이 O(L)로 성장하여 초기 레이어 정보가 매몰됨
Attention Residuals (AttnRes)	고정 단위 가중치 잔차 누적을 학습된 softmax 어텐션으로 대체하는 방법. 각 레이어가 이전 레이어 출력에 선택적으로 접근 가능
Pseudo-query	AttnRes에서 각 레이어에 할당된 학습 가능한 d차원 벡터 w_l. 이전 레이어 출력에 대한 어텐션 가중치를 결정하는 query 역할
Block AttnRes	L개 레이어를 N개 블록으로 묶고 블록 수준에서 어텐션을 수행하는 AttnRes의 효율적 변형. 메모리를 O(Ld) → O(Nd)로 감소
Depth Mixing Matrix	M ∈ R^{L×L} 행렬로, M_{i→l}이 레이어 l이 레이어 i의 출력에 부여하는 가중치. 모든 잔차 변형을 통합적으로 비교하는 프레임워크
Sequence-Depth Duality	시퀀스 방향의 반복(RNN)과 깊이 방향의 누적(잔차 연결)이 구조적으로 동일한 문제라는 관찰. AttnRes의 핵심 동기
Online Softmax	softmax 계산을 스트리밍 방식으로 수행하는 기법. Block AttnRes의 two-phase 추론에서 inter-block과 intra-block 결과를 효율적으로 병합하는 데 사용
Cross-Stage Caching	파이프라인 병렬에서 블록 표현을 로컬 메모리에 캐시하여 중복 전송을 제거하는 최적화 기법
Semiseparable Rank	구조화 행렬 M의 복잡도를 측정하는 지표. 표준 잔차(rank-L, 고정) → Highway(1-semiseparable) → mHC(m-semiseparable) → AttnRes(dense)
MoE (Mixture of Experts)	입력에 따라 일부 전문가 네트워크만 활성화하는 아키텍처. Kimi Linear의 기반 구조(48B 전체/3B 활성)
Kimi Linear	Moonshot AI의 MoE Transformer 아키텍처. KDA(Kimi Delta Attention)와 MLA(Multi-head Latent Attention)를 3:1 비율로 교차 배치

Juhyeon's Blog

탐색기

Attention Residuals