문샷 AI, ‘AttnRes’ 공개: 고정 잔차 연결을 어텐션으로 대체해 추론 성능 향상

딥러닝 모델을 조금이라도 깊게 쌓아본 사람이라면 잔차 연결(Residual Connection)의 위대함을 잘 알고 있을 것이다. 입력값을 그대로 출력에 더해주는 이 단순한 아이디어가 기울기 소실 문제를 해결했고, 오늘날 거대 트랜스포머의 기틀을 닦았다.

하지만 지난 10년 가까이 당연시되던 이 방식에는 고질적인 한계가 있었다. 모든 레이어의 출력을 동일한 가중치로 단순 누적하다 보니, 신경망이 깊어질수록 초기 레이어의 특징이 희석되고 PreNorm 환경에서 은닉 상태의 크기가 걷잡을 수 없이 커지는 현상이 발생했다.

문샷 AI(Kimi 팀)는 이 오랜 기본기를 재설계한 ‘AttnRes(Attention Residuals, 어텐션 잔차)’를 발표했다. 고정된 덧셈식 잔차 연결을 걷어내고, 깊이 축에 소프트맥스 어텐션을 적용하는 방식이다.

이전 레이어의 출력을 스스로 골라 담는다

기존의 잔차 구조는 비유하자면 보고서를 쓸 때마다 이전 초안들을 무조건 한 뭉치로 묶는 것과 같았다. 뒤로 갈수록 뭉치가 너무 두꺼워져 정작 첫 장에 적어둔 중요한 단서를 찾아보기 힘들어진다.

AttnRes의 원리는 단순하면서도 직관적이다. 동일 레이어 안에서 토큰 간 관계를 파악할 때 쓰던 어텐션을 레이어의 깊이 방향으로 확장한 것이다. 각 레이어마다 하나의 학습 가능한 의사 쿼리(pseudo-query)를 두고, 이전의 모든 레이어 출력을 대상으로 어텐션을 계산해 필요한 특징만 취사선택한다.

다만 모든 레이어가 매번 이전의 전 층을 돌아보면 메모리 사용량이 O(Ld)로 치솟아 실제 학습에 쓰기 어렵다. 문샷 팀은 이를 해결하기 위해 실용적인 절충안인 ‘블록 AttnRes(Block AttnRes)’를 내놓았다.

전체 레이어를 약 8개 안팎의 블록으로 나누고, 블록 내부에서는 기존의 가벼운 잔차 연결을 쓰되 블록 사이를 오갈 때만 어텐션을 거치도록 한 것이다. 이렇게 하면 메모리 복잡도가 O(Nd)로 대폭 줄어들어, 연산 부담은 최소화하면서 전체 어텐션 구조의 성능을 거의 그대로 가져올 수 있다.

7.5점 뛴 복합 추론과 1.25배 연산 효율

공개된 벤치마크 지표는 인상적이다. 블록 AttnRes를 적용한 모델은 동일한 학습 스텝에서 기존 베이스라인 모델이 연산량을 1.25배 더 투입해 얻어낸 손실값 수준을 달성했다. 구조만 살짝 손봤을 뿐인데 연산 자원을 25% 더 쥐어짠 셈이다.

특히 다단계 논리 추론과 코드 생성 분야에서 두드러진 도약이 나타났다. 고난도 과학 추론 벤치마크인 GPQA-Diamond에서 베이스라인의 36.9점에서 44.4점으로 무려 7.5점 상승했다. 코딩 능력을 측정하는 HumanEval에서는 59.1점에서 62.2점으로, MBPP는 72.0점에서 73.9점으로 올랐다.

범용 평가에서도 MMLU가 73.5에서 74.6, BBH가 78.0으로 개선되었으며, 중국어 벤치마크 C-Eval 역시 79.6에서 82.5로 올랐다. 내부적으로도 출력 크기가 무한정 커지는 것을 억제해 레이어 전반의 그래디언트 분포가 훨씬 고르게 유지되는 것이 확인됐다.

왜 중요한가

이번 성과는 모두가 의심 없이 받아들이던 인공신경망의 가장 기본적인 뼈대를 건드려 유의미한 효율을 입증했다는 점에서 영리하다. 최근 2년간 업계가 컨텍스트 윈도우 확장이나 토큰 어텐션 경량화에만 매달리는 사이, 문샷은 레이어 간의 수직적 연결선에 여전히 최적화할 공간이 남아있음을 보여줬다.

자체 파운데이션 모델을 사전 학습하는 팀 입장에서는 코드 몇 줄의 수정으로 연산량을 1.25배 아끼면서 GPQA-Diamond 같은 추론 성능을 크게 끌어올릴 수 있어 매력적인 선택지가 될 것으로 보인다.

다만 대규모 분산 학습 환경에서의 실질적인 오버헤드는 지켜볼 필요가 있다. 블록 구조로 메모리를 아꼈다고 해도, 수천 대 이상의 GPU에서 파이프라인 병렬화(Pipeline Parallelism)를 수행할 때는 레이어 간 어텐션이 통신 단계를 복잡하게 만들 가능성이 있다. 이 병목을 인프라 차원에서 어떻게 매끄럽게 풀어내느냐가 실전 도입의 관건이 될 것이다.