소프트맥스 선형 어텐션으로 글로벌 경쟁 회복

초록

본 논문은 선형 어텐션이 소프트맥스 정규화를 없애면서 발생하는 전역 경쟁 상실 문제를 해결한다. 멀티‑헤드 구조를 활용해 헤드 수준에서 소프트맥스 게이트를 적용함으로써, 토큰‑레벨이 아닌 의미‑슬롯 수준에서 ‘승자 독식’ 동역학을 재현한다. 이 설계는 O(L) 복잡도를 유지하면서도 쿼리 크기에 따른 주의 집중도와 헤드 간 선택성을 회복한다. 실험 결과 SLA가 RetNet, GLA, GDN 등 최신 선형 어텐션 모델을 일관되게 향상시킴을 보인다.

상세 분석

SLA는 기존 선형 어텐션이 갖는 두 가지 근본적 결함, 즉 ‘Magnitude Neglect’와 ‘Context Collapse’를 헤드 차원의 소프트맥스 경쟁으로 완화한다. 논문은 먼저 전통적인 소프트맥스 어텐션이 토큰 간 확률 질량을 공유함으로써 전역 경쟁을 강제하고, 이는 높은 신뢰도(큰 쿼리 노름)일수록 분포가 한두 토큰에 집중되는 ‘winner‑take‑all’ 효과를 만든다고 설명한다. 반면 선형 어텐션은 ϕ(Q)·ϕ(K)ᵀ 형태로 토큰 간 상호작용을 없애고, 쿼리 스케일링이 출력 크기만 늘릴 뿐 분포 형태를 바꾸지 못한다.

SLA는 멀티‑헤드가 이미 의미적 다양성을 제공한다는 점에 착안한다. 각 헤드를 ‘semantic slot’으로 보고, 헤드별 중요도를 G_Q와 G_K라는 스칼라 게이트로 계산한다. 이 게이트는 (QW_GQ)와 (KW_GK)를 헤드 차원에 softmax 함으로써 얻으며, 쿼리·키의 크기가 커질수록 특정 헤드에 확률 질량이 집중돼 엔트로피가 감소한다(정리 4.2). 따라서 헤드 수준에서 전역 경쟁이 재현되고, 이는 곧 토큰‑레벨 경쟁을 대체할 수 있는 충분히 거친 선택 메커니즘이 된다.

수식 (3)‑(7)에서 보듯, SLA는 기존 선형 어텐션의 ϕ(Q)·ϕ(K)ᵀ·V 연산에 G_Q·G_K라는 점곱 가중치를 삽입한다. 이는 읽기(read)와 쓰기(write) 두 단계의 게이트를 동시에 제공해, 키가 메모리에 저장될 때와 쿼리가 메모리를 조회할 때 각각 가장 관련성 높은 헤드만 활성화한다. 구현 측면에서 G_Q와 G_K는 토큰당 스칼라이므로 순환형(RNN‑like) 업데이트와 청크 단위 병렬 처리 모두에 거의 비용을 추가하지 않는다. 파라미터 증가도 W_GQ, W_GK 두 개의 작은 투사 행렬에 불과해 전체 모델 크기의 0.02% 수준이다.

이론적 분석에서는 (1) 쿼리 노름에 민감한 헤드 게이트를 통해 magnitude sensitivity를 복원하고, (2) 헤드 차원의 softmax가 λ→∞일 때 one‑hot 분포로 수렴함을 보이며 asymptotic winner‑take‑all 특성을 증명한다. 실험에서는 언어 모델링(Pile, WikiText)과 장문 컨텍스트(LongChat, Retrieval‑augmented QA)에서 SLA가 RetNet, GLA, GDN 대비 perplexity와 정확도에서 평균 5‑12% 향상을 기록한다. 특히 노이즈가 섞인 검색 시나리오에서 SLA는 잡음 억제 능력이 크게 개선돼, 전역 경쟁이 복원된 것이 실제 성능 향상으로 이어짐을 입증한다.

요약하면, SLA는 “헤드‑레벨 전역 경쟁”이라는 새로운 설계 패러다임을 제시함으로써, 선형 복잡도를 유지하면서도 소프트맥스 어텐션이 제공하던 선택적 집중 메커니즘을 효과적으로 재현한다. 이는 앞으로 초장문 처리, 메모리 효율이 중요한 대규모 모델, 그리고 검색 기반 시스템 등에 널리 적용될 수 있는 실용적이고 이론적으로도 탄탄한 접근법이다.