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

Summary

This paper makes a notable contribution towards fully-differentiable MoE LMs by proposing a couple of ways to help scale the "soft MoE" (i.e. fully-differentiable MoE) approach. Specifically, they propose: * doing routing decisions at the segment level and not the token level. They group tokens into segments, and make subsequent routing decisions based on the segment history (some kind of averaging of hidden representations of the segments) * a training strategy that groups semantically-similar segments, which they find leads to much better expert utilization The claims were empirically validated at the 0.3B and 1.5B parameter-level on both perplexity (LM) tasks as well as downstream applications. Compared to equivalently-sized dense models they show improvements

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

* a reasonable step towards scaling soft MoE approaches, with validation on the smaller end (<= 1.5B params), although some validation of 7B param models is done (but with no obvious improvements on downstream tasks).

Reasons to reject

* while it does seem clear that casual segment routing does better than prefix routing (and perhaps somewhat obvious?) I think the true comparison would be token-level vs. causal segment, which I believe is done in Figure 5. This Figure still shows that token-level routing is better. I'm not sure if the comparison with prefix routing really adds anything? * section 5.3: "suggesting that a fully differentiable architecture is more effective than a sparse MoE when using the same routing strategy" - this seems a bit bold of a claim, as it's unclear if the same behaviors would extend to token-level? * downstream improvement patterns can be explored further. For example, I noticed that for the 1.5B model the dense baseline was actually pretty good except for two specific text classification datasets (Amazon and Yelp), is there some reason why? * the gains on 7B param models are not as large, and in fact downstream performance is mixed. This should be highlighted upfront.

Questions to authors

* "fully-differentiated" --> "fully differentiable" * section 2.1: "exper" --> "expert" * "We leave how to extending Soft MoE on decoder language models as the future work" --> I don't understand this - hasn't this work shown a way to extend soft MoE from SMEAR to decoder language models?

Summary

The paper looks into the problem of improving the model quality of MoE during pre-training via fully differentiable architecture. MoE models, which rely on sparsely activated experts and token routing such as top-K gating, pose training stability challenges due to their non-differentiable nature. Recent work, SMEAR, introduced softly merged experts in the parameter space, making MoE fully differentiable. However, SMEAR was primarily tested during fine-tuning for classification tasks. This paper extends SMEAR to the pre-training stage and introduces additional optimizations, including segment routing and similarity-based batching to enhance expert specialization. Evaluation shows that the proposed method can lead to lower pretraining loss than standard MoE architecture.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

- The paper tackles an important problem, which is to improve the convergence and quality of MoE models. - The paper builds upon SMEAR, extending its benefits from fine-tuning to pre-training. Notably, the introduced causal segment routing allows training with soft experts while maintaining the autoregressive nature of language models. - Promising results on the tested MoE dataset and MoE architecture.

Reasons to reject

- The paper's technical novelty is limited. The core idea, e.g., the soft differentiable MoE was inherited from SMEAR, and the main difference is to extend SMEART from fine-tuning to pre-training. - Missing discussion and comparison of related work. The paper claims that prior work used similarity-based batching only for reasoning, not expert specialization. This is not true. HashLayer, for instance, studied clustered hashes, which routes similar tokens to the same experts. The paper also got contradictory results to HashLayer without giving any explanation. The authors of HashLayer actually found that routing similar tokens to the same expert hurts the quality of MoE. More careful examination is needed before claiming similarity-based routing helps expert specialization. - The paper does not do a thorough comparison of its approach with related work. For instance, prior work actually also addresses the differentiability of the top-K gating function, e.g, using gumble-softmax for gradient estimation and has been shown to be effective in stabilize training of MoE models (e.g., SwitchTransformer). A discussion and direct comparison of these methods would provide evidence in terms of how much benefit from a fully differentiable architecture in this work provides.

Questions to authors

Please discuss and compare with HashLayer. Please provide more justification of the benefits of fully differentiable MoE vs. prior work that provide gradient estimation for discrete gating functions.

Summary

This paper introduces 'Lory,' a novel mixture-of-experts pretraining approach. Lory features two innovative techniques: (1) a causal segment routing strategy that efficiently merges expert contributions while preserving the autoregressive properties of language models, and (2) a similarity-based data batching method that fosters expert specialization by grouping similar documents during training. The authors pretrain a series of Lory models on 150 billion tokens from scratch, incorporating up to 32 experts and 30 billion parameters. Compared to traditional dense models, the mixture-of-expert pretraining approach demonstrates superior performance in both perplexity and in-context learning evaluation.

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

This paper presents a novel mixture-of-experts pre-training solution. The experimental results demonstrate that the proposed solution outperforms traditional dense models.

Reasons to reject

The experimental setup requires improvement, particularly in terms of baseline comparison. Currently, the sole baseline is a traditional dense model. Given the availability of several off-the-shelf, pretrained, and instruction-tuned mixture-of-experts models, it would be beneficial for the authors to include these as additional baselines. Comparing the proposed solution against these models would provide a more comprehensive evaluation of its effectiveness and contextualize its performance within the current state of the art.

Summary

