BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts
Authors: Qizhen Zhang
Nikolas Gritsch
Dwaraknath Gnaneshwar
Simon Guo
David Cairuz
Bharat Venkitesh
Jakob Nicolaus Foerster
Phil Blunsom
Sebastian Ruder
Ahmet Üstün
Acyr Locatelli
Venue: NeurIPS 2024 poster (2024)
Keywords: Mixture of Experts
Large Language Models
Primary area: natural_language_processing
TL;DR: A simple but efficient improvement for upcycle dense models parameters and use them to initialise Mixture of Experts Models.

Abstract: Mixture of Experts (MoE) framework has become a popular architecture for large language models due to its superior performance compared to dense models. However, training MoEs from scratch in a large-scale regime is prohibitively expensive. Previous work addresses this challenge by independently training multiple dense expert models and using them to initialize an MoE. In particular, state-of-the-art approaches initialize MoE layers using experts' feed-forward parameters while merging all other parameters, limiting the advantages of the specialized dense models when upcycling them as MoEs.  We propose BAM (Branch-Attend-Mix), a simple yet effective improvement to MoE training. BAM makes full use of specialized dense models by not only using their feed-forward network (FFN) to initialize the MoE layers but also leveraging experts' attention weights fully by leveraging them as mixture-of-attention (MoA) layers. We explore two methods for upcycling MoA layers: 1) initializing separate attention experts from dense models including key, value, and query matrices; and 2) initializing only Q projections while sharing key-value pairs across all experts to facilitate efficient inference. Our experiments using seed models ranging from 590 million to 2 billion parameters show that our approach outperforms state-of-the-art approaches under the same data and compute budget in both perplexity and downstream tasks evaluations, confirming the effectiveness of BAM.

BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts
Qizhen Zhanglt
Nikolas 2ritsch
Dwaraknath Gnaneshwar
Simon Guo
David Cairuz
Bharat Venkitesh²
Jakob Foerster' Phil Blunsom 1.2
Sebastian Ruder
Ahmet Üstün Acyr Locatelli
1 University of Oxford 2 Cohere
Cohere For Al 4 Stanford University
Abstract
The Mixture of Experts (MoE) framework has become popular architecture for large language models due superior performance over dense models. However, training MoEs from scratch in large-scale regime is prohibitively expensive. Existing methods mitigate this by pre-training multiple dense expert models independently and using them to initialize MoE. This is done by using experts' feed-forward network (FFN) to initialize the MoE's 's's experts while merging other parameters. However, this method limits the reuse of dense model parameters to only the FFN layers, thereby constraining the advantages when "upcyeling" these models into MoEs. We propose BAM Frranch-Attend-Wee method that addresses this shortcoming. BAM makes full use specialized dense models by not only using their FFN to initialize the MoE layers but also leveraging experts' attention parameters fully by initializing them into a soft-variant of Mixture Attention (MoA) layers. We explore two methods for upcycling attention parameters: I) initializing separate attention experts from dense models including all attention parameters for the best model performance; and 2) sharing key and value parameters across all experts to facilitate for better inference efficiency. To further improve efficiency, we adopt a parallel attention transformer architecture to MoEs, which allows the attention experts and FFN experts to be computed concurrently. Our experiments on seed models ranging from 590 million to 2 billion parameters demonstrate that BAM surpasses baselines in both perplexity and downstream task performance, within the same computational and data constraints.
simple yet effective
Introduction
Large language models (LLMs) have demonstrated remarkable performance across wide range of tasks 4]. Often, LLMs use dense architecture, where models apply all parameters to every input during training and inference [2]3 Consequently, increasing model capacity results in increased computational cost. An empirical finding in the scaling laws of LLMs [5] indicates that model performance scales with model capacity. However, naively increasing model size achieve predicted improvements challenging due to 1) increased hardware requirements and associated costs; and 2) greater training instabilities such as gradient spikes during large-scale training runs (6). Mixture of Experts (7 have emerged as a popular solution to these issues [0 During training and inference, each input activates only a subset of the model parameters (often referred to experts), thereby decoupling computation cost from the total parameter count, allowing the total number of
'Fqual Mentorship
Work done Cohere
Corresponding authors: bachen.zhangee fing.xx.a..uk (ahmet, acyr] @cohere.ai
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

AIt
FFN
Data
Layer
Aft
AMT
Alt
Att
FFN
FFN
PPN
PFN
Att
FFN
Routerr
Router
Layer
FFN
Layer Norm
Data
Branch from dense model
Layer Norm
Continue pre-training dense experts their corresponding data domains
Initialiee mixture model's attention & FFN experts from dense experts
Figure 1: BAM operates in three phases. Different colors correspond to different expert domains, which indicates the pre-trained seed model. White indicates random parameter initialization, gradient color indicates parameter merging 1) Branching: Begin with a pre-trained dense seed model and create N copies of it. 2) Continued Pre-training: Continue to pre-train each copy independently on its own data mixture. This process yields specialized dense expert models. 3) Mixture Model Training: Utilize these specialized dense expert models to initialize both the FFN and attention experts of the mixture model. The router layers are initialized randomly. All other parameters are derived by averaging the corresponding layers in each of the dense experts. Note that BAM employs parallel attention transformer architecture that concurrently computes attention experts and FFN experts. The figure is loosely based on Figure 1 from Sukhbaatar et al.
parameters grow without additional computation. Moreover, MoEs empirically outperform dense models with equivalent computational requirements 13
However, training MoEs from scratch is expensive due larger compute requirements [12 7 and is difficult due training instabilities/8). To address this, recent works have explored more efficient alternatives by initializing MoEs using pre-trained dense models [14,1, followed by small number of MoE training steps. This approach offers several advantages: 1) By utilizing pre-trained model, we can leverage existing knowledge rather than starting from a random initialization; 2) the dense model used for initialization is smaller in total parameter count, making it faster and easier to train; 3) various pre-trained dense models are readily available in the open source community, which can be improved further without incurring additional pre-training cost [5 [. particular, Branch- Trxin-MiX [ initializes Initially, N copies of a pre-trained seed model are created. Each copy then further pre-trained on domain-specifie data for specialization. The MoE's FFN expert parameters initialized from these specialized dense models by directly copying their FFN parameters. It is important to note that the MoE's non-FFN parameters, such attention layers, are not converted to MoE layers. Since there are N dense models, each with its own set of attention parameters, BTX uses the average of the N attention parameters to initialize the MoE's attention module while the router parameters are initialized randomly. However, this approach under-utilizes the potential of the specialized dense models for two key reasons. 1) The expert knowledge embedded in the attention modules of each of these specialized models is not leveraged, despite their potential contribution performance enhancement; 2) Uniform averaging of parameters from distinet models, even though they originate from the same seed, can compromise performance due to averaging divergent specializations [17].
MoE with N FFN experts through three-step upcycling process.
subsequently
Motivated by these limitations, we propose Branch-Attend-M (BAM), a simple yet effective improvement to BTX that upcycles both FFN and attention parameters into expert layers. See Figure for an illustration of BAM. Our contributions are as follows:
To fully leverage the benefits of specialized attention experts and improve training stability, BAM uses variant of Mixture of Attention [MoA;[8] with soft routing where each token is assigned to every attention expert. We validate through ablation studies that this critical surpass baseline performance.

---

We propose two variants of BAM: I) For best model performance, BAM with KV Experts includes all attention parameters in the attention experts; 2) For enhanced inference efficiency, BAM with KV sharing shares the key and value parameters across all attention experts, similar to the method done in /Zhng et al.;
To alleviate the increased computations from soft-routing attention experts, BAM employs a parallel attention transformer architecture [9 [00 that allows the attention experts and FFN experts to be
computed in parallel, thereby improving throughput.
We validate our approach through experiments on seed models ranging from 590 million billion parameters. Our results demonstrate that BAM outperforms the BTX model in terms of both perplexity and downstream task performance across various domains under equivalent data and computational constraints, confirming the effectiveness of our proposed enhancements.
2 Related Work
Efficient initialization for MoEs Two recent works have studied how to use pre-trained dense models to initialize MoEs [4 I]. Sparse Upeycling [14] upcycles single dense model by simply copying over dense model parameters as MoE parameters. It replicates FFN parameters times into each FFN expert in sparse MoE layers. BTX (1] improves this approach by upcycling not just single dense seed model but also multiple specialized dense models branched from the seed model. This is advantageous because not only can the dense models be trained in parallel, but also they specialized in particular domains. However, unlike our method, both approaches only upcycle the FFN part of the dense (seed or specialized) models.
Alternative architecture for MoEs Alternative to the standard MoE architectures where only feed-forward blocks (FFNs) are used as MoE layers, recent work [8] 21 investigates the mixture- of-attention framework, which uses attention experts in addition to FFN experts. This approach has not gained widespread popularity because only achieves modest performance improvements while requiring additional engineering tricks for optimization [22). However, in the setting of upcycling specialized dense experts into an MoE, the usage of attention experts much more appropriate, as the attention parameters contain specialized domain knowledge that would otherwise not be accessible in the final MoE.
Alternatives Methods for Model Merging Recent work has tried to combine many openly avail- able dense models to create an improved model, via layerwise merging or stitching of pre-trained model weights (24). However, this approach creates has looked into dynamically adding new blocks to MoA whenever the model trained on a new domain, but uses random initialization for the new experts instead of utilizing existing models [25]
dense model, not sparse one. Other work
3 Background
3.1 Parallel Attention Transformers
The typical architecture for LLMs is based on stacking multiple blocks of the transformers (26]. The conventional transformer block comprises Multi-Heade Attention (MHA) module, commonly referred as the attention layer, followed by a residual connection and a feed-forward neural network (FFN). Recent works [19] introduced the parallel attention transformer, also known as parallel layers (2). Parallel attention transformers
variant of the original architecture which computes the outputs of the attention and the FFN layers in parallel using the same input. The final transformer output is the projected concatenation of the attention output and the FFN output with its residual connection (See Figure Processing the two layers in parallel increases the computational throughput without degrading the performance (20).
3.2 Multi-Headed Attention
In attention layer with h heads, the input is linearly projected to form the query, key, and value vectors through learned parameters. Each head applies the attention mechanism to its vectors inde- pendently to compute its attention outputs. The individual outputs from all h heads are concatenated
3

---

and linearly projected back to the dimensionality of the original model, ensuring consistent input and output dimensions throughout the network layers.
@WW(KW?S VW:
Head; = softmax
vdz
MHA(Q, K, V) = concatenate (head)..... head,) ) W"
where W W E Rdmaaxdd W.O E Rd xdv and Wo E hgd x dmaan dr is the key dimension.
3.3 Mixture of Experts
The Mixture of Experts [MoE: 9 model replaces the FFN layer by MoE layer in the conventional dense transformer. An MoE layer consists of a linear router, parameterized by WEFN couter E T de X
and a set of N FFN experts, denoted as (x) The router outputs the normalized router logits p(z) softmax (WFFN rourer z), where pe (z) corresponds the gate value for the i-th expert, FFN, The router assigns the input token representation x to the k experts with the highest gate values. Denoting the of selected top-k indices as T, the MoE layer's final output is the weighted sum of the selected experts' outputs, weighted by their respective gating values:
FFN Σ (x) iET
3.4 Mixture of Attention
Mixture of Attention [M 23 extends MoEs by also replacing the conventional attention layer with attention experts. Each MoA layer comprises of a linear router parameterized by Watn Watnrour router E Rddwa N and set of N attention experts, (MHA,(=)) The router outputs the normalized router logits g(x) softmax (Wam att reater z), with each gy (xx) represents the gate value for the j-th expert. Each attention expert is composed of query-projection experts { (W3)
NV
J==1 and the attention output
j- The key projection W and value projection W are shared among all experts for computation efficiency. The final output of the MoA layer is gate-value weighted sum of the computations from the top-k selected experts M:
N projection experts { [Wj} j=1'
MHAMaA = Σ (HA
iE
BAM: Branch-Attend-Mii
There are three phases to BAM:
1) Branching: We create copies of the model from a pre-trained dense model. We will refer to this dense model the seed model from here We a pre-trained seed model with a parallel attention transformer architecture. This is motivated by two reasons: 1) Parallel attention transformers offer better throughput compared to traditional transformers. with only minor performance degradation at smaller scales but no degradation at larger scales 20 2) Employing a dense seed model with a parallel attention transformer architecture allows the upcycled mixture model to inherit this parallel attention setup. This architecture enables the attention experts and FFN experts to be computed concurrently, thereby enhancing both training and inference efficiency.
2) Continued Pre-training: Each copy of the seed model undergoes independent pre-training on specialized data mixture, each tailored to different domains such as law, mathematies, and coding, This phase results in specialized dense expert models that have enhanced performance within their respective domains compared to the seed model. However, these models might perform wose in domains outside their specialization.
3) Mixture Model Training:
We initialize the mixture model using the specialized dense models from step 2. We also incorporate a copy the original pre-trained dense seed model (similar to the BTX approach), ensuring the seed model's general knowledge acquired in the pre-trate pre-training stage

---

is transferred to the mixture model. A key differentiator for BAM is it takes full advantage of the dense specialized models by leveraging both the FFN and attention layers. The parameters for each attention expert and FFN expert in the mixture model are directly inherited from their corresponding expert or seed dense model. Non-expert parameters in the mixture model, such as layer normalization, input and output embeddings, are initialized by uniformly averaging the corresponding parameters from the dense models. The router layers are randomly initialized. The upcycling of both FFN experts and attention experts is beneficial for two reasons: 1) It maximizes the utilization of specialized capabilities in the dense models; and 2) avoids averaging the attention parameters, which has shown to degrade model performance [17]. In the subsequent section, we will discuss in more detail the training methodologies of the mixture model.
4.1 Mixture Model Training: Attention Experts
The MoA architecture proposed by Zhang et al. uses only the query projection (W") and the attention output projection (W") as attention experts, while sharing the key-value projections among all experts to enhance efficiency. Like conventional MoEs, /Zhang et al. also employs top-k routing, which assigns input tokens to only & experts with the highest attention gate values.
To adapt MoA to our BAM framework, we propose two modifications to the original architecture. First, instead of the traditional top-k routing, BAM employs soft routing for the attention experts, where the attention router assigns each input token to all attention experts. We show in Section 7 that using soft-routing is crucial for achieving performance superior to the baseline. Second, we propose variant of MoA where that also incorporates the key-value projections (W* and W as part of the attention experts (i.e. we have one copy of W WT for each attention expert i). This is as oppose to than sharing them across all experts. These two modifications are motivated by three reasons:
1. Employing soft-routing and including KV experts allow us to fully exploit the capabilities of the specialized attention experts.
2. In transformer architectures, FFNs computations typically account for a larger portion of computational load than attention computations [27] 28]. By adopting the parallel attention transformer architecture, BAM enables the concurrent computation of both the attention experts module and the FFN experts module. In addition, we implement expert parallelism 77] where the individual FFNs and attention are computed in parallel across experts. Thus, we are partially "masking" the extra computations introduced by soft-routing and KV experts under the more expensive FFN experts computations.
3. Using soft-routing allows for more stable training dynamics since the router is longer discrete optimization problem. This alleviates the attention experts training from common MoE problems such as imbalanced router load 7
We discuss in Section 5.3 how evaluate BAM fairly comparing baselines.
4.2 Mixture Model Training Phase: FFN Experts
We use t routing of FFN experts for all experiments unless otherwise stated. For training stability, we use the auxiliary load balancing loss LLB 7 to encourage uniform load across all experts and the auxilliary router z-loss [8] to penalize large logits going into the router. Let be the number of tokens in the batch, fa be the fraction of tokens assigned to FFN expert i, and P; be the fraction of the expert i's gate value, then the two auxiliary losses are as follows:
LLB N Σ
S.P
(2)
B Σ
Cg (LogSumExp, (xt))²
The model's final loss function L LXT.. + MoE LNLL is the negative likeliioddooo loss, and and 3 are hyperparameters.
V (oLL + BL2)) is weighted sum, where
5

