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  foerster
phil  blunsom
sebastian  ruder
ahmet  üstün
acyr  locatelli
  cohere
cohere  for
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:12:17

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.

Main paper: 
 introduction: 
 Large language models (LLMs) have demonstrated remarkable performance across a wide range of tasks [2] [3] [4] . Often, LLMs use a 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 to achieve predicted improvements is 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 [MoEs; [7] [8] [9] have emerged as a popular solution to these issues [10] [11] [12] . During training and inference, each input activates only a subset of the model parameters (often referred to as experts), thereby decoupling computation cost from the total parameter count, allowing the total number of parameters to grow without additional computation. Moreover, MoEs empirically outperform dense models with equivalent computational requirements [7, 13] .However, training MoEs from scratch is expensive due larger compute requirements [12, 7] and is difficult due to 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 a small number of MoE training steps. This approach offers several advantages: 1) By utilizing a 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 [15, 16, 3] . In particular, Branch-Train-MiX [BTX; 1] initializes an MoE with N FFN experts through a three-step upcycling process. Initially, N copies of a pre-trained seed model are created. Each copy is then further pre-trained on domain-specific data for specialization. The MoE's FFN expert parameters are subsequently 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 as 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 to performance enhancement; 2) Uniform averaging of parameters from distinct models, even though they originate from the same seed, can compromise performance due to averaging divergent specializations [17] .Motivated by these limitations, we propose Branch-Attend-Mix (BAM), a simple yet effective improvement to BTX that upcycles both FFN and attention parameters into expert layers. See Figure 1 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 a variant of Mixture of Attention [MoA; 18] with soft routing where each token is assigned to every attention expert. We validate through ablation studies that this is critical to surpass baseline performance.• We propose two variants of BAM: 1) 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 Zhang et al..• To alleviate the increased computations from soft-routing attention experts, BAM employs a parallel attention transformer architecture [19, 20] 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 to 2 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. 
related work: 
 Efficient initialization for MoEs Two recent works have studied how to use pre-trained dense models to initialize MoEs [14, 1] . Sparse Upcycling [14] upcycles a single dense model by simply copying over dense model parameters as MoE parameters. It replicates FFN parameters N times into each FFN expert in sparse MoE layers. BTX [1] improves on this approach by upcycling not just a 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 are 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 [18, [21] [22] [23] investigates the mixtureof-attention framework, which uses attention experts in addition to FFN experts. This approach has not gained widespread popularity because it 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 is 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 available dense models to create an improved model, via layerwise merging or stitching of pre-trained model weights [24] . However, this approach creates a new dense model, not a sparse one. Other work has looked into dynamically adding new blocks to MoA whenever the model is trained on a new domain, but uses random initialization for the new experts instead of utilizing existing models [25] .3 Background 
parallel attention transformers: 
 The typical architecture for LLMs is based on stacking multiple blocks of the transformers [26] . The conventional transformer block comprises a Multi-Headed Attention (MHA) module, commonly referred to 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 [20] . Parallel attention transformers is a 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 1 ). Processing the two layers in parallel increases the computational throughput without degrading the performance [20] .In an 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 independently to compute its attention outputs. The individual outputs from all h heads are concatenated and linearly projected back to the dimensionality of the original model, ensuring consistent input and output dimensions throughout the network layers.EQUATIONwhereW q i , W k i ∈ R dmodel×d k , W v i ∈ R dmodel×dv and W o ∈ R hdv×dmodel , d k is the key dimension.The Mixture of Experts [MoE; 9] model replaces the FFN layer by an MoE layer in the conventional dense transformer. An MoE layer consists of a linear router, parameterized by W FFN router ∈ R dmodel×N , and a set of N FFN experts, denoted as{FFN i (x)} N i=1. The router outputs the normalized router logits p(x) = softmax (W FFN router • x), where p i (x) corresponds to the gate value for the i-th expert, FFN i . The router assigns the input token representation x to the k experts with the highest gate values. Denoting the set 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 MoE = i∈T p i (x)FFN i (x).Mixture of Attention [MoA; 18, 23] extends MoEs by also replacing the conventional attention layer with attention experts. Each MoA layer comprises of a linear router parameterized by W attn router ∈ R dmodel×N and a set of N attention experts, {MHA j (x)} N j=1 . The router outputs the normalized router logits g(x) = softmax (W attn router • x), with each g j (x) represents the gate value for the j-th expert. Each attention expert is composed of query-projection experts W q j N j=1 and the attention output projection experts W o j N j=1 . The key projection W k and value projection W v are shared among all experts for computation efficiency. The final output of the MoA layer is a gate-value weighted sum of the computations from the top-k selected experts M:MHA MoA = i∈M g i (x)MHA i (x). 
