SlimQwen: 대형 MoE 모델 pre-training에서의 structured pruning 및 knowledge distillation 탐구

대형 MoE LLM의 사전학습은 비용이 큰 문제다. 본 연구는 depth/width/experts를 모두 포함하는 구조적 pruning과 추후 학습 전략으로 프리training에서의 성능 회복을 분석하고, progressively pruning이 one-shot보다 더 안정적이고 효과적임을 보여준다. 또한 NTP KD와 LM loss의 결합과 MTP KD가 지식 집중형 태스크에서 이점을 준다.

MoE 아키텍처는 라우터가 선택한 전문가들만 활성화하는 구조로, 전체 파라미터의 대부분이 활성화되지 않는 특성을 가진다. 본 연구는 이 특성을 고려하여 depth/width/experts의 pruning을 수행하고, pruning 후 학습을 통해 잃은 지식을 KD 및 MTP KD로 재전이한다. 먼저 prune된 teacher MoE에서 target architectures로의 점진적 전이(두 단계의 progressive pruning)로 정보 손실을 최소화하고, partial-preservation merging으로 핵심 지식을 보존하면서도 불필요한 전문가를 효과적으로 합친다. 또한 언어 모델 손실과 지식 증류를 결합한 학습( LM loss + LKD)과 multi-token prediction distillation(MTP KD)을 도입해 지식 전달의 품질을 높이고, speculative decoding의 효율도 개선한다. 실험은 120B 및 400B 토큰의 토큰 수를 고정한 조건에서 다양한 벤치마크를 통해 프리training의 초기화, 전문가 압축 방법, 학습 전략의 효과를 확인한다. 마지막으로, 80A3B 크기의 교사 모델에서 23A2B 모델로의 압축이 실질적 속도/메모리 이점을 제공하는지 확인하고, progressive trajectory의 이점을 강조한다.

1. 전체 접근: MoE 기반 모델의 depth, width, and expert를 대상으로 구조적 pruning을 수행하고, partial-preservation merging으로 2단계 압축을 진행한다. 2) 구조적 pruning의 구체: Depth pruning은 Lkeep = {1, ..., L-N}, L~ = L-N으로 마지막 N층 제거; Width pruning은 RMSNorm 활성화의 평균값을 이용한 중요도 점수 I(k)norm를 계산하고 dt_dim을 보존한다. 3) Expert compression은 I_Freq, I_Soft, I_REAP를 이용한 전문가 중요도 산정과, 상위 전문가의 보존 및 남은 전문가의 가장 유사한 보존 대상과의 머징으로 구성한다. Partial-preservation merging은 상위 N~ 전문가의 절반을 보존하고, 나머지 절반은 merge bases에서 가장 비슷한 파트너와 합쳐 최종 Ẽ를 구성한다(Algorithm 1). 4) Distillation pretraining: MTP 모듈을 도입해 k번째 미래 토큰의 확률 분포를 예측하고, LM loss와 LKD를 사용하며, LMTP-LM, LMTP-KD를 포함하는 총 손실 L = (1 − λ) LLM + λ LKD + β((1 − λ)LMTP-LM + λLMTP-KD)로 학습한다. 5) Progressive pruning & distillation: 두 단계의 스케줄로, 40B tokens에서 초기화 후 남은 토큰으로 최종 구성을 완료하는 방식으로 학습한다. 6) 실험 설계: 80A3B 교사 모델에서 23A2B 학생 모델로의 압축 및 Continual pretraining(120B/400B 토큰) 수행, MMLU, MMLU-Pro, MMLU-Redux, BBH, GSM-8K, EvalPlus, C-Eval, CMMLU 등 벤치마크 평가를 포함한다.