---

Experiment Details
5.1 Experimental Setup
We validate BAM against the baseline model, BTX, through two sets of experiments that vary in scale. Detailed architecture specifications are shown in Table in the Appendix. For both sets of experiments, we train a mixture model upcycled from four dense experts: general pre-training seed model, along with specialized experts in code, mathematics, and law. Each expert undergoes a continued pre-training phase using 100 billion tokens. We use the A optimizer [29] with a weight decay of 0.1. To ensure training stability, the peak learning rate for the continued pre-training phase set to 50% that of the pre- vre-trainig phase.
Small-Scale Experiments Our first set of experiments uses a dense seed model with 590 million parameters, pre-trained 400 billion token with batch size of 1 million tokens. During the pre-training and mixture model training phases, we employ a learning rate schedule that warms up from 0 to 1.2e over 1500 steps, then undergoes a cosine decay to 10% of the peak learning rate.
Large-Scale Experiments The second of experiments utilizes a dense seed model with billion parameters, pre-trained on 750 billion tokens with batch size of 4 million tokens. During the pre-training phase, the learning rate schedule warms up from cosine-decays to 3% of the peak learning rate. In the mixture model training phase, we observes gradient spikes when using the same learning rate schedule as in the pre-training phase. To mitigate this, we lower the peak learning rate to le - 4, warm-up steps to 1000. And as before, cosine-decay to 10% of the peak learning rate.
2 over 2000 steps, then
5.2 Data Details
Below we describe the data details for both BAM and BTX experiments. In the continued pre- training phase, each dense expert trained on specialized data domain for 100 billion tokens. The data domains contain the following data:
Mathematics: Utilizes the MathGLM [30]. GSM8K 31 proof-pile-2 [5 and MathPile [32 datasets. Notably, we ensure that there no overlap between the GSM8K train and evaluation splits.
Code:
Comprises the Starcoder dataset
33
Law: Includes the lile-of-law 34 and HUPD 35 datasets.
Additionally, each data domain is augmented with 10% data sourced from Common Crawl. In our ablation experiments, we observed that incorporating portion of general-text data into the expert training mix enhances model performance performance as it allows the expert data distribution to more matching with the data distribution of the seed model.
In the mixture model training phase, the training data comprises an equal distribution of data from each domain (including the pre- vre-traning domain), with each domain contributing 25% to the overall mix.
5.3 Evaluation
We compare BAM to BTX, as well as against the original dense seed model and the dense expert models. Similar to the evaluations in Sukhbaatar et al. [1, BAM is evaluated in two settings:
Data-Matching (DM): We use exactly the same training data in terms of both content and quantity for BAM as used for the BTX baseline.
(CM): We train BAM with the same amount of computation resources as BTX is trained on, measured in TPU-core-days. In the small-scale experiments, all MoE training phases are allocated 25 TPU-core-days, while the large-scale experiments, the allocation increases to 305 TPU-core-days.
Compute-Matchhe

---

Pretrain
Code
Law
Math
Average
Base Model
23.98
8.62
8.92
14.12
13.91
Specialized Models
Code Dense Expert
39.50
3.71
10.83
8.77
15.70
Law Dense Expert
87.82
23.45
7.53
30.79
37.40
Math Dense Expert
57.61
6.78
10.43
5.67
20.12
Generalist Models
BTX (Baseline)
26.72
3.78
6.63
5.77
10.72
BAM DM (Expert KV), ours
24.02
3.45
6.04
5.31
9.70
BAM CM (Expert KV), our
24.84
3.53
6.21
5.43
10.00
BAM DM (Shared KV), ours
25.67
3.64
6.39
5.57
10.32
BAM CM (Shared KV), ours
26.00
3.72
6.50
5.63
10.46
Table I: Perplexity evaluation (L) of BAM versus BTX for small-scale experiments (using a seed
model with 590M parameters). Highlighted entries denote models that outperform the BTX baseline. BAM consistently outperforms the baseline under both compute matching and token matching regimes.
Pretrain
Code
Law
Math
Average
Base Model
9.83
2.41
3.86
3.71
3.33
Specialized Models
Code Dense Expert
15.39
2.18
5.22
4.34
3.91
Law Dense Expert
32.69
6.84
3.09
8.61
6.18
Math Dense Expert
20.32
3.20
5.11
3.20
3.84
Generalist Models
BTX (Baseline)
10.35
2.40
3.76
3.64
3.27
BAM DM (Expert KV), ours
10.11
2.36
3.66
3.55
3.19
BAM CM (Expert KV), our
10.19
2.37
3.69
3.57
3,21
BAM DM (Shared KV), ours
10.20
2.37
3.69
3.59
3.22
BAM CM (Shared KV), ours
10.28
2.38
3.72
3.61
3.24
Table Perplexity evaluation (4) of BAM versus BTX for large-scale experiments (using a seed
model of 2B parameters). Highlighted entries indicate models outperforming the BTX baseline.
BAM consistently surpasses BTX in all domains under both compute matching and token matching regimes.
In addition evaluating perplexity, we also asses large-scale models' performances with zero-shot and few-shot downstream tasks relevant to the expert domains
For each domain, report the average of task scores within that domain.
Math: For mathematical reasoning, we report the average performance on GSMSK (8-shot) [3]
and MATH (4-shot) 36).
Code: For code generation, evaluate on HumanEval (0-shot) 37 and MBPP (3-shot)
3S
Law: We use the International Citizenship Questions sub-task in LegalBench 39
World Knowledge: For general knowledge on facts, we report the performance on Natural
Questions (0-shot) 440 and TriviaQA (0-shot) 4
Reasoning: For reasoning abilities, we use ARC-Easy, ARC-Challenge 42,. SIQA 43 PIQA (44), WinoGrande
45). and QUAC 46 (all 0-shot).
General Knowledge: We use MMLU (5-shot)
test general language understanding.
'For LegalBench, we specifically selected the International Citizenship QA sub-task, as it was the only sub-task among those implemented in code-base that scored above random.

---

Math
Code
Law
Know.
Reason.
MMLU
Average
Seed Model
3.68%
9.41%
73.34%
21.33%
47.73%
34.13%
31.60%
Specialized Models
Math Dense Expert
4.92%
12.39%
68.21%
13.32%
46. 16.11% %%
34.29%
29.87%
Code Dense Expert
3.19%
18.80
21.49%
12.18%
44.29%
31.50%
21.91%
Law Dense Expert
3.05%
0.20%
88.80 %
10.41 %
44.08%
32.18%
29.79%
Generalist Models
BTX (Baseline)
3.86%
10.05%
81.85%
19.07%
47.36%
34.07%
32.71%
BAM DM
4.44%
12.83%
85.47%
19.89%
47.11%
34.42%
34.02%
(Expert KV), ours
BAM CM
4.34%
12.48%
82.79%
19.51%
47.43%
34.43%
33.50%
(Expert KV), ours
BAM DM
4.10%
11.76%
86.73%
19.48%
47.27%
34.55%
33.98%
(Shared KV), ours
BAM DM
3.65%
11.77%
80.98%
19.22%
47.56%
34.16%
32.89%
(Shared KV), ours
Table 3: Benchmark evaluations (1) for BAM versus BTX. Table shows large-scale experiments
(using seed model of 2B parameters). Highlighted entries indicate models outperform the BTX baseline. All BAM variants outperform BTX average.
6
Results
We evaluate two variants of BAM: BAM (Expert KV), where all attention parameters are utilized in the attention experts, and BAM (Shared KV), where key-value (KV) projections are shared among all attention experts. We report evaluations on both variants because each presents its own advantages. BAM (Expert KV) yields better model performance because leverages more specialization from expert dense models. Conversely, BAM (Shared KV) offers greater inference efficiency because of reduced memory requirements and the shared KV cache. A more detailed analysis on inference efficiency is provided in Section]
Perplexity Evaluations We present perplexity results for both dense models (seed and specialized dense models) and generalist models (BAM and BTX). Results from small-scale experiments are shown in Table [ and large-scale experiments in Table 2 In all experiments, BAM consistent outperforms BTX under both the compute and token-matching regimes. While dense expert models sometimes excel in their respective specialized domain, they significantly under perform compared to generalist MoE models all other domains.
Benchmark Evaluations See Table 3 for downstream tasks evaluations on large scale experiments where we show that all BAM variants outperform BTX on average. We omit the small-scale downstream results most of the performance was indistinguishable from random guessing, likely due to the limited scale the models.
7 Ablations
7.1
On the Importance of Upcycling Attention Experts
We ablate the importance of upcycling attention experts in BAM by comparing it to BTX, ensuring both models have matching total parameters (equivalent computations) and/or active parameters (equivalent model capacity). All ablation experiment use the same number of 100B tokens in the mixture model training phase. See Table the Appendix for detailed walk-through on calculating the number of total and active parameters of each model.

---

Matching the Number of Total Parameters with BTX To match the number of total parameters BTX with BAM, we simply increase the number of FFN experts while still using top-1 routing. All newly added FFN experts are initialized from the seed model's FFN. As shown in Table 4 BTX matches BAM's total parameter count by employing 6 FFN experts. Even when BTX's s total number of FFN experts is increased to 8, thereby exceeding the total parameter count of BAM, BAM continues to outperform BTX.
Total Params
Active Params
Code
Law
Math
Pre-train
BAM, ours
776M
663M
3.44
6.04
5.31
24.06
BTX Experts
700M
587M
3.78
6.63
5.77
26.72
BTX Experts
738M
587M
4.11
7.19
6,25
29.08
BTX Experts
776M
587M
3.99
7.03
6.10
28.50
BTX Experts
813M
587M
3.90
6.84
6.00
27.36
BTX Experts
851M
587M
3.94
6.92
6.05
27.79
Table 4: Perplexity (4) ablation studies for small-scale experiments assess how BAM's total parameter
count compares with BTX. To match BTX's total parameters with those of BAM, we incrementally
increase the number of FFN experts BTX from up to a maximum of The additional experts upcycled using the FFN parameters from the same dense seed model.
Matching the Number of Total & Active Parameters with BTX To match the number of active parameters with BAM, we simply increase the number of top-k from the usual top-1 to top-3 3 in every MoE layer of BTX. We that BTX matches total parameters and active parameters with BAM when using top-3 routing with a total of 6 experts in T 3 of the 6 experts are upeycled from the 3 expert dense models while the other are upcycled from the
copy of the dense seed model.
Despite matching parameters, BTX does not outperform BAM.
Total Params
Active Params
Code
Law
Math
Pre-train
BAM, ours
776M
662M
3.44
6.04
5.31
24.06
BTX top-3 w/ 6 Experts
776M
662M
3.55
6.23
5,48
24.13
Table Perplexity ablation studies on small-scale experiments compare the performance of BAM to BTX under conditions where both models have equivalent numbers of active parameters and total parameters. In addition utilizing six FFN experts, we have adjusted BTX's routing mechanism from top-1 I to top-3 routing. This change means that instead of activating just one FFN expert per token representation, three are now activated, effectively increasing the number of active parameters per input.
7.2 On the Importance of Using Soft-Routing for Attention Experts
We compare soft routing with the most commonly used sparse routing approaches (i.e. top-1 and top-2 routing) in BAM's attention experts layers. See Table for ablation results. All experiments are compute matched the 590M scale. Our results indicate that BAM with top-1 and top-2 attention experts routing does not show improvement over the baseline BTX on most domains. Conversely, BAM with soft-routing consistently outperforms BTX across all domains. This demonstrates that using soft routing in MoA layers is crucial and supports our decision to implement this approach.
Pretrain
Code
Law
Math
Average
BTX
26.72
3.78
6.63
5.77
10.72
BAM soft-routing MoA
26.00
3.72
6.51
5.64
10.47
BAM top-2 routing MoA
26.68
3.78
6.69
5.75
10.72
BAM 1 routing MoA
26.89
3.83
6.69
5.82
10.81
Table Perplexity ablation (4) of BAM compute matched with different attention experts routing methods.
7.3 Analysis on Inference Efficiency
We analyze the FL.OPs count per token 5 48 BTX versus BAM during inference. Table we provide the estimates for non-embedding parameter per Transformer layer, excluding non-linearities,

---

biases, and layer normalization as they are negligible. Our results are given for the small-scale model dimensions stated in Table in the Appendix. For the detailed FLOPs calculation steps, refer to Table 10 in the Appendix. Compared to the standard BTX, BAM consumes more FLOPPs due to the soft-routing attention experts. However, the increase FLOPss is partially mitigated by our implementation of a parallel attention transformer architecture and expert parallelism. In addition, when we compare FLOPs between BAM and the paraheter-matching variant of BTX, BAM contains approximately the same number of FLOPS while still performing better model-quality wise (see T).
For empirical inference latency, we measure the time taken to generate 16 tokens from a prompt of length 256, averaged over 10 runs on TPUs. We observe that BAM slightly slower (6.17s) than standard BTX (4.81s). The increased latency in BAM is likely attributable to increased memory pressure from the individual KV caches maintained by each attention expert, as attention mechanisms are typically memory-boung 49]. Motivated by this, the K K-ssarrin version of BAM reduces the inference latency from 6. 7s to 5.96s, as all attention experts share a single KV cache. In future work, can further optimize inference time to be closer to the standard BTX baseline (representative of standard MoEs) as our hardware and framework code is optimized for training rather than inference latency.
Contig
Params
FLOPs FFN
Method
nepperts
RLopk
Total
Active
Attention
Total
BAM
4
776M
663M
35,651,584
12,589,056
48,257,024
BTX
4
700M
587M
8,912,896
12,589,056
21,310,144
BTX
776M
663M
8,912,896
37,767,168
46,692,352
Table 7: Estimates for FLOPs per token during inference. Row 1 shows BAM with KV experts, row 2 shows the standard BTX, and row 3 shows a parameter-matching variant of BTX with 6 FFN experts & top-3 routing (refer Table[5jin the paper). We use a prompt with context length of 256.
8 Conclusion & Future Work
In this work, we introduce BAM (Branch-Attend-Mix)) a simple yet effective improvement for upcycling dense models into MoE models. In addition upcycling just the FFN parameters, also upcycle the dense model's s attention parameters into attention experts. To maximize the benefits of specialized dense models and enhance training stability, BAM uses variant of Mixture of Attention with soft routing whereby each token is assigned to every attention expert. Our ablation studies confirm that this approach is crucial for surpassing baseline performance levels. We propose two variants of BAM to address different needs: BAM with KV Experts, which includes all attention parameters in the attention experts for optimal model performance, and BAM with KV Sharing, which shares the key and value parameters across the attention experts to boost inference efficiency, To alleviate the increased computations from soft-routing attention experts, BAM employs parallel attention transformer architecture. This setup enables concurrent computation of the attention expert sand FFN experts. Finally, we validate our approach through experiments on seed models ranging from 590 million to 2 billion parameters. Our results consistently show that BAM outperforms the BTX model in both perplexity and downstream task performance across various domains. These outcomes validate the effectiveness of the enhancements proposed in BAM. Future work could focus on optimizing the training data mixture for BAM's three phases and improving the training framework to accelerate both training and inference processes.
Acknovledgements
The authors would like thank Siddhartha Rao Kamalakara for his helpful pointers on debugging the inference code. The authors thank Felipe Cruz for his tremendous help in training and evaluating the seed model used in the large scale experiments. The authors also thank Amr Kayid, Cécile Robert-Michon, Joon Kim, Milad Alizadehand, Louise Rust, and Autumn Moulder for their amazing support on GPU and TPU infrastructure. The authors are grateful to John Lin, Viraat Aryabumi, and Roman Castagné for providing valuable suggestions and insights on data and evaluation. QZ also thanks Joanna Yoo, Kuba Perlin, and Sid for their help and discussions on a previous related project on MoEs sometime ago. QZ thanks Chris Lu for his discussions and feedback on the project.
10

