어텐션 레지듀얼: 학습 가능한 깊이 방향 어텐션을 통한 잔차 연결의 일반화

기존 LLM의 잔차 연결은 모든 층의 출력을 단순히 동일한 비중으로 더하기만 하여 깊어질수록 초기 정보가 희석되는 한계가 있었다. 이 논문은 각 층이 이전 층들의 정보를 선택적으로 가져오는 '깊이 방향 어텐션'을 도입해, 동일 연산량 대비 더 높은 성능과 안정적인 학습을 가능하게 한다.

Transformer의 Residual Connection은 그래디언트 소실을 막는 고속도로 역할을 하지만, 수학적으로는 모든 이전 층의 출력을 동일한 비중(가중치 1)으로 더한다. 이로 인해 모델이 깊어질수록 초기 층의 정보가 누적된 값들에 묻혀버리는 PreNorm 희석 현상이 발생하며, 깊은 층일수록 자신의 영향력을 확보하기 위해 더 큰 출력값을 내놓아야 하는 불안정성이 생긴다.

본 논문은 시퀀스 데이터에서 특정 토큰을 선택하는 Self-Attention 원리를 깊이(층) 차원에 적용한다. 각 층은 고정된 합산 대신, 학습 가능한 의사 쿼리(Pseudo-query)를 사용하여 이전 모든 층의 출력 중 현재 층에 가장 필요한 정보를 선택적으로 수집한다. 이는 마치 도서관에서 모든 책을 한 권으로 합쳐 읽는 대신, 필요한 페이지들만 골라 읽는 것과 같다.

결과적으로 모델은 층 깊이에 상관없이 필요한 시점에 초기 정보를 명확하게 복원할 수 있게 된다. 실험 결과, 이 방식은 기존 방식보다 1.25배 더 많은 연산 자원을 투입한 모델과 대등한 성능을 보였으며, 48B 규모의 Kimi Linear 모델에서도 모든 벤치마크 성능을 일관되게 개선했다.

Full AttnRes는 각 층 l의 입력 h_l을 이전 모든 층 출력 v_i의 가중합으로 계산한다. 이때 가중치 α는 각 층마다 정의된 학습 가능한 벡터 w_l을 쿼리로, 이전 층들의 출력을 키로 사용하는 Softmax 연산을 통해 결정된다. [학습 가능한 벡터 w_l과 이전 층 출력 v_i를 입력으로] → [내적 연산 후 지수 함수를 씌워 전체 합으로 나누는 Softmax 연산을 수행해] → [0에서 1 사이의 가중치 α를 얻고] → [이 값은 현재 층이 이전 특정 층의 정보를 얼마나 중요하게 참조할지를 결정하는 지표가 된다.]

대규모 학습 효율을 위해 제안된 Block AttnRes는 L개의 층을 N개의 블록으로 나눈다. 블록 내부에서는 기존처럼 단순 합산을 수행하고, 블록 간에만 AttnRes를 적용하여 관리해야 할 상태 수를 줄인다. [L개 층을 N개 블록으로 그룹화] → [블록별 요약 벡터 생성] → [요약 벡터 간 어텐션 수행] 과정을 통해 메모리와 통신 복잡도를 O(Ld)에서 O(Nd)로 낮춘다.

인프라 측면에서는 파이프라인 병렬화 시 발생하는 중복 통신을 막기 위해 교차 스테이지 캐싱(Cross-stage caching)을 도입한다. 또한 추론 시에는 2단계 계산 전략을 사용하여, 블록 간 어텐션을 먼저 배치 처리하고 블록 내 계산을 순차적으로 수행한 뒤 온라인 Softmax로 병합함으로써 지연 시간을 최소화한다.

Scaling Law 실험 결과, 모든 모델 크기에서 AttnRes가 Baseline(PreNorm)보다 낮은 손실값을 기록했다. 특히 Block AttnRes는 Baseline 대비 약 1.25배의 연산량 이득(Compute Advantage)을 제공하는 것으로 나타났다.

48B 파라미터(활성 3B) 규모의 Kimi Linear 모델을 1.4T 토큰으로 사전 학습한 결과, MMLU(+1.1), GSM8K(+0.7), GPQA-Diamond(+7.5) 등 주요 벤치마크에서 Baseline을 능가했다. 특히 다단계 추론이 필요한 수학 및 코드 작업에서 개선 폭이 컸다.

학습 역학 분석에서 AttnRes는 깊이에 따른 은닉 상태 크기(Output Magnitude) 증가를 억제하고 그래디언트 노름을 전 층에 걸쳐 균일하게 분포시키는 효과를 보였다. 이는 깊은 모델의 학습 안정성을 높여준다.

AttnRes는 잔차 연결을 깊이 차원의 Softmax 어텐션으로 공식화한다. 각 층 l은 고유한 학습 가능 파라미터 w_l을 가지며, 이를 쿼리로 사용하여 이전 층들의 출력 v_i와 어텐션 스코어를 계산한다. 커널 함수 ϕ(q, k) = exp(qᵀ RMSNorm(k))를 사용하여 값의 폭주를 막는다.

Block AttnRes는 층을 S = L/N 크기의 블록으로 분할하여 메모리 I/O를 최적화한다. 블록 n의 요약 벡터 b_n은 내부 층 출력의 단순 합으로 정의되며, 외부 어텐션은 이 b_n들을 대상으로 수행된다. 이는 수학적으로 행렬 분해(Matrix Decomposition)를 통한 저순위 근사와 유사한 효과를 낸다.

시스템 구현 시 파이프라인 병렬화(Pipeline Parallelism) 환경에서 O(Ld)의 통신 오버헤드를 피하기 위해 스테이지 간 증분 블록만 전송하는 캐싱 메커니즘을 사용한다. 추론 시에는 의사 쿼리가 입력 독립적이라는 점을 활용해 블록 간 어텐션을 행렬 곱셈으로 배치화하여 처리한다.

저자들은 잔차 연결의 변종들을 깊이 혼합 행렬(Depth Mixing Matrix) M의 관점에서 분석한다. 표준 잔차는 하삼각 행렬(All-ones lower-triangular)인 반면, AttnRes는 데이터에 의존적인 밀집 행렬(Dense Matrix)을 형성하여 어텐션 싱크(Attention Sink) 현상 등을 깊이 차원에서도 재현한다.

Block AttnRes에서 블록 크기 S가 커질수록 성능이 Baseline에 가까워지는 경향이 있으며, 현재 하드웨어 제약상 Full AttnRes의 O(Ld) 통신 비용을 완전히 극복하기 위해서는 향후 상호 연결(Interconnect) 기술의 발전이 필요하다.