Same Architecture, Different Capacity: Optimizer-Induced Spectral Scaling Laws

Optimizer의 기하학이 FFN width의 스펙트럴 용량 활용 방식에 직접적 영향을 미친다. 같은 Transformer 아키텍처와 동일한 손실 목표라도 AdamW와 Muon/NorMuon 사이에 hard-soft 스펙트럼 분포 차이가 관찰되며, 이는 optimizer–architecture 공동 설계의 필요성을 시사한다.

출발점: FFN은 Transformer에서 매개변적 비중이 큰 구성요소이지만, 표현 차원의 차원 축을 통해 학습된 정보의 차원을 측정하기 어렵다. 본 연구는 FFN 표현의 eigenspectra를 Soft/Hard spectral rank로 추정하고, FFN 폭 D가 커질 때 실제로 얼마나 많은 spectral capacity가 활용되는지를 측정한다. 해결 원리: 프리-액티베이션(z)과 포스트-액티베이션(a)을 각각 프로브로 삼아 공분산 스펙트럼 C를 계산하고, λ_i를 고유값으로 가져와 p_i = λ_i/∑λ_j로 정규화한다. Rényi 엔트로피(Hα) 기반의 Rα(p)로 다양한 α에서 용량을 추정하며, α=1(soft rank)과 α=2(hard rank)을 두 축으로 사용한다. 입력으로 z의 스펙트럼과 nonlinearity 이후의 스펙트럼 간의 재주입(reinjection) 비율도 함께 측정한다. 토큰 빈도별로 HEAD/MID/TAIL로 구간화하여 빈도 분포에 따른 용량 배분을 확인한다. 결과: AdamW는 HARD-rank에서 sublinear한 βhard를 보이는 경향이 강하고, Muon/NorMuon은 MID/TAIL에서 βhard가 거의 선형에 가까운 증가를 보이며, soft-rank와 hard-rank 간의 차이(A1,2)도 크게 달라진다. 이 차이는 아키텍처의 영향보다 옵티마이저 기하학의 영향이 더 크다는 것을 시사한다.

입력값과 연산 흐름을 명확히 정의한다. 입력: hn ∈ R^D를 FFN 각 층에서 얻고, zℓ,t = W(ℓ)in xℓ,t, aℓ,t = φ(zℓ,t)로 비선형 전이한다. [어떤 값을 입력으로] → [covariance C 계산] → [Eig값 λ_i 산출] → [정규화된 고유값 p_i 계산] → [Rα(p) = (∑i p_i^α)^(1/(1−α)) 계산] → [Soft Rank R1 = exp(−∑p_i log p_i), Hard Rank R2 = (∑λ_i)^2/(∑λ_i^2) 계산] → [βsoft/βhard를 D에 대해 선형 회귀로 추정] → [A1,2 = βsoft − βhard 정의]를 수행한다. 프로브 포인트: pre-activation(z)와 post-activation(a)에서 각각 C를 계산하고, nonlinearity를 거친 재주입 비율 ρα = Rpostα/Rpreα를 측정한다. Token-frequency 기반 stratification: HEAD, MID, TAIL로 구분하여 각 regime에서 Rα를 추정한다. Scaling 법칙은 R(D) ∝ D^β로 모델링하고, β는 FFN width 증가에 따른 용량 확장의 효율을 나타낸다. 실험 설정은 GPT-2 160M/350M 디코더-스타일 모델에서 FFN 폭 D를 mdmodel으로 확장하고, AdamW, Muon, NorMuon, Dion(r=1/2,1/16) 등의 옵티마이저를 비교한다. 6k/12k 스텝에서 perplexity와 함께 spectal-scaling 지표를 수집하고, 3-4개의 주된 Ablation과 scale-up 실험을 수행한다.

주요 결과는 다음과 같다. (1) Soft rank는 폭 증가에 대해 βsoft ∈ [0.66, 1.01]로 견고하게 증가하는 반면, Hard rank는 βhard ∈ [0.29, 0.82]로 옵티마이저에 크게 의존한다. AdamW는 HARD-rank에서 가장 낮은 βhard를 보이고(약 0.29), Muon/NorMuon은 MID/TAIL에서 거의 선형에 근접하는 βhard를 보이며(약 0.93~1.04), Dion의 낮은 업데이트 랭크(r=1/16)는 TAIL에서 βhard ~0.40으로 AdamW과 비슷한 수준으로 떨어진다. (2) Head/MID/TAIL 토큰 regime에서 MID, TAIL에서 차이가 가장 큼. (3) Extended AdamW 훈련에서도 hard-rank scaling이 감소하며, perplexity가 개선되어도 spectral capacity의 구조는 유지되지 않는다. (4) Dion update-rank(r=1/2 vs 1/16)에 따라 hard-rank 축이 감소하고 soft-rank 축은 비교적 견고하게 유지된다. (5) RoPE 제거(NoPE) 시 Muon은 HEAD hard-rank를 증가시키고 NorMuon은 감소시키는 등 optimizer 간의 차이가 재배치된다. (6) Attention-rank 증가와 같은 아키텍처 개입보다 optimizer의 기하학 변화가 더 큰 spectral-scaling 차이를 내고, optimizer 간의 최적 매칭도 아키텍처에 따라 달라진다.

전체 아키텍처는 modded-nanoGPT 구성을 사용하며, FFN은 두 개의 linear projection과 squared-ReLU를 포함한다. 프리-액티베이션(z)과 포스트-액티베이션(a)을 각각 프로브로 삼아 Covariance Spectrum을 계산하고, λ_i를 고유값으로 취해 p_i = λ_i/∑λ_j로 정규화한다. Rényi 엔트로피 Hα(p)와 Rα(p)로 spectrum의 집중도를 평가하며, α=1(R1, soft rank)과 α=2(R2, hard rank)을 주요 축으로 삼아 βsoft, βhard를 D에 대해 선형 회귀로 추정한다. HEAD/MID/TAIL로 토큰 빈도 층을 나누어 Regime별 스케일링을 측정하고, R(D) ∝ D^β를 가정해 β를 구한다. 두_probe 포인트(z, a) 간의 차이는 프리-액티베이션 기하학과 포스트-액티베이션 재배치를 구분하는데 활용한다. Dion(r=1/2,1/16)의 update_rank를 통해 업데이트 랭크가 하드-capacity에 미치는 영향을 제어한다. 4.6 RoPE 제거(NoPE) 시 optimizer 간의 spectral 재배치를 확인하고, 350M 스케일에서 NoPE가 각 Regime에 미치는 영향을 분석한다. 학습 설정은 160M/350M에서 FFN width를 D = m·dmodel 형태로 확장하며, optimizer로 AdamW, Muon, NorMuon, Dion(r=1/2, r=1/16)을 사용한다.

본 연구는 160M/350M 파라미터 규모에 한정되며 1B+ 스케일에서 Frontier-scale 일반화를 추가로 확인할 필요가 있다. 또한 Dense FFN 아키텍처에 한정되며 다른 아키텍처나 다른 옵티마이저 계열에서의 일반화 여지가 있다. 또한, 우리의 측정은 FFN의 latent 공간 활용도에 대한spectral-지표에 의존하며, downstream task 성능과의 직접적 매핑은 추가 연구가 필요하다.