1. The paper proposes a novel method to train fully differentiable MoE models for autoregressive tasks. Concretely it proposes 1.1 A segment level routing tasks that computes the expert mixing weights for a segment based on the (aggregate) hidden representation from the previous segment, and uses that to softly mix the experts in parameter space 1.2 Leveraging similarity based data batching, similar to In-context pretraining [1] to avoid expert under-utilization, that subsequently results in better domain specialization of experts 2. The authors demonstrate superior performance on both perplexity as well as downstream tasks compared to an (active) parameter matched dense model setup across a series of compute budgets. They also demonstrate competitive performance compared to token level hard-routing baselines (Expert Choice [2]) 3. Furthermore, they also demonstrate that the proposed use of similarity based batching benefits MoE training more than it's random sampling counterpart, and that strong performance gains are observed even for dense models (with the dense sim-batch model outperforming the MoE random batch baseline). 4. The paper also demonstrates that the proposed method achieves good domain level expert specialization without any domain level supervision. [1] Shi, Weijia, et al. "In-Context Pretraining: Language Modeling Beyond Document Boundaries." arXiv preprint arXiv:2310.10638 (2023). [2] Zhou, Yanqi, et al. "Mixture-of-experts with expert choice routing." Advances in Neural Information Processing Systems 35 (2022): 7103-7114.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

1. The proposed method allows the previously proposed SMEAR approach to be used for autoregressive modeling with the segment based routing approach, potentially allowing for fully differentiable MoE models to be scaled. 2. The discussions on parallelism, especially Appendix H is very informative, and in my opinion forms for a good discussion point for scaling fully differentiable MoE models compared to the hard decision based token routed MoE models 3. The authors demonstrate the utility of ICL pre-training for MoE model training, showing improved gains compared to random batching. While the results are preliminary, that forms for a good avenue of exploration for training MoE models

Reasons to reject

My concerns with the paper are as follows: 1. A fundamental mismatch between training and test objective: While the paper does discuss the train-test mismatch in Appendix G3, one thing to note is that for most of the downstream tasks, the distribution between the prompt-lengths and the generation lengths is not very varied. Furthermore, without an ablation on the segment length, it's hard to understand at what point the segment length is too large for the routing function to be learned properly. In my opinion, some experiments to test this hypothesis are important. For example: 1.1 Controlling for prompt-length: By averaging the hidden representations for generating the FFN weights, intuitively there would be a diffusion of information. If that is indeed the case, then a generative needle in a haystack ([1]) or a multi-key NIAH ([2]) would be interesting to test if this diffusion indeed causes issues in learning the routing weights 1.2 Another controlled experiment that can help measure this is via perplexity measures: specifically, by fixing a prompt length and completion length to a fixed value, and then measuring the perplexity on the completion as a function of the segment length for a trained model. The difference in perplexity should be a good indication on how much impact does the train-test mismatch actually have. 1.3 Another test that is important to understand the (potential) limitations for this approach would be to have generative tasks that are long (i.e the generations are >> segment length). Intuitively, with a larger generation, the model should progressively adapt to use different experts based on the prompt + generation content, but that would not happen based on the proposed inference methodology. It would be good to quantify this limitation if possible. 2. Comparison with token level MoE routing: Based on the experiments presented in the paper, it is hard to understand what advantages would having the soft MoE approach have compared to vanilla token based routing. 2.1 If the hypothesis is that token based routing requires advanced training techniques like Expert Choice to be better than the proposed method, it would be good to demonstrate the proposed methods superiority compared to a vanilla top-k based routing approach 2.2 If the hypothesis is that the proposed routing achieves better OOD generalization because of expert specialization, it would be good to demonstrate that. Concretely, in the current paper, it is hard to disentangle if the expert specialization comes from the sim-batch data approach or from the proposed routing (or rather, if the token routed models were also trained in the sim-batch setup, would they also result in expert specialization). [1] Kamradt, G. "Needle in a Haystack–pressure testing LLMs." (2023). [2] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv preprint arXiv:2404.06654 (2024).

Questions to authors

1. The expert choice with segment level routing model (EC, segment-level in Figure 5) is barely better than the dense model (Figure 3). What is the intuition on the poor performance ? 2. For the first segment routing, why not use a BOS segment of 1 token length, and use that for priming the mixture weights for all subsequent generations ? Otherwise, you might be limiting the model from learning from the first segment. For example, if the model used absolute position embeddings, then if I understand correctly, the embeddings for the first segment would never get trained. 2. [Minor] It would be good to have the notation for e_{i} to be consistent between Equation (1), (2) and (3). Maybe consider changing Eqn (1) to be Top-k over (e_i) ?

Ethics concerns details

N/A

Rebuttal Acknowledgment

Reviewer tkvL acknowledges the rebuttal and thanks the authors for explaining why comparisons were not made to off-the-shelf, pre-trained, and instruction-tuned MoE models. However, the explanation does not fully address my concerns regarding the positioning of the proposed MoE solution within current MoE research. Therefore, I will maintain my original score. Reviewer tkvL

Post-rebuttal comment

The rebuttal addressed my concerns. I raised my score from 6 to 7.