[Paper] 깊이 혼합 Attention

Source: arXiv - 2603.15619v1

개요

이 논문은 Mixture‑of‑Depths Attention (MoDA) 라는 새로운 어텐션 원시 연산을 제안합니다. MoDA는 각 트랜스포머 헤드가 현재 레이어의 토큰 표현뿐만 아니라 이전 레이어에서 생성된 키‑밸류(KV) 쌍에서도 정보를 끌어올 수 있게 합니다. 모델이 깊어짐에 따라 “얕은” 신호를 보존함으로써, MoDA는 매우 깊은 대형 언어 모델(LLM)에서 흔히 발생하는 신호‑감쇠 문제를 완화합니다. 저자들은 또한 FlashAttention‑2만큼 빠르게 동작하는 하드웨어 친화적인 구현을 소개하여, 실제 학습 파이프라인에서 이 아이디어를 실용적으로 사용할 수 있게 만들었습니다.

Key Contributions

• Depth‑aware attention: 각 헤드는 현재 레이어의 KV 쌍과 이전 레이어에서 선택된 KV 쌍의 혼합에 and 주의를 기울인다.
• Efficient kernel: 맞춤형 알고리즘이 MoDA에 내재된 비연속 메모리 접근을 해결하여 64 K 시퀀스 길이에서 FlashAttention‑2 처리량의 **97.3 %**를 달성한다.
• Empirical gains: 10개의 검증 언어‑모델 벤치마크에서 MoDA는 평균 0.2 포인트의 퍼플렉시티 감소를 보이며, 10개의 다운스트림 작업에서는 성능을 2.11 % 향상시킨다.
• Low overhead: 이 개선은 FLOPs가 **3.7 %**만 증가한 수준으로, 대규모 학습에 적용 가능하다.
• Norm‑type insight: 실험 결과 MoDA는 일반적인 pre‑norm 구성보다 post‑norm 트랜스포머 블록에서 더 잘 작동한다.
• Open‑source release: 전체 구현 및 학습 스크립트가 GitHub에 공개되어 커뮤니티 채택을 장려한다.

방법론

전통적인 트랜스포머 레이어는 현재 레이어에서만 생성된 KV 쌍을 사용해 어텐션을 계산합니다. 깊이가 증가함에 따라 초기에 포착된 유용한 패턴이 새로운 잔차 업데이트와 반복적으로 섞이면서 결국 사라지게 됩니다. MoDA는 이를 다음과 같이 해결합니다:

1. KV 히스토리 수집 – 순전파 과정에서 각 레이어는 자신의 KV 텐서를 압축 버퍼에 저장합니다.
2. 혼합 선택 – 특정 헤드에 대해 경량 게이팅 네트워크가 현재 레이어 KV와 구성 가능한 수의 이전 레이어 KV 텐서들의 가중합 중 어느 정도 비중을 둘지 결정합니다.
3. 어텐션 계산 – 혼합된 KV 세트를 표준 스케일드‑닷‑프로덕트 어텐션 루틴에 입력합니다. KV 텐서가 레이어별로 연속적으로 저장되기 때문에, 저자들은 비용이 많이 드는 임의 메모리 접근 없이 필요한 슬라이스를 스트리밍하는 커널을 설계했습니다.
4. 정규화와의 통합 – 저자들은 프리‑노름(잔차 전에 레이어‑노름)과 포스트‑노름(잔차 후에 레이어‑노름) 트랜스포머 블록 모두를 실험했으며, 후자가 MoDA의 깊이‑혼합과 더 잘 시너지를 이룬다는 것을 발견했습니다.

그 결과는 기존의 멀티‑헤드 어텐션 모듈을 그대로 대체할 수 있는 드롭‑인 형태이며, 단일 플래그 하나로 켜고 끌 수 있습니다.

결과 및 발견

• Perplexity: 10개의 검증 코퍼스에서 일관된 0.2 포인트 감소는 언어 모델링 품질이 향상되었음을 나타냅니다.
• Downstream tasks: 질문 응답, 요약, 코드 생성과 같은 작업에서 성능 향상이 관찰되며, 이는 보존된 얕은 신호가 더 강력한 다운스트림 추론으로 이어짐을 시사합니다.
• Speed: 맞춤형 커널이 64 K 컨텍스트 윈도우에서 FlashAttention‑2의 처리량의 97.3 %에 도달하여 추가 메모리 조회가 병목이 되지 않음을 확인했습니다.

Ablation studies는 (i) 두 개 이상의 이전 레이어를 혼합하면 수익이 감소하고, (ii) 게이팅 메커니즘이 필수적이며—KV 히스토리를 단순히 연결하면 성능이 저하된다는 것을 확인합니다.

실용적 시사점

• 추가 학습 비용 없이 더 깊은 LLM – MoDA가 초기 레이어 표현이 덮어쓰이는 것을 방지하므로, 개발자는 깊이 스케일링(예: 48‑layer에서 96‑layer 모델) 을 진행하면서도 학습 안정성을 유지할 수 있다.
• 긴 컨텍스트 응용 – 효율적인 커널이 매우 긴 시퀀스(64 K 토큰)에서도 잘 작동하여, 컨텍스트 길이가 중요한 검색 기반 생성, 문서 수준 QA, 코드베이스 분석 등에 MoDA가 매력적이다.
• 플러그‑앤‑플레이 업그레이드 – MoDA가 표준 어텐션 블록을 대체하므로, 기존 코드베이스(Hugging Face Transformers, DeepSpeed, FlashAttention)를 최소한의 리팩터링으로 도입할 수 있다.
• 하드웨어 친화적 – FlashAttention에 근접하는 속도로 인해, 맞춤형 ASIC 없이 현재 GPU/TPU 스택에 배포할 수 있어 산업 채택 장벽이 낮아진다.
• 파인튜닝 가능성 – 적당한 FLOP 오버헤드로 인해 MoDA를 사용한 대규모 사전 학습 모델의 파인튜닝이 실현 가능하며, 처음부터 재학습하지 않고도 다운스트림 성능 향상을 기대할 수 있다.

제한 사항 및 향후 연구

• 스케일업 검증 – 실험은 1.5 B 파라미터 모델에 한정되어 있으며, 메모리 대역폭 제약이 더 엄격한 10 B 이상 파라미터 LLM에서 MoDA가 어떻게 동작하는지는 아직 확인되지 않았다.
• 메모리 사용량 – 여러 레이어의 KV 텐서를 저장하면 다소 있지만 무시할 수 없는 메모리 오버헤드가 발생하며, 메모리 제한이 있는 장치에서 매우 긴 컨텍스트를 처리할 때 제한 요소가 될 수 있다.
• 게이팅 복잡도 – 현재 게이팅 네트워크는 단순하며, 보다 표현력이 풍부한 게이팅(예: 동적 라우팅 또는 학습된 깊이‑주의 정책)은 결과를 더욱 향상시킬 수 있지만 연산량이 증가한다.
• 노름 상호작용 – 포스트‑노름이 더 잘 작동하지만, 그 근본적인 이유는 완전히 분석되지 않았다; 보다 깊은 이론적 분석이 보다 원칙적인 아키텍처 선택을 안내할 수 있다.

향후 연구 방향으로는 MoDA를 인코더‑디코더 아키텍처에 확장하고, 학습 중 적응형 깊이‑혼합 스케줄을 탐색하며, 희소성 또는 전문가 혼합 패러다임과 결합하여 깊이와 폭 스케일링을 동시에 추진하는 것이 있다.

저자

• Lianghui Zhu
• Yuxin Fang
• Bencheng Liao
• Shijie Wang
• Tianheng Cheng
• Zilong Huang
• Chen Chen
• Lai Wei
• Yutao Zeng
• Ya Wang
• Yi Lin
• Yu Li
• Xinggang Wang

논문 정보

• arXiv ID: 2603.15619v1
• 분류: cs.CL, cs.AI
• 발행일: 2026년 3월 16일
• PDF: PDF 다운로드