대규모 MoE 학습률을 단계적으로 확장하는 방법

왜 중요한가

대규모 MoE는 전체 파라미터를 크게 늘리면서도 토큰당 활성 파라미터를 제한할 수 있지만, 모델 폭과 학습 토큰 수가 바뀔 때마다 learning rate를 다시 탐색해야 합니다. 이 연구는 MoE에 맞춘 μP로 작은 모델의 learning rate를 폭이 큰 모델로 옮기고, 짧은 토큰 구간의 결과를 log-log 회귀해 10T-token 설정까지 확장합니다. 155B total, 17B active parameters 모델을 10T tokens로 사전학습하면서 3.85 × 10⁻⁴ learning rate를 사용했고, 회귀 적합도 R² = 0.95와 안정적인 loss 궤적으로 대규모 탐색 비용을 줄일 근거를 확보했습니다.

핵심 기여

MoE에 맞춘 μP 학습률 전이

MoE의 hidden dimension과 total expert 수를 함께 키우는 확장 경로에 μP를 맞췄습니다. router와 expert FC1은 matrix-like, expert FC2는 vector-like 파라미터로 분류해 초기화와 learning rate scaling을 다르게 적용합니다. 0.6B total, 0.3B active proxy에서 찾은 설정이 2.2B, 8B, 30.7B total 모델로 전이됐습니다.

토큰 예산 기반 두 단계 예측

짧은 학습 구간에서 여러 learning rate를 실행하고 EMA checkpoint의 validation loss를 비교해 각 토큰 예산의 최적값을 찾습니다. 이후 최적 learning rate와 token budget을 log-log 선형 회귀해 관측하지 않은 긴 학습 구간으로 외삽합니다. 10T-token 목표에 필요한 최적 learning rate를 3.85 × 10⁻⁴로 추정했습니다.

155B MoE 사전학습으로 검증

155B total, 17B active parameters의 foundation MoE를 10T tokens로 처음부터 사전학습했습니다. proxy 실행 비용보다 약 98배 큰 목표 학습에서 loss spike 없이 안정적인 학습 궤적을 얻었습니다. MMLU-Pro 기준으로 비교 대상 일부보다 비슷하거나 낮은 추정 학습 연산에서 높은 정확도를 기록해 설정의 실용성을 뒷받침했습니다.

핵심 아이디어 이해하기

일반적인 hyperparameter sweep은 모델 크기와 token budget을 각각 바꿔 최적 learning rate를 찾아야 하므로, 대규모 MoE에서는 실험 조합이 급격히 늘어납니다. 특히 MoE는 hidden dimension뿐 아니라 total expert 수와 active expert 수의 관계도 바뀌어 dense 모델용 전이 규칙이 그대로 맞는지 불확실합니다. 이 논문은 먼저 μP의 초기화와 파라미터별 learning rate scaling으로 폭 확장에 따른 업데이트 크기를 맞춥니다.

μP를 적용하면 작은 proxy에서 얻은 최적 learning rate가 hidden dimension과 total expert 수를 키운 모델에서도 유지됩니다. 입력은 폭이 다른 MoE 모델과 동일한 후보 learning rate이고, 처리 과정은 matrix-like와 vector-like 파라미터에 서로 다른 scaling을 적용한 학습이며, 출력은 각 모델에서 가장 낮은 loss를 만드는 learning rate입니다. SP에서는 모델 폭에 따라 최적값이 이동했지만 μP에서는 0.6B proxy의 선택이 2×, 4×, 8× 확장 모델로 전이됐습니다.

다음으로 학습 토큰 수가 늘어날 때 최적 learning rate가 조금씩 낮아지는 관계를 추정합니다. proxy를 decay 없이 stable phase까지 학습하고 EMA로 가중치를 합친 뒤, 10B-token 간격 checkpoint의 validation loss와 log learning rate 사이에 2차 곡선을 맞춥니다. 곡선의 꼭짓점이 해당 token budget의 최적값이 되며, 이 값들을 다시 token budget과 log-log 선형 회귀해 10T-token learning rate를 계산합니다.