---

References
[1] Sainbayar Sukhbaatar, Olga Golovneva, Vasu Sharma, Hu Xu, Xi Victoria Lin, Baptiste Rozière, Jacob Kahn, Daniel Li, Wen tau Yih, Jason Weston, and Xian Li. Branch-train-mix: Mixing ex- pert Ilms into a mixture-of-experts 1mm, 2024. URL AtRpsttp::/ h //axivv..vvvv//.a org/abs//////ooooooeo.ooooeeeaae 2rrg/3bb/2403... 2403. 07816
[2] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling, In International Conference on Machine Learning, pages 2397-2430. PMLR, 2023.
[3] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint M.Xi::3307.092888 2023.
[4] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vosse Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
[5] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint wrXiv:2001.0831, 2020.
[6] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2255.01066 2022.
[7] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23 (120):1-39, 2022.
[8] 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. arXiv preprint 22Xiv:2202.08900 2022.
[9] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate mixture-of-experts layer. arXiv preprint 1701:0700 2017.
[10] DeepSeek-AL. Deepseek-v22 A strong, economical, and efficient mixture-of-experts language model, 2024.
[11] Albert 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 arXiv:2401 240.000000000888, 2024.
[12] Trevor Gale, Deepak Narayanan, Cliff Young, and Matci Zaharia. Megablocks: Efficient sparse training with mixture-offeexpett Proceedings of Machine Learning and Systems, 5, 2023.
[13] Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski. Maciej Pióro. Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóždż, Piotr Sankowski, et al. Scaling laws for fine-grained mixture of experts. arXiv preprint arXiv:2402.0783 077.777777777 2024.
[14] Aran Komatsuzaki, Joan Puigcerver, James Lee- Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcyeling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:22 22. 2022.
[15] Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Albert Q Jiang, Jia Deng, Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics. arXiv preprint arXiv:2310.10631 2023.
11

