소프트맥스 대신 다항식으로 보는 트랜스포머 어텐션

초록

이 논문은 트랜스포머에서 소프트맥스 어텐션이 성공한 이유가 확률 분포 형성이라기보다 어텐션 행렬의 Frobenius 노름을 암묵적으로 정규화하는 데 있다고 주장한다. 이를 바탕으로 행렬 정규화 효과를 제공하는 다항식 활성화 함수를 제안하고, 이론적 분석과 광범위한 실험을 통해 정규화만으로도 소프트맥스와 동등하거나 더 나은 성능을 얻을 수 있음을 보인다.

상세 분석

논문은 먼저 기존 트랜스포머에서 소프트맥스가 갖는 “비음성, 행 정규화, 희소성”이라는 세 가지 전통적 특성이 실제 성능에 필수적인지 의문을 제기한다. 저자는 소프트맥스가 행별로 1로 정규화되면서 어텐션 행렬의 Frobenius 노름이 √N(여기서 N은 시퀀스 길이) 이하로 제한된다는 정리를 제시한다(정리 4.1). 이 정리는 역전파 시 그래디언트가 과도하게 커지는 것을 방지하고, 학습 안정성을 보장한다는 점에서 핵심적인 역할을 한다.

그 다음 저자는 동일한 정규화 효과를 다항식 활성화 ϕ(x)=x^p에 스케일링 계수 1/√N을 곱해 구현한다. 정리 4.2와 그에 따른 코롤라리 4.3은 XQKᵀXᵀ의 p제곱 행렬에 대해 기대값 차원에서 Frobenius 노름이 O(N)이고, 1/√N 스케일링을 적용하면 O(√N)으로 감소함을 보인다. 이는 소프트맥스와 동일한 성장률을 갖는다. 또한 정리 4.4·4.5는 Q와 K에 대한 미분(그래디언트) 역시 동일한 스케일링으로 제한될 수 있음을 증명한다.

실험에서는 ViT‑Tiny 모델에 ϕ(x)=x³와 ϕ(x)= (1/√N)x³을 적용해 시퀀스 길이(N=256,64,16,8)별 최적 스케일을 탐색한다. 결과는 N이 커질수록 스케일 k가 √N에 비례해야 최적 성능을 얻는다는 이론적 예측과 일치한다. 또한 동일 조건에서 소프트맥스와 비교했을 때, 스케일링된 다항식은 어텐션 행렬과 그 Jacobian의 Frobenius 노름을 소프트맥스 수준으로 낮추어 학습 안정성을 확보한다. 최종 Top‑1% 정확도 실험에서는 1/16·x³이 소프트맥스와 거의 동등하거나 약간 앞서는 성능을 보이며, 스케일이 없는 x³은 크게 떨어진다.

핵심 인사이트는 어텐션 가중치가 확률 분포를 이루는 것이 아니라, 행렬 노름을 적절히 제한하는 것이 트랜스포머 학습의 핵심 안정성 요인이라는 점이다. 따라서 비음성, 정규화, 희소성 같은 전통적 제약을 완전히 포기하고도, 스케일링된 다항식 활성화만으로 충분히 강력한 어텐션 메커니즘을 설계할 수 있다. 이는 향후 경량화, 하드웨어 친화적 구현, 혹은 새로운 정규화 기법을 탐색하는 연구에 중요한 설계 지침을 제공한다.