최종적으로 10.8B total, 3.3B active proxy의 약 500B-token 결과에서 10T-token 목표값을 추정하고 155B total, 17B active 모델에 적용합니다. 3.85 × 10⁻⁴ 설정으로 10T-token 사전학습을 진행했을 때 loss spike가 나타나지 않았고, MMLU-Pro와 여러 영어·다국어·수학·코드 평가에서 경쟁력 있는 결과를 얻었습니다. 따라서 작은 proxy 실행을 폭과 토큰 축에 순차적으로 전이하는 방식이 전체 규모의 반복 sweep을 대체하는 경로가 됩니다.

방법론

첫 단계는 MoE 구조에 μP를 적용하는 과정입니다. 모델 폭을 키울 때 hidden dimension, attention head 수, total expert 수를 함께 늘리고 active expert 수와 MoE intermediate dimension은 고정합니다. Embedding과 bias는 vector-like로, FFN·Attention·router·expert FC1은 matrix-like로 분류해 초기화와 learning rate scaling을 다르게 적용하며, depth는 고정해 불안정한 depth scaling을 피합니다.

MLA에서는 query와 key-value의 low-rank projection dimension을 폭 확장 중에도 고정합니다. 이 차원이 해당 up-projection의 fan-in으로 남기 때문에 그 행렬의 learning rate scaling factor는 1이 됩니다. Muon optimizer와 WSD scheduler를 사용하고, batch size는 학습률 전이 대상에서 제외한 채 GPU 처리량을 높이는 값으로 고정하거나 일정 시점 이후 조정합니다.

두 번째 단계는 짧은 token budget에서 최적 learning rate를 측정하는 과정입니다. 각 후보 learning rate로 proxy를 stable phase까지 학습하고, 약 2B-token 간격으로 EMA 가중치를 갱신한 뒤 10B-token 간격 checkpoint를 만듭니다. EMA smoothing factor α = 0.6을 사용해 최근 업데이트에 큰 가중치를 주면서 초기 불안정성을 누그러뜨리고, 각 checkpoint를 고정 validation batch로 평가합니다.

각 token budget에서 얻은 validation loss를 log-transformed learning rate의 2차 함수로 적합합니다. 포물선의 꼭짓점은 해당 token budget에서 loss가 가장 낮은 learning rate이며, 입력은 후보 learning rate와 EMA checkpoint의 loss이고 출력은 추정 최적값입니다. 예를 들어 10.8B proxy의 255B 이후 안정화된 지점들을 log-log 선형 회귀에 넣어 10T-token 목표값 3.85 × 10⁻⁴를 산출했으며, 회귀의 R² = 0.95가 관측값과 직선 관계의 일치도를 나타냅니다.

주요 결과

폭 전이 실험에서 base proxy는 0.6B total, 0.3B active parameters이고, 비교 모델은 2.2B total, 0.7B active, 8B total, 1.5B active, 30.7B total, 3.6B active로 확장됐습니다. SP에서는 폭이 커질수록 training loss가 가장 낮은 learning rate가 이동했지만, μP에서는 proxy의 최적 learning rate가 각 확장 모델에서 일관되게 유지됐습니다. 이 결과는 hidden dimension과 total expert 수를 함께 확장하는 경로에서 learning rate 전이가 작동함을 뜻합니다.

토큰 축 실험에서는 5.6B total, 1.8B active proxy와 20.7B total, 3.8B active held-out 모델을 각각 40B, 60B, 80B, 100B tokens까지 학습했습니다. token budget이 늘어날수록 최적 learning rate는 소폭 하락했고, 두 모델의 validation-loss 포물선은 비슷한 곡률과 꼭짓점 위치를 나타냈습니다. 따라서 μP가 폭 축의 전이를 유지하는 동안 token budget 축에서는 별도의 감소 추세를 회귀할 수 있습니다.

10T-token 예측에서는 10.8B total, 3.3B active proxy를 약 500B tokens까지 학습하고, batch size가 바뀐 뒤 동역학이 안정된 255B 이후의 자료만 회귀에 사용했습니다. log-log 선형 회귀는 R² = 0.95를 기록했고 10T-token 학습의 최적 learning rate를 3.85 × 10⁻⁴로 추정했습니다. 목표 모델의 전체 사전학습은 proxy 실행 합계보다 약 98배 많은 연산을 요구했지만, 전체 규모 sweep 없이 안정적인 설정을 얻었습니다.

