Papers·2일 전
DOT-MoE: 밀집 모델을 MoE로 변환하는 미분 가능 최적 수송 프레임워크 — 파라미터 50% 감소, 성능 90% 유지
DOT-MoE는 사전 학습된 밀집 Transformer의 FFN을 MoE로 변환할 때 뉴런 할당을 미분 가능 최적 수송(DOT) 문제로 공식화합니다. Sinkhorn-Knopp 반복으로 전문가 용량 제약을 강제하고 Straight-Through Estimator로 이산 할당과 라우팅을 end-to-end 학습합니다. 여러 벤치마크에서 구조적 프루닝·랜덤 분할 대비 우수하며, 활성 파라미터 50% 감소에도 밀집 모델 성능의 90%를 유지합니다.
DOT-MoE는 밀집 모델을 희소 MoE로 변환하는 새로운 프레임워크로, 뉴런 할당을 최적 수송 문제로 모델링해 기존 휴리스틱 대비 일관된 성능 향상을 보입니다.
핵심 결론
- 성능 — 활성 파라미터 50% 감소 조건에서 밀집 모델 성능의 90% 유지 — 구조적 프루닝·랜덤 분할·클러스터링 기법보다 일관되게 우수.
- 벤치 — LLaMA, GPT-Neo 등 여러 아키텍처와 언어 모델링·다운스트림 태스크에서 검증.
방법
- 최적 수송 — FFN 뉴런을 전문가에 할당하는 문제를 balanced transport 문제로 보고, differentiable Sinkhorn-Knopp 반복으로 해결.
- End-to-end — Straight-Through Estimator를 이용해 이산적인 뉴런-전문가 할당과 토큰-전문가 라우팅 정책을 동시에 학습.
- 기존 휴리스틱(랜덤 분할, k-means 클러스터링)과 달리 할당이 데이터와 라우터에 의해 최적화되는 점이 핵심입니다.
한계·조건
- 변환 비용 — 사전 학습된 밀집 모델을 변환하는 과정에서 추가 학습이 필요하며, Sinkhorn 반복으로 인한 계산 오버헤드가 있습니다.
- 희소성 — 50% 활성 파라미터 감소는 고정 — 더 높은 희소성에서의 동작은 보고되지 않았습니다.
- 코드 — GitHub 공개 예정 (현재 abstract만 공개).
편집자 한 줄
MoE 변환 접근은 학습 안정성과 추론 효율 사이의 트레이드오프를 줄여주는 흐름인데, DOT-MoE는 할당 최적화를 수리적으로 깔끔하게 정식화한 점이 인상적입니다.
- #mixture-of-experts
- #optimal-transport
- #model-compression
- #efficiency
Udbhav Bamba