어텐션 싱크가 만든 네이티브 MoE와 헤드 붕괴 해결

초록

본 논문은 대형 언어 모델에서 첫 토큰에 과도하게 집중되는 “어텐션 싱크” 현상이 사실상 헤드별 게이트 역할을 수행해 자연스러운 전문가‑전문가(MoE) 구조를 형성한다는 이론을 제시한다. 이 구조가 특정 헤드에만 의존하는 “헤드 붕괴”를 야기함을 밝히고, 싱크 가중치를 명시적 게이트로 전환하고 로드‑밸런싱 손실을 추가하는 “Sink‑Aware Training”을 제안한다. 실험 결과, Vanilla, Sink, Gated 세 종류의 어텐션 모두에서 헤드 활용 균형이 개선되고 전반적인 성능이 상승한다.

상세 분석

논문은 먼저 어텐션 싱크가 발생하는 메커니즘을 수식적으로 분석한다. Vanilla Attention에서는 softmax 정규화 때문에 전체 가중치 합이 1이 되도록 강제되며, 실제 의미가 없는 잉여 가중치는 첫 토큰에 몰린다. 첫 토큰의 value 벡터는 학습 과정에서 거의 0에 수렴하므로, 이 가중치는 출력에 거의 영향을 주지 않는다. Sink Attention은 이 잉여 가중치를 학습 가능한 스칼라 ‘sink’ 파라미터로 대체해 첫 토큰 대신 0값을 갖는 가상의 토큰에 흡수한다. Gated Attention은 각 헤드마다 sigmoid 기반 게이트를 도입해 직접적인 라우팅을 수행한다.

핵심 통찰은 “sink에 할당된 가중치 = 1‑Gating”이라는 관계이다. 즉, A_sink이 클수록 해당 헤드의 실제 출력 가중치는 (1‑A_sink)만큼 감소한다. 이는 곧 헤드별 게이트 값 G = 1‑A_sink와 동일하며, MoE에서 전문가를 선택하는 라우팅 메커니즘과 일치한다. 따라서 Vanilla와 Sink Attention 모두 암묵적인 MoE 구조를 내포하고 있음을 증명한다.

이러한 암묵적 MoE는 전문가 붕괴(Expert Collapse)와 유사한 “헤드 붕괴” 현상을 초래한다. 실험에서는 특정 소수의 헤드가 지속적으로 높은 활성도를 보이며, 나머지 헤드는 거의 사용되지 않는다. 이는 모델이 제한된 표현력에 의존하게 만들어 긴 컨텍스트에서의 정밀한 토큰 선택을 방해한다.

헤드 붕괴를 완화하기 위해 논문은 두 단계의 해결책을 제시한다. 첫째, sink 가중치를 명시적 게이트로 재해석하고, 이를 목표 함수에 포함시켜 모든 헤드가 균등하게 라우팅되도록 한다. 둘째, 로드‑밸런싱 손실 L_balance = Σ_h (π_h – 1/H)^2 (π_h는 헤드 h의 평균 사용 비율, H는 헤드 수) 를 추가한다. 학습 초기에는 전체 헤드에 균등한 라우팅을 강제하고, 사전 학습된 모델을 파인튜닝할 때는 이미 활성화된 상위 헤드를 고정하고 나머지 헤드에만 균형 손실을 적용해 기존 지식을 보존한다.

효율성 측면에서 구현은 Flash Attention의 log‑sum‑exp 값을 재활용해 추가 연산 없이 LSE와 sink 파라미터만으로 게이트와 로드‑밸런싱을 계산한다. 실험에서는 0.6B, 1B, 2B 규모 모델을 Vanilla, Sink, Gated 각각에 적용하고, LLaMA3와 Qwen3 같은 대형 사전 학습 모델에도 파인튜닝을 수행했다. 모든 설정에서 헤드 사용률이 고르게 분포하고, Perplexity, LongBench, MMLU 등 다양한 벤치마크에서 평균 1‑2% 포인트의 성능 향상이 관찰되었다. 특히 긴 컨텍스트(>2k 토큰)에서 Sink와 Gated Attention이 Vanilla 대비 주목할 만한 개선을 보였으며, 헤드 붕괴를 완화한 모델은 양자화와 프루닝에도 더 강인한 특성을 나타냈다.

결론적으로, 어텐션 싱크는 단순한 버그가 아니라 내재된 MoE 라우팅 메커니즘이며, 이를 명시적으로 제어하고 로드‑밸런싱을 도입함으로써 헤드 붕괴를 해결하고 전반적인 모델 효율성과 성능을 동시에 끌어올릴 수 있음을 입증한다.