💻 computer science

Multi-Head Attention Residuals

본 논문은 서로 다른 특징 서브스페이스가 다중 라우팅 헤드를 통해 깊이 이력(depth history)에 독립적으로 어텐션을 수행할 수 있게 함으로써 서브스페이스 불일치 문제를 해결하고, 표준 트랜스포머와 비교하여 다양한 모델 규모에서 검증 손실 및 다운스트림 태스크 성능을 유의미하게 향提升하는 파라미터 프리(parameter-free) 구조적 수정 방식인 멀티 헤드 어텐션 레지듀얼(Multi-Head Attention Residuals, MHAR)을 소개한다.

원저자: Cheng Luo, Zefan Cai, Junjie Hu

게시일 2026-07-31

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Cheng Luo, Zefan Cai, Junjie Hu

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 다층 도서관을 짓고 있다고 상상해 보세요. 새로운 층이 추가될 때마다 새로운 지식의 방이 생겨납니다. 인공지능의 세계에서 이 '도서관'들은 트랜스포머(Transformers)라고 불리며, 우리가 사용하는 많은 스마트한 도구들의 두뇌 역할을 합니다. 이야기나 수학 문제를 이해하기 위해서, 이 도서관은 열 번째 층에서 작업하는 동안 첫 번째 층에서 무슨 일이 있었는지 기억해야 합니다. 전통적으로 이 도서관들은 모든 층을 연결하는 하나의 좁은 복도를 가지고 있었습니다. 열 번째 층의 방이 과거를 되돌아보고 싶을 때, 그 방은 바로 아래층의 방만을 엿볼 수 있습니다. 그 방은 아래층이 첫 번째 층부터 지금까지의 모든 내용을 완벽하게 요약했을 것이라고 믿어야만 합니다. 이는 마치 영화 전체를 기억하기 위해 바로 앞에 앉아 있는 사람에게만 지금까지 무슨 일이 있었는지 물어보는 것과 같습니다. 만약 그 사람이 세부 사항 하나를 놓쳤다면, 당신도 그것을 놓치게 됩니다.

이를 해결하기 위해, 연구자들은 최근 "어텐션 레지듀얼(Attention Residuals)"이라는 새로운 아이디어를 시도했습니다. 단순히 아래층을 바라보는 대신, 도서관의 모든 방에 이전의 어떤 층이라도 볼 수 있는 마법의 망원경을 준 것입니다. 하지만 여기에는 함정이 있었습니다. 망원경에는 단 하나의 렌즈만 있었다는 점입니다. 방 안의 모든 세부 사항들—모든 색상, 모든 소리, 모든 숫자—은 어느 층에 집중할지 결정하기 위해 동일한 단 하나의 렌즈를 통해 보아야 했습니다. 이는 마치 서로 완전히 다른 취향을 가진 친구들(한 명은 액션 영화를 좋아하고, 다른 한 명은 다큐멘터리를, 또 다른 한 명은 만화를 좋아하는 식)이 모두 하나의 채널을 함께 보기 위해 합의하도록 강요하는 것과 같았습니다. 그들은 타협해야 했고, 종종 그 타협은 아무도 자신이 정말 필요로 하는 것을 보지 못하게 만드는 결과를 초낳았습니다.

이 논문은 영리한 업그레이드인 **멀티 헤드 어텐션 레지듀얼(Multi-Head Attention Residuals, MHAR)**을 소개합니다. 연구진은 영화관이 여러 장르를 동시에 보여주기 위해 여러 개의 스크린을 갖추고 있는 것처럼, 도서관의 망원경도 여러 개의 렌즈를 가져야 한다는 것을 깨달았습니다. 그들은 그 하나의 공유된 렌즈를 여러 개의 더 작고 독립적인 렌즈로 나누었습니다. 이제, 방의 "액션 영화" 부분은 액션이 가득한 층들을 볼 수 있고, "다큐멘터리" 부분은 역사적인 층들을 볼 수 있으며, 이 과정에서 서로 다투거나 타협할 필요가 없습니다.

