EMO, 문서 단위 MoE 사전학습

AllenAI가 EMO를 공개했다. 1B active / 14B total parameter 규모의 MoE를 1T tokens으로 사전학습했고, 8개 활성 expert / 128개 total expert 구조를 썼다.

문서 경계를 약한 감독 신호로 사용해 같은 문서의 토큰이 동일한 expert pool 안에서만 라우팅되도록 만들었다. 표준 MoE처럼 토큰이 표면적 패턴으로 흩어지지 않고, 문서 단위의 의미 모듈이 형성되게 한 설계다.

• 글로벌 load balancing으로 expert 쏠림을 줄이고 문서 내부 라우팅 일관성을 유지했다.
• 문서별 pool 크기는 고정하지 않고 랜덤 샘플링해 다양한 subset 크기에 대응하게 했다.
• EMO의 라우팅 클러스터는 Health, News Reporting, US Politics, Film & Music 같은 도메인으로 묶였고, 표준 MoE는 Prepositions, Proper Names, Copula Verbs, Definite Articles 같은 표면 특징에 쏠렸다.

전체 벤치마크에서는 표준 MoE와 비슷한 성능을 유지했고, 모든 expert를 함께 쓰면 일반 목적 모델 성능도 유지했다. 하지만 expert를 줄여도 성능 저하가 작아 선택적 사용에 유리했다.

• 25% expert subset에서는 평균 약 1%p 하락했다.
• 12.5% subset에서는 평균 약 3%p 하락에 그쳤다.
• 같은 조건의 표준 MoE는 subset이 줄수록 급격히 무너져 작은 subset에서는 거의 랜덤 수준까지 떨어졌다.
• 소량의 검증 데이터나 1개 예시만으로도 작업별 모듈 선택이 가능했고, Easy-EP와도 결합됐다.

함께 공개된 것은 EMO 전체 모델, 같은 데이터로 학습한 standard-MoE baseline, 그리고 training code다.