PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, Yang You

arXiv:2402.01739v2cs.CLcs.AIcs.DCcs.LG

TL;DR

MoE LLMs can reduce computation pressure, but the behavior and specialization of their experts remain insufficiently understood. OpenMoE trains and releases reproducible open-source MoE models while analyzing their routing, finding positive effectiveness and three routing phenomena: Context-Independent Specialization, Early Routing Learning, and Drop-towards-the-End.

  • Problem

    LLMs are computationally expensive, while the specialization and routing behavior of MoE experts remain poorly understood.

  • Method

    The paper trains and releases reproducible open-source decoder-only MoE LLMs and conducts an in-depth analysis of their routing mechanisms.

  • Results

    The study verifies MoE-based LLM effectiveness and identifies Context-Independent Specialization, Early Routing Learning, and Drop-towards-the-End.

  • Takeaways & Limitations

    The findings and proposed solutions provide insights for future developers of open-source MoE LLMs.

  • Takeaways & Limitations

    Tokens appearing later in sequences face higher drop risk when experts reach capacity, with the issue more severe in instruction-tuning data.

Abstract

from arXiv · show

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.

1 Introduction

OpenMoE addresses the computational cost of LLMs by releasing open-source MoE models while investigating routing mechanisms and advanced training strategies. The study also analyzes routing limitations and uses its findings to propose directions for future MoE development.

  • Releasing OpenMoE: OpenMoE releases open-source MoE LLMs spanning debugging, pretrained, and chat variants, including an 8B-parameter model activating about 2B parameters per token and trained on over 1T tokens.The released series includes OpenMoE-Base/16E, OpenMoE-8B/32E, and OpenMoE-8B/32E-Chat.
  • Exploring Advanced Training Strategies: OpenMoE explores training beyond conventional practice by using up to 52.25% code during early pre-training and investigating the UL2 objective.These strategies extend beyond text-dominated data and next-token prediction training.
  • Studying MoE Routing In-depth: The work studies sparse routing because MoE must assign each token to a few experts to keep computation fixed while scaling trainable parameters.The introduction identifies sparse expert selection as crucial to MoE functionality and notes that its behavior remains insufficiently studied.
  • Studying MoE Routing In-depth: Tokens later in a sequence face greater dropping risk when experts reach capacity, with the problem worsening on instruction-tuning data whose domain differs from pre-training data.Early pre-training establishes balanced token-assignment strategies that may not transfer effectively to instruction-tuning scenarios.
  • Rethinking Our Mistakes and Proposing Potential Solutions: The authors use their empirical findings and routing visualizations to reassess sub-optimal project decisions and propose potential solutions for future MoE development.They frame sharing mistakes and lessons from this early open-source effort as an important contribution alongside successful strategies.

2 Designing OpenMoE

OpenMoE combines multilingual tokenization, sparse interleaved MoE architecture, UL2 pre-training, and limited supervised fine-tuning to build scalable, instruction-capable models. Its design emphasizes stable routing, balanced expert utilization, and reduced routing overhead.

  • Tokenization: The model uses a 256K umT5 tokenizer to support low-resource languages and out-of-vocabulary tokens through multilingual vocabulary and byte fallback.The tokenizer is chosen over smaller-vocabulary alternatives and older tokenizers without byte fallback.
  • Architecture and routing: OpenMoE uses token-choice top-2 routing, with each expert implemented as an FFN rather than a complete Transformer.The gating function selects K=2 experts from E using sparse softmax routing, enabling conditional computation.
  • Architecture and routing: Residual MoE keeps one fixed FFN active for every token, while interleaving MoE layers every 4 or 6 Transformer layers reduces routing overhead.MoE is used every 4 layers in OpenMoE-Base/16E and OpenMoE 34B/32E, and every 6 layers in OpenMoE-8B/32E.
  • Training objectives: Training combines load-balance loss with router z-loss to improve expert parallelism and stabilize routing by penalizing large gating logits.The z-loss encourages smaller absolute logits and reduces round-off errors in MoE layers.
  • Training objectives: OpenMoE explores UL2, combining span corruption and prefix language modeling, with more PrefixLM and only 20% low-mask-ratio training.The configuration prioritizes PrefixLM for zero-shot and in-context learning, while low mask ratios avoid slower learning from fewer output tokens.
  • Supervised fine-tuning: OpenMoE applies supervised fine-tuning to 58K GPT-4-generated conversations from WildChat to enhance instruction following and compare MoE behavior before and after SFT.The project limits SFT data to GPT-4 instruction-response pairs because of late-stage computation constraints.

