[Kimi Team] — KAttention Residuals: Attention Residuals(AttnRes)를 통한 가중치 기반 깊이 방향 정보 통합 기법

배경 및 개요

현대적인 거대 언어 모델(LLM)에서 표준적으로 사용되는 PreNorm 기반 잔차 연결(Residual Connections)은 모든 레이어의 출력을 동일한 가중치(유닛 가중치)로 누적하여 전달하는 방식을 취합니다. 그러나 이러한 균등 집계 방식은 모델이 깊어짐에 따라 은닉 상태(Hidden-state)의 크기가 $O(L)$로 제어되지 않고 성장하게 만들며, 결과적으로 각 레이어의 상대적인 기여도를 희석(Dilution)시키는 문제를 야기합니다. 특히 초기 레이어의 정보가 깊은 층으로 갈수록 매몰되어 선택적으로 회수하기 어렵고, 실증적으로 많은 레이어가 최소한의 손실만으로도 프루닝(Pruning)될 수 있다는 사실은 기존 잔차 연결의 비효율성을 시사합니다.

본 논문에서는 이러한 한계를 극복하기 위해 Attention Residuals (AttnRes)를 제안합니다. 이는 고정된 누적 방식 대신, 이전 모든 레이어 출력에 대한 소프트맥스 어텐션(Softmax Attention)을 적용하여 각 레이어가 입력에 따라 이전 표현들을 선택적으로 집계할 수 있게 합니다. 이는 순환 신경망(RNN)이 시퀀스 차원에서 겪었던 병목 현상을 어텐션 메커니즘이 해결했던 것과 유사한 논리를 깊이 차원에 적용한 것입니다.

시스템 아키텍처 및 주요 결과

해당 모델은 크게 Full AttnRes와 대규모 학습 및 추론 효율성을 고려한 Block AttnRes 두 가지 변체로 구성됩니다.

• 성능 향상: 스케일링 법칙(Scaling Law) 실험 결과, AttnRes는 모든 모델 크기 및 연산 예산 범위에서 베이스라인을 지속적으로 상회했습니다. 특히 Block AttnRes는 베이스라인 대비 1.25배 더 많은 연산량을 사용해 학습한 모델과 대등한 손실(Loss) 값을 달성했습니다.
• 학습 역학 최적화: PreNorm 베이스라인에서는 깊이에 따라 출력 크기가 단조 증가하는 반면, Block AttnRes는 블록 경계에서 선택적 집계가 이루어짐에 따라 출력 크기가 경계 내에서 주기적으로 유지되는 양상을 보입니다. 또한, 고정된 잔차 가중치로 인해 초기 레이어에 경사도(Gradient)가 집중되던 현상이 완화되어, 깊이 전반에 걸쳐 경사도가 더 균일하게 분포하게 됩니다.
• 추론 효율성: 인프라 최적화를 통해 Block AttnRes는 전형적인 추론 워크로드에서 2% 미만의 지연 시간(Latency) 오버헤드만을 발생시키며 실용적인 교체 가능성을 입증했습니다.

데이터 및 평가

본 연구는 Kimi Linear 아키텍처(총 48B 파라미터, 활성 3B 파라미터)에 AttnRes를 통합하여 1.4T 토큰으로 사전 학습을 진행한 후 평가되었습니다.

• 평가 데이터셋: MMLU, GPQA, GSM8K, HumanEval 등 언어 이해, 수학, 코드 생성 분야의 다양한 벤치마크가 사용되었습니다.
• 주요 성과: Block AttnRes는 모든 평가 과제에서 베이스라인과 대등하거나 이를 능가했습니다. 특히 다단계 추론이 요구되는 GPQA-Diamond(+7.5) 및 Minerva Math(+3.6), 코드 생성 과제인 HumanEval(+3.1)에서 괄목할 만한 성능 향상을 보였습니다. 이는 깊은 층이 이전 층의 표현을 선택적으로 추출하고 구축할 수 있는 능력이 복합적인 태스크 수행에 유리함을 시사합니다.

방법론 및 기술적 세부 분석

1. Attention Residual 메커니즘

기존 잔차 연결의 수식 $h{l} = h{l-1} + f{l-1}(h{l-1})$을 확장하여, 레이어 $l$의 입력을 다음과 같이 정의합니다:

여기서 $\alpha$는 각 레이어별 학습 가능한 의사 쿼리(Pseudo-query) $w_{l}$을 이용해 계산된 소프트맥스 어텐션 가중치입니다.

2. Block Attention Residuals (Block AttnRes)

모든 이전 레이어를 참조하는 Full AttnRes는 대규모 분산 학습 시 $O(Ld)$의 메모리 및 통신 오버헤드를 발생시킵니다. 이를 해결하기 위해 레이어를 $N$개의 블록으로 나누는 Block AttnRes를 도입합니다.

• 블록 내 누적: 블록 내 레이어 출력들은 단순 합산으로 요약됩니다 ($b{n} = \sum f{j}(h_{j})$).
• 블록 간 어텐션: 각 레이어는 $N$개의 블록 수준 요약본과 토큰 임베딩만을 대상으로 어텐션을 수행하여 통신 복잡도를 $O(Nd)$로 줄입니다. 실증적으로 $N \approx 8$ 정도의 설정만으로도 Full AttnRes 성능의 대부분을 회수할 수 있었습니다.

3. 인프라 설계 및 최적화

• 교차 스테이지 캐싱 (Cross-stage Caching): 파이프라인 병렬화 시 동일 물리적 장치가 처리하는 여러 가상 스테이지 간에 블록 표현을 캐싱하여 중복 통신을 제거했습니다. 이를 통해 통신 비용을 $V$배(가상 스테이지 수) 향상시켰습니다.
• 2단계 연산 전략 (Two-phase Strategy): 의사 쿼리가 입력 독립적이라는 특성을 활용하여, 한 블록 내 모든 레이어의 쿼리를 배치화(Batching)하여 한 번의 메모리 읽기로 처리하고(Phase 1), 블록 내 의존성은 순차적으로 처리하며 온라인 소프트맥스로 병합합니다(Phase 2).
• 메모리 효율적 프리필링: 긴 문맥 입력 시 블록 표현을 시퀀스 차원으로 샤딩(Sharding)하여 개별 장치의 메모리 점유율을 대폭 낮추었습니다.

시사점

본 기술은 모델 아키텍처의 깊이-너비 트레이드오프를 재편합니다. 아키텍처 스윕(Sweep) 결과, AttnRes는 동일 연산 예산 내에서 기존 베이스라인보다 더 깊고 좁은 모델 설정에서 최적의 성능을 발휘하는 경향을 보였습니다. 이는 AttnRes가 깊이 방향의 정보 흐름을 더 효과적으로 활용할 수 있음을 의미하며, 향후 하드웨어 제약이 완화됨에 따라 더 미세한 블록 크기나 Full AttnRes를 채택함으로써 추가적인 성능 향상을 기대할 수 있습니다.

References:

1. TECHNICAL REPORT OF ATTENTION RESIDUALS (https://github.com/MoonshotAI/Attention-Residuals)

Tags:

LargeLanguageModels #ResidualConnections #AttentionMechanism #DeepLearning #ScalingLaws #ModelOptimization #거대언어모델 #잔차연결 #어텐션메커니즘 #딥러닝 #스케일링법칙 #모델최적화