155B total, 17B active foundation MoE는 10T-token Stage 1 사전학습 동안 loss spike 없이 안정적인 loss 궤적을 유지했습니다. 6T-token 시점에는 데이터 혼합 비율을 English 45%에서 22.5%, Math/STEM 12.5%에서 27.5%, Code 27.5%에서 25.0%, Multilingual 15.0%에서 25.0%로 조정해 저대표 도메인 비중을 높였습니다. MMLU-Pro 기준 비교에서 dots.llm1과 GLM-4.5-Air보다 비슷하거나 낮은 추정 학습 연산으로 높은 정확도를 기록해 Pareto frontier에 위치했습니다.

기술 상세

μP 적용은 파라미터의 폭 확장 방식에 따라 초기화와 업데이트 크기를 조정하는 데서 출발합니다. 하나의 무한 확장 차원을 가진 Embedding, bias, expert FC2는 vector-like로 분류하고 μP initialization을 적용합니다. 두 개의 무한 확장 차원을 가진 FFN, Attention, router, expert FC1은 matrix-like로 분류하며 μP initialization과 함께 fan-in 비율에 따른 learning rate scaling을 적용합니다.

이 논문은 MLA의 low-rank query와 key-value projection dimension을 고정하고 attention head dimension도 고정합니다. hidden dimension이 커질 때 head 수를 비례적으로 늘리며, total expert 수도 함께 늘리지만 token당 active expert 수는 고정합니다. 그 결과 total scale, active scale, sparsity가 독립 축이 아니라 하나의 결합된 확장 경로가 되며, 개별 Expert의 fan-in과 fan-out 변화는 폭 scaling 규칙 안에서 처리됩니다.

학습률 탐색은 WSD scheduler의 장기 설정과 연결되지만 proxy 단계에서는 decay를 조기에 넣지 않습니다. 조기 decay는 각 실행에서 마지막 checkpoint 하나만 유효하게 만들고 loss 추정에 편향을 줄 수 있으므로, stable phase의 θ를 EMA로 합쳐 여러 시점의 평가값을 만듭니다. EMA는 현재 가중치와 이전 EMA 가중치를 α = 0.6 비율로 결합하며, 약 2B-token마다 갱신한 결과를 10B-token 간격으로 평가합니다.

각 token budget에서 후보 learning rate별 validation loss를 log learning rate에 대한 2차 곡선으로 맞춥니다. 입력은 learning rate와 고정 validation batch의 loss이고, 처리 과정은 log 변환 후 곡선 적합과 꼭짓점 계산이며, 출력은 해당 token budget의 최적 learning rate입니다. 이후 최적값의 log와 token budget의 log 사이에 선형 회귀를 적용해 10T-token 같은 미관측 horizon으로 외삽하며, 255B 이후 자료에서 R² = 0.95를 얻었습니다.

실험은 NVIDIA H200 GPU와 내부 Megatron-LM fork에서 Muon optimizer를 사용해 수행했습니다. 155B total, 17B active 모델의 학습 연산은 active parameters를 N, token 수를 D로 두고 6ND로 추정했으며, Stage 1 데이터 혼합은 6T tokens에서 도메인별 비율을 변경했습니다. 최종 모델은 MMLU, MMLU-Pro, BBH, Global-MMLU, MATH, GSM8K, MBPP, HumanEval으로 평가됐고, 동일한 평가 harness로 비교 모델을 측정했습니다.

한계점

저자들은 이 연구가 MLA와 Muon optimizer를 사용하는 MoE 구조에 초점을 맞췄다고 명시합니다. 다른 MoE 구조와 optimizer에서도 같은 전이 규칙이 유지되는지는 추가 실험이 필요합니다. 따라서 제안된 절차의 적용 범위는 현재 실험 조건에 한정됩니다.

목표 learning rate 3.85 × 10⁻⁴가 정말 최적인지 확인하는 전체 규모 sweep은 계산 비용 때문에 수행하지 못했습니다. 안정적인 loss 궤적과 benchmark 결과가 효과성을 뒷받침하지만, 목표 모델에서 여러 learning rate를 직접 비교한 확정적 증거는 아닙니다. 이 점은 예측값의 최적성을 판단할 때 고려해야 합니다.

확장 실험은 width와 sparsity를 함께 키우는 경로를 사용했습니다. 따라서 성능과 learning rate 전이에서 sparsity 자체가 만든 효과와 width 변화가 만든 효과를 분리하지 못했습니다. sparsity 축만 독립적으로 확장하는 통제 실험이 남은 과제입니다.

