가중치 정규화 학습의 진짜 핵심은 유효 학습률 스케줄
원제 Your LR Schedule May Not Be Principled: Effective Learning Rate Scheduling in Weight-Normalized Training
추천 11댓글 2
Key Point
Weight Decay 기반 학습에서 최고 성능을 달성하려면 명목 학습률이 아닌 유효 학습률 궤적을 직접 최적화해야 한다는 실질적 교훈을 제시한다.
핵심 요약
- Weight-Normalized(WN) 학습이 기존 Weight Decay(WD)보다 사전학습 손실을 크게 줄이는 것으로 보고되고 있으나, 그 개선은 정규화 자체보다는 유효 학습률(ELR) 궤적이 주된 원인임을 밝혔다.
- WN과 WD의 ELR 궤적을 일치시키면 훈련 손실 동역학이 대부분 겹쳐지므로, 스케줄 설계는 명목 학습률(LR)이 아닌 ELR을 직접 대상으로 해야 한다.
- 이 원칙을 바탕으로 이론적으로 유도된 1/t 감쇠에 선형 보정을 더한 AIR(Affine Inverse-square-Root decay) 스케줄을 제안하며, 유한 훈련 기간 동안 목표 최고점과 말기 LR을 만족하도록 설계했다.
- 128 B-토큰 Mixture-of-Experts 사전학습 실험에서 AIR은 코사인 스케줄 대비 훈련 손실을 약 0.01 상대 감소시켰고, 다운스트림 성능과 하이퍼파라미터 강건성에서도 우수했다.
- 말기 ELR 스케줄링이 핵심 이득의 원천임을 확인했으며, 공유된 말기 스케줄이 초기 LR 선택의 영향을 약화시킨다.
- 같은 ELR 궤적을 WD에 적용해도 동일한 훈련 동역학을 재현하므로, ELR이 가중치 정규화를 넘어 원칙적인 스케줄링의 기본 대상임을 확립했다.