Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training
Authors: Zexuan Zhong
Mengzhou Xia
Danqi Chen
Mike Lewis
Venue: COLM (2024)
Keywords: Mixture of Experts, Language Modeling, Language Model Pre-training
TL;DR: We introduce Lory, a fully-differentiable MoE architecture, designed for autoregressive language model pre-training

Abstract: Mixture-of-experts (MoE) models facilitate efﬁcient 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 (Muqeeth et al., 2023), which softly merges experts in the parameter space. Nevertheless, its effectiveness was only demonstrated in downstream ﬁne-tuning on classiﬁcation tasks. In this paper, we present Lory, a novel 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 efﬁciency 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. We pre-train a series of Lory models from scratch on 150B tokens, with up to 32 experts and 30B (1.5B active) parameters. Experimental results show signiﬁcant performance gains over parameter-matched dense models in 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. We further demonstrate that the trained experts capture domain-level specialization without supervision. Our work highlights the potential of fully-differentiable MoE architectures for language model pre-training and advocates future research in this area.

Published as a conference paper at COLM 2024
Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training
Zexuan Zhong', Mengzhou Xia", Danqi Chen', Mike Lewis
'Princeton University, M Al
{zzhong, mengzhou, dangic]@cs, heap:j/cs.ppiiicnteeee princeton. edu, www.llewws@eetta.ccom
Abstract
Mixture-of-experts (MoB) models facilitate efficient scaling; however, training the routing network introduces the challenge of optimizing a non-differentiable, discrete objective. Recently, fully differentiable MoE architecture called SMEAR was proposed (Muqeeth et al., 2023), which softly merges experts in the parameter space. Nevertheless, its effectiveness has only been demonstrated in downstream fine-tuning on classification tasks. In this work, we present Lory, novel 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, and (2) a menilarity-baseer data batching method that encourages expert specialization by grouping similar documents in training instances. We pre-train series of Lory models from scratch on 150B tokens, with up to 32 experts and 30B (1.5B active) parameters. Experimental results show significant performance gains over parameter-matched dense models in both perplexity (+13.9%) and variety of downstream tasks (+1.5%-11.1%)) Despite segment-level routing, Lory models achieve competitive performance compared to state-of-the-arr MoE models with token-level routing. We further demonstrate that the trained experts capture domain-level specialization without supervision. Our work highlights the potential of fully differentiable MoE architectures for language model pre-training and advocates for future research in this area.
perplexity
Introduction
Mixture-of-expeee (MoE) architectures with sparse activation enable the scaling of model sizes while maintaining high training and inference efficiency (Lepikhin et al., 2021; Fedus et al., 2022; Du et al., 2022; Zoph et al., 2022; Lewis et al., 2021; Zhou et al., 2022; Jiang et al., 2024; Xue et al., 2024; Shen et al., 2024). However, training the routing network (or router) in MoE architectures introduces the challenge of optimizing a non-differentiable, discrete objective (Shazeer et al., 2017; Zoph et al., 2022). Various techniques, such as switch routing (Fedus et al., 2022), top-k expert-choice routing (Zhou al., 2022), and linear programming (Lewis et al., 2021), have been developed to address this challenge, often requiring carefully designed balancanni objectives (Fedus et al., 2022) or introducing additional complexity in assignment algorithms (Lewis et al., 2021; Roller et al., 2021).
Recent research exploring fully differentiable MoE architectures to overcome training challenges. SMEAR (Muqeeth et al., 2023), for instance, softly merges experts by averaging their parameters rather than activating the top-k experts. However, SMEAR's effectiveness has only been shown in small-scale fine-tuning experiments on downstream classification tasks (Wang et al., 2018).
In this work, we propose Lory', the first approach that scales such fully differentiable MoE architectures to autoregressive language model pre-training. Unlike text classification tasks,
l Lory tribe of parrots known for their rainbow-like colors, reflecting the spirit of 'soft' MoF.

---

