We would like to thank R fKKY for their time and the feedback provided. We appreciate any critical insights and will take them into account as we continue refining our work.

We want to address the points raised one by one:

The methodology part of the paper is too short.

We provide a comprehensive description of our method in “Background” (Section 2), and “Adaptive Router for Upcycling Specialized Experts as MoE” (Section 3). Additionally, we also provided pseudo code for the router showing its details in Appendix A. Together with the method, we also provide very detailed explanation of our experiments in Section 4 which helps the reader understand our method. We believe that these sections clearly describe our methodology and demonstrate our experimental setup. We are happy to extend any missing point or unclear description if R fKKY points out so. We kindly would like to note that length of the methodology sections should not be considered as a criteria for rejecting an ICLR paper.

Some methods mentioned in the paper have been explored in previous research. For example, “shared expert” has been discussed in Deepseekmoe, and the adaptation to new domains (expert merging) has been discussed in Mergekit.

Thanks for your comments. We would like to clarify that our paper’s main novel contribution is the novel router that leverages information about the training domain of specialized experts to enable better MoE upcycling from separately trained dense models.

For the use of shared experts, we already cite the original paper where shared experts are trained end-to-end which is different from our setting. It is worth mentioning that the use of shared experts is new in our upcycling setting. Previous work such as BTX [1], used seed models MLP’s as separate experts but does not utilize it as a shared expert. In our experiments, we find that using the seed model’s MLP as the shared expert is particularly beneficial when adapting a new domain expert into Nexus (i.e. extending nexus with new domain expert). We believe that this is novel and demonstrates a new approach compared to the baseline in upcycling setting.

For the merging of the attention weights, we follow BTX [1] by simply averaging parameters and clearly mention this in the manuscript. We are happy to cite Mergekit as well in the related works in the context of merging.

[1] Sukhbaatar, S., Golovneva, O., Sharma, V., Xu, H., Lin, X. V., Rozière, B., Kahn, J., Li, D., tau Wen-Yih, Weston, J., & Li, X. (2024). Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM. https://arxiv.org/abs/2403.07816

Lack of comparison baselines: This paper lacks baseline models. There are already many open-sourced MoE networks for comparison, e.g., Mixtral [3], Deepseekmoe [1]. The author should compare these models, in terms of training time, data, performance, etc.

Thanks to the suggestions by R sqxx and R Ykss, we extended the baseline comparisons with data-matched seed models, where we show Nexus outperforms dense baselines even when they are further trained much longer with the data of all the experts.

However, comparing Nexus with Mixtral [2], or DeepseekMoE [3] would not be fair due to multiple reasons:

• Both models are much larger than the models that we trained in this paper. Mixtral includes 46.7B and DeepseekMoE has 16.4B total parameters.
• DeepseekMoE is trained for 2T tokens which is a much larger token count compared to our experiments. Also, they do not disclose the detailed information for the training datasets. Similarly, Mixtral does not disclose any information about their training data and number of training tokens.

Therefore, we believe that comparing models trained in this paper with these MoEs would not be fair.

As the main baseline, we compared Nexus with the BTX [4]-style standard MoE models that we upcycle from the same dense experts, using the exact same dataset. In this way, we aimed to isolate the benefits of our proposed router and showcase that Nexus router outperforms the linear router for upcycling MoE from specialized experts without any additional cost. In particular, when extending an upcycled MoE with a new domain, our experiments (Section 5.2, Figure 3) show that Nexus enables clearly better performance for the new domain with a lightweight fine-tuning.

As we address their comments, we would like to ask R fKKY to consider re-evaluating their review.

As the discussion period is nearing its end, we wanted to ask R fKKY if there are any follow-up points we can clarify. We have responded to all concerns raised, in addition to taking a few days to run the additional experiments necessary to demonstrate our points, all of which are incorporated in the appendix of the manuscript. If there are no further points of clarification regarding the manuscript, our methodology, our extended baselines, and the difference between our method and existing expert sharing and merging approaches, we would kindly ask that reviewer R fKKY considers increasing their score to reflect this.

We would like to thank R sqxx for the detailed and thoughtful feedback. We appreciate that the reviewer acknowledges that “This work shows that improving experts specialization can be helpful and can generalize the new domains.”

We want to address the points raised one by one:

Nexus introduces a dependency on the domain associated to each sample, which does not exist in the baselines

We politely want to point out that there seems to be a misunderstanding of the behavior of Nexus during inference.

• Nexus does not require data annotation during deployment, instead it utilizes a learned router to choose which expert to activate for which input.
• The router weights are calculated by a shallow MLP hypernetwork. The inputs to this network are the embeddings of each of the domains (not only the domain corresponding to the input sample!). The embeddings are constant during the whole training and inference process, and are computed before the training starts by embedding a representative sample of each training data fold once.
• At inference, not only the inputs to the MLP, but also the weights of the MLP are constant, which also leads to constant outputs. These can be precomputed and hardcoded into the router layer, which makes Nexus during inference absolutely identical to a corresponding vanilla MoE model in terms of architecture and number of total/active parameters. The distinguishing feature is that if an additional expert is added, the Nexus hypernetwork can compute the routing weights for this additional expert, instead of learning it from scratch.
• This means that Nexus does not have an advantage over the baseline and does not receive an additional input signal.