3 Training OpenMoE

OpenMoE’s training design and MoE scaling improve cost-effectiveness and performance across several benchmarks, while revealing limitations from UL2 saturation, code-heavy training, and multi-turn token dropping. The results support MoE’s scalability but indicate that training objectives and data composition require careful adjustment.

  • Ablation Study: The MoE approach, UL2 objective, increased code-data emphasis, and LLaMA tokenizer each contribute positively to zero-shot TriviaQA performance.The LLaMA-tokenizer model outperforms the umT5-tokenizer model, despite the possible drawback of a larger vocabulary.
  • Ablation Study: Code-domain training is easier to optimize: models reach higher accuracy and lower loss, with over 80% token-prediction accuracy on GitHub despite the model’s small size.The paper attributes this partly to long-tail token frequencies, including frequent newline and tab tokens.
  • Training Objective: UL2 accelerates early learning but saturates later, prompting OpenMoE-8B/32E to switch to CasualLM after 390K steps and OpenMoE-34B/32E to use UL2 for 25K steps.The later accuracy drop after 780B tokens reflects the more difficult CasualLM objective and less-easy code data, while the easy-to-hard curriculum may remain useful.
  • Benchmark Results: On 13 LM-Evaluation-Harness tasks, OpenMoE and TinyLLaMA underperform OpenLLaMA, potentially because initially high code-data sampling harms text-dominated benchmarks.The authors nevertheless consider OpenMoE’s scores acceptable and identify this data-mixture issue for further discussion.
  • Benchmark Results: OpenMoE strongly outperforms baselines in single-turn MTBench, especially on coding, but performance drops on the second turn and yields worse multi-turn results.The paper attributes this degradation probably to token dropping.

4 Analyzing OpenMoE

OpenMoE’s routing is primarily shaped by token identity and position rather than broad domain or coding-language semantics, with specialization emerging early and remaining largely context-independent. Capacity limits then create a Drop-towards-the-End effect, while language- and task-level specialization appears in finer-grained inputs.

  • Domain and Language Specialization: OpenMoE shows little expert specialization across domains or coding languages, including low-resource Assembly and Blitzmax versus Java and Python.Domain subsets are mostly uniformly distributed across experts, with only slight preferences such as E21 for code and E10 for books.
  • Domain and Language Specialization: Natural-language routing exhibits clearer specialization: Chinese favors E5 and E16, while Japanese and Korean favor E14.These findings come from the TED-Parallel-Corpus analysis across 12 languages.
  • Task and Position Specialization: MT-Bench conversation routing shows similar finer-grained specialization, especially for math data, likely because math tasks contain more special tokens.Position IDs also exhibit specialization, with consecutive positions preferring similar experts such as E10 and E19.
  • Context-Independent Specialization: Token IDs strongly determine routing: identical tokens favor only a few experts despite appearing in diverse contexts, demonstrating Context-Independent Specialization.Experts also cluster tokens with similar low-level semantics, such as “can,” “will,” and “would.”
  • Early Routing Learning: Routing preferences are almost totally overlapped across intermediate checkpoints, indicating that token-to-expert assignments are learned and fixed very early in pre-training.Supervised fine-tuning does not significantly change the Drop-towards-the-End issue, consistent with early-fixed routing behavior.
  • Drop-towards-the-End: Later tokens are more likely to be dropped when earlier tokens fill expert capacity, producing a Drop-towards-the-End effect in decoder-only MoE models.The issue arises from predefined per-expert capacity and autoregressive processing, and is observed in Mixtral and Deepseek-MoE as well as OpenMoE.

5 Rethinking OpenMoE