Published as a conference paper at COLM 2024
Lory: Fully Differentiable Eexture-of-Expertt for Autoregressive Language Model Pre-training
Causal segment routing
McE layer
Oulput of the MoE layer x
Merged FFN
Merged FFN
Merged
layer
Training instance: simvlor docs
doc
doe
2FN
FFN
PPN
FFN
FFN
FFN'S
4FN
FFN
FFNN
FFN
The Fidlss Medil is ariz
Hluh Ane Morbé
Segmant ( d)
Segmant 2 x
Segment x
Similarity-base data batching
Imput the MrE layer (L d
Figure 1: We propose Lory, a fully differentiable MoE architecture designed for autoregressive language models based expert merging (Section 2.2). We introduce two key techniques to train Lory: First, we propose the causal segment routing strategy, which conducts expert merging at the segment level and preserves the autoregressive property of language models. Second, we use the similarity-based data batching method to construct training instances, which steers the experts toward specializing in specific domains or topics.
which only require routing each input sequence to different experts, language modeling makes predictions for each input token. Performing token-level routing is prohibitively expensive because the computational cost of merging operations scales linearly with the number of experts.
Lory is based two key techniques (Figure 1). First, we propose causal segment routing: for a sequence of input tokens, we split them into multiple segments of fixed length and use the previous segment to determine the router's weights and calculate the merged expert for the subsequent segment. During inference, we can simply use the prompt to make single routing decision throughout the generation process. This segment-level routing strategy preserves the autoregressive nature of language models while keeping the merging operations efficient. However, since the text data for pre-training language models usually concatenates random sets of documents, we find that such routing can lead to scenarios in which experts are not sufficiently specialized. To address this, we propose our second technique: Similarity-basee data batching for MoE training, which groups semantically similar documents to form consecutive segments. This idea has been recently proposed to train LMs to better reason across document boundaries (Shi et al., 2024), and we find that it leads to more effective training of expert routing.
We pre-train series of Lory models from scratch under training budget of 150B tokens, with 0.3B and 1.5B active parameters, and 8, 16, or 32 experts (corresponding to 6.8B and 29.5B full parameters at most; see Table 4). Experimental results show that our Lory models significantly outperform equal-sized dense models trained with the same amount of data, achieving performance gains in both perplexity (+ 13.9%) and wide range of downstream tasks, including commonsense reasoning (+3.7%), reading comprehension (+3.3%), closed- book QA (+1.5%), and text classification (+11.1%). Interestingly, despite using segment- level routing, Lory achieves competitive performance compared to state-of-the-art MoE models with token-level, non-differentiable discrete routing (Zhou et al., 2022). Our analysis further shows that the trained experts capture domain-level specialization without any supervision, making it distinct from previous MoE LMs with token-level routing, which only exhibit local patterns uniformly distributed across different domains (Xue et al., 2024; Jiang et al., 2024). Together, we present the first fully differentiable MoE model suitable for language model pre-training and demonstrate its effectiveness at scale. We hope our work sheds light on the potential of fully differentiable MoE architectures in cultivating specialized experts and encourages continued exploration in this research field.
2

---

Published as a conference paper at COLM 2024
2 Preliminaries
2.1 Sparsely-activated MoE
Transformer-based MoE language models typically substitute feed-forward network (FFN) layers with sparsely activated MoE layers (Shazeer et al., 2017; Fedus et al., 2022; Zoph et al., 2022). Assume an MoE layer consists of E expert FFNs, each parameterized as
(FN((::1)),.... ...., FFN((;;E), where the function FFN Rd - Rd defines a single expert module. For each token in sequence, an MoE layer takes the hidden representation hy E Rd as the input and computes its output 0x e Rd by sparsely activating experts in this layer and aggregating the outputs through weighted sum:
0x = Dey ej (F;(),;;;)), where ej = Top-k(Softmax
(1)
The routing weight e; for the i-th expert is determined by routing network router R, which takes ny as input and calculates the weight for each expert. In practice, to achieve sparsity and computational efficiency, only one (Fedus et al., 2022) experts with the highest routing weights are activated at each MoE layer. The weights of the remaining experts are set to 0 (i.e., e; 0), eliminating the need to compute (Fx 8;) and effectively deactivating the i-th expert.
(Lepikhin al., 2021)
2.2 Fully Differentiable MoE Architectures via Expert Merging
The primary challenges in training sparsely activated MoE models arise from the difficulty in training discrete routers. A promising direction is to design fully differentiable MoB architectures that do not rely on extra loss formulations for stabilized training. A recent model architecture (Muqeeth al., 2023) demonstrates the feasibility of this approach by computing weighted average of all expert FFNs in the parameter space (Matena & Raffel, 2022; Wortsman et al., 2022), thereby creating a "merged FFN." " Given an input x and its corresponding routing weights e, the output ox of merged FFN is computed as:
0x = (]x Σ e 0)), where = Softmax((((x))):
(2)
However, naively extending this approach to autoregressive language models, which would require computing the merged FFN for each token in a sequence, is infeasible as the computational costs of merging operations scale linearly with the number of experts. SMEAR (Muqeeth et al., 2023) has only been evaluated for downstream fine-tuning on text classification tasks, making routing decisions based pooling representation of the entire
D " )); Such operations disrupt the autoregressive
input sequence, i.e, e; Softmax (R( property in language model pre-training. In this work, we address these challenges by developing fully differentiable MoE. architecture suitable for autoregressive language modeling and pre-training such models at scale.
3 Our Approach: Lory
Lory is an approach for pre-training fully differentiable MoE language models (Figure 1). The core technique that enables Lory to be fully differentiable is expert merging (Muqeeth et al., 2023, see details in Section 2.2). To make it computationally feasible, propose a causal segment routing method that merges experts only once per segment, effectively reducing the number of merging operations (Section 3.1). We also propose data batching strategy that groups semantically similar texts, which is crucial for the effective training of the segment-level router (Section 3.2).
Notations. We denote an input sequence of L tokens as X = segment size T, we divide the input sequence into N [L/T]] segments, denoted as
XT.). Considering
3

---

Published as a conference paper at COLM 2024
S2,......, SN. We use R to denote the routing network (parameterized linear layer), which computes the weights for expert merging. Let ft represent the hidden representation of the token The parameters of the i-th expert FFN are denoted by e,.
3.1 Efficient Expert Merging via Causal Segment Routing
Challenges. An intuitive way to reduce the computational cost is to use segment- routing instead of token-level routing, which can reduce the number of merging operations from to N. However, simply using the current segment to compute the routing weights can cause information leakage.
Training design. We propose causal segment routing to effectively route information across segments in an autoregressive manner." This method merges FFNs in an MoE layer based on the previous segment's 's information and uses it to process the current segment. Specifically, given a training instance X that consists of L tokens (e-g., = 4096), we split the training instance into N segments, each containing T (e-g., T = 256) consecutive tokens. For the k-th segment Sk, where k > 1, we compute the average of the hidden representations of its preceding segment S denoted as f421-- Using the average hidden representation allows the model to adapt to prompts of varying lengths during inference. hx-1 is then utilized to determine the routing weights, resulting in a merged expert
1 T
h
=
Σ x
hx, = Softtmxx (R(h 1)), 0= -= 0,.
(3)
We then use the merged expert to process all the tokens in the current segment Sk, i.e., Or = FF;(),;;)) Yx E Sx- This approach ensures that the routing decisions made by the model are based exclusively on data from preceding positions. For the first segment S₁, the representation of the segment itself is used to compute the merging weights for its own FFN. To prevent information leakage, we implement a stop-gradient operation on R R(i)). As demonstrated in Appendix B, merging experts the segment level incurs minimal overhead compared to the training of dense models.
Prompt-only routing during inference. During inference, we begin with a given prompt and make a single routing decision per layer based on the average hidden representations of the prompt. This routing decision determines merged FFN, which is used consistently throughout the entire generation process. It is important to note that this inference process is as simple and computationally efficient as that of dense models.?
3.2 Similarity-based Data Batching
The standard practice of pre-training LMs is to randomly concatenate documents to construct training instances with fixed length. This approach can lead to under-specialized experts because tokens within adjacent segments may come from very different and irrelevant documents.
To mitigate this issue, we employ a similarity-based data batching technique inspired by Shi et al. (2024), which sequentially concatenates similar documents to construct training instances. This method encourages high similarity between adjacent segments, enabling the experts to specialize in different domains or topics.
We use Contriever (Izacard et al., 2022) to measure document similarity and apply a greedy search algorithm to concatenate similar documents to form batches (see Appendix C).
PPeeddccode of the causal segment routing strategy can be found in Appendix A.
Inn Appendix G.2, we compare the prompt-only routing strategy to using the causal segment routing strategy that faithfully follows the training design, and find they do not lead to significant differences. We also discuss the potential of converting Lory to sparsely MoF models for memory- efficient inference.
4

---

Published as a conference paper at COLM 2024
Although our data batching technique is similar to that of Shi et al. (2024), our goal is different. While they focus on enhancing language models' reasoning across document boundaries, find this approach particularly effective for fostering expert specialization during MoE model training.
4 Experiments
In this section, we evaluate Lory by training a series of language models from scratch. We first describe the experimental setups (Section 4.1) and then present the results (Section 4.2).
4.1 Setups
Models. We evaluate our approach by training decoder-only Transformer models which consist of 0.3B and 1.5B active parameters.' 4 For each FFN layer in the Transformer model, we replace it with MoE layers with E E {8, 16, 32} experts with exactly the same architecture." Appendix D shows the configuration of model architectures as well as the total parameter count. We follow LLaMA (Touvron et al., 2023a) and use SwiGLU (Shazeer, 2020) the activation function in FFNs. We use the same tokenizer as the LLaMA models (Touvron et al., 2023a;b). All models are trained with 4096-token context window. In the causal segment routing strategy, we set the length of each segment to be T 256.
Training details. We employ the A optimizer (Loshchilov & Hutter, 2019) with B1 = 0.9 and B2 = 0.95 and use a learning rate of 2e-4 with a cosine learning rate scheduler. All models with a batch size of 1 million tokens. We employ the data parallelism with the
ZeRO optimization (Rajbhandari et al., 2020) for distributed training. At the beginning of training, we train parameter-matcher dense model and duplicate the FFN layers as initialization of the MoE model. In our experiments, we use the first 5% training steps as the warmup to initialize the MoE weights. We find that without warmup training, there may be more experts under-utilized (see Appendix G.3 for an ablation study). We also apply linear warmup the learning rate scheduler for the first 5% training steps. We train our models with up to 64 A100 GPUs
Training datasets. We randomly sample subset (150 billion tokens) of the Commoncrawl dataset (Wenzek et al., 2019) for training. Using the similarity-base data batching method from Shi et al. (2024), construct all training instances (see Appendix C for details).
Evaluation datasets. We evaluate all the models on language modeling tasks by measuring the perplexity of trained models on held-out evaluation datasets sampled from arXiv, Books, Wikipedia, C4 (Raffel et al., 2020), and Python code (a Python subset of Github). Each evaluation dataset contains 1K samples, each of which consists of 4096 tokens.
We also evaluate models in downstream tasks with in-context learning (Brown et al., 2020), including common sense reasoning: BoolQ (Clark et al., 2019), PIQA (Bisk et al., 2020), SIQA (Sap et al., 2019), HellaSwag (Zellers et al., 2019), WinoGrand (Sakaguchi et al., 2020); reading comprehension: RACE (Lai al., 2017), ARC (Clark al., 2018)); closed- book QA: Natural Questions (Kwiatkowski et al., 2019), TriviaQA (Joshi et al., 2017); and text classification: AGNews (Zhang al., 2015), SST-2 Socher al. (2013), Amazon and Yelp (Zhang et al., 2015), FEVER (Thorne et al., 2018), MRPC (Dolan & Brockett, 2005). For text classification tasks, we follow the evaluation setup of Min al. (2022); for the rest of tasks, we follow the same setup as Touvron et al. (2023b).
Here, "active parameters" refers to the size of the model after merging each MoF layer. 'n Appendix E, we additionally conduct experiments on 7B dense model and 7B/4E MoE model without using similaritybbasee data batching. Due to the limited computing resources, we not able to train 7B models on the similarity-based batched dataset.
5

