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 및 후처리 학습 전략을 대규모 continual pretraining에서 평가하고, pruning 초기화가 scratch 대비 우수한 초기 성능을 제공함을 보인다. 또한 다양한 expert compression 방법 간 차이가 크지 않음을 확인한다.

부분 보존 기반 expert merging 전략

target expert 수를 half로 절반 보존하고, 남은 절반은 merging하여 성능 저하를 완화하는 partial-preservation merging을 제안하고, 주요 벤치마크에서 일관된 개선을 보였다.

결합적 학습 목표의 설계

NTP KD와 LM Loss의 결합이 pure KD보다 지식집중형 벤치마크에서 우수한 회복을 보이고, MTP KD가 추가적으로 일관된 이득을 준다.

진보적 pruning 및 distillation의 이점

깊이 우선/너비 우선/공동(progressive) 세 가지 스케줄 중 어떤 경로가 최적의 트레이닝 곡선을 제공하는지 분석하고, progressive pruning이 one-shot 대비 더 나은 최적화 경로를 제공함을 제시한다.

실용적 압축: 23A2B로의 약 4× 축소

Qwen3-Next-80A3B를 23A2B로 압축해 광범위 벤치마크에서 경쟁력 있는 성능을 달성하고, 대규모 MoE 압축의 compute 효율성에 관한 실용 가이드를 제시한다.

핵심 아이디어 이해하기

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의 이점을 강조한다.

방법론

전체 접근: 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에서 가장 비슷한 파트너와 합쳐 최종 Ẽ를 구성한다(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 등 다수의 태스크를 포함한다.

실무 활용

구조적 pruning과 KD를 결합한 MoE 압축을 통해, pretraining 규모에서의 파라미터 효율성과 추론 효율을 동시에 향상할 수 있다. Progressive pruning 전략은 one-shot 대비 더 안정적인 지식 이동을 가능하게 한다.

제약된 컴퓨트 리소스에서의 대형 MoE 모델 배포
모델 서버의 메모리 제약이 있는 환경에서의 파생 모델 생성
추론 속도 향상을 위한 MoE 축소 및 지식 증류 적용
지식 중심 태스크에서의 파인튜닝 효율성 개선

코드 공개 여부: 비공개

키워드

structured pruningknowledge distillationmixture-of-expertspretrainingprogressive pruningpartial-preservation mergingmulti-token prediction distillation

용어 해설

MoE(혼합-전문가)(MoE (Mixture-of-Experts)): — MoE는 입력에 따라 다수의 Expert 중 일부를 선택적으로 활성화해 계산 자원을 절약하고, 모델 규모 확장을 가능하게 하는 아키텍처이다. 각 Expert는 독립적인 파라미터를 가지며, 라우터는 토큰별로 활성화할 Expert들을 결정한다. 이와 같은 구조는 파라미터 효율성과 연산 집중화의 균형에 의존한다.
가지치기(pruning): — 네트워크의 일부 모듈(예: Layer, Hidden Dimension, Expert)을 제거해 파라미터 수를 줄이는 방법으로, 구조적 pruning은 실행 중인 아키텍처에서 직접 속도 향상과 메모리 감소를 유도한다.
지식 증류(knowledge distillation): — 교사 모델의 확률 분포를 학생 모델에 전달해 학습을 돕는 방법으로, KD는 프루닝 후 손실 보완에 자주 활용된다.
멀티토큰 예측(Multi-Token Prediction): — 다음 토큰 하나가 아니라 다수의 미래 토큰을 예측하도록 하는 distillation 모듈로, backbone 학습 및 speculative decoding의 효율성을 높인다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

Top‑100 캐시와 Fused Chunked KL로 VRAM 비용을 크게 낮춘 지식 증류

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 기사

Top‑100 캐시와 Fused Chunked KL로 VRAM 비용을 크게 낮춘 지식 증류

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