BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts

The Mixture of Experts (MoE) framework has become a popular architecture for large language models due to its superior performance over dense models. However, training MoEs from scratch in a large-scale regime is prohibitively expensive. Existing methods mitigate this by pre-training multiple dense expert models independently and using them to initialize an MoE. This is done by using experts' feed-forward network (FFN) to initialize the MoE'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"upcycling"these models into MoEs. We propose BAM (Branch-Attend-Mix), a simple yet effective method that addresses this shortcoming. BAM makes full use of 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 of Attention (MoA) layers. We explore two methods for upcycling attention parameters: 1) 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.

Paper

Similar papers

Peer review

Reviewer ScVn6/10 · confidence 5/52024-06-14

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.

Strengths

* The paper is well put together with clear figures and illustrations of methodology. * 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. * 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

See above.

Rating

6

Confidence

5

Soundness

1

Presentation

2

Contribution

2

Limitations

Yes.

Reviewer MjaL5/10 · confidence 4/52024-06-28

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.

Strengths

1. The results demonstrate that BAM outperforms BTX under the same data and compute budget in both perplexity and downstream task evaluations. 2. 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.

Questions

Why not include the baseline of standard upcycling?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

1

Limitations

Open-sourcing the implementation and models would significantly enhance the contribution of this paper.

Reviewer MjaL2024-08-08

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.

Authorsrebuttal2024-08-08

Regarding Open-sourcing

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

Authorsrebuttal2024-08-10

NeurIPS Guidlines

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

Reviewer MjaL2024-08-11

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

Authorsrebuttal2024-08-11

Response to Reviewer Comments

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

Authorsrebuttal2024-08-12

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

Reviewer 4yUX6/10 · confidence 4/52024-07-15

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.

Strengths

1. The paper proposes a single fix (mixture-of-attention) to address the limitations of parameter averaging in previous work 2. Performance gains are somewhat small, but consistent across scale and benchmarks 3. 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 ?

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 ?

Rating

6

Confidence

4

Soundness

4

Presentation

3

Contribution

2

Limitations

Yes

Reviewer ScVn2024-08-10

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

Reviewer 4yUX2024-08-12

Re: rebuttal

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.

Authorsrebuttal2024-08-12

Reply

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

Reviewer 4yUX2024-08-12

Re: rebuttal

Thank you for the clarification.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC