기술 요약: UMoE (도메인 특화 학습을 위한 모든 전문가의 잠재력 해방)

1. 문제 정의

혼합 전문가(Mixture-of-Experts, MoE) 모델은 추론 비용을 비례적으로 증가시키지 않고도 용량을 확장할 수 있는 능력 덕분에 프런티어 대규모 언어 모델(LLM)의 표준 아키텍처가 되었습니다. 그러나 도메인 특화 사후 학습(예: 지도 미세 조정 또는 SFT) 과정에서 결정적인 비효율성이 발생합니다.

사전 학습된 MoE 모델은 혼합 도메인 사전 학습에 의해 형성된 전문가 풀을 상속받습니다. 특정 타겟 도메인 하에서 이 풀은 이질적입니다. 즉, 상당 부분의 전문가 집단이 낮은 "도메인 정렬 살리언스(domain-aligned saliency)"를 보이며, 이는 이들이 적지 않은 라우팅을 받음에도 불구하고 타겟 작업에 거의 기여하지 못함을 의미합니다. 표준 SFT는 이러한 상속된 전문가들을 제자리에서 업데이트하지만, 전문가 슬롯을 재할당하지는 않습니다. 결과적으로, 고정된 모델 용량의 일부가 타겟 도메인과 제대로 정렬되지 못한 채 남아 있어 미세 조정 과정의 효과를 제한하게 됩니다.

핵리 질문은 다음과 같습니다: 모델의 아키텍처나 추론 예산을 변경하지 않고도, SFT 전에 전문가 풀을 더 도메인에 정렬된 초기화 상태로 재구성할 수 있는가? 즉, 더 많은 전문가 슬롯이 타겟 도메인 역량을 학습하는 데 유용하게 쓰이도록 할 수 있는가?

2. 방법론: UMoE

저자들은 SFT 이전에 전문가 풀을 재구성하는 예산 보존형 파이프라인인 **UMoE (Unlocking MoE Capacity)**를 제안합니다. 이 방법은 세 가지 뚜렷한 단계로 작동하며, 원래의 전문가 수($E$), 활성 전문가 수($K$), 파라미터 수 및 추론 비용을 그대로 유지합니다.

1단계: 도메인 정렬 전문가 프루닝 (Domain-Aligned Expert Pruning)

모델은 먼저 타겟 도메인에 가장 관련성이 낮은 전문가들을 식별하고 제거합니다.

• 교정(Calibration): 타겟 도메인 데이터의 하위 집합($C \subset D$)을 사용하여 전문가를 점수화합니다.
• 살리언스 지표: 저자들은 전문가의 활성화 강도를 라우터 게이트 가중치와 결합한 REAP(Lasby et al., 2026) 살리언스를 활용합니다. 레이어 $l$에서의 전문가 $e$에 대한 살리언스 $S(l, e)$는 전문가가 선택된 토큰들에서의 게이트 가중치 적용 출력 노름(output norm)의 조건부 평균입니다.
• 조치: 이 살리언스 점수를 기준으로 (레이어당) 하위 50%의 전문가를 프루닝합니다.

2단계: 예산 보존형 확장 (Budget-Preserving Expansion)

원래의 모델 크기를 복구하고 고정된 추론 예산을 유지하기 위해, 프루닝된 슬롯을 남은 "생존" 전문가들로부터 다시 생성합니다.

• 초기화: 각 해제된 슬롯은 라운드 로빈 방식으로 유지된 "부모" 전문가의 가중치(MLP 및 라우터 행)를 복사하여 초기화됩니다.
• 섭동(Perturbation): 대칭성을 깨고 전문화를 유도하기 위해, 부모와 새로 생성된 전문가 모두를 독립적으로 섭동시킵니다. 파라미터 텐서 $W$에 대해, 섭동된 가중치 $W'$는 다음과 같이 계산됩니다:
  $$W' = W + \sigma s(W) \epsilon$$
  여기서 $\epsilon \sim \mathcal{N}(0, I)$이고, $s(W)$는 파라미터 그룹의 경험적 표준 편차이며, $\sigma = 0.05$입니다.
• 설계 원칙: 결정적으로, 섭동은 부모와 새로 생성된 전문가 모두에게 대칭적으로 적용됩니다. 저자들의 절제 연구(Ablation study)에 따르면, 한쪽 방향 섭동(새로 생성된 전문가에만 적용)은 라우팅 비대칭과 성능 저하를 초nd합니다. 보완적인 라우팅 및 기능적 비대칭을 제공하기 위해 라우터 행과 MLP 가중치 모두에 섭동을 적용합니다.

3단계: 도메인 SFT

이제 도메인 정렬된 전문가 풀로 초기화된 재구성된 모델은 타겟 도메인 데이터에 대해 표준 전체 파라미터 지도 미세 조정(Full-parameter SFT)을 거칩니다. 별도의 도메인별 하이퍼파라미터 튜닝은 필요하지 않으며, 모든 도메인에 동일한 레시피가 적용됩니다.

3. 주요 기여

1. 전문가 풀 재구성: 사전 학습된 MoE를 미세 조정하기 전에 특정 타겟 도메인에 맞게 전문가 구성을 재정렬하여, "중복된" 용량을 "도메인" 용량으로 전환하는 새로운 방법을 제시했습니다.
2. 예산 보존: 이 방법은 원래의 아키텍처, 파라미터 수, 추론 비용을 엄격히 유지하므로, 표준 SFT 파이프라인의 드롭인 교체(drop-in replacement)가 가능합니다.
3. 대칭 섭동 전략: 부모 전문가와 새로 생성된 전문가 모두를 대칭적으로 섭동시키는 것이 균형 잡힌 라우팅 역학을 유지하고 새로 생성된 절반의 유용성 붕괴를 방지하는 데 필수적임을 확인했습니다.
4. 실증적 검증: 두 가지 MoE 아키텍처(Qwen3-30B-A3B 및 Qwen3.5-35B-A3B), 다섯 가지 다양한 도메인(수학, 코드, 과학, 도구 사용, 에이전트 코딩) 및 12개 벤치마크에 대한 포괄적인 평가를 수행했습니다.

4. 실험 결과

UMoE는 평가된 모든 설정에서 Direct SFT(재구성 없이 사전 학습된 모델을 직접 미세 조정)보다 일관되게 우수한 성능을 보였습니다.

• 수학 도메인:
  • Qwen3-30B-A3B (OpenMathReasoning 데이터)에서, UMoE는 6개 경진대회 벤치마크의 평균 정확도를 +3.40 포인트 향상시켰으며, AIME 2024에서는 최대 +6.56 포인트의 향상을 보였습니다.
  • 이미 Qwen3-30B-A3B-Thinking 모델(82.81 vs 81.06)을 능가하는 강력한 자체 내부 수학 코퍼스에서도, UMoE는 평균을 84.17(+1.36 포인트)로 높여 강력한 베이스라인에 대해서도 견고함을 입증했습니다.
  • Qwen3.5-35B-A3B에서, UMoE는 평균 +1.67 포인트를 개선했습니다.
• 교차 도메인 성능:
  • 코드: LiveCodeBench v5/v6에서 +2.73 / +0.85 포인트 향상.
  • 과학: GPQA-Diamond에서 +3.10 포인트 향상.
  • 도구 사용: $\tau$-bench / $\tau^2$-bench에서 +1.20 / +1.40 포인트 향상.
  • 에이전트 코딩: SWE-bench Verified에서 상당한 +6.0 포인트 향상 (pass@1이 16.2%에서 22.2%로 상승).
• 데이터 스케일링: SFT 데이터 크기(0.5B에서 4.1B 토큰까지)를 변화시킨 실험에서, UMoE의 성능 이점은 모든 데이터 규모에서 지속되고 양의 값을 유지함을 보여주었습니다.

5. 메커니즘 분석

본 논문은 UMoE가 작동하는 이유에 대해 세 가지 핵심 통찰을 제공합니다:

1. Direct SFT의 중복성: Direct SFT 모델에 대한 분석 결과, 이 모델은 라우팅된 전문가 연산의 상당 부분(약 42%)을 낮은 살리언스를 가진 전문가 집단에 할당하고 있음을 발견했습니다. Direct SFT 모델에서 이 집단을 사후 프루닝하더라도 성능 저하가 미미히(-0.12 포인트)했는데, 이는 이 용량이 기능적으로 중복되었음을 확인시켜 줍니다.
2. UMoE의 결정적 용량: 반면, 학습된 UMoE 모델을 64개의 전문가로 프루닝하면 상당한 성능 저하(-5.00 포인트)가 발생합니다. 이는 UMoE가 이전의 중복된 전문가 풀의 절반을 다운스트림 성능에 결정적인 용량으로 성공적으로 전환했음을 나타냅니다.
3. 최적화 역학: UMoE는 재구성으로 인해 Direct SFT보다 약간 높은 학습 손실(loss)에서 시작하지만, 더 낮은 최종 학습 손실(0.396 vs 0.404)로 수렴합니다. 이는 더 효율적인 최적화 궤적을 시사합니다.
4. 오류 수정: 사례 연구(예: AIME 2026 Problem 25)는 UMoE가 Direct SFT가 전파하는 특정 추론 오류(예: 각의 이등분선 정리에서의 분수 약분 오류)를 수정하여 최종적으로 정답을 도출하는 것을 보여줍니다.

6. 의의 및 주장

본 논문은 UMoE가 MoE 모델의 도메인 적응을 위한 단순하고, 예산을 보존하며, 보편적으로 적용 가능한 접근 방식이라고 주장합니다. 그 의의는 다음과 같습니다:

• 효율성: 추가적인 파라미터나 추론 비용 없이 기존 모델 내의 잠재된 용량을 끌어냅니다.
• 일반성: 단일 고정 레시피를 사용하여 다양한 도메인(추론, 코딩, 에이전트 작업)과 서로 다른 모델 아키텍처 전반에서 작동합니다.
• 견고성: 베이스라인 SFT가 이미 강력한 경우와 훈련 데이터 규모가 커지는 경우에도 성능 향상을 제공합니다.

저자들은 상속된 전문가 용량을 재구성함으로써 모델이 고정된 컴퓨팅 예산을 타겟 도메인의 특정 요구 사항에 더 잘 맞출 수 있으며, 이를 통해 "죽은" 전문가를 도메인 특화 성능에 대한 "활발한" 기여자로 변모시킬 수 있다고 결론짓습니다.