대형 언어 모델에서 scale vectors는 크기는 미미하지만 효과는 크다

Normalization 레이어는 학습의 안정성을 제공하지만 scale vectors의 역할은 불분명했다. 본 연구는 scale vectors가 expressivity를 늘리기보다 Pre-Norm 아키텍처에서 self-amplifying preconditioning으로 학습을 가속한다는 이론적·실험적 근거를 제시한다. 또한 Input-Norm과 Output-Norm 간의 weight decay 차이를 보이고, heterogeneity, placement, reparameterization의 세 가지 경량 개선으로 다양한 설정에서 터미널 손실을 낮추고 확장성을 향상시킨다.

단락1: RMSNorm의 scale vector γ는 Pre-Norm에서 선형 맵 앞에 곱해져 표현력에 미치는 직접적 영향은 없을 수 있지만, γ가 다양한 채널에 걸쳐 선형 변환의 학습 다이나미를 변화시켜 최적화 속도를 좌우한다. 단일 표준 설계는 입력 채널의 균형을 완전히 보장하지 못하므로 학습 동역학이 불균형해질 수 있다. 이를 self-amplifying preconditioning으로 해석할 수 있다. 단락2: Input-Norm vs Output-Norm의 차이는 wd의 효과에도 반영된다. Input-Norm는 최적화에 도움을 주는 반면 Output-Norm는 표현력에 직접 관여하므로 wd에 민감하다. 이로써 wd의 적용 여부를 구분해야 한다. 단락3: HG/Placement/ 재매개변화는 모두 scale vector를 매개로 한 미세한 좌표계 조정으로, 파라미터 수는 작지만 학습 곡선의 기울기에 큰 차이를 만든다. 단락4: OR/ER 재매개변화는 모듈 간의 배치와 채널별 방향성에 따라 효과가 달라진다는 점에서, 기존의 단일 벡터 설계보다 더 미세한 최적화 경로를 제공한다. 단락5: 통합 시나리오로 HG+DP+OR+IWD를 하나의 프리컨디셔너 체계로 바라볼 수 있으며, 이로 인해 0.12B~2B 사이 Dense/MoE 모델에서 터미널 손실 감소와 더 나은 스케일링 법칙이 관찰된다.

설정: RMSNorm 기반의 Pre-Norm Transformer에서 scale vector γ를 고려한다. 입력 x에 Norm(x)를 적용한 뒤 γ로 스케일링하여 W에 전달하는 구조를 가정한다. [입력 →Norm(X) → γ ⊙ Norm(X) → 선형 변환]의 구성에서 γ의 존재 자체가 표현력을 증가시키지 않는다고 가정하되, 학습 다이나믹스에 영향이 있음을 보인다. 2.2의 Theorem 2.2에 따르면, γ를 도입한 모델의 손실 Lf(t)는 wd 없이도 Lg(t)보다 작아지며, 이는 self-amplifying preconditioner Pf,j(t) = γj(t)^2 I_c + w_fj w_fj^⊤에 의해 실현된다. 2.3의 Theorem은 Input-Norm에 wd를 적용하면 Hessian 샤프니스를 억제해 학습 속도를 높이고, Output-Norm에 wd를 적용하면 표현력이 감소해 성능이 저하된다고 보여준다. 3.1의 Theorem 3.1은 DP가 표준 설계보다 즉시 가속화될 수 있음을 보인다. 3.2의 Theorem 3.2는 OR 재매개변화가 표준 설정보다 우수하다고 주장한다. 3.4는 HG/Placement/OR/ER를 Unified Preconditioning으로 해석하는 일반적 프레임을 제시한다. 실험은 0.12B~2B 규모의 Llama 계열(Dense)과 LlamaMoE(MoE)에서 수행되며, 다양한 옵티마이저 AdamW/Muon 및 wsd 스케줄과 토큰 예산에서 평가한다.

4.1 Step-by-step 검증에서 HG로 공유되는 scale vector보다 브랜치별 scale vector가 성능을 개선하고, DP/ DNP의 배치가 훈련 속도에 추가 이점을 제공한다. OR/ER 재매개변수를 적용하면 최적화 속도가 빨라지며, OR+IWD 조합은 가장 낮은 터미널 손실을 달성한다. 4.2 unified 전략은 0.12B~2B 규모의 Dense 및 MoE 모델에서 baselines 대비 일관되게 터미널 손실을 감소시키고, 확장성 면에서도 더 우수한 스케일링 법칙을 보인다. 4.3 Muon 및 wsd 설정과의 호환성에서도 성능 이점이 확인되며, Dense-0.12B에서 0.0156, Dense-0.5B에서 0.0167, MoE-0.5B에서 0.0187만큼 터미널 손실이 감소한다. 4.4 Overhead 분석에서 파라미터 증가율은 O(d)로 매우 작고, Dense-1B 모델에서 wall-clock은 baseline 대비 1.04×, 메모리 사용은 1.01×로 나타난다. 파라미터 총합은 baseline 대비 대략 7.85×10^-5 증가한다.

1. 전체 아키텍처 구조: RMSNorm를 포함하는 Pre-Norm Transformer에서 각 RMSNorm(X; γ) 뒤에 선형 변환이 위치하며, γ가 Norm(X)의 출력에 채널별로 곱해진다. 2) 핵심 메커니즘: scale vector γ를 도입하면 Af(t) = Wf(t) diag(γ(t))가 되고, 각 열의 업데이트는 γj(t)^2 I_c + wf_j(t)wf_j(t)^⊤에 의해 가중치가 조정된다. 3) Prior work 대비 차별점: scale vector의 필요성을 실험적으로 입증하고, expressivity 증가가 아니라 최적화 효과를 강조한다. 4) 구현/학습 세부사항: 0.12B~2B 규모의 Llama 계열 및 LlamaMoE에서 AdamW, Muon, wsd 스케줄로 평가. 5) 이론적 분석: Theorem 2.2, Theorem 2.3, Theorem 3.1, Theorem 3.2 및 Unified 재매개변형의 프린시펄 제시.