OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models

To help the open-source community have a better understanding of Mixture-of-Experts (MoE) based large language models (LLMs), we train and release OpenMoE, a series of fully open-sourced and reproducible decoder-only MoE LLMs, ranging from 650M to 34B parameters and trained on up to over 1T tokens. Our investigation confirms that MoE-based LLMs can offer a more favorable cost-effectiveness trade-off than dense LLMs, highlighting the potential effectiveness for future LLM development. One more important contribution of this study is an in-depth analysis of the routing mechanisms within our OpenMoE models, leading to three significant findings: Context-Independent Specialization, Early Routing Learning, and Drop-towards-the-End. We discovered that routing decisions in MoE models are predominantly based on token IDs, with minimal context relevance. The token-to-expert assignments are determined early in the pre-training phase and remain largely unchanged. This imperfect routing can result in performance degradation, particularly in sequential tasks like multi-turn conversations, where tokens appearing later in a sequence are more likely to be dropped. Finally, we rethink our design based on the above-mentioned observations and analysis. To facilitate future MoE LLM development, we propose potential strategies for mitigating the issues we found and further improving off-the-shelf MoE LLM designs.

한 줄 요약

OpenMoE는 650M~34B 파라미터 규모의 오픈소스 MoE 언어 모델로, 라우팅 메커니즘 분석과 개선 전략 제시를 통해 MoE 연구를 촉진한다.

핵심 기여도

OpenMoE-8B/32E 모델이 OpenLLaMA-3B, TinyLLaMA-1.1B보다 유사한 성능을 내며 더 낮은 학습 비용을 기록.
MoE 라우팅 분석에서 "Context-Independent Specialization", "Early Routing Learning", "Drop-towards-the-End" 세 가지 현상을 발견.
UL2 학습 목적과 코드 데이터(52.25%)를 활용한 학습 전략을 제시.
OpenMoE-8B/32E 모델의 5개 중간 체크포인트를 공개하여 연구 재현성을 높임.

핵심 아이디어

OpenMoE는 Mixture-of-Experts(MoE) 기반 언어 모델의 개방성과 재현성을 높이기 위해 설계되었다. 기존 MoE 연구는 주로 비공개 데이터와 모델로 이루어져 있었으나, OpenMoE는 1조 토큰 이상의 데이터로 학습한 650M~34B 규모의 모델을 완전히 오픈소스로 공개함으로써 커뮤니티의 접근성을 확대했다. 특히, MoE의 핵심 구성 요소인 라우팅 메커니즘을 분석하여, 토큰 ID 기반의 라우팅이 문맥과 무관하다는 점을 밝혔다. 이는 토큰이 학습 초기에 고정된 전문가에 할당되고, 이후 변경되지 않는다는 "Early Routing Learning" 현상을 유발한다. 또한, 시퀀스 끝에 위치한 토큰이 전문가 용량 초과로 삭제되는 "Drop-towards-the-End" 문제를 발견하여, 이는 대화형 작업에서 성능 저하를 초래할 수 있음을 분석했다. 이러한 발견은 MoE 모델의 설계 및 학습 전략 개선에 중요한 통찰을 제공한다.

기술적 접근법

**모델 아키텍처**: Decoder-only ST-MoE 구조를 사용. 각 MoE 레이어는 최대 32개의 전문가(Fully Connected Layer)로 구성.
**토크나이저**: umT5 토크나이저(256K vocab)를 사용하여 다국어 및 저자원 언어 지원을 강화.
**라우팅 메커니즘**: Top-2 선택 방식을 적용. 입력 토큰은 2개의 전문가에 할당되며, 이는 `softmax(f(x))`를 통해 결정됨.
**Residual MoE**: MoE 레이어와 고정 FFN 레이어를 병렬로 사용. OpenMoE-8B/32E는 6개 레이어마다 MoE를 적용.
**손실 함수**: Cross-Entropy Loss 외에도 Load Balance Loss와 Router Z-Loss를 추가하여 라우팅 안정성과 전문가 간 부하 균형을 유지.
**학습 전략**: UL2 학습 목적을 사용하며, 코드 데이터 비중을 52.25%로 설정.

주요 결과

OpenMoE-8B/32E는 OpenLLaMA-3B, TinyLLaMA-1.1B와 유사한 성능을 보이며, 학습 비용은 더 낮음.
MT-Bench에서 OpenMoE-8B/32E-Chat은 단일 턴 대화에서 두 밀집 모델을 크게 앞선다.
코드 데이터에서 토큰 예측 정확도가 80% 이상 달성.
라우팅 분석에서 토큰이 문맥과 무관하게 전문가에 할당되는 현상이 확인됨.

의의 및 한계

OpenMoE는 MoE 기반 언어 모델의 재현성과 개방성을 확보한 첫 사례로, 커뮤니티 연구에 중요한 기초 자료를 제공한다. 특히, 라우팅 메커니즘 분석은 MoE 모델의 내부 동작을 이해하고 개선하는 데 기여한다. 그러나, 모델 성능이 기대보다 높지 않으며, 일부 설계 선택(예: 과도한 데이터 혼합)이 부적절했다는 점이 지적된다. 또한, 대규모 모델에서의 라우팅 분석 결과의 일반화 가능성은 추가 연구가 필요하다.

실용적 활용

OpenMoE는 저비용으로 대규모 언어 모델을 학습하고, 라우팅 메커니즘을 분석하려는 연구자들에게 유용하다. 특히, 코드 데이터를 다루는 산업(예: 소프트웨어 개발, 자동화 테스트)에서 활용 가능하며, 대화형 AI 개발에도 기여할 수 있다.