As mentioned in the manuscript, the domain labels are only required during the training time, however, this labeling can also be simply done through an unsupervised domain clustering similar to c-BTM [1]. We think that assuming the availability of the labels in training time is realistic since large pre-training datasets either include data sources or can be pre-classified for such labels.

In case we misunderstood the intention of this point raised by the reviewer, we politely ask for clarification on this issue.

[1] Gururangan, S., Li, M., Lewis, M., Shi, W., Althoff, T., Smith, N. A., & Zettlemoyer, L. (2023). Scaling Expert Language Models with Unsupervised Domain Discovery. https://arxiv.org/abs/2303.14177

What are the training, inference, and memory costs of Nexus compared to the baselines? The paper does not include a complexity analysis.

We thank reviewer R sqxx for suggesting this addition, and have added Table b) to the global response (also in the new Appendix B, Table 3), comparing the number of total and activate parameters for Nexus and all baselines, which allows comparing the memory and compute cost during training and inference.

At 470M scale, both Nexus and the baseline MoE with linear router include approximately 1.3B total parameters, however, they only use 606M active parameters at inference time. During the training, Nexus additionally updates 15M additional parameters, which correspond to the parameters of the hypernetwork to learn router weights from the domain embeddings. After the hypernetwork is trained, its constant output (as the domain embedding inputs are constant) can be calculated once and stored, afterwards these additional parameters have zero impact on the memory and compute cost of Nexus.

Similarly, at 2.8B scale, Nexus and the baseline have around 9B total and 4.3B active parameters, and Nexus updates 85M additional parameters for the router hypernetwork during training, which have no influence on the inference behavior.

Comparison with the seed model is not entirely fair as it is only trained on the SlimPajama dataset and not on the domain specific data.

We thank reviewer R sqxx for raising this valid point. We have ran the proposed additional baselines and included the results in Table a) in the global response (also in the new Appendix F, Table 6 in the paper). The experiments show that the 2.8B Nexus model outperforms the continued pretraining baseline clearly, while also providing the benefits of extensibility with new experts for easy adaptation to new domains. For more details, we refer to the global response.

The second new baseline proposed by the reviewer, which uses an “oracle model” to always pick the correct expert, is an interesting idea, but unfortunately would require information about the correct domain for a token which is not available to Nexus. While the training data does have domain labels which could be used, this approach would make inference on other datasets impossible, e.g. on downstream task evaluations or in a chat setting.

As the discussion period is nearing its end, we wanted to ask R sqxx if there are any follow-up points we can clarify. We have responded to all concerns raised, in addition to taking a few days to run the extensive experiments necessary to demonstrate our points, all of which are incorporated in the appendix of the manuscript. If there are no further points of clarification regarding the manuscript, domain dependencies, seed model baselines, or the complexity of Nexus compared to the baselines, we would kindly ask reviewer sqxx to consider increasing their score to reflect this.

The comparison to related work is inaccurate, BTM and c-BTM can be evaluated sparsely

We thank R Ykss for pointing this out. We revised the related work with a note about how BTM and c-BTM can be used sparsely.

However, even though both BTM and c-BTM could be used with top-k, this comes with unique drawbacks. Firstly, each expert in BTM / c-BTM retains a separate set of attention weights for all experts, which results in substantially higher memory requirements and an increased total/active parameter count compared to Nexus (e.g. 5 times the 2.8B seed model for c-BTM = 13.7B total parameters, instead of 9.1B total parameters for Nexus with 5 experts). Particularly, having a separate KV cache for each expert would impact inference times significantly at scale. Secondly, while both BTM and c-BTM allow top-k computations, they are restricted to using the same experts ensemble for all tokens in the sequence, while vanilla MoE and Nexus offer token-wise routing. For BTM, calculating the ensemble weights even requires a forward pass on the prompt for every expert, which breaks sparsity, and for c-BTM, a separate embedding model has to be trained and evaluated during inference. To summarize, as we propose a new MoE architecture, we mainly compare our method with linear MoE architectures like BTX [1] that are directly comparable in terms of memory and compute requirements.

[1] Sukhbaatar, S., Golovneva, O., Sharma, V., Xu, H., Lin, X. V., Rozière, B., Kahn, J., Li, D., tau Wen-Yih, Weston, J., & Li, X. (2024). Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM. https://arxiv.org/abs/2403.07816

In Experiment 5.2, it might be more suitable to expand the parameters of the linear router's weights after expanding the experts rather than resetting them. The poor performance of the linear router at 200M tokens in Figure 3 (left) might be due to the randomly initialized router being undertrained.

We thank R Ykss for this suggestion, and will include this additional experiment in the camera-ready version of the paper.