연구팀은 1억 개의 파라미터를 가진 작은 모델부터 10억 개의 파라미터를 가진 거대한 모델에 이르기까지 다양한 크기의 도서관을 구축하여 이 아이디어를 테스트했습니다. 그들은 작은 도서관의 경우 기존의 단일 렌즈 시스템이 잘 작동한다는 것을 발견했습니다. 하지만 도서관이 커지고 복잡해질수록 단일 렌즈는 병목 현상이 되었으며, 실제로 기존의 복도 시스템보다 AI를 더 나쁘게 만들었습니다. 그러나 새로운 멀리 렌즈 시스템은 도서관이 커질수록 점점 더 좋아졌습니다. 가장 큰 규모에서, 새로운 방식은 표준 모델에 비해 AI의 혼란(검증 손실로 측정됨)을 약 0.063 감소시킨 반면, 기존의 단일 렌즈 방식은 오히려 상황을 0.140만큼 악화시켰습니다.

또한 연구진은 이 업그레이드가 더 큰 망원경을 만들거나 도서관에 더 많은 무게를 더할 필요가 없음을 보여주었습니다. 그것은 단지 기존의 유리를 더 똑똑하게 배치하는 방법이었습니다. 그들은 또한 도서관이 추가적인 탐색 작업으로 인해 느려지지 않도록 특별하고 매우 빠른 도구들을 만들었으며, 이를 통해 속도를 원래 설계와 거의 비슷하게 유지했습니다. 마지막으로, 그들은 이 기술이 이미 만들어진 거대한 80억 파라미터 규모의 도서관에서도 작동하며, 처음부터 다시 만들 필요 없이 수학 문제를 풀고 까다로운 과학 질문에 답하는 능력을 높여준다는 것을 증명했습니다. 핵심적인 발견은 AI 모델이 더 넓어지고 복잡해짐에 따라, 그들이 무엇을 기억할지에 대해 서로 동의하지 않게 되며, 여러 개의 독립적인 방식으로 과거를 돌아볼 수 있게 하는 것이 그 잠재력을 끌어올리는 열쇠라는 점입니다.

기술 요약: 멀티 헤드 어텐션 레지듀얼 (Multi-Head Attention Residuals, MHAR)

1. 문제 정의

표준 트랜스포머는 단일 가산형 레지듀얼 스트림( $h_\ell = h_{\ell-1} + f_\ell(h_{\ell-1})$ )을 통해 깊이(depth)를 따라 정보를 전파하며, 이 과정에서 각 서브레이어는 직전의 상태만을 읽습니다. 이 방식은 이전 레이어에서 계산된 유용한 특징들이 누적 합계(running sum) 내에 유지되기는 하지만, 개별적으로 접근(addressable)할 수는 없다는 한계가 있습니다.

어텐션 레지듀얼(Attention Residuals) (Kimi, 2025)는 모든 이전 서브레이어의 출력을 메모리로 취급함으로써 이를 완화하며, 각 서브레이어가 전체 깊이의 이력을 학습된 소프트맥스 가중치 조합으로 검색할 수 있도록 합니다. 그러나 이 메커니즘은 구조적 병목 현상을 겪습니다. 즉, 깊이 이력을 스코어링하기 위해 단일 공유 라우팅 쿼리( $q \in \mathbb{R}^d$ )를 사용한다는 점입니다. 결과적으로, $d$ 개의 모든 특징 좌표는 단일 분포 $\alpha$ 를 통해 동일한 비율로 동일한 레이어들로부터 정보를 읽어야만 하는 상황에 놓입니다.

논문은 이것이 **"강제된 타협(forced compromise)"**을 초래한다고 주장합니다. 모델의 너비(width)가 증가함에 따라, 서로 다른 특징 서브스페이스(feature subspaces)들은 어떤 역사적 레이어가 가장 관련 있는지에 대해 서로 다른 선호도를 가질 가능성이 높습니다. 단일 쿼리는 이러한 상이한 선호도를 동시에 만족시킬 수 없으며, 이로 인해 서로 다른 깊이 선호도가 하나의 읽기(read)로 붕괴됩니다. 이러한 타협의 비용은 모델의 너비와 함께 증가하며, 대규모 모델에서 단일 쿼리 라우팅을 해롭게 만들 수 있습니다.