bam: branch-attend-mix: 
 There are three phases to BAM:1) Branching: We create N copies of the model from a pre-trained dense model. We will refer to this dense model as the seed model from here on. We use 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 a specialized data mixture, each tailored to different domains such as law, mathematics, 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.We initialize the mixture model using the specialized dense models from step 2. We also incorporate a copy of the original pre-trained dense seed model (similar to the BTX approach), ensuring the seed model's general knowledge acquired in the 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) It 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.The MoA architecture proposed by Zhang et al. uses only the query projection (W q ) and the attention output projection (W o ) 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 k 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 a variant of MoA where that also incorporates the key-value projections (W k and W v ) as part of the attention experts (i.e. we have one copy of W k i , W v i 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 [7] 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 no longer a 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 on how to evaluate BAM fairly comparing to baselines.We use top-1 routing of FFN experts for all experiments unless otherwise stated. For training stability, we use the auxiliary load balancing loss L LB [7] to encourage uniform load across all experts and the auxilliary router z-loss [8] to penalize large logits going into the router. Let B be the number of tokens in the batch, f i be the fraction of tokens assigned to FFN expert i, and P i be the fraction of the expert i's gate value, then the two auxiliary losses are as follows:EQUATIONThe model's final loss function L = L NLL + ∀ MoE Layers (αL LB + βL z ) is a weighted sum, where L NLL is the negative log-likelihood loss, and α and β are hyperparameters.5 Experiment Details 
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 8 in the Appendix. For both sets of experiments, we train a mixture model upcycled from four dense experts: a 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 AdamW optimizer [29] with a weight decay of 0.1. To ensure training stability, the peak learning rate for the continued pre-training phase is set to 50% that of the pre-training phase.Small-Scale Experiments Our first set of experiments uses a dense seed model with 590 million parameters, pre-trained on 400 billion token with a 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 -4 over 1500 steps, then undergoes a cosine decay to 10% of the peak learning rate.The second set of experiments utilizes a dense seed model with 2 billion parameters, pre-trained on 750 billion tokens with a batch size of 4 million tokens. During the pre-training phase, the learning rate schedule warms up from 0 to 1e -2 over 2000 steps, then 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 1e -4, warm-up steps to 1000. And as before, cosine-decay to 10% of the peak learning rate.Below we describe the data details for both BAM and BTX experiments. In the continued pretraining phase, each dense expert is trained on a specialized data domain for 100 billion tokens. The data domains contain the following data:• Mathematics: Utilizes the MathGLM [30] , GSM8K [31] , proof-pile-2 [15] , and MathPile [32] datasets. Notably, we ensure that there is no overlap between the GSM8K train and evaluation splits.• Code: Comprises the Starcoder dataset [33] .• Law: Includes the pile-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 a 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-training domain), with each domain contributing 25% to the overall mix.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.• Compute-Matching (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 in the large-scale experiments, the allocation increases to 305 TPU-core-days. BAM consistently surpasses BTX in all domains under both compute matching and token matching regimes.In addition to evaluating perplexity, we also asses large-scale models' performances with zero-shot and few-shot downstream tasks relevant to the expert domains1 .For each domain, we report the average of task scores within that domain.• Math: For mathematical reasoning, we report the average performance on GSM8K (8-shot) [31] and MATH (4-shot) [36] .• Code: For code generation, we evaluate on HumanEval (0-shot) [37] and MBPP (3-shot) [38] .• 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) [40] and TriviaQA (0-shot) [41] .• 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) [47] to test general language understanding.  
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 it 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 7.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 1 , 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 in 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 as most of the performance was indistinguishable from random guessing, likely due to the limited scale of the models. 
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 9 in the Appendix for a detailed walk-through on calculating the number of total and active parameters of each model.To match the number of total parameters in 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 5 : 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 to utilizing six FFN experts, we have adjusted BTX's routing mechanism from top-1 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.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 6 for ablation results. All experiments are compute matched at 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. We analyze the FLOPs count per token [5, 48] for BTX versus BAM during inference. In Table 7 , 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 9 in the Appendix. For the detailed FLOPs calculation steps, refer to Table 10 in the Appendix. Compared to the standard BTX, BAM consumes more FLOPs due to the soft-routing attention experts. However, the increase in FLOPs 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 parameter-matching variant of BTX, BAM contains approximately the same number of FLOPS while still performing better model-quality wise (see Table 5 ).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 is slightly slower (  
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 to upcycling just the FFN parameters, we also upcycle the dense model's attention parameters into attention experts. To maximize the benefits of specialized dense models and enhance training stability, BAM uses a 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 a 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. 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. 
 
 Appendix: For LegalBench, we specifically selected the International Citizenship QA sub-task, as it was the only sub-task among those implemented in our code-base that scored above random. 


Decision: 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.

Reviews: review: paper_summary: paper_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. 

The 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 
Experiments are conducted on both small and large-scale base LLMs, using 3 expert domains (code, law, and math). 

The authors demonstrate that overall, BAM outperforms the BTX baselines over both data and training compute budgets, for both in-distribution perplexity and common benchmarks.
strength_weakness: strength_weakness: 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 ?
questions: 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 ?
2. Could you expand on why you add 10% of text from Common Crawl in every domain ?
3. 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 ?
4. For the second ablation, was the model trained with top-3 ? What is the performance of top-1 model at that scale ?
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: questions: Why not include the baseline of standard upcycling?
limitations: limitations: Open-sourcing the implementation and models would significantly enhance the contribution of this paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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.
* 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.
* 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. 
* 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.
questions: questions: See above.
limitations: limitations: Yes.

Comments: title: Re: rebuttal
Reply
Re: rebuttal
Response to Reviewer Comments
NeurIPS Guidlines
Regarding Open-sourcing
comment: Thank you for the clarification.
Dear reviewer,

As the rebuttal period ends in just one day, we wish to gently request your feedback above once more.

We would appreciate it if you could share any remaining valid concerns you have regarding our work to justify your rating of 5, which has been advised to be used sparingly. We would be happy to try to address any remaining concerns you have.

Thank you again for all your time and feedback on our work,

Authors of 19862
Dear reviewer,

Thank you for your feedback and questions.

> Given that each added experts is copied from the base seed model and then finetuned on its own (unique?) data mixture, how exactly did you train 8 experts if there are only 4 data domains?
Here's the 3 stages for obtaining this compute-matched variant of BTX:
1) We start with the dense seed model and create three copies  (this step is identical to BAM)
2) Each of these three models is independently trained on its specific data domains. The three domains were law, math, and code. (This step is also identical to BAM)
3) We initialized three of the MoE’s FFN experts using the three expert dense models. The remaining MoE experts are all initialized from the same FFN from the seed dense model. This approach is similar to [1], where if the number of different dense models available is less than the number of experts, the same dense model's FFN is replicated multiple times to upcycle into multiple experts. Following this initialization, the entire MoE model is trained (including the expert parameters) for a small number of steps, just like BAM.

