Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free

Gating mechanisms have been widely utilized, from early models like LSTMs and Highway Networks to recent state space models, linear attention, and also softmax attention. Yet, existing literature rarely examines the specific effects of gating. In this work, we conduct comprehensive experiments to systematically investigate gating-augmented softmax attention variants. Specifically, we perform a comprehensive comparison over 30 variants of 15B Mixture-of-Experts (MoE) models and 1.7B dense models trained on a 3.5 trillion token dataset. Our central finding is that a simple modification-applying a head-specific sigmoid gate after the Scaled Dot-Product Attention (SDPA)-consistently improves performance. This modification also enhances training stability, tolerates larger learning rates, and improves scaling properties. By comparing various gating positions and computational variants, we attribute this effectiveness to two key factors: (1) introducing non-linearity upon the low-rank mapping in the softmax attention, and (2) applying query-dependent sparse gating scores to modulate the SDPA output. Notably, we find this sparse gating mechanism mitigates 'attention sink' and enhances long-context extrapolation performance, and we also release related $\href{https://github.com/qiuzh20/gated_attention}{codes}$ and $\href{https://huggingface.co/QwQZh/gated_attention}{models}$ to facilitate future research.

한 줄 요약

Gated Attention을 통해 SDPA 출력에 시그모이드 게이트를 적용한 결과, 모델 성능과 훈련 안정성이 향상되며 'attention sink' 현상을 제거할 수 있었다.

핵심 기여도

15B MoE와 1.7B dense 모델에서 30개 이상의 Gated Attention 변형을 비교하여, SDPA 출력에 헤드별 시그모이드 게이트 적용이 가장 효과적임을 밝힘.
헤드별 게이트 적용으로 PPL 0.2 감소, MMLU 2점 상승, 훈련 안정성 향상 및 손실 스파이크 감소.
게이트가 'attention sink' 현상을 제거하고, RULER에서 10점 이상 긴 문장 추론 성능 향상.
코드와 모델을 공개하여 향후 연구를 지원.

핵심 아이디어

기존 어텐션 메커니즘은 선형 변환에 제한되어 표현력이 낮고, 토큰 초기부에 과도한 가중치가 집중되는 'attention sink' 문제가 발생한다. 본 연구는 SDPA 출력에 헤드별 시그모이드 게이트를 적용하여 두 가지 핵심 효과를 도출했다. 첫째, 게이트가 SDPA의 저랭크 선형 변환에 비선형성을 도입하여 모델 표현력을 높인다. 둘째, 입력에 따라 희소한 게이트 점수를 적용하여 어텐션 출력의 희소성을 유도함으로써 attention sink를 제거하고, 긴 문맥 처리 능력을 향상시킨다. 이는 기존의 라우팅 기반 게이트와는 별도로, 게이트 자체의 내재적 효과임을 실험적으로 입증했다.

기술적 접근법

**모델 아키텍처**: 15B MoE (15A2B, 128개 전문가, top-8 softmax 게이팅), 1.7B dense 모델.
**게이트 적용 위치**: G1 (SDPA 출력), G2 (value projection), G3 (key projection), G4 (query projection), G5 (최종 dense 출력).
**게이트 유형**: 헤드별/공유, 요소별/헤드별, 덧셈/곱셈.
**데이터셋**: 3.5T 토큰, 다국어, 수학, 일반 지식 포함.
**하이퍼파라미터**: AdamW 최적화기 기본값, 배치 크기 및 학습률은 각 실험별로 조정.
**추가 기술**: Group Query Attention (GQA), z-loss, fine-grained experts, global-batch LBL.

주요 결과

**성능 향상**: SDPA 출력에 헤드별 시그모이드 게이트 적용 시, PPL 0.2 감소, MMLU 2점 상승.
**훈련 안정성**: 손실 스파이크 감소, 더 높은 학습률 허용.
**attention sink 제거**: SDPA 출력에 희소 게이트 적용 시, 초기 토큰의 과도한 가중치 집중 현상이 사라짐.
**긴 문맥 처리**: RULER에서 10점 이상 긴 문장 추론 성능 향상.
**추가 지표**: GSM8k, HumanEval, C-eval, CMMLU 등 다양한 벤치마크에서 개선.

의의 및 한계

본 연구는 게이트가 단순히 라우팅 기능을 넘어서, 어텐션 표현력과 훈련 안정성에 직접적인 영향을 미친다는 점을 실증적으로 밝혔다. 특히, 희소 게이트가 attention sink를 제거하고 긴 문맥 처리 능력을 향상시킨다는 점은 기존 연구에서 다루어지지 않았던 중요한 통찰이다. 그러나 본 연구는 특정 데이터셋(3.5T 토큰)에서만 실험되었으며, 다른 도메인이나 더 큰 모델에서의 일반화 가능성은 추가 연구가 필요하다. 또한, 게이트가 도입하는 계산 비용은 낮지만, 대규모 배치 처리 환경에서는 미세한 성능 저하가 발생할 수 있다.

실용적 활용

Gated Attention은 대규모 언어 모델의 훈련 안정성과 추론 성능을 동시에 향상시키므로, 특히 긴 문맥 처리가 필요한 챗봇, 문서 요약, 코드 생성 등에 유용하게 활용될 수 있다. 또한, attention sink 문제를 제거함으로써 모델이 초기 토큰에 과도하게 의존하는 현상을 줄일 수 있어, 보다 균형 잡힌 어텐션 분포를 필요로 하는 응용 분야에도 적합하다.