---

[16] Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, al. Code llama: Open foundation models for code. arXiv preprint arXiv:2308. 12950, 2023.
[17] Margaret Li, Suchin Gururangan, Tim Dettmers Mike Lewis, Tim Althoff, Noah A Smith, and Luke Zettlemoyer. Branch-train-merge:: Embarrassingly parallel training of expert language models, arXiv preprint arXii:2208.033333 03306, 2022.
[18] Xiaofeng Zhang. Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. arXiv preprint arXiv:2210.05114, 2022.
[19] Ben Wang and Aran Komatsuzaki. GPT-J-6B: 6 Billion Parameter Autoregressive Language Model. http:///ggtuuuu...ooonoiggggoooo wttp:/u...oooo/iaa//a//
transf jax, May 2021.
[20] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):
1-113, 2023.
[21] Hao Peng. Roy Schwartz, Dianqi Li, and Noah A Smith. mixture of h 1 heads is better than h heads. arXiv preprint 1.Xiv:2005.06537 2020.
[22] Róbert Csordás, Piotr Piękos, and Kazuki Irie. Switchhead: Accelerating transformers with mixture-of-experts attention. arXiv preprint arXiv:2312.07998, 2023.
[23] Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. Jetmoe: Reaching Ilama2 performance with 0.1 m dollars. arXiv preprint 24Xx:2244..04113, 2024.
[24] Takuya Akiba, Makoto Shing, Yujin Tang, Qi Sun, and David Ha. Evolutionary optimization of model merging recipes, 2024.
[25] Yikang Shen, Zheyu Zhang, Tianyou Cao, Shawn Tan, Zhenfang Chen, and Chuang Gan. Moduleformer: Modularity emerges from mixture-of-expprts, 2023.
[26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[27] Longteng Zhang, Xiang Liu, Zeyu Li, Xinglin Pan, Peijic Dong, Ruibo Fan, Rui Guo, Xin Wang, Qiong Luo, Shaohuai Shi, al. Dissecting the runtime performance of the training, fine-tuning, and inference of large language models, arXiv preprint arXiv:2311.03667, 2023.
(28) Andrej Karpathy. I fixed the transformer diagram, 2023. URL URR.httppe://e.com/aaapppp #tatuss/6668116772255560022
[29] Iyya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101,200 2017.
[30] Zhen Yang. Ming Ding, Qingsong Lv. Zhihuan Jiang, Zehai He, Yuyi Guo, Jinfeng Bai, and Jie Tang. Gpt can solve mathematical problems without calculator. arXiv preprint 230..031,, 2023.
[31] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14663 2021.
[32] Zengzhi Wang. Rui Xia, and Pengfei Liu. Generative ai for math: Part mathhieee: A billion- tcken pretraining corpus for math. arXiv preprint arxiv:2312.171220 2023.
[33] Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint 2r11v22305.06661, 2023.
12

---

[34] Peter Henderson, Mark Krass, Lucia Zheng, Neel Guha, Christopher D Manning, Dan Jurafsky, and Daniel Ho. Pile of law: Learning responsible data filtering from the law and 256gb open- source legal dataset. Advances in Neural Information Processing Systems, 35:29217-29222 2022.
[35] Mirac Suzgun, Luke Melas- Mariazi, Suproteem Sarkar, Scott D Kominers, and Stuart Shieber. The harvard uspto patent dataset: large-scale, well-structured, and multi-purpose corpus of patent applications. Advances in Neural Information Processing Systems, 36, 2024.
[36] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021.
[37] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr. Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei. Sam McCandlish, Ilya Sutskever, and Wojeiech 7aremba. Evaluating large language models trained on code. 2021.
[38] Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang. Carrie Cai, Michael Terry, Quoe Le, and Charles Sutton. Program synthesis with large language models, 2021.
[39] Neel Guha, Julian Nyarko, Daniel Ho, Christopher Ré, Adam Chilton, Alex Siohlas-Wood Austin Peters, Brandon Waldon, Daniel Rockmore, Diego Zambrano, al. Legalbench: collaboratively built benchmark for measuring legal reasoning in large language models. Advances in Neural Information Processing Systems, 36, 2024.
[40] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming- Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: benchmark for question answering research. Transactions of the Association of Computational Linguistics, 2019.
[41] Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension, 2017.
[42] 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 reason- ing challenge. ArXiv, abs/1803.055577, 2018. httthhtss///ppp. www.nticccooooaaaaaaamaaaaaaaaaa http:pp://ppp..smmaiiicccccoalat org/
CorpusID:392283
[43] Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiga: Com- monsense reasoning about social interactions, 2019.
[44] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.
[45] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019.
[46] Eunsol Choi, He Hc, Mohit lyyer, Mark Yatskar, Wen tau Yih, Yejin Choi, Percy Liang, and Luke Zettlemoyer. Quac Question answering in context, 2018.
[47] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding, 2021.
13

---

[48] Alexandre Muzio, Alex Sun, and Churan He. Seer-moe: Sparse expert efficiency through regularization for mixture-of-ex ex-eeersss.. arXiv preprint arXiv:244..05089, 2024.
[49] Sehoon Kim, Coleman Hooper, Thanakul Wattanawong, Minwoo Kang, Ruohan Yan, Hasan Genc, Grace Dinh, Qijing Huang, Kurt Keutzer, Michael W Mahoney, et al. Full stack optimization of transformer inference: a survey, arXiv preprint arXiv:2302.44017, 2023.
14

---

A Model Architecture Details
Large Scale
Small Scale
Embedding dimension
2304
1024
FFN dimension
18432
4096
# Heads
18
8
# KV heads
6
8
Vocabulary Size
256000
256000
Activation Function
swiglu
swiglu
# layers
18
6
Positional Embedding
rope Yes
rope
Share Input & Output Embedding
No
Seed Model Parameters
Billion
590 Million
Table 8: Model architecture details for large and small scale experiments
Dense
BAM
BAM (KV sharing)
BTX top-1
BTX top-3
layernorm Block
1024
1024
1024
1024
1024
attn_or / Block
1,048,576
4,194,304
4,194.304
1,048,576
1,048,576
gkv_proj Block
3,145,728
1,258,2912
4,194.304
3,145,728
3,145,728
ffn_exp Block
4,194,304
4,194,304
4,194,304
4,194,304
12,582,912
ffn ffn_red/ Block
2,097,152
2,097,152
2,097,152
2,097,152
6,291,456
router / Block
0
4096
4096
4096
4096
input embedding
262,144,000
262,144,000
262,144,000
262, 144,00
262, 1400000
output embedding
262,144,000
142,00
362,144,000
262, 144,000
262,144,000 144,000
layernorm
1024
1024
1024
1024
1024
Active Params
587,209,728
662,731,776
612,400,128
387,234,304
662731776
Total Params
587,209,728
776,002,560
738,253,824
700,480,512
700,480,512
Table 9: Active and total parameter counts for small scale ablation experiments, All transformer parameters are reported as per transformer block.
15

---

B Inference Efficiency Analysis
Ops
BTX
BAM
Params
Heaperts dmodel Memperts moddel Neaperts to terrrtt@ooon moddet Sattnn 6 de madel
Attention Router
FLOPs
Params
Amodel
Attention: QKV
Attention
FLOPs
6d model
Params
Attention: Mask
FLOPs
2n eeserta Det model nezperta model d Me model
Params
datta dmodel 2d 2dmoodl
Attention: Projection
FLOPs
Params
Reaperta@modd
nemperts dmmddel
FFN Router
FFN
FLOPs
Mepperts Gomdel
Mespert S
Params I.neepeert Amodeddys
L. Amodeddss
FFN
FLOPs 3n moodddd
Table 10: Detailed breakdown of parameters and FLOPs per token in forward pass (inference) for each MoE layer, comparing the standard BTX and BAM with KV experts. Note that BAM uses soft-gating MoA which utilizes all the attention experts, increasing total attention FLOPs and introduces a second router gating network for attention experts.
16

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes].
Justification:
We support our claims with experiments and analysis in Sections 6 and 7.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes].
Justification: We discuss this Section 7.3 and Section 8.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA].
Justification: No proofs are needed.
4. Experimental Result Reproducibility
Question: [Yes].
Justification: We provide hyper-parametent and model architecture details in Section 5 and the Appendix.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No].
Justification: We cannot provide the code and data due to intellectual property agreements.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes].
Justification: See Section 5 for experimental settings, and appendix for architecture details.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No].
Justification: is infeasible to provide error bars due to the large scale compute intensive nature of pretraining experiments.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes].
Justification: See Section 5.3 for compute descriptions.
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes].
17