2. 방법론: 멀티 헤드 어텐션 레지듀얼 (MHAR)

저자들은 파라미터나 상당한 연산량 추가 없이 '강제된 타협'을 제거하는, 어텐션 레지듀얼의 엄격한 일반화 모델인 **멀티 헤드 어텐션 레지듀얼 (MHAR)**을 제안합니다.

핵심 메커니즘

단일 쿼리 $q$ 대신, MHAR은 라우팅 쿼리를 $H$ 개의 서브스페이스당 헤드로 재구성하여 각 헤드가 깊이 이력에 대해 독립적인 소프트맥스를 갖도록 합니다.

재구성(Reshaping): 단일 쿼리 벡터 $q \in \mathbb{R}^d$ 를 $(H, d/H)$ 형태의 텐서로 재구성합니다.
독립적 라우팅(Independent Routing): 각 헤드 $h$ 에 대해, 모델은 해당 헤드에 대응하는 소스 벡터의 슬라이스를 기반으로 별도의 깊이 분포 $\alpha^{(h)}$ 를 계산합니다.
블록 대각 읽기(Block-Diagonal Read): 검색된 혼합물은 블록 대각 형태를 띱니다. 즉, 각 서브스페이스는 깊이 이력을 독립적으로 읽습니다.
복구(Recovery): $H=1$ 로 설정하면 기존의 단일 헤드 어텐션 레지듀얼과 정확히 일치하게 복구됩니다.

구현 세부 사항

파라미터 효율성: 이 변환은 재구성(reshape) 작업이므로, 단일 헤드 베이스라인 대비 추가 파라미터가 전혀 없습니다.
연산 효율성: 연산 오버헤드는 미미하며, 단일 대형 소프트맥스 대신 $H$ 개의 작은 소프트-맥스를 수행하는 수준입니다(길이 $\le 2L+1$ ).
시스템 최적화: 깊이 라우팅이 메모리 집약적(memory-bound)이라는 점( $O(L^2 B T d)$ 트래픽 필요)을 인식하여, 저자들은 **융합된 Triton 커널(fused Triton kernels)**을 제공합니다. 이 커널들은 RMSNorm과 결론된 온라인 소프트맥스를 수행하여 키(key)와 스택된 소스의 중간 매테리얼라이제이션(materialization)을 제거합니다. 이를 통해 라우팅 시간을 torch.compile 대비 $2\times$ , eager 참조 구현 대비 $6.4\times$ 단축하여, 훈련 처리량을 베이스라인의 $0.55\text{--}0.88\times$ 수준까지 끌어올렸습니다.

하이퍼파라미터 설정

라우팅 헤드의 수 $H$ 는 모델의 KV 헤드 수와 동일하게 설정됩니다. 저자들은 이것이 모델의 소비 단위(consumption granularity)와 라우팅의 정밀도를 일치시키는, 하이퍼파라미터 조정이 거의 필요 없는 최적의 기본값임을 발견했습니다.

3. 주요 기여

MHAR 아키텍처: 단일 쿼리 깊이 라우팅의 강제된 타협을 제거하기 위해 라우팅 쿼리를 독립적인 헤드로 분할하는, 파라미터가 필요 없는 수정 방식입니다.
효율적인 커널: 깊이 라우팅의 메모리 오버헤드와 지연 시간을 크게 줄여 대규모 적용을 가능하게 하는 융합된 Triton 커널을 제공합니다.
항등성 보존 변환(Identity-Preserving Conversion): "델타 어텐션 레지듀얼(delta attention residuals)"을 사용하여(기존 스트림을 대체하는 대신 라우팅된 항을 추가함) 손실 급증 없이 사전 학습된 모델에 MHAR을 이식하는 방법을 제시합니다.
메커니즘적 통찰: 훈련된 쿼리에 대한 직접적인 프로빙(probing)을 통해, 서브스페이스 간의 불일치가 너비와 함께 증가하며 이것이 단일 헤드 라우팅 성능 저하의 주된 원인임을 확인했습니다.

4. 실험 결과