Section 5 distills design lessons from OpenMoE’s analysis and proposes future improvements, while noting uncertainty from expensive ablations and dependence on model size and data quality. Recommendations address code-data composition, tokenizer selection, efficient MoE architecture, and early instruction-data mixing to improve routing balance.

  • Overview: The section presents these recommendations as the main takeaways from the authors’ mistakes and analytical insights, intended to guide future practitioners.The authors explicitly acknowledge that they lack precise answers for some design choices and report their findings without reservation.
  • Code Composition: The authors suggest around 30% code, as used in OpenMoE-34B/32E, because over 50% code may harm text-task abilities.They caution that the precise conclusion may depend strongly on model size and data quality, and that large-scale ablations are expensive.
  • Tokenizer Selection: Quantitative tokenizer analysis finds umT5 much better than LLaMA on the multilingual dataset, motivating smarter tokenizer selection despite output-layer computation overhead.Table 10 compares the tokenizers using vocabulary activated and the umT5/LLaMA token-count ratio across subsets from different datasets.
  • More Efficient MoE Architecture: The authors propose removing the trainable router after warmup, parallelizing Transformer and FFN computation, and overlapping attention computation with MoE all-to-all communication.These changes target hardware utilization; parallel Transformer computation can enable overlap without performance drop when scaling up [9].
  • Load Balance and Drop-towards-the-End: Mixing instruction-following data into pre-training warm-up is proposed to improve load balance and alleviate Drop-towards-the-End, which matters on multi-turn MT-Bench.The rationale is that MoE routing behavior is learned and fixed early in pre-training.

6 Conclusion · Appendix · A Frequent Asked Questions

The work demonstrates the effectiveness and reproducibility of open MoE-based LLMs while identifying three routing phenomena and proposing solutions for future developers. It aims to improve the open-source community’s understanding of MoE models and provides point-to-point answers to potential frequently asked questions.

  • 6 Conclusion: The study verifies the effectiveness of MoE-based LLMs in the post-ChatGPT stage.
  • 6 Conclusion: The released models are fully reproducible because the authors disclose all details along with open-sourced code and data.
  • 6 Conclusion: The analysis identifies Context-independent Specialization, Early Routing Learning, and Drop-towards-the-End as important MoE routing phenomena.
  • 6 Conclusion: The authors acknowledge mistakes in their MoE design and propose possible solutions for future developers.
  • 6 Conclusion: The work is intended to help the open-source community better understand MoE models.
  • A Frequent Asked Questions: The appendix lists potentially frequently asked questions with point-to-point answers.

A.1 Why not show the token specialization of the checkpoints at the warmup stage? … A.4 Why not use AdamW optimizer and Cosine Learning Rate Schedule?

The appendix explains checkpointing, comparison, architectural, and optimization choices in OpenMoE, emphasizing reproducibility, training efficiency, and stability constraints. It also notes that routing was found to become fixed earlier than expected, while AdamW caused frequent instability.

  • A.1 Why not show the token specialization of the checkpoints at the warmup stage?: Routing was learned and fixed earlier than expected, but checkpoints were retained only every 200B tokens because storage was limited.This checkpoint cadence prevented showing token specialization during the warmup stage.
  • A.2 Why not compare with advanced open MoE models like Mixtral and DeepSeek-MoE?: OpenMoE was announced and released over four months before Mistral and more than four months before DeepSeek-MoE.
  • A.2 Why not compare with advanced open MoE models like Mixtral and DeepSeek-MoE?: Unlike models trained on in-house data, OpenMoE is fully transparent, with all details and code disclosed for training a comparable model from scratch.
  • A.3 Why not use MoE upcycling?: MoE is more efficient for training than inference because large batch sizes induce better parallelism.
  • A.3 Why not use MoE upcycling?: Building MoE on dense LLMs is a faster construction route, but not more efficient long term; distilling MoE into a dense model [53] may help when performance loss is small.
  • A.4 Why not use AdamW optimizer and Cosine Learning Rate Schedule?: OpenMoE used Adafactor with the Inverse Square Root learning-rate schedule following ST-MoE, rather than AdamW and cosine decay.
  • A.4 Why not use AdamW optimizer and Cosine Learning Rate Schedule?: AdamW was rejected because it frequently produced unstable NaN losses and would require substantial hyperparameter sweeping under limited computational resources.

A.5 Why not use better and larger datasets?

The project began when few open-source pre-training datasets were available, limiting data scale and quality. Since then, larger and more carefully cleaned datasets have emerged, suggesting that future improvements in training data could substantially improve LLM performance.

  • Data availability: In May 2023, only a few open-source pre-training datasets were available when the project launched.This constrained the available choices at the start of the project.
  • Data availability: Soldaini et al. released 3T carefully cleaned tokens, while Computer released 30T tokens in total.These examples illustrate subsequent gains in the scale and quality of open-source pre-training data.
  • Implications: The authors expect training on better future data to improve LLM performance generally by a large margin.This frames dataset quality and scale as an important avenue for future improvement.

B Hyper-parameters · C Related Work · C.1 Before OpenMoE