---

Justification: We don't 't violate any code of ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA].
Justification: Readers can refer to previous pretraining papers for positive and negative societal impacts.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We do release any models or data.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes].
Justification: We cite all existing datasets in Section 5.2.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]-
Justification: The paper does not release new assets.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA].
Justification: The paper does not involve crowdsourcing nor research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA].
Justification: The paper does not involve crowdsourcing research with human subjects.
18


Reviews:
[{"review_id": "JYkrrlg31q", "reviewer": "Reviewer_4yUX", "summary": "The paper extends previous work (BTX), which combines different expert LLMs into an MoE model by a) building FFN experts,  b) averaging the remaining parameters, and c) training a router over the experts. The authors propose BAM, which uses Mixture of Attention (MoA) to consolidate the different attention modules across experts, rather than parameter averaging. \n\nThe authors make some modifications on the underlying architecture as well, using soft routing across all attention experts, and employing parallel FFN / Attention blocks to absorb the cost of this more expensive attention layer. They also propose \nExperiments are conducted on both small and large-scale base LLMs, using 3 expert domains (code, law, and math). \n\nThe authors demonstrate that overall, BAM outperforms the BTX baselines over both data and training compute budgets, for both in-distribution perplexity and common benchmarks.", "questions": "1. I am not sure I fully understood the different BAM attentions. BAM with KV Sharing is essentially the standard MoA implementation, akin to multi-query attention. For BAM with KV experts, you run the full attention module, and perform a weighted average of the individual attention outputs ?\n2. Could you expand on why you add 10% of text from Common Crawl in every domain ?\n3. For the first ablation, how exactly did you add new experts ? Did you add new domains, or did you split domain(s) into multiple experts ?\n4. For the second ablation, was the model trained with top-3 ? What is the performance of top-1 model at that scale ?", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "1. The paper proposes a single fix (mixture-of-attention) to address the limitations of parameter averaging in previous work\n2.  Performance gains are somewhat small, but consistent across scale and benchmarks\n3.  The authors do a good job at taking into account both training data and training compute into their analysis to properly evaluate their method.", "weaknesses": "1. It's unclear to me what the additional cost (both in flops and in latency) that is incurred by having to go through all Attention Experts. This raises questions on the scalability of the method to more expert domains. Can the authors compare FLOP / latency at inference time for both BAM and BTX ?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer,\n\n\nThank you for your valuable feedback! We will recap your concerns/questions and address them one by one as follows\n\n> Question 1: Can the authors compare FLOP / latency at inference time for both BAM and BTX\n\nPlease see our global response for an analysis on inference efficiency. TLDR: \n\n1) BAM contains more FLOPs per token than the standard BTX. But under the optimal implementation, many of BAM’s FLOPs occurred due to attention experts can be “swept under the rug” by the parallel attention architecture we employ, as well as expert parallelization [1]. \n\n\n2) When we compare FLOPs between BAM and the parameter-matching variant of BTX, BAM contains approximately the same number of FLOPS while performing better model-quality wise (see Table 6 in the paper). \n\n\n3) We also empirically ablate on inference latency, and show that our preliminary implementation of BAM inference is slower than BTX. But as our paper mainly focuses on training gains rather than inference, our implementation is not optimal but can be greatly optimized for future work.\n\n\n[1] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity\n\n\n> Question 2: Clarifications on BAM’s KV sharing architecture\n\nFor BAM with KV Experts, each expert runs its full attention module independently. Afterward, we perform a weighted average of the outputs from these attention experts to produce the final output. \n\nBAM with KV Sharing operates similarly to the BAM with KV experts, but with a key distinction: the key and value parameters are shared among all experts to enable a single kv computation for the whole MoA layer which leads to decrease in compute and memory requirement. The aggregation remains the same, where a weighted average of the individual attention experts is used to obtain the final output.\n\nWe acknowledge the need for clearer explanations of these mechanisms and will provide further details in the camera-ready version of the paper!\n\n> Question 3: Why you add 10% of text from Common Crawl in every domain?\n\nIn our ablation experiments, we observed that incorporating a portion of general-text data into the training mix enhances the expert dense model’s performance since it allows a portion of the data matching with the data distribution of the seed model. We specifically chose Common Crawl because it is one of the largest genera-text datasets we have available, which allows us to minimize seeing repeated tokens during training. In future work, one can further explore what is the most optimal data mixture to use.\n\n> Question 4: For the first ablation, how exactly did you add new experts ? Did you add new domains, or did you split domain(s) into multiple experts.\n\nAll new experts are initialized from the same seed dense model. We do this because we assume we are working with a fixed set of existing seed and expert dense models. \n\n> Question 5: For the second ablation, was the model trained with top-3 ? What is the performance of top-1 model at that scale.\n\n\nIn the second ablation study, BAM uses soft-routing for 4 attention experts and top-1 routing for 4 FFN experts. In contrast, BTX uses top-3 routing for 6 FFN experts. We selected this exact setup for BTX so that it matches the number of total parameters as well as active parameters with the BAM experiments shown in Tables 2 and 3 in the paper. These two runs were token matched, meaning they were trained on the exact same training data.\n\nWe did not test BTX with top-1 routing for 6 FFN experts. This decision was based on the observation that, in Mixture of Experts, performance typically improves as the 'k' value increases in top-k routing due to the increased number of active parameters. Therefore, we anticipated that top-1 routing would underperform compared to top-3 under these conditions.\n\n\nWe want to thank the reviewer again for their feedback. If you believe we have addressed your concerns adequately, we would greatly appreciate it if you could consider raising the score. Otherise, please let us know any remaining concerns and we will provide further details.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uLetBirz1U", "reviewer": "Reviewer_MjaL", "summary": "The paper proposes BAM (Branch-Attend-Mix), a novel approach to improve the training of Mixture of Experts (MoE) models by fully leveraging the parameters of pre-trained dense models. The authors introduce a method to initialize both feed-forward network (FFN) and attention layers from specialized dense models, enhancing the efficiency and performance of MoE training. The approach is evaluated on language models ranging from 590 million to 2 billion parameters, demonstrating superior performance in terms of perplexity and downstream task evaluations compared to existing methods.", "questions": "Why not include the baseline of standard upcycling?", "limitations": "Open-sourcing the implementation and models would significantly enhance the contribution of this paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "1. The results demonstrate that BAM outperforms BTX under the same data and compute budget in both perplexity and downstream task evaluations.\n2. The paper is exceptionally well-written, with a clear and logical flow that makes it easy to follow.", "weaknesses": "The technical novelty of this paper is relatively low, primarily because it combines existing methods, namely Branch-Train-Mix (BTX) and Mixture of Attention, without introducing any significant contributions. While the integration of these approaches is executed well, the paper does not offer significant new insights or methodologies beyond this combination.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer,\n\nThank you for your valuable feedback! We will recap your concerns/questions and address them one by one as follows:\n\n> Concern #1: Novelty of this work\n\nWe respectfully disagree with the assessment that the combination of BTX and Mixture of Attention is trivial. We list some of the non-triviality below:\n\n\n1) **Validation of Concept**: It cannot be simply assumed that combining Mixture of Attention with BTX would automatically yield superior results. Our hypothesis required ablation, which we undertook and presented in Section 7 of the paper. This testing confirmed that our approach does indeed enhance performance over simply scaling up the number of feed-forward experts.\n\n2) **Challenges**: Naively combining the mixture of attention approach and BTX does not work. In fact, initial attempts using the typical top-1 and top-2 routing for the attention experts performed worse than the baseline BTX across most domains (See Table 3 in the PDF for these ablations). As a solution, we propose to use a soft-routing variant of the original Mixture of Attention that effectively addressed this.\n\n3) **Efficiency Improvements**: To ensure that our model, BAM, operates efficiently, we adopted a parallel transformer architecture so that the attention experts and FFN experts can be computed in parallel. This hide away some of the additional computation costs incurred from adding the attention experts. Please see the global response for additional ablations on flops and latency for BAM. This adaptation was necessary to manage the increased complexity introduced by the combined methodologies.\n\nIn the camera-ready version of our paper, we will further emphasize these non-trivial aspects of our research to clarify the contributions of our work. We urge the reviewer to reconsider the reject score due to “a lack of novelty”. We believe that a simple, effective, and “well-executed” solution should rather be applauded.\n\n> Question #2: Why not include the baseline of standard upcycling?\n\nWe didn’t include it because it has already been established that BTX is a stronger baseline than sparse upcycling [1], and that we want to focus our limited compute on the main comparison and ablation showcasing the validity of our approach. \n\t\n> Concern #3: Not open-sourcing\n\nWe understand the value of open-sourcing to the research community. However, our decision not to release the codebase and models is mainly due to the training framework and models are proprietary to the organization this work was done at. Despite these limitations, we have provided detailed descriptions of the hyperparameters, model architecture, and methods within our paper. This level of detail should enable readers to replicate our results using existing open-source LLM training codebases. Note that many significant works in the field, such as ones we build off of like PALM [2] and Switch Transformers [3] have also not been open-sourced but are still considered important contributions to the community. We are committed to contributing meaningfully to the field while navigating these complexities.\n\n[1] Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM, 2024\n\n[2] PaLM: Scaling Language Modeling with Pathways, JMLR, 2023\n\n[3] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, JMLR, 2022\nWe want to thank the reviewer again for their feedback. If you believe we have addressed your concerns adequately, we would greatly appreciate it if you could consider raising the score. Otherise, please let us know any remaining concerns so we can provide further details.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M1BoR9ArlD", "reviewer": "Reviewer_ScVn", "summary": "This paper proposes a method for upcycling specialized dense models into mixture models. The upcycling is applied to both FFNN and Attention, within a parallel-attention transformer architecture. The evaluation shows their method’s superiority upon a baseline that only upcycles the FFNN layers.", "questions": "See above.", "limitations": "Yes.", "rating": 6, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "* The paper is well put together with clear figures and illustrations of methodology.\n* The improvement by applying MoA upon baseline (i.e., BTX) is well-motivated.", "weaknesses": "* Table 3 shows the superiority on average of BAM over BTX. However, the specialized models seem not that specialized in their corresponding tasks, which makes the conclusion less convincing. For example, the “Law Dense Expert” beats the “Math Dense Expert” on the mask task.\n* This work employs a parallel-attention transformer architecture which is not quite popular. Since the authors state it “increases the computational throughput without degrading the performance”, it would be better to conduct more ablation experiments to further understand its influence on performance and efficiency.\n* The MoA requires “additional engineering tricks for optimization” as the authors said, it would be better to conduct further ablations studies by comparing the performance with BTX under the same inference throughput budget. \n* It would be better to understand the role of soft routing, instead of vanilla sparse routing, in the MoA layers either from empirical or theoretical perspectives.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer,\n\n\nThank you for your valuable feedback! We will recap your concerns/questions and address them one by one as follows.  \n\n> Concern #1: Specialized models seem not that specialized in their corresponding tasks, making the conclusion less convincing.\n\nThank you for highlighting this issue. We observe this unspecialized behavior in two areas of our original submission’s results:\n\n1) **In the 590M models' downstream evaluations**: We recommend interpreting the abilities of the 590M models primarily through the perplexity results and disregarding the downstream tasks results. The performance of these models on downstream tasks is very noisy and models are close to random guessing in most tasks due to their small size. We initially included these results for readers who might be curious about downstream evaluations at this scale, but for clarity, we will remove them from the camera-ready copy.\n2) **In the 2B models' downstream evaluations**: Specifically, the \"Math Dense Expert\" outperforms the \"Law Dense Expert\" on the law task. Upon further examination, we found that the LSAT benchmark [1] we used for evaluating law downstream performance primarily tests general reasoning abilities [1] rather than legal knowledge. Consequently, we re-evaluated all models using subtasks from LegalBench [2], an appropriate benchmark for law tasks. As shown in Table 2 of the updated PDF (see global response), all dense experts now perform best within their respective domains. Our conclusion that BAM outperforms baselines remains robust.\n\nTLDR: By excluding the noisy downstream evaluations of small-scale experiments, our results indeed show that specialized models do specialize in their corresponding tasks. Our conclusion holds true through these results.\n\n\n[1] From LSAT: The Progress and Challenges of Complex Reasoning, 2021\n\n\n[2] LegalBench: A Collaboratively Built Benchmark for Measuring Legal Reasoning in Large Language Models, 2023\n\n\n> Concern #2: Parallel-attention transformer architecture is not quite popular […]\n\nBoth PALM [3] (5000+ citations since 2023) and the open-source project GPT-J [4] (6000+ stars on GitHub since 2021) employ parallel-attention transformers, indicating that this architecture is indeed “quite popular”. Section 2 of the PALM paper already discussed findings on the efficiency of this architecture. We have cited both of these works in our paper and will further highlight these works in the related-works section. We apologize for any confusion caused by referring to this architecture as \"parallel-attention transformers\" in our paper instead of \"parallel-layers\" as used in the PALM paper. We will correct the naming of this architecture in the camera-ready version.\n\nGiven the established findings from these popular works, it should not be expected that we need to reinvent their results through additional ablation experiments. Our focus remains on building upon these validated architectures to explore new avenues of research.\n\n[3] PaLM: Scaling Language Modeling with Pathways, Journal of Machine Learning Research 2023 \n\n\n[4] GPT-J, 2021 https://github.com/kingoflolz/mesh-transformer-jax\n\n\n> Concern #3: Comparing the performance with BTX under the same inference throughput budget.\n\nPlease see our global response for an analysis on inference efficiency both theoretically and empirically! \nTLDR: \n1) BAM contains more FLOPs per token than the standard BTX. But under the optimal implementation, many of BAM’s FLOPs occurred due to attention experts can be “swept under the rug” by the parallel attention architecture we employ, as well as expert parallelization [5]. \n\n2) When we compare FLOPs between BAM and the parameter-matching variant of BTX, BAM contains approximately the same number of FLOPS while performing better model-quality wise (see Table 6 in the paper). \n\n3) We also empirically ablate on inference latency, and show that our preliminary implementation of BAM inference is slightly slower than BTX. But as our paper mainly focuses on training gains rather than inference, our implementation is not optimal but can be greatly optimized for future work.\n\n[5] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity\n\n> Concern #4:  [Authors should] understand the role of soft routing, instead of vanilla sparse routing, in the MoA layers either from empirical or theoretical perspectives.\n\n\nTo address this concern, we have included additional ablation studies comparing soft routing with the most commonly used sparse routing approaches (i.e. top-1 and top-2 routing) in BAM’s attention experts layers. These experiments were conducted at the 590M scale, with all models given the same amount of compute. These results are in Table 3 of the PDF (see global rebuttal for PDF).  Our results indicate that BAM with top-1 and top-2 attention experts routing does not show improvement over the baseline BTX on most domains. Conversely, BAM with soft-routing consistently outperforms BTX across all domains.This demonstrates that using soft routing in MoA layers is crucial and supports our decision to implement this approach.\n\nWe want to thank the reviewer again for their feedback. If you believe we have addressed your concerns adequately, we would greatly appreciate it if you could consider raising the score. Otherise, please let us know any remaining concerns so we can provide further details.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear reviewers,

Thank you again for your valuable feedback! 

In the attached PDF, we have provided updated experiments and additional ablations. 

In addition, see blow for additional analysis on inference efficiency. 

# Inference Arithmetic

We analyze the parameter and arithmetic for BTX vs. BAM during inference (the forward pass). Table 1 below shows the estimates for "FLOPs per Token" per Transformer layer. We use the common FLOPs counting methodology employed [1] and [2], where we exclude non-linearities, biases, and layer normalization which are negligible. 

**Table 1**
| **Operation**                 |           | **BTX**               | **BAM**                        |
|-------------------------------|-----------|-----------------------|--------------------------------|
| **Attention Router**          | Params    | -                     | $n_{experts}d_{model}$         |
|                               | FLOPs     | -                     | $2n_{experts}d_{model}$        |
| **Attention: QKV**            | Params    | $d_{model}3d_{attn}$  | $n_{experts}d_{model}3d_{attn}$|
|                               | FLOPs     | $6d_{model}^2$        | $6n_{experts}d_{model}^2$      |
| **Attention: Mask**           | Params    | -                     | -                              |
|                               | FLOPs     | $2n_{ctx}d_{model}$   | $2n_{experts}n_{ctx}d_{model}$ |
| **Attention: Projection**     | Params    | $d_{attn}d_{model}$   | $n_{experts}d_{attn}d_{model}$ |
|                               | FLOPs     | $2d_{model}^2$        | $2n_{experts}d_{model}^2$      |
| **FFN Router**                | Params    | $n_{experts}d_{model}$| $n_{experts}d_{model}$         |
|                               | FLOPs     | $2n_{experts}d_{model}$| $2n_{experts}d_{model}$       |
| **FFN***                      | Params    | $1.5n_{experts}d_{model}d_{ff}$ | $1.5n_{experts}d_{model}d_{ff}$|
|                               | FLOPs     | $3n_{topk}d_{model}d_{ff}$ | $3n_{topk}d_{model}d_{ff}$|





# Theoretical and Empirical Anlysis
## Theoretical FLOPS Counting

Using the arithmetic formulation above, Table 2 examines the inference FLOPs on the small-scale experiments (see Appendix A). Row 1 shows BAM with KV experts, row 2 shows the standard BTX, and row 3 shows a parameter-matching variant of BTX with 6 FFN experts & top-3 routing (refer to Table 6 in the paper). We use a prompt with context length of $256$.

**Table 2**
| Method                  | **$n_{experts}$** | **$n_{topk}$** | Total Param | Attention FLOPs | FFN FLOPs  | Total FLOPs |
| ----------------------- | ------------- | ---------- | ----------- | --------------- | ---------- | ----------- |
| BAM                     | 4             | 1          | 776M        | 35,651,584      | 12,589,056 | 48,257,024  |
| BTX (Standard)| 4             | 1          | 700M        | 8,912,896       | 12,589,056 | 21,510,144  |
| BTX (Param Matching)    | 6             | 3          | 776M        | 8,912,896       | 37,767,168 | 46,692,352  |



## BAM Efficiency Optimization
Compared to BTX (standard), BAM consumes more FLOP due to the soft-routing attention experts. To address the increase in FLOPs, we implement the following optimizations:
1) All attention experts are computed in parallel via expert parallelism [3].
2) The attention experts and FFN experts are computed in parallel rather than sequentially (see Figure 1 in the paper), using this parallel reformulation from [4]: $$y = x + MLP(LayerNorm(x)) \textbf{+} Attention(LayerNorm(x))$$ 

