단일 헤드 어텐션 고차원 일반화와 스펙트럼 스케일링 법칙

초록

본 논문은 단일 헤드 결합 어텐션 레이어를 고차원(입력 차원 → ∞) 설정에서 분석한다. 무작위 행렬 이론, 스핀 글라스 이론, 근사 메시지 전달(AMP) 기법을 이용해 훈련 및 테스트 오류, 보간·복구 임계점, 그리고 학습된 쿼리·키 행렬의 특이값 분포를 정확히 예측한다. 결과는 실제 트랜스포머에서 관찰되는 저차원 붕괴, 벌크 변형, 고립된 스펙트럼 외삽을 설명하고, 목표 함수가 파워‑러 법칙 스펙트럼을 가질 때 단계적 스펙트럼 회복을 통해 파워‑러 스케일링 법칙이 나타남을 보인다.

상세 분석

이 연구는 고차원 극한(d → ∞, n ∝ d², p ∝ d)에서 단일 헤드 결합 어텐션을 정확히 해석한다는 점에서 혁신적이다. 입력 시퀀스를 독립적인 가우시안 토큰으로 가정하고, 소프트맥스 행렬 A_W(x) 를 키·쿼리 행렬 W에 묶어 표현한다. 손실 함수는 ℓ₂ 정규화가 포함된 제곱 손실이며, 목표 함수는 주어진 스펙트럼을 가진 S₀ = W₀W₀ᵀ/√(p₀d) 형태의 저차원 행렬이다. 저차원 행렬의 고유값 분포 µ₀가 마르첸코‑파스토(또는 일반적인 파워‑러) 형태라고 가정하면, 근사 메시지 전달(AMP) 프레임워크를 통해 학습 과정은 6차원 스칼라 변수(Σ, m, q, ĤΣ, Ĥm, Ĥq)의 최적화 문제로 귀결된다. 이 변수들은 각각 목표‑학습자 겹침, 학습자 노름, 손실 헤시안 역전파 등을 의미한다. 최적화 식(8)은 자유 에너지 형태이며, 이를 수치적으로 해결하면 훈련 오류와 테스트 오류가 정확히 예측된다. 특히, 식(13)은 테스트 오류가 목표와 학습된 행렬 사이의 소프트맥스 차이의 평균 제곱으로 수렴함을 보여준다.

스펙트럼 분석에서는 학습된 W Wᵀ/√(pd) 의 특이값 분포가 두 개의 벌크와 하나의 δ‑피크(영점)로 구성된다는 점을 증명한다. 저차원(랭크 p₀) 성분은 고립된 외삽 고유값으로 나타나며, 이는 실제 트랜스포머에서 관찰되는 “스펙트럼 아웃라이어”와 일치한다. 또한, 샘플 수 α = n/d² 가 증가하면 벌크가 두 개로 분리되는 현상이 나타나며, 이는 학습이 저차원 구조와 고차원 잡음 성분을 순차적으로 회복한다는 물리적 직관과 부합한다.

목표 함수가 파워‑러 스펙트럼(λ_i ∝ i^{-γ})을 가질 경우, AMP 동역학은 강한 모드부터 약한 모드까지 순차적으로 복구한다. 이 과정에서 학습 손실은 단계적 감소를 보이며, 각 단계는 특정 스펙트럼 구간의 회복에 대응한다. 결과적으로 전체 성능은 데이터·파라미터·연산량에 대한 파워‑러 스케일링 법칙(e_test ∝ N^{-β})을 따르게 된다. 이러한 스케일링은 기존의 “lazy” 선형화 이론이 설명하지 못하는 비선형 표현 학습을 포착한다는 점에서 의미가 크다.

마지막으로, 실험에서는 d=100~200 정도의 중간 규모 시뮬레이션을 수행해 이론적 예측과 수치적 결과가 거의 일치함을 확인하였다. 특히, 마르첸코‑파스토 목표와 파워‑러 목표 모두에서 특이값 분포와 테스트 오류가 이론과 일치했으며, 그래디언트 기반 최적화(Adam)도 전역 최소점에 수렴함을 보여준다. 전체적으로 이 논문은 고차원 무작위 행렬 이론과 AMP를 결합해 어텐션 메커니즘의 스펙트럼 구조와 일반화 성능을 최초로 정량화한 작업이라 할 수 있다.