MGUP: 확률적 최적화(Stochastic Optimization)를 위한 모멘텀-그래디언트 정렬 업데이트 정책

효율적인 최적화(Optimization)는 대규모 언어 모델(Large Language Models)을 학습시키는 데 필수적입니다. 레이어 내부의 선택적 업데이트(Intra-layer selective updates)에 대한 연구가 진행되어 왔으나, 수렴 보장(Convergence guarantees)을 유지하면서도 미세한 제어(Fine-grained control)를 가능하게 하는 일반적인 메커니즘은 여전히 부족한 실정입니다. 이러한 격차를 해소하기 위해, 우리는 선택적 업데이트를 위한 새로운 메커니즘인 \textbf{MGUP}를 제안합니다. \textbf{MGUP}는 각 반복(Iteration)마다 선택된 고정된 비율의 파라미터에는 더 큰 단계 크기(Step-sizes)를 적용하고, 나머지 파라미터에는 더 작지만 0이 아닌 단계 크기를 적용함으로써 표준 모멘텀 기반 최적화 도구(Momentum-based optimizers)를 강화합니다. 거의 {plug-and-play} 모듈로서, \textbf{MGUP}는 AdamW, Lion, Muon과 같은 최적화 도구와 원활하게 통합됩니다. 이를 통해 \textbf{MGUP-AdamW}, \textbf{MGUP-Lion}, \textbf{MGUP-Muon}과 같은 강력한 변형 모델들을 생성할 수 있습니다. 표준적인 가정 하에, 우리는 확률적 최적화(Stochastic optimization)에서 (가중치 감쇠(Weight decay)가 없는) \textbf{MGUP-AdamW}에 대한 이론적 수렴 보장을 제공합니다. MAE 사전 학습(Pretraining), LLM 사전 학습, 그리고 다운스트림 미세 조정(Downstream fine-tuning)을 포함한 다양한 작업에 걸친 광범위한 실험을 통해, 우리의 \textbf{MGUP} 강화 최적화 도구들이 기존의 기본 최적화 도구들에 비해 우수하거나 더 안정적인 성능을 달성함을 입증합니다. 우리는 효율적인 레이어 내부 선택적 업데이트를 위한 원칙적이고 다재다능하며 이론적 근거를 갖춘 전략을 제공하여, 대규모 모델의 학습을 가속화하고 안정화합니다. 코드는 https://github.com/MaeChd/MGUP 에서 공개적으로 사용할 수 있습니다.