---

Published as a conference paper at COLM 2024
6.00 (eeen)
6.30-E Bary)
1.5 jaeve)
1.5022 (un)]
Model
arXiv
Books
Wiki
C4
Python
2.5
6.20n1e (Len))
(Lar))
233
1.567EE Jues
0.38
8.4
18.0
10.3
13.8
15.2
Perklodd
0.38/8E
7.4
13.3
12.5
0,3B/16E
7.2
15.7
9.1
13.1
12.2
0.38/32E
7.2
2.3
2.1
1.50
6.6
13.6
7.8
10.7
10.4
1.58/8E
1.5B/16E
6.0
12.4
7.1
10.6
8.9
23
1.0
Figure 2: Left: training curves (log perplexity) of models with different sizes and experts. Right: Perplexity of trained models different evaluation sets (arXiv, Books, Wikipedia, C4, and Python). We include the detailed model configurations and sizes in Appendix D.
Commonsense Reasoning
Reading Comprehension
Model
PIQA
SIQA
BoolQ
HellaSwag
WinoGrande
RACE-m
RACE-h
ARC-e
ARC-c
0.3B
65.8
42.7
44.6
34.6
51.2
41.7
30.9
51.5
21.3
0.3B/SE
67.5
41.2
34.8
54.4
43.1
31.4
52.4
0.3B/16B
67.2
44.1
56.6
34,9
54.1
43.9
31.1
54.8
24.9
0.3B/32R
68.2
43.0
58.0
34.7
53.4
42.7
32.0
57.4
26.3
1.5B
71.2
45.0
54.0
43,9
60.9
50.1
36.7
65.0
31.0
1.5B/SE
72.1
45.2
43.6
63.7
51.2
36.5
66.3
1.5B/16B
71.3
56.0
43.7
61.5
51.7
37.3
66.3
32.7
1.5B/32E
72.1
47.1
59,9
43,8
61.9
51.5
32.4
66.7
32.7
Closed-book QA
Text Classification
Avg
Model
NQ
TQA
AGNews
Amazon
SST-2
Yelp
Fever
MRPC
0.3B
4.7
8.8
30.3
53.6
54.6
66.0
47.6
62.0
41.8
0.38/8F
38.4
52.3
54.6
62.6
56.6
59.0
42.7
0.3B/165 16E
6.0
10.2
36.3
75.6
53.3
64.0
57.0
65.0
45.8
0.3B/32R
5.3
10.2
47.3
64.0
55.3
73.3
55.7
56.0
46.0
1,5B
7.6
23.8
64.0
65.3
80.0
58.6
59.0
66.7
51.9
1.5B/8E
7.3
24.2
94.0
80.0
57.0
64.0
1.5B/16E
7.3
25.6
61.6
78.3
84.6
93.6
57.3
63.6
55.1
1.5B
7.0
25.4
62.3
94.7
85.0
95.3
56.3
66.7
56.5
Table 1: We compare the Lory MoE models with the parameter-matched dense models on downstream tasks, including commonsense reasoning, reading comprehension, closed-book QA, and text classification.
4.2 Main Results
Training efficiency and convergence. Figure (left) shows the training loss curves of the dense model and our MoE models with different model sizes. First, find that with the same amount training tokens, our models clearly achieve better training loss compared to the dense model baseline. For the 0.3B and 1.5B models, the same level of loss with fewer than half of the training tokens. This indicates that our approach achieves much better performance with the same training compute (see analysis of additional FLOPs from MoE layers in Appendix B). We also observe that when using more experts, we are able to gain more improvement.
models with 32 experts achieve
Language modeling. We evaluated the models on language modeling tasks and found that MoE models consistently outperformed the dense baseline, significantly reducing perplexity across all domains. For instance, the 0.3B/32 B2F model improved perplexity by 13.9% on Books compared to the 0.3B dense model. Notably, improvements were most pronounced in test domains distinct from the training data (e.g., Python), indicating strong expert specialization, which is further explored in Section 5.4).
Downstream tasks. Table shows the model performance on downstream tasks. We observe significant performance across all tasks. For example, our 0.3B/32E model achieves

---

