Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training

Overview

• Mixture-of-Experts (MoE) models can efficiently scale, but training the router network is challenging
• Recently, a fully-differentiable MoE architecture called SMEAR was proposed, but its effectiveness was only shown in downstream fine-tuning tasks
• This paper introduces Lory, the first approach that scales such architectures to autoregressive language model pre-training

Plain English Explanation

Mixture-of-Experts (MoE) models are a type of AI model that use a collection of specialized "experts" to handle different parts of a task. This can help the model scale efficiently. However, training the part of the model that decides which expert to use (the "router" network) is challenging because it involves optimizing a non-differentiable, discrete objective.

Recently, researchers proposed a new MoE architecture called SMEAR that avoids this issue by "softly" merging the experts in the parameter space. But SMEAR's effectiveness was only demonstrated on classification tasks, not on the more complex task of training large language models from scratch.

This paper introduces "Lory", a new approach that adapts the SMEAR architecture to enable efficient pre-training of large language models with multiple experts. Lory uses two key techniques:

1. A "causal segment routing strategy" that efficiently merges experts while preserving the autoregressive nature of language models.
2. A "similarity-based data batching method" that encourages the experts to specialize by grouping similar documents together during training.

The researchers show that Lory models trained on 150 billion tokens can outperform dense (non-MoE) models with the same number of parameters on both perplexity and a variety of downstream tasks. Importantly, Lory achieves this while still being competitive with state-of-the-art MoE models that use a more complex token-level routing approach.

Technical Explanation

Mixture-of-experts (MoE) models facilitate efficient scaling; however, training the router network introduces the challenge of optimizing a non-differentiable, discrete objective. Recently, a fully-differentiable MoE architecture, SMEAR, was proposed, which softly merges experts in the parameter space; nevertheless, its effectiveness was only demonstrated in downstream fine-tuning on classification tasks.

In this paper, the authors present Lory, the first approach that scales such architectures to autoregressive language model pre-training. Lory introduces two key techniques:

1. A causal segment routing strategy that achieves high efficiency for expert merging operations while preserving the autoregressive nature of language models.
2. A similarity-based data batching method that encourages expert specialization by grouping similar documents in training instances.

The researchers pre-train a series of Lory models on 150B tokens from scratch, with up to 32 experts and 30B (1.5B active) parameters. Experimental results show significant performance gains over parameter-matched dense models on both perplexity (+13.9%) and a variety of downstream tasks (+1.5%-11.1%). Despite segment-level routing, Lory models achieve competitive performance compared to state-of-the-art MoE models with token-level routing. The authors further demonstrate that the trained experts in Lory capture domain-level specialization without supervision.

Critical Analysis

The paper presents a novel and promising approach for scaling MoE architectures to language model pre-training. The proposed techniques, such as causal segment routing and similarity-based data batching, appear to be effective in enabling efficient expert merging and specialization.

However, the paper does not delve into the potential limitations or caveats of the Lory approach. For example, it would be valuable to understand the computational and memory overhead associated with the proposed routing and batching methods, especially as the number of experts scales up. Additionally, the paper does not discuss the potential challenges in interpreting and understanding the learned expert specialization, which could be an important consideration for certain applications.

Furthermore, the paper focuses on the performance of Lory on standard benchmarks, but it would be interesting to see an analysis of the model's behavior and robustness in more challenging or adversarial settings. Toward Inference-Optimal Mixture-of-Experts for Large Language Models provides a useful perspective on considering the inference-time properties of MoE models, which could be a relevant avenue for future research.

Overall, the Lory approach represents a significant advancement in the field of scaling MoE architectures to large language models, and the results are highly promising. However, further exploration of the method's limitations and robustness would strengthen the work and provide a more comprehensive understanding of its potential and practical implications.

Conclusion

This paper introduces Lory, a novel approach for scaling fully-differentiable Mixture-of-Experts (MoE) architectures to the pre-training of large autoregressive language models. Lory employs a causal segment routing strategy and a similarity-based data batching method to enable efficient expert merging and specialization, respectively.

The experimental results demonstrate that Lory models can outperform parameter-matched dense models on both perplexity and a variety of downstream tasks, while remaining competitive with state-of-the-art MoE models that use more complex token-level routing. Importantly, the trained experts in Lory are shown to capture domain-level specialization without explicit supervision.

This work highlights the potential of fully-differentiable MoE architectures for language model pre-training and advocates for further research in this area. By addressing the challenges of scaling MoE models to this domain, the Lory approach opens up new avenues for building more efficient and specialized large language models, with potential implications for a wide range of natural language processing applications.