저자들은 top-k routing으로 Expert마다 할당되는 token 수가 달라질 수 있어 Expert별 effective batch size와 gradient noise scale도 달라질 가능성을 향후 과제로 남겼습니다. Expert별 learning rate를 적용하려면 학습 중 routing 변화와 데이터 혼합에 따른 Expert별 batch 차이를 함께 처리해야 합니다. 이를 대규모 학습에 효율적으로 구현하는 방법은 아직 다뤄지지 않았습니다.

실무 활용

대규모 MoE 사전학습에서 모델 폭과 token budget을 모두 바꿔가며 learning rate를 반복 탐색하는 비용을 줄이는 절차입니다. 작은 proxy를 먼저 학습하고 μP로 폭과 Expert 수의 변화에 전이한 뒤, EMA checkpoint와 log-log 회귀로 긴 학습 구간의 설정을 정합니다.

100B 이상 MoE를 사전학습할 때 전체 모델을 여러 learning rate로 반복 실행하지 않고 proxy의 최적값을 확장할 수 있습니다. hidden dimension과 total expert 수를 목표 모델에 맞춰 키우되 active expert 수와 intermediate dimension을 고정하는 설계에 적합합니다. 목표 학습 전에 loss 곡선과 validation 결과로 전이된 설정의 안정성을 점검할 수 있습니다.
수조 개 token을 사용하는 장기 사전학습에서 직접적인 decay 실험을 반복하지 않고 EMA checkpoint를 이용해 여러 token budget의 loss를 한 번에 얻을 수 있습니다. 각 token budget의 최적 learning rate를 산출한 뒤 log-log 회귀로 목표 horizon의 값을 계산합니다. 논문 사례에서는 255B 이후의 안정화된 proxy 자료로 10T-token learning rate를 추정했습니다.

코드 공개 여부: 비공개

키워드

Mixture-of-Experts(MoE)Maximal Update Parameterization(μP)Multi-head Latent Attention(MLA)Muon optimizerhyperparameter transferscaling lawlearning rate extrapolation

용어 해설

Mixture-of-Experts: — 여러 Expert 중 일부만 토큰마다 선택해 계산하는 구조입니다. 전체 파라미터는 크게 유지하면서 토큰당 활성 파라미터와 추론 연산량을 제한합니다. Expert 수를 늘리면 계산량을 비례적으로 늘리지 않고 모델 용량을 확장할 수 있습니다.
Maximal Update Parameterization: — 모델 폭이 커져도 파라미터 업데이트의 크기가 일관되도록 초기화와 학습률을 조정하는 방법입니다. 작은 Proxy 모델에서 찾은 최적 학습률을 더 큰 모델로 옮길 수 있게 만듭니다. 이 논문에서는 MoE의 폭과 Expert 수를 함께 확장하는 데 사용합니다.
Multi-head Latent Attention: — Attention의 Key-Value 정보를 저차원 잠재 공간으로 압축하는 구조입니다. 입력의 Key-Value 상태를 작은 표현으로 저장한 뒤 필요한 Attention 계산에 사용해 KV cache 메모리를 줄입니다. 이 논문에서는 모든 폭 확장 MoE 모델에 적용합니다.
Muon optimizer: — 신경망 가중치를 갱신하는 최적화 알고리즘으로, 논문에서는 기존 AdamW와 비교되는 학습 도구로 사용합니다. 각 학습 단계에서 gradient를 바탕으로 파라미터를 업데이트해 수렴을 돕습니다. MoE의 μP 기반 학습률 전이 실험에 사용됩니다.
scaling law: — 모델 크기나 학습 토큰 수의 변화와 성능 또는 최적 설정 사이의 관계를 수식으로 근사하는 방법입니다. 이 논문은 토큰 예산과 최적 learning rate의 관계를 log-log 선형 회귀로 적합합니다. 짧은 학습 구간의 결과로 10T-token 학습 설정을 추정하는 데 쓰입니다.

대규모 MoE 학습률을 단계적으로 확장하는 방법

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

314B MoE 언어 모델 Motif 3의 구조와 학습 시스템

MoE 학습용 토폴로지 인식 Expert Replica 배치

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

314B MoE 언어 모델 Motif 3의 구조와 학습 시스템

MoE 학습용 토폴로지 인식 Expert Replica 배치