디리치렛 라우팅으로 구현한 차별화된 전문가 선택과 기여 분배

초록

DirMoE는 전문가 선택을 베르누이 변수, 선택된 전문가들의 기여 비율을 디리치렛 분포로 모델링한 완전 미분 가능한 라우터이다. Gumbel‑Sigmoid과 암시적 재파라미터화 기법을 이용해 선택·기여 과정을 모두 역전파 가능하게 만들고, λ라는 스파스성 조절 파라미터로 기대 활성 전문가 수와 기여 집중도를 독립적으로 제어한다. ELBO 기반 학습과 단계적 온도·농도 스케줄링을 통해 기존 Top‑k+Softmax 라우터보다 높은 전문화와 안정성을 달성한다.

상세 분석

DirMoE는 기존 MoE 라우터가 Top‑k+Softmax을 사용해 이산적인 전문가 선택과 연속적인 확률 할당을 하나의 Softmax 단계에 얽어두는 문제를 근본적으로 해소한다. 논문은 라우팅을 “spike‑and‑slab” 사전으로 정의하고, spike 부분을 베르누이(π_i)로, slab 부분을 디리치렛(α)로 분리한다. 선택 마스크 ˜z는 Gumbel‑Sigmoid(temperature τ_z)으로 샘플링해 연속적인 근사 이진값을 얻으며, τ_z를 점진적으로 감소시켜 탐색‑수렴 과정을 제어한다. 선택된 마스크에 조건부로 정의된 디리치렛 파라미터 α_i = λ·