Q1. 프리트레이닝에서 prune된 MoE가 scratch 초기화보다 우수하다: 120B 토큰 조건에서 Pruned + KD Loss의 평균값은 73.45로 Random Init + KD Loss의 61.66보다 +11.79 포인트 높다. Pruned 모델은 교사 성능의 86.5%를 회복하며(73.45 vs 82.68), 모델 규모는 3.4× 감소했다. 학습 곡선은 LM Loss가 더 낮게 수렴한다. Q2. one-shot expert compression의 차이는 작고, partial-preservation merging이 주요 벤치마크에서 일관된 개선을 보인다. Q3. 23A2B로의 압축에서 NTP KD와 LM Loss의 결합은 지식 집중 태스크에서 우수하며, MTP KD를 추가하면 추가 이득이 있다. 표 3에서 NTP KD + LM Loss + MTP Loss + MTP KD 구성이 MMLU 75.67, MMLU-Pro 51.19, MMULU-Redux 74.37, BBH 72.29, GSM-8K 83.17, EvalPlus 69.30, C-Eval 80.67, CMMLU 80.95를 달성한다. 표 4는 전-훈련 단계에서 MTP KD가 다중 토큰 생성의 수용률(acc4)을 증가시키며, SFT 단계에서도 비슷한 경향을 보인다. 표 5에서 progressive pruning은 One-stage보다 벤치마크 성능을 향상시키고, 특히 MMLU/ MMLU-Redux에서 큰 개선을 보인다. 특히 Depth-first SlimQwen은 MMLU 77.39, MMLU-Redux 78.01을 달성한다. A11의 효율성 분석에서 SlimQwen-23A2B는 peak memory 43.30 GB로 대조군 156.56 GB 대비 약 3.6× 감소, vLLM 백엔드에서 Decoding Throughput이 210.87 Toks/s로 큰 개선을 보인다. HF 백엔드에서도 Prefill Latency가 0.44 s로 감소하고 Decoding Throughput이 6.55 Tok/s로 증가한다.

MoE 아키텍처의 핵심 수식은 MoE(x) = sum_{e=1}^{nrouted} ze(x)Expert_e(x) + sum_{s=1}^{nshared} zs(x)Expert_s(x)이다. Router는 TopK(z(x), k)로 활성 Expert를 선택하고, Shared Gate는 σ(xwsh)로 공유 Expert를 선택한다. Depth pruning은 Lkeep = {1,...,L−N}, L~ = L−N으로 마지막 N층을 제거한다. Width pruning은 Mean(RMSNorm(X))를 이용해 각 hidden dimension의 중요도 I^(norm)(k) = sum{i=0}^L Mean(RMSNorm(X))i(:,k)을 계산해 dt 차원의 상위 값을 남긴다. Expert pruning/merging은 I_Freq, I_Soft, I_REAP를 통해 중요도를 산출하고, top-N 비율의 Expert를 보존, 남은 Expert를 Sbase의 가장 유사한 Partner에 머지한다(식 (8)). Partial-preservation merging은 보존 Expert의 수를 ⌊N' / 2⌋로 정하고, 나머지 절반은 merge bases에서 선택된 Expert와 머지한다. MTP Distillation은 Embedding 및 OutHead를 공유하는 모듈로, i번째 입력 토큰 ti에 대해 depth k에서 ti+k 토큰의 확률 분포 p^k_i+k를 예측한다. LMTP-LM = (1/D) sum{k=1}^D -(1/(T−k) ∑{i} log p^k_i+k[t_i+k]); LMTP-KD = (1/D) sum{k=1}^D ∑_{v} q_i+k[v] log p^k_i+k[v]. 총 손실은 L = (1−λ)L_LM + λL_KD + β((1−λ)LMTP-LM + λLMTP-KD)이다. Progressive pruning과 distillation은 Two-stage로 40B tokens에서 초기화하고, 남은 360B tokens에서 재 pruning 및 재학습을 수행한다. 실험 데이터는 120B 및 400B 토큰 설정에서 교사-학생 구조로 평가되며, MMLU 계열 벤치마크, BBH, GSM8K, CMMLU 등 다수의 태스크를 포함한다.