Published as a conference paper at COLM 2024
average performance improvement of +3.7% in common sense reasoning, +..3% in reading comprehension, +1.5% in reading comprehension, and +11.1% in text classification.
5 Analysis and Ablation Studies
In this section, we conduct ablation studies and analysis understand the significance of each component of our approach. In Appendix G, we provide additional insights, demonstrating that (1) during the inference of downstream tasks, whether the entire input prompt is routed once or each segment is routed individually does not result in substantial differences in the tasks we evaluated; and (2) warmup training is crucial for achieving high expert utilization, particularly when training MoE models with large number of experts.
5.1 Importance of Causal Segment Routing
We compare our causal segment routing strategy with an alternative prefix routing strategy for training. In prefix routing, expert merging is performed only once for each sequence based on the first segment. The merged FFN then used to process the rest of the sequence without further updates. Figure shows that using only prefix for routing leads to much worse performance compared to causal segment routing. These results highlight the importance of using every segment to provide strong training signals for routers.
2.5
2.5
0.10 prrte rarign
[annen] icsses segoen sortnn)
2.5
0.30 barT))
2:30RE peintertrg
SE inand
paid NIN
Iememmmn (MM Over Dhrret)
0.08
12nd bick))
Em3
2.4
0.00
2.3
0.04
2.2-
0.02
21
2.1
60
150
0.00
50
150
Billon
Figure 3: Training curves of causal segment routing and prefix routing. The latter is straightforward segment-level routing strategy that uses the first segment to route the entire input.
Figure 4: Left: Training curves of similarity- based data batching (sim batch) or the standard random batching (rand batch). Right: Training loss difference between Lory and dense model when using different batching strategies. Lory leads to a larger loss improvement over the dense model when using similarity-base data batching.
5.2 Importance of Similarity-based Data Batching
To investigate the importance of similarity-based data batching, we compare the performance improvement of MoE models over dense models with and without this batching method. Figure 4 (left) shows the training loss of dense (0.3B) and MoE models with eight experts (0.3//80 /E) using similarity-batched (sim batch) and randomly-batchee (rand batch) data. MoE models consistently outperform dense models in both setups. However, the loss improvement (i.e., the difference in loss between dense and MoE models) is much larger with similarity-based batching, and this effect is amplified with more training data (Figure 4 (right)). These results strongly support the importance of similarity-based batching for effectively training our MoE model.
Comparison to Existing MoE Models
We compare our approach with Expert Choice (EC) (Zhou et al., 2022), a state-of-the- art MoE method that ensures balanced load during training by having each expert select top-k inputs according to the routing weights. During inference, we route each token into top-k experts to avoid to leverage global information for routing.
7

---

Published as a conference paper at COLM 2024
We consider two variants of EC MoE models, both with a capacity factor of to match the computation of our MoE models. First, we train sparse EC MoE model using our segment routing strategy, where each expert selects top segments and processes all tokens within those segments. This variant allows us to directly compare our expert-merging strategy with the expert choice method while using the same segment-level routing approach. Second, we consider the original EC setting with token-level routing to provide an end-to-end comparison with state- of-the-art MoE models using the same amount of training computation. Figure shows the training loss curves. We observe that Lory (blue curve) significantly outperforms segment-level EC (orange curve) with the same routing setting, suggesting that a fully differentiable architecture is more effective than a sparse MoE when using the same routing strategy. Comparing Lory with the token-level EC model (green curve), we find that Lory achieves competitive results despite using segment-level routing and not requiring any advanced training techniques. These results highlight the significant potential of Lory.
2.6
0.36IBE (Lory)
0,383E (FC, segment-ivell
2.5
0.36/8E (EC, Xakan-lore)
Log Parpleaty
2.4
2.3
2.1
50
100
150
Biliin takens
Figure Comparison with the state-of- the-art MoE training technique Expert Choice (EC) with a segment-level or token-level routingg. For both EC models, we use the capacity factor of 1 with the same amount of FLOPs as our training method for the fair comparison.
Table 2 shows model perplexity on held-out evaluation sets. The token-level EC model outperforms ours on C4, likely due to its similarity to the training set (Commoncrawi), However, on arXiv, Books, and Wikipedia, EC performs similarly or slightly worse. Notably, our model excels on the Python evaluation (12.5 vs. 13.6 perplexity), suggesting segment-level routing can be particularly effective for out-of-domain data (i.e., Python in CommonCrawl), Our analysis in Section 5.4 shows that segment-level routing models are indeed able to learn experts that are specialized in specific domains (e.g,, Python code), potentially helping models achieve high performance in less frequent domains.
Model
arXiv
Books
Wiki
C4
Python
0.3B/88 (Lory)
7.4
16.0
9.2
13.3
12.5
0.3B/ /BE (EC, seg-level)
7.9
17.6
10.5
14.1
14.1
0.3B/ /BE (EC, tok-level)
7.5
17.0
9.2
12.8
13.6
Table 2:
Perplexity of our trained MoE model and EC models on evaluation sets. We instantiate EC methods with our segment-level routing and the original token-level routing.
5.4 Expert Utilization and Specialization
Utilization: How many experts actively utilized? One potential issue of training MoE models is the models may collapse to dense models because most experts are under-utilized (e.g., some some experts have never been activated). In Appendix G.1, we show although without using any auxiliary loss on load balancing, Lory is able to achieve high expert utilization, preventing the MoE models from collapsing to dense models.
Specialization: What do experts learn? In order to study the expert specialization, we investigate the averaged routing weights at different layers of the 0.3B/ /8E model, different domains (Books, arXiv, Python, and Wikipedia). Figure shows the routing weights at layer O, 11, and 23 (the first, middle, and last layer) of the 0.3B/ model. 6 First, we find that there exists clear domain-level expert specialization in our trained MoE models, even though no additional domain-level supervision is used during training. For instance, expert 7 at layer 11 is specialized to process inputs in the arXiv domain. We also observe that routing weights on arXiv and Python code are more similar compared to Books and Wikipedia, likely because LaTex code and Python code are dissimilar to natural language.
6n Appendix F, we show the averaged routing weights all layers of the 0.3B/8F model.
8

---

Published as a conference paper at COLM 2024
Books
Python
Wikipedia
Layer
Layer 11
Layer 23
Averaged Weights
0.2
0.1
0.0
0
0
Expert ID
Figure 6: Averaged routing weights at layer (0, 11, 23) of the 0.3B/8E model different domains (Books, arXiv, Python, Wikipedia), We observe that the experts in our MoE models learn domain-level specialization, especially at middle and higher layers.
Second, experts at the middle or high layers are more specialized in specific domains, while the routing weights at lower layers are similar and flat across domains.
worth noting that our learned experts behave differently from those of prior token-level MoE models, where shallow token-level specialization is observed. For example, some experts are specialized for a specific type of word (e.g, punctuations, articles), and few deep semantic features captured by the learned routers (Jiang al., 2024; Lewis et al., 2021; Zoph et al., 2022; Shazeer et al., 2017; Xue et al., 2024). Our models learn domain-level specialization, which attribute to the segment-level routing strategy used during training. This strategy allows routers to capture global semantic features beyond the token level. The complementary nature of features captured by segment/ sentence-leeel and token-level routing strategies suggests the possibility of combining them to build even stronger models, and we leave it for future work.
Discussion of Parallelism Strategies for Scaling Lory
In our experiments, we utilized data parallelism with ZeRO optimization (Rajbhandari et al., 2020) to reduce memory usage. Data parallelism involves either maintaining a copy of the model parameters on each device or using fast communication between devices to synchronize parameters with ZeRO optimization. As we scale up the Mixture of Experts (MoE) models to sizes exceeding 100 billion parameters, the substantial increase in MoE parameters can lead to a bottleneck due to the significant amount of parameter communication required.
To address this problem, expert parallelism is utilized for sparsely MoE models, where experts are distributed across different devices and operate independently (Lepikhin et al., 2021; Fedus et al., 2022). Unlike traditional MoE models, which use a top-k discrete routing network to dispatch inputs access to all expert parameters for computations.
subset of experts, our soft-routing MoE models require
To overcome this challenge, we can employ expert-aise model parallelism strategy, as shown in Figure (left). This involves partitioning the MoE layers across multiple devices by sharding all experts along with their hidden dimension. Consequently, identical dimensions from different experts are allocated to the same device. We shard all the experts along with the hidden dimension and distribute the same shard of different experts one device. This expert-wise model parallelism enables us to communicate only hidden activations instead of model parameters, which is more efficient since hidden activations do not scale with the number of experts. Furthermore, for non-MoE components, such as attention layers, we can employ data parallelism to minimize communication overhead. Figure 7 (right) illustrates combination of data parallelism and model parallelism to
training setup in which facilitate training at scale. We plan to explore the experiments for scaling our approach to very large models in future work.

---

Published as a conference paper at COLM 2024
Expert perallllam
Device
Devie'
Dewie
FTN
FFN
Data parallelism
Model parallelism (expert-wise)
Model parallelism (expert-wise)
Ataenian
Dewic
Dewce
Device
PPNT
Device
Attention
Alle
FFNZ
Device
Attention
FFN1
FFNZ
FFN3
: PINT
FTM
Fna
Figure 7: Left: Comparison between expert parallelism and expert-wise model parallelism. Right: Using a mixed strategy of data parallelism and expert-wise model parallelism. We use data parallelism for non-MoE components (e.g., attention) and model parallelism for MoE layers.
Related Work
Mixture of Experts. Sparsely activated MoE models (Shazeer et al., 2017) have been proposed to demonstrate the potential of massively scaling up model sizes. GShard (Lepikhin et al., 2021) adapts the sparse MoE architecture into Transformer models and achieves strong results on machine translation. Recent work has extended to general language models (Fedus et al., 2022; Zoph et al., 2022; Jiang al., 2024; Dai et al., 2024; Zhou et al., 2022; Du et al., 2022; Artetxe et al., 2021; Xue et al., 2024). Traditional MoE models are trained to route given inputs one or a few specialized expert modules, which introduces a non-differentiable, discrete decision-learning problem. These existing models are trained with the top-1 or top-2 routing strategy on carefully designed load balancing objective (Lepikhin et al., 2021; Fedus et al., 2022; Zoph et al., 2022), or employ complicated assignment algorithms to distribute inputs (Lewis al., 2021; Roller et al., 2021; Zhou et al., 2022). Training MoE models has been shown to be difficult, facing the issues of training instability, expert under-specialization, poor training efficiency (Zoph et al., 2022).
Our approach enables end-to-end gradient pack-propagation using fully differentiable MoE architectures. While SMEAR (Muqeeth et al., 2023) and Soft MoË (Puigcerver et al., 2024) also achieve this, SMEAR is limited to text classification tasks with an encoder backbone, and Soft MoE is only evaluated on vision tasks. Lory is the first to scale this architecture for autoregressive language model pre-training. Extending Soft MoE to decoder language models is left for future work.
Similarity-based data batching. There exists research that applies similar data batching method during training. In-context pre-training (Shi et al., 2024) groups relevant documents together to encourage language models to leverage long-range contexts and improve the results of in-context learning and retrieval augmentation. Zhong et al. (2022) batch documents with high lexical similarity to collect more positive pairs in contrastive learning framework to provide stronger training signals. Despite sharing similar idea, the goal of our data batching method is to avoid routing irrelevant documents together, which may hurt the expert specialization.
8 Conclusion
In this paper, we introduce Lory, fully differentiable MoE model specifically designed for autoregressive language model pre-training. Through extensive experiments, we demonstrate that Lory significantly outperforms its dense counterparts in both language modeling and downstream tasks. Our findings reveal that the trained experts in Lory are highly specialized and adept at capturing domain-level information. Future research directions include scaling up Lory further (Section 6), integrating token-level and segment- level routing, and developing efficient decoding methods tailored for Lory.
10

---

Published as a conference paper at COLM 2024
Acknowledgements
We appreciate the valuable comments and feedback from members of the Princeton NLP group. We also thank Weijia Shi for the help with the experiments and discussions related to the similarity-based data batching method.
Ethics Statement
This paper presents novel approach for training sparse large language models. It is crucial to acknowledge that, similar to existing language models, those trained using our method may have comparable societal consequences. For instance, language models can generate factually inaccurate outputs and unethical content, posing risks of spreading misinformation and perpetuating harmful stereotypes or biases (Zhao et al., 2024). Furthermore, malicious users may extract the training data used to develop these models, leading to potential privacy and licensing issues (Carlini et al., 2021). We recognize these potential negative consequences and advise those employing our approach to remain vigilant about these risks when developing powerful language models.
References
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan lyer, Ramakanth Pasunuru, et al. Efficient large scale language modeling with mixtures of experts. ar preprint ar arXie:2112.1068, 2021.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Conference on Artificial Intelligence (AAAI), 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert- Herbert-Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ulfar Erlingsson, al. Extracting training data from large language models. In USENIX Security Symposium, 2021.
Christopher Clark, Kenton Lee, Ming- Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In North American Chapter of the Association for Computational Linguistics (NAACL), 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arii preprint arx:: Att::180300557,, 2018.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint ar Xii:2401.06066 2024.
William B. Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In International Workshop on Paraphrasing (IWP), 2005.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning (ICML), 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. In The Journal of Machine Learning Research (JMLR), 2022.
11

---

Published as a conference paper at COLM 2024
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. In Transactions Machine Learning Research (TMLR), 2022.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, al. Mixtral of experts. arXiv preprint ar arXi::2401.04m: 2024.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 7(3):535-547, 2019.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Association for Computational Linguistics (ACL), 2017.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. In Transactions of the Association of Computational Linguistics (TACL), 2019.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. In Empirical Methods in Natural Language Processing (EMNLP), 2017.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. In International Conference Learning Representations (ICLR), 2021.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning (ICML), 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019.
Michael S Matena and Colin A Raffel. Merging models with fisher- averaging. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCL: Learning to learn in context. In North American Chapter of the Association for Computational Linguisties: Human Language Technologies (NAACL-HLT), 2022.
Mohammed Muqeeth, Haokun Liu, and Colin Raffel. Soft merging of experts with adaptive routing. Xiv preprint ar XrX::22366003744 2023.
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In International Conference on Learning Representations (ICLR), 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. In The Journal of Machine Learning Research (JMLR), 2020.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2020.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. In Advunces in Neural Information Processing Systems (NeurIPS), 2021.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In Conference on Artificial Intelligence (AAAI), 2020.
12

---

Published as a conference paper at COLM 2024
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. In Empirical Methods in Natural Language Processing (EMNLP), 2019.
Noam Shazeer. Glu variants improve transformer. arxiv preprint arXiv:2002.052022 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture- of-experts layer. In International Conference on Learning Representations (ICLR), 2017.
Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. JetMoE: Reaching Ilama2 performance with 0.1 m dollars. Xrv preprint ar xrxo:2404.07444 2024.
Weijia Shi, Sewon Min, Maria Lomeli, Chunting Zhou, Margaret Li, Victoria Lin, Noah A Smith, Luke Zettlemoyer, Scott Yih, and Mike Lewis. In-context pretraining: Language modeling beyond document boundaries. In International Conference on Learning Representations (ICLR), 2024.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In Empirical Methods Natural Language Processing (EMNLP), 2013.
James Thorne, Andreas Vlachos, Christos Christodoulopoulls, and Arpit Mittal. FEVER:
large-scale dataset for fact extraction and verification. In North American Chapter of the Association for Computational Linguistics (NAACL), 2018.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint ar 10071130,.13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arv preprint arXiv:2307.09288 2023b.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE: A multi-lask benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations (ICLR), 2018.
Guillaume Wenzek, Marie-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Armand Joulin, and Edouard Grave. Cenet: Extracting high quality monolingual datasets from web crawl data. In International Conference on Language Resources and Evaluation (LREC), 2019.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo- Lopes, Ari Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning (ICML), 2022.
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. OpenMoE: An early effort on open mixture-of-experts language models. arXiv preprint X01,22201,,,,,,,,0,0 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Association for Computational Linguistics (ACL), 2019.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In Advances Neural Information Processing Systems (NeurIPS), 2015.
Wenting Zhao, Xiang Ren, Jack Hessel, Claire Cardie, Yejin Choi, and Yuntian Deng. Wildchat: 1m chatgpt interaction logs in the wild. arxiv preprint ar Xie:2405.004700, 2024.
13

---

Published as a conference paper at COLM 2024
Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmentation. In Empirical Methods in Natural Language Processing (EMNLP), 2022.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. In Advances in Neural Information Processing Systems (NeurlPS), 2022.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. ST-MoE: Designing stable and transferable sparse expert models. ar Xiv preprint ar arXiv:2202.08906 2022.
14

---

Published as a conference paper at COLM 2024
A Pseudocode of Causal Segment Routing
We include a pytorch-style pseudocode of our proposed causal segment routing strategy in Algorithm 1.
Algorithm 1 Pseudocode of causal segment routing.
B: batch size (nunber training instances)
# L: length of each training instance
d: hidden dinension
E: number experts
# T: length of each segnent
R: routing network (input: hidden rep, output: routing weights)
input x x: input tensor (BxLxd)
// number of segrents seg_x = (B T, d) split into segnents
sample
representation each segnent (BNKE)
repr = mean(seg_x, din=t)
routing results (not causal) (BNXE) softnex(R(repr), dime-1)
routing results for the first segnent firrsst e.view(B, N. E)C:, 0]
nake causal routing results (shift roll(e, 1)
by
# set routing results of the first segnent e.view(B, N,
instance
e[s, 0] stop_grad(e_first) # assign w/ stop gradient vieee(BNN,, E)
HoE FFN forward with expert weights e foe_efn(see_x,, seg.y: BAN x
back to the instance view vi.
return y
moe_ffri ffn: compute the merged expert and process the input (equation 2).
B Computational Overhead of Routing and Merging
Here we investigate the computational overhead of our Moe models compared to the dense counterpart. We consider an MoE layer and an input tensor x consisting of tokens and dimensions: x L x d. We assume that the model uses SwiGLU as the activation function in FFNs and it up-projects the input x to d' o'ddimensional activations in FFNs. In this case, processing the input on an FFN requires roughly 6 x L x d x d' FLOPs (there are two up projections and one down projections in SwiGLU-based FFNs). The overhead of soft-routing MoE comes mainly from the merging operation. Suppose that there are F. experts and that the model makes routing decision for every segment of tokens (equivalently, there are L/T routing decisions). Each merging operation on E experts takes 6 x E d x d' FLOPs (we compute three merged matrices). Therefore, the total overhead will be 4 6 E d x d' FLOPs. This indicates that compared to a dense FFN layer, an MoE layer with E experts E
requires T more FLOPs, compared to the dense counterpart. In our experiments, we set = 256; this suggests that using E experts introduces 3. more computations and using E = 32 experts introduces 12.5% more computations at the FFN layers. It is worth noting that the computations from FFN layers are only a subset of the full model computations, so 3.1% is overhead upperbound when measuring experiments, our most straightforward implementation leads to a 15% or 28% slowdown of training efficiency when using 8 or 32 experts (Table 3).
full models. In our
15

---

Published as a conference paper at COLM 2024
Model
Throughput (tokens/s/gpu)
0.3B
29,000
0.3B/8E
24,500
0.3B/10 16E
22,900
0.3B/32E
20,800
Table 3: Training throughput (tokens/s) gpu) of our MoE models and the dense counterpart. Our
implementation based data parallelism with the ZeRO optimization (Rajbhandari al., 2020).
C Details of Similarity-based Data Batching
We adapt the pipeline of in-context pre-training (Shi et al., 2024) in our approach. Given of documents D, for each document d E D, we first use Contriever (Izacard et al., 2022) to retrieve top-k most similar documents N(d). The similarity between the document dj and d, is defined the cosine similarity of their Contriever embeddings, i.e,, (d; cos(C(d;), (dd)), where denotes the Contriever encoder model. We implement an efficient approximate nearest-neighbors search based on the FAISS library (Johnson et al., 2019). Then, we sort all the documents according to the similarity and construct training instances by batch consecutive documents. We use the same greedy algorithm as Shi et al. (2024). We start from a single document and repeatedly add the document that has the highest similarity value and has not been added to the list; we restart the process with a new document if all documents that connected to the last document of the list are selected. We repeat this process until there are no documents left.
Model
params
N
D
Thaad
0.3B
0.3B
0.3B/8E
1.8B
24
1024
16
0.3B/ 33B/166 /16E
3.5B
0.3B/32E
6.8B
1.5B
1.5B
1.5B/8t /EE
7.8B
48
1536
24
1.5B/ / 16E
15.0B
1.5B/32E 29.5B
Table 4: Model architectures and sizes used in our experiments. For MoE models, we replace each FFN layers with MoE layer. kE (e.g., "16B" in "0.3B/16E") represents the architecture in which each FFN layer is replaced with MoE layer of experts. N: number of layers; D: hidden dimension of the model; Whead: number of attention heads.
D Model Configurations
In our experiments, we employ Lory to train decoder-only models which consists of effective parameters of 0.3B and 1.5B. For each FFN layer in the Transformer model, we replace with MoE layers with E (E E {8, 16, 32}) experts with exactly the same architecture. Table shows the configurations of model architectures.
E Experiments on 7B models
Experimental Setups. We conduct experiments on 7B architecture. Table 5 shows the configuration of the model architectures. We train a dense 7B model and 7B/4E MoE model. For the 7B models, we follow LLaMA2 (Touvron et al., 2023b) and use combination of several corpora as the training set. We down-sample the full training set to a subset of 200B tokens for 7B models. Due to limited resources, we only conduct experiments on randomly batched training data for 7B models and do not apply the similarity-based batching method.
16

---

Published as a conference paper at COLM 2024
Model
Mparams
N
D
Whead
7B
7B
32
4096
7B/4E
19.7B
Table 5: Model architectures and sizes used in our 7B experiments. For MoE models, replace each FEN layers with a MoE layer. N: number of layers; D: hidden dimension of the model; Thead number of attention heads.
Model
PIQA
SIQA
BoolQ
HellaSwag
1.5B/EE (prompt)
72.1
45.2
62.0
43.6
1.58/8B (segment)
72.1
45.6
60.2
43,9
1.BB/1E (prompt)
71.3
45.0
56.0
43.7
1.5B/168 16E (segment)
72.9
45.4
55.2
43.6
Model
Wino
NQ
TQA
Avg
1.5B/8E (prompt)
63.7
7.3
24.2
45,4
1 15B/88 (segment)
61.8
7.3
24.4
45.1
1.6B/16E/// 1.5B/ 166 (prompt)
61.5
7.3
25.6
44.4
1.5B/161 16E (segment)
62.4
7.6
25.5
44.7
Table 6: Downstream performance of using different inference methods. We study two routing strategy for inference, prompt: we make the routing decision once on the entire input prompt; segment: we re-route and get new merged FFNs every segment.
Language Modeling Results. We show the training loss curves in Figure 8 and the perplexity on held-out evaluation sets in Table 7. We find that compared to the 0.3B and 1.5B models (see Section 4), the improvement of the 7B/ /4E model is less significant. We think it because (1) the similarity-based batching method is not applied in this case, making the experts under-utilized; (2) we only use four experts in the MoE model. We leave the experiments with the similarity-base batching method on MoE models with more experts as future work.
2.2
7B
7BME
2.1
Log Perplexity
2.0
1.9
1.8
1.7
100
150
Billion of tokens
Figure Training curves (log perplexity) of the 7B dense model and the 7B/4B MoË model. Note that when training the 7B/4E model, we do not apply the similarity-bass batching method.
Performance Downstream Tasks.
Table shows the performance of the models on downstream tasks. We find that although the similarity-based batching method is not used when training the 7B/ 4E model, it still achieves clearly better results on various tasks compared to the dense 7B model. This further suggests the effectiveness of our causal routing strategy.
17

---

Published as a conference paper at COLM 2024
Model
arXiv
Books
Wiki
C4
Python
7B
2.3
9.1
5.9
8.0
2.3
7B/4E
2.2
8.7
5.7
7.7
2.2
Table 7: Perplexity of trained models on different evaluation sets (arXiv, Books, Wikipedia, C4, and Python). Note that when training the 7B/4E model, we do not apply the similarity- batching method.
Commonsense Reasoning
Reading Comprehension
Model
PIQA
SIQA
BoolQ
HellaSwag
WinoGrande
RACE-m
RACE-h
ARC-e
ARC-c
7B
76.9
50.2
65.2
52.6
66.2
55.3
40.5
73.0
38.5
7B/4E
77.7
50.1
67.6
54.8
67.3
57.0
41.3
73.5
39.6
Closed-bond QA
Text Classification
Avg
Model
NQ
TQA
AGNews
Amazon
SST-2
Yelp
Fever
MRPC
7B
17.3
42.5
80.6
94.3
92.7
98.3
53.7
67.0
62.5
7B/4E
18.8
44.5
81.7
95.7
93.1
96.7
57.7
69.7
63.9
Table We compare the 7B/ MoE models trained with our routing strategy (without using the similarity-bant -bseed batching method) with the parameter-matched dense models on downstream tasks, including commonsense reasoning, reading comprehension, closed-book QA, and text classification.
Expert Specialization: Full Results of Routing Weights
Figure 11 shows the routing weights all layers of the 0.3B/8E model. It clearly shows that our MoE models are able to learn domain-level specialization.
G More Analysis and Ablation Studies
900
# Expert Aclivated
300
1.5
1.58/16E
1.5B/32E
100
150
Billion of tokens
Figure 9: We show how many experts are actively utilized at least once every 10 training steps. We define an expert is activated if the weight is larger than 2, where E denotes the number of experts at each MoE layer. Our models have 48 layers; therefore, the 1.5B/8E, 8E, 1.5B/16E, 1.5B/32E models have 384, 768, 1536 experts in total, respectively.
G.1 Expert Utilization
Here, investigate the expert utilization of our MoE models. We define expert that is activated for given input if the routing weight larger than (e.g., larger than 25% there are 8 experts each layer). In Figure 9, we plot the number of experts that are activated at least once among training steps (consisting of 10M tokens and 40K segments) when training 1.5B MoE models. We see that after the warmup phase at the beginning, the expert utilization quickly increases. 1.5B/SE and 1.5B/16E models have quickly utilized most of the experts; while the expert utilization of the 1.5B/32 model continues to increase until the end of the training and it able to activate about 900 experts among 1536 experts at the end of training. This indicates that approach is able to prevent the MoE models from
18

---

Published as a conference paper at COLM 2024
Model
PIQA
SIQA
BoolQ
HellaSwag
1.5B/8K (prompt)
72.1
45.2
62.0
43.6
1.58/8K (segment)
72.1
45.6
60.2
43.9
1.5B/166 16E (prompt)
71.3
45.0
56.0
43.7
1.5B/ 16E (segment)
72.9
45.4
53.2
43.6
Model
Wino
NQ
TQA
Avg
1.5B/8E (prompt)
63.7
7.3
24.2
45.4
1. 1.5B//B (segment)
61.8
7.3
24.4
45.1
1.5B/ 16E (prompt)
61.5
7.3
25.6
44.4
1.5B/ 16E (segment)
62.4
7.6
25.5
44.7
Table 9: Downstream performance using different inference methods. We study two routing strategy for inference. prompt: we make the routing decision once on the entire input prompt; segment: re-route and get new merged FFNs every segment.
collapsing to dense models and achieves high expert utilization. However, when training with large number of experts, achieving high expert utilization is more challenging.
G.2 Inference Methods
Comparison to segment-level routing during inference. During inference of downstream tasks, by default, we take the task input prompt as the input of the routers in each layer and make the routing decision once. This inference method enables the decoding process to be simple and achieves low latency, since after encoding and routing the input, we do not need to use the routers again the rest generation can be run in (merged) dense model. As such an inference method introduces train-test discrepancy, we study the method that routes every segment as we do during training. As shown in Table 9, routing the input once or routing each segment does not make substantial differences in the downstream tasks we evaluate. Due to simplicity and efficiency, use the entire prompt as the routing input and perform routing only once.
Warrup)
1.0032 (no waruup)
2.2
2.1
900
Experts Adivaled
600
300
1 158/32 (9% wantup)
1.58/3E (no wammp?
0.
50
100
150
Billion of tokens
Figure 10: Training curves and expert utilization of employing warmup phrase or not. We find without a warmup phrase, training leads to worse MoE model (top) and worse expert utilization (bottom).
Converting to sparse models for efficient inference. In our MoE models, merge FFN experts for each segment. This can lead to a large memory usage during inference when the model processes a large inference batch, because the parameters of a merged FFN per
19

---

Published as a conference paper at COLM 2024
batch are cached in the GPU memory. One possibility to alleviate this issue is to fine- tune the trained models with a hard-decision routing mechanism (e.g, top-k routing) after the pre-training stage. This method transitions the models with soft routers to ones with hard routers, significantly reducing the memory usage during inference. We leave further investigation this direction future work.
G.3 Warmup Training
At the beginning of training (i.e., the first 5% training steps), train dense LM with the same configuration before training the MoE model. We initialize the MoE layers by duplicating the FFN layers of the warmup trained model. We find that this warmup phase is crucial for achieving high expert utilization especially when there is large number of experts. Figure 10 visualizes the training loss curves and expert utilization of the 1.5B/32E model (with or without warmup training). As shown in the figure, without the warm-up phrase, the model achieves worse performance and much fewer experts are utilized.
20

---

Published as a conference paper at COLM 2024
Python
Wilipedia
Layer
Layer
Layer
0.4
0.2
0.4
0.2
0.0
Layer
Layer
Layer
Layer
Layer
0.4
Averaged Routing Weight
0.2
0.0
Layer
Layer
0.4
0.2
0.0
Layer 15
Layer 16
Layer 17
0.4
0.2
0.4
0.2
0.0
Layer
Layer
Layer
Expert ID
Figure 11: Averaged routing weights at all the layer of the 0.3B/8E model on different domains (Books, arXiv, Python, Wikipedia). We observe that the experts in our MoE models learn clear domain-level specialization.
21


Reviews:
[{"review_id": "4riuS1YxRy", "reviewer": "Reviewer_benk", "summary": "1. The paper proposes a novel method to train fully differentiable MoE models for autoregressive tasks. Concretely it proposes\n\n1.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\n\n1.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\n\n2. 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])\n\n3. 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).\n\n4. The paper also demonstrates that the proposed method achieves good domain level expert specialization without any domain level supervision.\n\n[1] Shi, Weijia, et al. \"In-Context Pretraining: Language Modeling Beyond Document Boundaries.\" arXiv preprint arXiv:2310.10638 (2023).\n\n[2] Zhou, Yanqi, et al. \"Mixture-of-experts with expert choice routing.\" Advances in Neural Information Processing Systems 35 (2022): 7103-7114.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments! We are glad that the reviewer sees great potential in training large-scale fully-differentiable MoEs. We address the concerns below.\n\n**Mismatch b/w training and testing**\n\nThanks for the excellent suggestions! In our paper, we showed that a simple inference strategy works well for a wide range of tasks, for both classification and generation. We agree that the inference strategy requires further investigation for long-form generation. However, a comprehensive study demands careful designs and evaluation, and due to computational constraints, we’re unable to delve deeper into this during rebuttal. We’ll explore these suggestions in future work.\n\n**“Expert specialization comes from the sim-batch data approach or from the proposed routing”?**\n\nIn Appendix G.1, we compare Lory and EC, both trained with sim-based batching. We show that Lory is better at modeling OOD areas such as coding than EC. This should sufficiently show that the specialization does not fully come from the batching strategy, but also due to Lory’s architectural designs. We didn’t do a full comparison on Lory between with or without sim-based batching, and we’ll add discussions around it. \n\n**Comparison with token level MoE**\n\nIn Figure 5, we show Lory performs similarly to EC. We’d have liked to compare our method with more approaches, e.g., gumbel-softmax. However, due to computational limitations, we only got to experiment with one baseline and opted for EC, which has shown to be the SoTA MoE architecture. \n\n**Poor performance of segment-level EC**\n\nOur hypothesis is that EC is less effective in segment-level routing. EC can lead to the token-dropping issue; while in segment-level EC, it may drop the whole segment completely, which may even more severely hurt the performance. \n\n**“The first segment would never get trained”**\n\nSorry about any confusion, but this is not true. The stop-grad is only applied on the merged FFN for the first segment, to avoid gradients from the first segment to update the FFNs and routers, thereby avoiding information leakage. However, the first segment will still be used to update the attention layers, embedding layers, etc.\n\n**“Why not use a BOS segment?”**\n\nGood point! We indeed conducted a similar experiment, where we make the first segment’s routing weights as trainable parameters. We find that doing so leads to slightly worse performance.\n\n**Notion for e_{i}**\n\nWe will revise Eq (1) to make it consistent.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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.\n\n2. 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\n\n3. 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:\n\n1. A fundamental mismatch between training and test objective:\nWhile 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:\n\n1.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\n\n1.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.\n\n1.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.\n\n2. Comparison with token level MoE routing:\nBased 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. \n\n2.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\n\n2.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).\n\n[1] Kamradt, G. \"Needle in a Haystack–pressure testing LLMs.\" (2023).\n\n[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 ?\n\n2. 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. \n\n2. [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) ?", "extra_scores": {"ethics_concerns_details": "N/A"}}, {"review_id": "xs2Do4oW0K", "reviewer": "Reviewer_tkvL", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback. We are encouraged that the reviewer finds our pre-training solution novel and our experimental results strong!\n\n**Compare with off-the-shelf, pre-trained, and instruction-tuned MoE models**\n\nOur goal is to explore model architectures and optimization of MoE, instead of building a state-of-the-art pre-trained MoE. Our models are only trained on 150B tokens, which is far less compared to what has been used in state-of-the-art pre-trained models (trillions of tokens). In addition, instruction-tuning of MoE is orthogonal to our contributions and is out of scope of our research.\n\nIn Section 5.3, we do make a comparison between Lory and Expert Choice (EC) under the same training setup. We show that Lory performs comparably to EC.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "This paper presents a novel mixture-of-experts pre-training solution.\n\nThe 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.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TlBYqzQgJb", "reviewer": "Reviewer_ToMC", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments! We are glad that the reviewer found that our paper tackles an important problem and our results are promising. We address the reviewer’s concerns below.\n\n**“The paper’s technical novelty is limited”, “the main difference is to extend SMEAR from fine-tuning to pre-training”**\n\nWe argue that it is important to explore new model architectures/objectives for LM pre-training in the current landscape of LLM research and development. It’s non-trivial to extend SMEAR to pre-training, due to the causal nature of generative models. To our best knowledge, we are the first to enable fully-differentiable MoE architectures for LM pre-training, which is achieved by the two techniques (causal segment routing and similarity-based data batching) proposed in our paper. \n\n**Discuss and compare with HashLayer. “The paper also got contradictory results to HashLayer without giving any explanation.”**\n\nWe disagree that our results contradict HashLayer's results. HashLayer works on token-level routing, and suggests that we should route clustered tokens to different experts. Our work uses segment-level routing, which routes adjacent tokens in a segment, which are usually topically similar to the same router. Conceptually, clustered tokens do not equate tokens in the same context window. And it’s hard to make direct comparisons between the two. We will add more discussions around it in our next revision!\n\n**Benefits from a fully differentiable architecture compared to prior work on gradient estimation**\n\nWe argue that the contribution of Lory is complementary - it is fully differentiable and easy to optimize. We think these solutions should be encouraged, and we hope to see future work further explore such paradigms, instead of resorting to the complicated optimization methods. In addition, gradient estimation methods still suffer from common MoE training issues such as token dropping and unbalanced loading, and are usually paired with an auxiliary loss. This complicates the training paradigm and makes it harder to tune. \n\nWe note that in Section 5.3, we compare Lory with Expert Choice (EC), which is a state-of-the-art method designed to address the token dropping and unbalanced loading issues. Comparing Lory with all the other MoE architectures (e.g., gumbel softmax, switch transformer) would be computationally infeasible for us, and we’d like to defer it to future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper tackles an important problem, which is to improve the convergence and quality of MoE models. \n- 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. \n- 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. \n\n- 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. \n\n- 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.\n\nPlease provide more justification of the benefits of fully differentiable MoE vs. prior work that provide gradient estimation for discrete gating functions.", "extra_scores": {}}, {"review_id": "2X9jIsMSp7", "reviewer": "Reviewer_bqcN", "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:\n\n* 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)\n* a training strategy that groups semantically-similar segments, which they find leads to much better expert utilization\n\nThe 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", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments and for pointing out the typos! We are encouraged that the review views our contribution as a step towards scaling soft MoE approaches! We address the reviewer’s concerns below.\n\n**Causal segment routing vs token-level routing**\n\nIn general, token-level routing models have greater expressiveness compared to segment-evel routing models, because of the more fine-grained routing strategy. We did preliminary experiments with segment-level routing and token-level routing with Lory, and we observed that token-level routing does perform better. However, because token-level routing is computationally expensive, as it requires merging experts for every token, we didn’t get to run a full experiment. Instead, we compare Lory (segment routing) with EC (token routing) directly, and present results in Figure 5. We acknowledge that this will be an interesting point to show, and we’d like to add more discussions in our next revision. \n\nAs for prefix routing, we consider it as a simpler alternative to causal segment routing. However,  it does not perform well. This justifies the additional considerations in Lory’s architectural design.\n\n**Downstream improvement**\n\nThe performance on text classification tasks have higher variance compared to the other tasks we used (e.g., Commonsense reasoning). But we note that our 1.5B Lory models *do* clearly outperform the 1.5 dense model on all evaluation tasks (e.g., commonsense reasoning, reading comprehension).\n\n**Gains on 7B models**\n\nWe are sorry about the confusion. Our 7B experiments are not complete unfortunately, as it only uses casual segment routing but not similarity-based data batching. We were not able to carry out full runs due to limited computational resources, and expect much bigger gains if we consider both techniques. We have mentioned it in the paper; we are also fine with taking out this result if the reviewer feels strongly about it. That is also the reason we left 7B results in the appendix instead of main paper.\n\n**“Hasn't this work shown a way to extend soft MoE from SMEAR to decoder language models?”**\n\nHere, Soft MoE is a particular architecture proposed in [1], and it is different from SMEAR. Soft MoE softly merges token representations instead of FFN weights. Hence, our techniques cannot be directly applied to Soft MoE for decoder model training, and we leave it as future work\n\n[1] Puigcerver et al., 2024. “From Sparse to Soft Mixtures of Experts”", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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? \n* 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? \n* 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?\n* 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\"\n* section 2.1: \"exper\" --> \"expert\"\n* \"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?", "extra_scores": {}}]

Author rebuttal:

Meta-review: Lory is a fully-differentiable MoE architecture, designed for autoregressive language model training. The reviewers agree to accept the submission. It's a good exploration given the context of large-scale MoE systems for LM training. The reviewers ask to include more discussions about limitations and controlled experiments.
Decision: Accept