Stable Rank 복원으로 대규모 언어 모델 훈련 불안정성 방지

초록

본 논문은 대규모 언어 모델(L​LM) 사전학습 중 급격한 그래디언트 폭발을 일으키는 두 가지 현상, 즉 가중치 행렬의 stable rank 급락과 인접 레이어 Jacobian의 정렬 증가를 발견한다. 이를 이론적으로 결합하면 깊이에 따라 그래디언트 노름이 지수적으로 성장함을 증명한다. 해결책으로 제안된 MSign 옵티마이저는 일정 주기마다 가중치 행렬에 matrix‑sign 연산을 적용해 stable rank를 복원하고, 실험적으로 5M‑3B 규모 모델에서 훈련 실패를 방지하면서 7 % 이하의 연산 오버헤드만 발생함을 보인다.

상세 분석

이 논문은 LLM 사전학습에서 흔히 보고되는 “gradient explosion” 현상을 미세하게 추적하여 두 가지 핵심 메커니즘을 규명한다. 첫 번째는 weight matrix의 stable rank가 급격히 감소하는 현상이다. stable rank는 ‖W‖_F²/‖W‖_2² 로 정의되며, 이는 스펙트럼 에너지가 소수의 상위 특이값에 집중됨을 의미한다. 실험에서는 5M‑parameter NanoGPT 모델을 µP 스케일링 하에 훈련시, 실패 직전 단계에서 대부분의 레이어에서 stable rank가 급락하는 것을 관측하였다. 두 번째는 인접 레이어 Jacobian 사이의 정렬(alignment)이 증가한다는 점이다. 정렬은 상위 오른쪽 특이벡터와 다음 레이어의 상위 왼쪽 특이벡터 사이의 코사인 유사도로 정의되며, 값이 1에 가까워질수록 두 Jacobian이 같은 방향으로 작용한다. 정렬이 높아지면 행렬 곱셈에서 발생하는 일반적인 “취소 효과”(cancellation) 가 억제되어, 각 레이어의 스펙트럴 노름이 곱해지듯 전파된다.

이 두 현상을 결합한 이론적 분석은 다음과 같다. 먼저, stable rank가 낮아지면 고정된 Frobenius norm 하에서 연산자 노름 ‖W‖_2 = ‖W‖_F / √srank(W) 가 증가한다(정리 4.4). 따라서 각 레이어 Jacobian의 노름 M이 커진다. 동시에 정렬 a가 1에 가깝게 되면 전체 Jacobian ‖J_total‖_2 ≥ (a·M)^L (정리 4.2) 로, a·M > 1 인 경우 깊이에 따라 지수적으로 성장한다. 이때 체인 룰에 의해 가중치에 대한 그래디언트는 J_total 과 직접적으로 연결되므로, 그래디언트 노름도 폭발한다(정리 4.8). 논문은 이 과정을 “positive feedback loop”이라 부르며, stable rank가 감소하면 Jacobian 노름이 커지고, 이는 다시 정렬을 강화해 더 큰 stable rank 감소를 야기한다는 순환 구조를 제시한다.

해결책인 MSign 옵티마이저는 매 P 스텝(기본값 100)마다 가중치 행렬에 matrix‑sign 연산을 적용한다. 구체적으로, W = U S Vᵀ 의 SVD를 구한 뒤 sign(W) = U Vᵀ 로 모든 비영특이값을 1로 만든다. 이렇게 하면 행렬의 열·행 공간은 보존되면서 singular value가 모두 1이 되므로 stable rank가 최대치(=rank) 로 복원된다. 이후 원래의 Frobenius norm을 다시 스케일링해 원래의 전체 에너지 수준을 유지한다. 이 연산은 주로 어텐션 출력 프로젝션에 적용되며, MLP 레이어에만 적용해도 효과가 제한적임을 실험적으로 확인하였다.

실험에서는 NanoGPT‑5M, Sigma‑40M, LLaMA‑1B, LLaMA‑MoE‑3B 네 가지 아키텍처에 대해 표준 Adam 기반 훈련과 MSign을 비교하였다. 표준 설정에서는 학습률 6e‑4 정도에서 급격한 loss spike와 gradient explosion이 발생했지만, MSign을 적용하면 stable rank가 일정 수준 이상 유지되고 Jacobian 정렬도 억제돼 훈련이 안정적으로 진행되었다. 연산 비용은 추가 SVD와 sign 연산을 포함해 전체 훈련 시간 대비 5‑7 % 수준에 머물렀다. Ablation study에서는 (1) 어텐션 출력 프로젝션에만 적용했을 때 충분히 안정성을 확보하고, (2) 정렬을 직접 정규화하는 방법(예: orthogonal regularizer)보다 비용 효율이 높음을 보여준다.

강점으로는 (①) 실험‑이론‑실제 적용을 일관되게 연결한 점, (②) 기존 옵티마이저와 학습 스케줄을 그대로 사용하면서 간단히 플러그인 형태로 적용 가능하다는 실용성, (③) 다양한 규모와 아키텍처에 대한 검증을 제공한 점을 들 수 있다. 한계점은 (①) SVD 연산이 GPU 메모리와 연산량을 추가로 요구해 매우 큰 모델(>10 B)에서는 구현 최적화가 필요하고, (②) stable rank 복원을 위한 주기적 강제가 실제 최적화 경로에 어떤 미세한 영향을 미치는지에 대한 정량적 분석이 부족하다는 점이다. 또한, 정렬 증가 원인이 학습 데이터 분포, 레이어 정규화, 혹은 학습률 스케줄과 어떻게 상호작용하는지에 대한 탐구가 더 필요하다.

향후 연구 방향은 (1) 저비용 근사 SVD(예: 랜덤화된 SVD)와 결합해 초대규모 모델에 적용, (2) stable rank와 Jacobian 정렬을 동시에 모니터링하고 자동으로 P 를 조절하는 적응형 MSign, (3) 다른 구조(예: 비-Transformer, CNN)에서 동일 현상이 나타나는지 검증, (4) 정렬 억제를 위한 대안적 정규화(예: 스펙트럼 정규화)와의 시너지 효과를 탐색하는 것이 될 것이다.