As a response to these very helpful comments, we have slightly clarified the sections in the paper touched by these issues, and added Tables 3, 5, and 6, and Figures 9 and 10 to the Appendix. We are happy to provide any further information if needed.

We would like to thank R Ykss for their thoughtful and constructive feedback. The comments have been very helpful in improving the clarity and depth of our work.

We want to address the points raised one by one:

The paper does not explain the necessity of composing domain-specific experts into experts

We would like to address the subpoints of this issue by providing additional evidence and resolving some misconceptions:

The authors could consider adding results from training a dense model on all data (including initial and subsequent data) to better illustrate the significance of converting dense models to MoE.

We thank the reviewer for the suggestion of this baseline, which can provide additional context to the performance of the Nexus method. We ran the experiment and report the results in the global response, Table a) (also included in the new Appendix F, Table 6). The 2.8B Nexus model outperforms the continued pretraining baseline clearly, while also providing the benefits of extensibility with new experts for easy adaptation to new domains. For more details, we refer to the global response.

[The authors could consider] showing the performance of each expert model on the corresponding domain task

As a response to the helpful suggestion, we systematically evaluated each expert on each domain and updated our paper and include a table of expert performance by domain in Appendix E, Table 5.

The number of tokens used to train the seed model (25B/40B) is similar to the number of tokens used for subsequent domain-specific training, the seed model may not have converged. [The authors could consider adding] experiments based on a stronger pretrained foundation model

Regarding the number of tokens used to train the seed model: the 470M/2.8B seed models were not, as assumed by R Ykss, trained on 25B/40B tokens, but instead both models were pretrained on 750B tokens, which at this scale is sufficient for the model to converge. Section 4.1 describes the training process. We have slightly updated it to make our training setting more easily understandable.

Figures 4 and 5 lack baseline comparisons when explaining specialization.

We thank R Ykss for this helpful feedback, and have ran the necessary analysis. The paper has been updated with Figures 9 and 10 in Appendix H, which show baselines for Figures 4 and 5, respectively. When comparing the routing distribution of Nexus with the baseline linear router, we find that for the training domains, both models achieve good specialization with high probability to the experts own domains. However, for the new domain (Code), Nexus demonstrates significantly higher specialization (~0.7 vs ~0.3) which also correlates with higher performance in this new domain.

The performance of the code domain decreases after the code expert is merged into the MoE model. Do experts from other domains have similar issues?

R Ykss is indeed correct in that the Nexus model with merged code expert initially performs worse than the code expert on code generations tasks. This is not surprising as the code expert is solely specialized in this domain with a high performance on the code task, but has lower performance on the other tasks (“catastrophic forgetting” as seen in the new Appendix E, Table 5). The advantage of Nexus is to adapt to new domains without regression on old domains. Furthermore, the merging procedure itself averages the non-FFN parameters of the models to merge, and we observe that this can temporarily decrease the performance on some tasks. However, after a small amount of training the model recovers. For the other domain experts, this effect is not visible because other domains do not directly correlate with any particular downstream tasks, except the Book expert, which slightly outperforms the final Nexus on MMLU (39.6 vs 39.4)

The comparison to MoE baselines is not entirely reasonable. The Nexus router has more parameters than the linear router

We thank R Ykss for pointing out the mismatch in total parameters in the models. However, we would like to clarify an important point, as the MLP in the Nexus router is subtly different from using an MLP in place of the linear router projection: The additional MLP in Nexus acts as a hypernetwork with the role of projecting domain embeddings to expert embeddings. Therefore output of this layer (i.e. hypernetwork) is used as the weights for the linear router. Most importantly, during inference, the MLP parameters do not change anymore, and neither do the inputs (=dataset embeddings). This means we can precompute the constant output of the MLP, and “hardcode” this as router weights of the model. This means that the architecture and parameter count of Nexus during inference are exactly identical to a vanilla MoE model. Finally, this additional layer adds less than 1% additional parameters compared to the overall parameter count of the model (see Table b) in the global response).

Furthermore, we present an overview of total and active parameters, from which memory need and FLOPs can be inferred, as requested by reviewer sqxx.

Table b): During training, Nexus has less than 1% additional parameters compared to the vanilla MoE, leading to no measurable increase in step times. During inference, they are even exactly equal, as the output of passing the constant dataset embeddings (1 per expert) through the frozen MLP hypernetwork can be precomputed and hardcoded as the router weights. Nevertheless, Nexus keeps the advantage of being able to flexibly recompute the router weights if a new expert is added.

In the individual responses to reviewers, we provide additional new data and clarification.

We also update the following parts of the paper:

• New Table 3 in Appendix B, showcasing the parameter counts of Nexus and baselines
• New Table 5 in Appendix E, showcasing the individual performance of each expert on each domain
• New Table 6 in Appendix F, showcasing the performance of continual training of the seed models
• New Figures 9 and 10 in Appendix F, showcasing the routing distribution of the linear MoE router, as compared to the existing Figures 4 and 5 for Nexus.
• Minor changes to clarify some sections

We look forward to engaging in meaningful discussion and welcome any additional questions at any point in time.