We hope this clarifies the question, and please let us know if anything remains unclear. We will also add these clarification details to the next version of the paper. 

In addition, we have conducted two additional ablation studies plus updated our existing results with more appropriate benchmark tasks during the rebuttal phase. These improvements are detailed in our response to Reviewer ScVn, and we believe these efforts have significantly strengthened the paper quality. If you're happy with our rebuttals, the authors would really appreciate if you might consider raising the score! :)

Thank you for your time,

Authors of Submission19862 

[1]  Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints, 2022
Thank you for addressing most of my concerns, and for clarifying key points. 

Regarding my second to last point, I apologize if my question was not clear. Given that each added experts is copied from the base seed model and then finetuned on its own (unique?) data mixture, how exactly did you train 8 experts if there are only 4 data domains ?

Overall I am happy with the clarifications provided by the authors.
Dear reviewer,

Thank you for engaging with our rebuttal.

We would appreciate it if you could clarify any lingering valid concerns you have regarding our work, which might justify your rating of 5. We note that the NeurIPS reviewer guide advises using a borderline score of 5 sparingly.

We believe that the inability to open-source our code (due to constraints beyond our control) cannot justify a low score like 5. Particularly since 1) As mentioned above, NeurIPS reviewer guidelines does not hold open-sourcing as a requirement 2) Our paper is neither a dataset nor a benchmark study whose primary contribution is the code.

