Scalable Multi-Domain Adaptation of Language Models using Modular Experts

1. The baselines compared in this work are limited to standard LoRA and SFT. The concept of integrating modular domain expert is somewhat similar to prior papers on MoE and LoRA-MoE architecture (e.g., [1][2][3]), while the difference is MoDE applies transformer block as expert instead of feed-forward networks or LoRA block. The paper may need to add discussions of these differences to clarify its novelty. Furthermore, including other MoE or LoRA-MoE architectures as baselines would help demonstrate the superiority of the proposed approach.

2. Many low-level details are omitted, such as training and test datasets, hyperparameters of baseline methods and the two-stage training procedure. Including these details would help understand the method and facilitate future research endeavors.

3. Although there is no consensus, models >10b parameters are generally considered LLMs. The backbone pretrain model in this paper is only 1.5b, it would be better to try multiple models with larger parameter sizes and apply other PLMs such as Phi and LLaMA to demonstrate scalability.

4. In Table 1, the baseline methods are trained on the mixed dataset, whereas MoDE is trained on Math and Code, respectively, and then on the mixed data. This seems unfair as it has more training steps.

[1] LoRAMoE: Alleviate World Knowledge Forgetting in Large Language Models via MoE-Style Plugin

[2] MOELoRA: An MOE-based Parameter Efficient Fine-Tuning Method for Multi-task Medical Applications

[3] MixLoRA: Enhancing Large Language Models Fine-Tuning with LoRA-based Mixture of Experts

1. The selection probability of the gating function is unclear to me. How does the gating function select experts that represent different tasks when an input from certain task is fed into the model? It is crucial to understand how can the MoE model use the domain-specific knowledge from different experts.

2. Since the MoE will significantly increase the parameter size, e.g., in Table 1. It is unfair to compare the performance of MoE model which has more parameters with those have fewer parameters.

3. The time cost in training a MoE model is very important. Althoght this work uses some model-parallel partitioning to accelerate the training. It is also unclear the time and memory cost within the proposed method.

4. Besides, the analysis of how the relationships between the number of experts and the top-k selection affects the performance is also miss in this work. Which has strong relationship to the efficiency and effectiveness of a MoE approach.

My current ratings are based on the following weaknesses, and I am open to hearing from the authors if I'm wrong and re-evaulate the paper.

• Limited catastrophic forgetting verification. Only the English corpus is utilized for verifying the effectiveness of dealing with catastrophic forgetting (Code + Math → English). More setups would help understand the effectiveness. e.g., Code → Math + English, Math → Code + English, English → Code + Math, Code + English → Math, Math + English → Code.
• Potential unfair comparisons. Since MoDE is not sparsely activated, there are more available parameters participating in the computation during inference. It is unclear whether the performance improvement comes from the parameter scaling or the proposed method.
• Unclear settings of the pretrained 1.5B model and the multi-domain datasets.
  • The baseline performance of the pretrained 1.5B model is missing, and it is not clear how LoRA and full-param SFT would improve the performance. Do you train models from scratch on specific datasets or use a pretrained language model? If the 1.5B model is pretrained already, why not present its performance in Table 1? If the model is pretrained from scratch by your own, why not utilize a smaller pretrained model like LLaMA 3.0B or Phi? Are they too good to diminish the performance improvements?
  • It is not clear whether these datasets are seen during the pre-training stage.
  • Besides, if you utilize open-sourced datasets, please consider disclose details about the datasets. e.g., specific subsets of a open-sourced pre-training corpus.
• The paper proposes a similar idea of ModuleFormer and BTX, but they are not discussed or compared as baselines.
  • Shen, Y., Zhang, Z., Cao, T., Tan, S., Chen, Z., & Gan, C. (2023). Moduleformer: Learning modular large language models from uncurated data. arXiv preprint arXiv:2306.04640.
  • Sukhbaatar, S., Golovneva, O., Sharma, V., Xu, H., Lin, X. V., Rozière, B., ... & Li, X. (2024). Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM. arXiv preprint arXiv:2403.07816.
  • Maybe sparse upcycling should also be a baseline: Komatsuzaki, A., Puigcerver, J., Lee-Thorp, J., Ruiz, C. R., Mustafa, B., Ainslie, J., ... & Houlsby, N. (2022). Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:2212.05055.

1. The paper disregards a large body of work such as [1], [2], [3], [4], [5] that are extremely relevant and similar to the proposed method in many aspects. Some of these works propose parameter-efficient MoE architecture, and others propose modular MoEs where experts are trained separately.

2. Connected to 1st point, the proposed MoDE architecture only compared with full model fine-tuning and simple LoRA methods while there are multiple baselines that should be considered such as BTX [1], BAM [2] or MoE-LoRA methods.

3. For both full model fine-tuning and LoRA, only one variant where the backbone model is fine-tuned using the combined dataset (Code+Math) is presented. Another set of comparisons where the backbone model is fine-tuned per domain separately, and also fine-tuned using combined dataset but including English data as well required.

4. It is unclear why the paper does not include any benchmark on Code, Math, and general language understanding in English. Only next-token prediction per domain is presented however, these results may not correlate with downstream benchmarks in the same degree.

5. For the selected MoDE configuration, the additional number of parameters for only 2 domains is 0.789B which corresponds to approximately 50% of the pretrained parameter count (0.789/1.583). This is quite a high number of parameters when the number of new domains gets higher.

[1] Li, Margaret, et al. "Branch-train-merge: Embarrassingly parallel training of expert language models." arXiv preprint arXiv:2208.03306 (2022). [2] Zhang, Qizhen, et al. "BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts." arXiv preprint arXiv:2408.08274 (2024). [3] Zadouri, Ted, et al. "Pushing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning." arXiv preprint arXiv:2309.05444 (2023). [4] Wu, Xun, Shaohan Huang, and Furu Wei. "Mixture of lora experts." arXiv preprint arXiv:2404.13628 (2024). [5] Dou, Shihan, et al. "Loramoe: Revolutionizing mixture of experts for maintaining world knowledge in language model alignment." arXiv preprint arXiv:2312.09979 4.7 (2023).