Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

발행일
출처
arXiv
논문 번호
970
분야
Machine Learning
arXiv 번호
2608.20061

작은 대리 모델에서 찾은 최적 학습률을 모델 폭과 토큰 예산 두 축으로 옮겨, 거대 MoE 사전학습의 학습률을 전면 탐색 없이 예측하는 방법을 제안한 논문이다.

이 논문은 한마디로 거대 MoE(전문가 혼합, 여러 전문가망 중 일부만 골라 쓰는 구조) 모델의 최적 학습률을 작은 모델에서 옮겨와 맞히는 2단계 방법이다. 모델 크기와 토큰 예산이 함께 커지면 학습률을 하나씩 시험해 보는 전면 탐색은 계산비용 때문에 사실상 불가능하다. 1단계는 μP(최대 업데이트 파라미터화)를 MLA 어텐션과 Muon 옵티마이저를 쓰는 MoE에 맞게 정식화해, 폭을 키워도 최적 학습률이 그대로 옮겨가게 만드는 것이다. 2단계는 짧은 토큰 예산에서 구한 최적값들을 로그-로그 공간의 직선으로 맞춰 훨씬 긴 학습 구간까지 밖으로 늘려 잡는 것이다. 저자들은 이렇게 예측한 값으로 총 155B, 활성 17B 파운데이션 모델을 10조 토큰까지 처음부터 사전학습해 안정적인 손실 곡선을 얻었다고 보고한다. arXiv 코멘트와 PDF 머리글에는 COLM 2026 학회 논문으로 실렸다고 적혀 있다.

핵심 요약

  • 1단계로 μP를 MLA 어텐션과 Muon 옵티마이저를 쓰는 MoE에 맞게 정식화했고, 기본 대리 모델(총 0.6B, 활성 0.3B)에서 찾은 최적 학습률이 폭을 2배, 4배, 8배로 키운 모델(총 30.7B, 활성 3.6B까지)에서도 그대로 옮겨감을 보였다.
  • 2단계는 감쇠 구간 없이 안정 구간에서 학습을 끊고 EMA(지수이동평균) 가중치로 여러 체크포인트를 뽑은 뒤, 255B에서 502B 토큰 구간의 최적 학습률을 로그-로그 선형회귀로 맞춰 더 긴 구간으로 늘려 잡는 방식이다.
  • 이 회귀의 적합도는 R²=0.95였고, 10조 토큰 학습에 대한 최적 학습률로 3.85×10⁻⁴을 예측했다.
  • 예측한 학습률로 총 155B, 활성 17B 파운데이션 모델을 10조 토큰까지 처음부터 학습했으며, 목표 학습의 계산량이 대리 실험 전체 합의 약 98배라서 탐색을 생략해 아끼는 비용이 크다.
  • 다만 검증 범위가 MLA와 Muon 조합에 한정되고, 배치 크기는 일부러 다루지 않았으며, 희소성을 폭과 함께 키운 탓에 희소성 축만의 효과는 떼어내지 못했고, 예측 학습률이 정말 최적인지 전면 탐색으로 확인하는 것도 계산상 불가능하다고 저자들이 밝혔다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.

원문 보기 (새 탭에서 열림)