하나의 학습률이 모두에게 맞지는 않는다: LLM을 위한 헤비테일(Heavy-Tail) 가이드 기반 계층별 학습률

학습률 (Learning rate) 설정은 현대 딥러닝 (deep learning)의 근본적인 측면입니다. 모든 계층에 동일한 학습률을 적용하는 기존의 관행은 Transformer의 구조적 이질성 (structural heterogeneity)을 간과하며, 이는 대규모 언어 모델 (LLMs)의 백본 (backbone)으로서의 효과를 잠재적으로 제한할 수 있습니다. 본 논문에서 우리는 개별 Transformer 계층에 서로 다른 학습률을 할당하는 적응형 방식인 계층별 학습률 (Layerwise Learning Rate, LLR)을 소개합니다. 우리의 방법은 가중치 상관 행렬 (weight correlation matrices)의 경험적 스펙트럼 밀도 (empirical spectral density, ESD)를 특징짓고 헤비테일성 (heavy-tailedness)을 정량화하는 헤비테일 자기 규제 (Heavy-Tailed Self-Regularization, HT-SR) 이론에 근거합니다. 헤비테일성이 약한 계층에는 훈련을 가속화하기 위해 더 큰 학습률을 할당하고, 헤비테일성이 강한 계층에는 더 작은 학습률을 할당합니다. 이러한 방식으로 학습률을 맞춤화함으로써, LLR은 계층 간의 균형 잡힌 훈련을 촉진하여 더 빠른 수렴 (convergence)과 향상된 일반화 (generalization)를 이끌어냅니다. 다양한 아키텍처 (LLaMA부터 GPT-nano까지), 옵티마이저 (AdamW 및 Muon), 그리고 파라미터 규모 (60M-1B)에 걸친 광범위한 실험을 통해 LLR이 최대 1.5배의 훈련 속도 향상을 달성하고 베이스라인을 능가하며, 특히 평균 제로샷 정확도 (zero-shot accuracy)를 47.09%에서 49.02%로 높인다는 것을 입증했습니다. LLR의 주요 장점은 낮은 튜닝 오버헤드 (tuning overhead)입니다. 즉, 균일한 베이스라인 (uniform baseline)으로부터 거의 최적에 가까운 학습률 (LR) 설정을 직접 전달받을 수 있습니다. 코드는 https://github.com/hed-ucas/Layer-wise-Learning-Rate 에서 확인할 수 있습니다.