In Table 2, row 3, we see that each attention expert is significantly less FLOP intensive than each FFN expert. Thus, with the parallel transformer block, we are computing the FLOP-expensive and FFN expert(s) alongside the less FLOP-expensive and smaller attention expert(s), all while using expert parallelism and effectively overlapping computations. This way we are effectively overlapping compute.

## Empirical Inference Latency
For empirical inference latency, we generate 16 tokens given a prompt with $256$ length, averaged over 10 runs on TPU v4s; we note BAM slightly is slower than BTX: 
    BAM with KV experts: `6.17s`
    BTX (standard): `4.81s`.
    
Although we overlap computation with the parallel transformers architecture and expert parallelism, BAM's longer inference time is likely casued by increased memory pressure from attention experts (each expert has a KV cache) as attention is memory-bound [5]. Motivated by this, the KV-sharing version of BAM reduces the inference latency from `6.17s` to `5.96s`, as all attention experts share one KV cache. We believe we can further optimize inference time to be close to BTX baseline (representative of standard MoEs) as our hardware and framework code is optimized for training rather than inference latency.



## TLDR
Takeaways:
1) BAM contains more FLOPs than BTX, but under optimal implementation, many of BAM’s FLOPs due to attention experts can be "swept under the rug" by the parallel attention architecture and expert parallelism.
2) When we compare FLOPs between BAM and the parameter-matching variant of BTX, BAM contains approximately the same number of FLOPS while performing better model-quality wise (see Table 6 in the paper). 
3) We also empirically ablate on inference latency, and show that our preliminary inference implementation of BAM is slower than BTX. But as our paper mainly focuses on training gains rather than inference, our primary implementation is not optimal but can be greatly optimized for future work.



# Reference
[1] Scaling Laws for Neural Language Models, 2020

[2] SEER-MoE: Sparse Expert Efficiency through Regularization for Mixture-of-Experts, 2024.

[3] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, 2021

[4] PaLM: Scaling Language Modeling with Pathways, 2022.

[5] Full Stack Optimization of Transformer Inference: a Survey, 2023.

Meta-review: The paper looks into methodologies for training high quality MoE models. In particular, the paper proposes BAM, a simple yet effective parameter upcycling method, which extends prior work Branch-Train-Mix (BTX) while also leveraging Mixture of Attention (MoA) to consolidate attention modules across experts. The paper also explores additional design choices, such as soft routing and parallel attention. Overall, evaluation shows promising results on small to large scale LLMs across different domains. 

How to train high quality MoE is a timely topic, and the paper has made some interesting observations, which shall help the community improve the training of MoEs. The writing of the paper is also of high quality, and the reviews are overall positive. Therefore, I recommend acceptance to this paper.
Decision: Accept (poster)