사전 학습 (처음부터 학습)

FineWeb-Edu 데이터셋을 사용하여 세 가지 규모(100M, 350M, 1B)로 학습:

검증 손실(Validation Loss): MHAR은 표준 트랜스포머 및 단일 헤드 어텐션 레지듀얼보다 일관되게 우수한 검증 손실을 보였습니다.
- 100M: 베이스라인 대비 $-0.049$ nats 개선.
- 350M: 베이스라인 대비 $-0.080$ nats 개선.
- 1B: 베이스라인 대비 $-0.063$ nats 개선.
스케일링 동작: MHE의 이점은 규모가 커질수록 단일 헤드 라우팅과의 격차가 벌어집니다. 단일 헤드 라우팅은 100M에서는 도움이 되지만($-0.039 $), 350M($ +0.055 $) 및 1B($ +0.140$)에서는 오히려 해로운 것으로 나타났습니다.
연산 등가성: MHE는 베이스라인 대비 $1.3\text{--}1.5\times$ 의 연산 등가 이득을 달성했습니다.

중간 학습 (8B 모델)

역량 중심 코퍼스에서 어닐링(annealing)된 8B 모델(Marin-8B)을 사용하여 제어 그룹과 매칭된 스케줄로 테스트:

항등성 보존: 변환 과정에서 최적화 충격(optimization shock)이 발생하지 않았습니다. 손실 곡선은 처음 100 스텝 동안 베이스라인과 일치했습니다.
다운스트림 이득: MHE는 추론 작업에서 통계적으로 유의미한 향상을 보였습니다:
- GSM8K: $+3.2$ 포인트 (paired McNemar $p=0.004$ ).
- GPQA: $+3.1$ 포인트 (paired McNemar $p=0.038$ ).
- MMLU, MATH 및 코드 작업에서의 이득은 이 특정 설정에서 통계적으로 유의미하지 않았습니다.

절제 연구 (Ablation Studies)

헤드 수: 최적의 헤드 수 $H$ 는 KV 헤드의 수를 따릅니다. 덜 훈련된 환경에서는 더 많은 헤드가 항상 도움이 되며, 수렴된 환경에서는 $H=KV$가 곡선의 '무릎(knee)' 지점이 됩니다.
너비 vs 깊이: 훈련된 쿼리를 프로빙한 결과, "너비 불일치(width disagreement, 서브스페이스 간의 상이한 라우팅 선호도)"가 모델 너비에 따라 단조 증가함을 확인하였으며, 이는 강제된 타협 가설을 입증합니다.
헤드 정렬: 라우팅 헤드를 어텐션 KV 그룹에 정렬하는 것은 임의의 서브스페이스 라우팅에 비해 측정 가능한 이득을 주지 않았습니다. 이는 이점의 근거가 특정 구조와의 정렬보다는 헤드의 다중성(multiplicity)에 있음을 시사합니다.

5. 의의 및 주장

본 논문은 멀티 헤드 어텐션 레지듀얼이 트랜스포머가 깊이 이력에 접근하는 방식의 근본적인 아키텍처적 한계를 해결한다고 주장합니다. 서로 다른 특징 서브스페이스가 독립적으로 서로 다른 레이어에 주목할 수 있게 함으로써, MHAR은 단일 쿼리 라우팅이 실패하는 지점에서 효과적으로 스케일링됩니다.

저자들은 다음을 강조합니다:

이 개선은 추가된 파라미터나 연산량 때문이 아니라, 구조적 병목 현상을 제거했기 때문에 발생합니다.
융합된 커널 덕분에 메모리 제한적인 깊이 라우팅을 대규모로 적용할 수 있어 실용적입니다.
방식이 다재다능하여, 처음부터 학습하는 경우와 항등성을 보존하며 중간 학습을 진행하는 경우 모두에 적용 가능합니다.
멀티 헤드 라우팅의 필요성은 모델 너비의 직접적인 결과입니다. 모델이 커질수록 단일 쿼리의 "강제된 타협"은 리스크가 되며, 따라서 MHAR은 깊이 인지형(depth-aware) 아키텍처를 위한 필수적인 요소가 됩니다.