The appendix specifies OpenMoE’s configurations and training choices, including a cost-conscious optimizer and a noted head-dimension limitation, then situates the model among prior MoE language, vision, routing, and scaling efforts.

  • B Hyper-parameters: OpenMoE’s model configurations define hidden size, MoE-layer layout, FFN size, attention dimensions, layer count, total parameters, and activated parameters.“Every 4” denotes one MoE layer per four Transformer blocks; activated-parameter accounting covers parameters used to process each token.
  • B Hyper-parameters: OpenMoE-8B/32E uses head dimension 128, although dimension 64 may provide a better cost-effectiveness trade-off for its 2B activated Transformer parameters.Parameter counts simplify accounting by including trainable attention- and FFN-layer parameters, which constitute most Transformer-block parameters.
  • B Hyper-parameters: OpenMoE’s training hyper-parameters are summarized in a dedicated configuration table.The supplied passage identifies the table but does not provide its individual values.
  • B Hyper-parameters: Adafactor replaces AdamW because its memory efficiency enables less model parallelism and more data parallelism, making training cheaper under the reported setup.Adafactor performs slightly worse than AdamW for equal training steps, while the cost gap depends strongly on hardware and model size.
  • C.1 Before OpenMoE: Earlier MoE language-model work progressed from recurrent MoE layers by Shazeer et al. [42] to Transformer-scale expert parallelism in GShard [27] and open-source Switch Transformer.Switch Transformer used an encoder-decoder architecture and trained on C4.
  • C.1 Before OpenMoE: Prior work also scaled MoE training through larger datasets and efficient architectures, including Artetxe et al. [3], GaLM, and Brainformer.Artetxe et al. [3] used 112B tokens, GaLM used 1.6T tokens, and Brainformer searched MoE attributes such as layer interleaving and capacity.
  • C.1 Before OpenMoE: MoE extensions beyond language include ViT-MoE for scalable vision Transformers, WideNet for parameter efficiency, and SoftMoE [35] for stabilized soft routing.SoftMoE [35] applies soft token selection while retaining efficiency and improving routing-gradient stability.

C.2 After OpenMoE

After OpenMoE’s release, partially open-sourced MoE models such as Mixtral [22] and Deepseek-MoE [12] achieved significantly better final results. However, their use of in-house training data prevents understanding how those results were obtained, highlighting OpenMoE’s fully open-sourced nature.

  • C.2 After OpenMoE: Table 13 provides a timeline of open-sourced MoE LLM releases, sorted by model release date and recording pre-training dataset size.
  • C.2 After OpenMoE: Mixtral [22] and Deepseek-MoE [12] were released after OpenMoE as partially open-sourced MoE models with significantly better final results.
  • C.2 After OpenMoE: Because these later models use in-house data, their training process and resulting performance cannot be fully understood.

D BigBench-Lite Results

Table 14 reports detailed BigBench-Lite results and clarifies that BIG-G-sparse 8B is an MoE model with 60B total parameters.

  • Table 14 provides detailed BigBench-Lite results.
  • BIG-G-sparse 8B is an MoE model with 60B parameters in total.

E Routing Decision Standard Deviation · F Top Token Selection by Experts

Routing decisions vary more across token IDs than position IDs, indicating predominantly context-independent routing. The section also catalogs the top tokens selected by each expert.

  • E Routing Decision Standard Deviation: Figure 15 reports routing-decision standard deviation across different position IDs.The figure examines how routing variability changes by sequence position.
  • E Routing Decision Standard Deviation: Figure 16 reports routing-decision standard deviation across token IDs, restricting analysis to IDs with more than 128 tokens.Extremely low-resourced token IDs are excluded because they always have large standard deviations.
  • E Routing Decision Standard Deviation: Token IDs that never appeared also exhibit routing-decision variance.This observation is included in the token-ID analysis.
  • E Routing Decision Standard Deviation: Token IDs show larger routing-decision standard deviations than position IDs, indicating that routing is more strongly associated with token identity.This comparison is reported from Figures 15 and 16.
  • E Routing Decision Standard Deviation: Most token IDs have relatively large routing-decision standard deviations, supporting context-independent routing for most tokens.The analysis interprets the observed variability across token IDs as evidence of context-independent routing.
  • F Top Token Selection by Experts: Table 15 lists the top tokens selected by each expert.The provided passage identifies the table’s scope but does not provide the token entries.
Loading 2402.01739v2…