In case you might have missed it, we have conducted two additional ablation studies plus updated our existing results with more appropriate benchmark tasks during the rebuttal phase. These improvements are detailed in our response to Reviewer ScVn.  In addition, we have incorporated your feedback to enhance the clarity significance of our work’s contributions in the forthcoming version of the paper. We believe these efforts have significantly strengthened the paper quality.

With these considerations, we urge you to reconsider and honor your previous indication of “raise the score to above 6”. If not, please let us know your remaining valid concerns.

Thank you for your time,

Authors of 19862
Dear authors,   
Thank you for your rebuttal and reminder of the policy. I'd like to increase my score to 5.
Dear Reviewer,

Thank you once again for your insightful comments and feedback on our work.

As the rebuttal period ends in just three days, we wish to gently remind you that, according to NeurIPS reviewer guidelines, giving a paper a rejection score if it's not open-sourced is against the rules. The NeurIPS submission checklist clearly states: **"While we encourage the release of code and data, we understand that this might not always be possible. Papers cannot be rejected simply for not including code."**

We kindly ask that you follow these guidelines and consider adjusting the score to "above 6" you previously committed to.

Thank you for your consideration,

Authors of Submission 19862
Thank the authors for their answers to my concerns. They have adequately answered every question I raised to my satisfaction, and therefore, I will increase my rating to 6. 

I wish this paper could be open-sourced, although it does not affect my scoring regarding NeurIPS policy. :)
Dear reviewer,

Thank you for your feedback and the positive remarks.

> I agree the combination of BTX and Mixture of Attention is meaningful and I appreciate the author's efforts

We are pleased that our rebuttal has clarified the contributions of our work. We will ensure that the final version of the paper reflects the detailed explanations provided during the rebuttal phase.

> However, if it's not open-sourced (I will raise my score to above 6 if it's open-sourced), I don't think this paper meets the bar of NeurIPS.

We appreciate your support for open-sourcing, which we also firmly believe benefits the research community. However, due to intellectual property constraints, which are beyond the authors' control, we are unable to release the source code at this time. To aid reproducibility, we have meticulously detailed the model architecture and training setup in Appendix A and Section 5 of our paper, respectively.

We also want to highlight that open-sourcing is not a "bar" for acceptance into NeurIPS main track.
Many other LLM pre-training papers have been published at NeurIPS maintrack without open-sourcing such as Chinchilla [1], and the community has found their contributions to be important (1241 citations since 2022).

Furthermore, the NeurIPS submission checklist explicitly state that the absence of open-source code should not be a ground for rejection. We quote the following directly from the NeurIPS checklist:

>"**Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code…**"

We hope this addresses your concerns, and we thank you for your consideration.

[1] Training Compute-Optimal Large Language Models 2022, https://proceedings.neurips.cc/paper_files/paper/2022/file/c1e2faff6f588870935f114ebe04a3e5-Paper-Conference.pdf
I agree the combination of BTX and Mixture of Attention is meaningful and I appreciate the author's efforts. However, if it's not open-sourced (I will raise my score to above 6 if it's open-sourced), I don't think this paper meets the bar of NeurIPS. I suggest the authors resubmit the paper to another venue.