긴 문맥 일반화를 위한 적응형 희소 어텐션

초록

이 논문은 기존 트랜스포머의 소프트맥스 기반 어텐션이 길어질수록 토큰 간 확산과 표현 붕괴를 일으키는 문제를 지적하고, α‑entmax를 이용한 동적 희소 어텐션이 이러한 한계를 극복함을 보인다. 특히 온도 파라미터를 학습하도록 확장한 Adaptive‑Scalable Entmax(ASEntmax)를 제안해, 희소와 밀집 어텐션 사이를 자동으로 전환하면서 긴 컨텍스트에서도 높은 정확도와 낮은 퍼플렉시티를 달성한다.

상세 분석

본 연구는 트랜스포머 모델에서 어텐션 가중치를 계산할 때 소프트맥스가 모든 토큰에 비제로 확률을 할당함으로써 발생하는 세 가지 근본적 문제—대표성 붕괴, 오버‑스쿼싱, 그리고 어텐션 분산—를 체계적으로 분석한다. 저자들은 α‑entmax(α>1) 가 정확히 0인 토큰을 만들 수 있는 희소 변환임을 이용해, 토큰 수가 증가해도 중요한 토큰에 집중된 확률 질량을 유지한다는 ‘Non‑Vanishing Attention Property’를 정리하고, 이를 정리 1로 수학적으로 증명한다. 또한, 어텐션 분산을 엔트로피 정규화 형태로 정의하고, 소프트맥스는 로그 n에 비례해 엔트로피가 증가해 완전 분산(complete dispersion) 상태에 도달하지만, α‑entmax는 지원 집합 크기 s가 전체 길이 n보다 훨씬 작을 경우 엔트로피가 O(log s)로 제한돼 ‘Concentration Resilience’를 보인다는 정리 1을 제시한다.

표현 붕괴 측면에서는, α‑entmax 기반 어텐션이 지원 집합이 s≪n일 때 토큰 간 L1 거리 차이가 일정 수준 이상 유지됨을 보이며, 이는 Proposition 2의 ‘Preserved representations’ 부분에서 수학적으로 입증된다. 오버‑스쿼싱 완화는 그래디언트 경로 수가 O(nL)에서 O(sL)으로 감소함으로써, 장거리 의존성을 학습할 때 그래디언트 소실을 크게 억제한다는 점에서 실질적인 이점을 제공한다.

하지만 고정된 α와 온도 파라미터만으로는 매우 긴 시퀀스에서 로그 n에 비례해 로그잇(logits) 범위가 확대돼 어텐션이 과도하게 피크(piky)하거나 반대로 너무 넓게 퍼질 위험이 있다. 이를 해결하기 위해 저자들은 온도 스케일링을 로그 길이와 곱하는 형태의 학습 가능한 파라미터 β, γ, δ를 도입한 ASEntmax를 설계한다. 식 (8)‑(9)에서 보듯, 각 헤드마다 입력 특성에 기반해 β와 γ를 동적으로 조정함으로써, 길이가 증가함에 따라 로그잇의 스케일을 적절히 보정한다. 실험적으로는 Gaussian 로그잇 가정 하에 γ<0이면 로그 n에 비례한 스케일링이 로그잇 범위 증가를 상쇄해 일정한 희소성을 유지한다는 이론적 근거를 제시한다.

실험 결과는 두 가지 축을 중심으로 강력히 뒷받침된다. 첫째, 합성 ‘Multi‑query Multi‑token Associative Recall’ 벤치마크에서 ASEntmax는 64 토큰으로 학습한 모델이 65K 토큰까지 1000배 길이 외삽에서도 95.3% 정확도를 기록, 기존 Scalable‑Softmax와 Adaptive‑Temperature 기법을 크게 앞선다. 둘째, 대규모 언어 모델링(예: WikiText‑103, OpenWebText)에서 ASEntmax는 8배 긴 훈련 컨텍스트에서도 퍼플렉시티 감소와 토큰 회수 정확도 상승을 보이며, 짧은 컨텍스트에서는 소프트맥스와 동등하거나 약간 우수한 성능을 유지한다. 또한, 어텐션 히트맵 분석에서 α‑entmax 기반 모델은 토큰 간 집중도가 높고, 그래디언트 흐름 측정에서는 sL 규모의 경로가 실제로 유지되는 것을 확인한다.

결론적으로, α‑entmax 자체가 제공하는 희소성은 긴 문맥 처리 시 어텐션 분산을 억제하고, 표현 붕괴와 오버‑스쿼싱을 완화한다. ASEntmax는 이러한 희소성을 동적으로 조절함으로써, 토큰 수가 급증하는 상황에서도 모델이 중요한 정보를 놓치지 않으며, 필요 시 밀집 어텐션으로 전환할 수 있는 유연성을 제공한다. 이는 앞으로 수십만 토큰까지 확장 가능한 LLM 설계에 핵심적인 설계 원칙이 될 가능성을 시사한다.