SoftMoE: LLM을 위한 Mixture-of-Experts의 소프트 미분 가능 라우팅 (Soft Differentiable Routing)
요약
SoftMoE는 기존 MoE의 미분 불가능한 이산적 top-k 라우팅 문제를 해결하기 위해 소프트 미분 가능 라우팅 방식을 제안합니다. 이를 통해 레이어별 전문가 할당을 최적화하며, 적은 전문가 활성화로도 기존 희소 MoE와 대등하거나 뛰어난 성능을 보여줍니다.
핵심 포인트
- 이산적 top-k 라우팅을 소프트 top-k LapSum 완화 방식으로 대체
- 경사 기반 최적화를 통해 전문가 라우팅 학습 가능
- 전역 예산 제약을 통해 레이어별 최적의 전문가 용량 할당
- 후반부 레이어에서 더 많은 전문가를 활성화하는 특성 발견
- 자기회귀 모델링과 호환되며 효율적인 파라미터 확장 가능
희소 Mixture-of-Experts (MoE) 아키텍처는 top-$k$ 라우팅을 통해 전문가(expert)의 작은 하위 집합만을 활성화함으로써, 고정된 추론 예산 내에서 LLM 파라미터를 확장할 수 있게 해줍니다. 이는 인과성(causality)을 보존하고 자기회귀(autoregressive) 언어 모델에 적합하지만, 이산적인 (discrete) top-$k$ 연산자는 미분 불가능하여 입력당 활성화되는 전문가의 수를 고정하게 만들고 결과적으로 계산의 비효율적인 사용을 초래합니다. 우리는 이산적 라우팅을 절단된 소프트 top-$k$ LapSum 완화(truncated soft top-$k$ LapSum relaxation)로 대체하여, 전문가 라우팅의 경사 기반 최적화(gradient-based optimization)를 가능하게 하는 SoftMoE를 제안합니다. 나아가 우리는 레이어당 평균 활성 전문가 수를 매개변수화하고 전역 예산 제약(global budget constraint)을 부과하여, 모델이 레이어 전반에 걸쳐 전문가 용량을 할당하는 방법을 학습할 수 있도록 합니다. SoftMoE는 자기회귀 모델링과 완전히 호환되며, 훨씬 더 적은 수의 전문가를 활성화하면서도 언어 모델링 및 다운스트림 태스크에서 희소 MoE와 대등하거나 더 나은 성능을 달성합니다. 특히, 학습된 할당은 매우 불균일하며, 후반부 레이어일수록 더 많은 전문가를 활성화합니다. 소스 코드는 공개적으로 사용 가능합니다$^\dagger$.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기