Soup-of-Experts: Pretraining Specialist Models via Parameters Averaging

Dear reviewer,

We thank you for your work reviewing our paper and for your comments. We are glad to read that "The experimental designs are well conceived", and we propose "a novel method that is both theoretically grounded and practically relevant".

We now try to address your concerns.

• "the authors focus solely on next-token prediction as the evaluation metric."

We have now trained bigger 350M models with 32 experts, and evaluated the corresponding models on MMLU, arc_easy and arc_challenge. The reported accuracy is here (https://anonymous.4open.science/r/rebuttal_soup-E8F9/accs.pdf); we see that the gains obtained in loss also transfer to accuracy.

• "the method is not directly compared to other merging techniques."

See answer to rev.No4u where we compare SoEs to model merging.

• "The authors do not adequately compare their approach to existing expert merging techniques (e.g., methods detailed in [0] and [1])"

We will add the following discussion to the paper. [0] propose a variant of mixture of experts, where instead of mixing activations, they mix expert weights. However, the routing mechanism is still that of an MoE, based on input samples. On the contrary, our routing mechanism is based on input domains. For instance, one cannot use [0] to instantiate a small model, while it is one of the main features of SoEs.

[1] propose a routing mechanism to merge a bank of Lora weights. Similarly to the previous work, this cannot instantiate a frozen model of small size in a flash, while SoEs can.

• "Several key papers that provide additional context to the proposed method were not discussed in the paper. These include"

Thanks for these useful references, we will discuss the most relevant in the paper.

• "The experiments are conducted on relatively small architectures, and given the computational resources available, it would be beneficial to evaluate the approach on larger models (e.g., at least 1B-parameter models) to better assess scalability."

As discussed above, we have now trained 350M models with 32 experts.

• "It is unclear whether the number of tasks used in the expert merging experiments aligns with those used in the Soup-of-Experts experiments, leading to ambiguity about the direct comparability of the approaches."

We have compared the two approaches (see above). We want to highlight that the both methods are not used in the same way; SoE is a pretraining technique, while merging is used on fine-tuned models.

• "model soup is also a specific algorithm [...] It would be beneficial for the authors to clarify these differences when discussing the model soup concept."

Indeed, we will clarify in the text that we simply mean model averaging when speaking of model soups (i.e. we only consider uniform soups using the Model soups' paper terminology).

• "The captions are overly long"

We will simplify the captions.

• "How does your method compare to hypernetwork-based and other mixture-of-experts"

As explained above and in table1, SoEs and MoEs serve different purposes: MoEs cannot instantiate a small model, and they are not aware of the input domain distribution. We will clarify.

• "How does the individual performance of a model instantiated via the Soup-of-Experts framework compare to that of a task-specific fine-tuned model on in-distribution data?"

We believe this is the experiment in fig.5, where we compare SoEs to fine-tuned models.

• "Is there a universal guideline for selecting the prior used for the dataset weights sampler?"

In our experiments, we used ad-hoc priors that are completely blind to the downstream tasks that are eventually going to be addressed. We still obtain very good performances compared to regular pre-training. Incorporating knowledge about the downstream distribution of tasks is an exciting research direction.

• "Given that the number of dataset-specific experts in the baseline (e.g., 64) is much lower than the number of domains used in the meta-soup (4096), doesn’t this disparity affect performance?"

We also tried training models with 64 domains, but the performances where worse than with 4096 domains. The routing mechanism learns to combine those domains in a meaningful way that makes sense for the experts.

• "While training 4096 individual models may be computationally prohibitive, the difference in exposure could naturally lead to a higher generalization capacity. Could you provide more explanation on this issue?"

We believe that these models would be over-specific, and have no general knowledge at all. We conducted that experiment with 64 individual models in fig.4, and those models are all very poor generalists. Such a method would therefore perform poorly compared to SoE.

We thank you again for your detailed review which help us improve the paper ! We hope that our comments have resolved your concerns.

Dear reviewer,

We thank you for your review and for helping us improve the paper. We are happy to read that "The experiment design and analysis are appropriate.", that "this paper delivers good results", and that "the proposed method is easy to understand".

We now try to address your concerns.

• " The explanation in Section 2.5 regarding the negligible overhead cost for training the group of experts is unclear. [...] Additionally, both the forward computation and backpropagation scale linearly with the number of experts."

This is an important remark, thanks ! We will greatly clarify this part with the following analysis: We compare training a model of size d with generic pre-training and training a Soup of Experts with n experts of size d. We let b be the mini-batch size. The training loop of models consists of A) computing gradients and B) updating the parameters with Adam. For the generic model, the cost of A) scales roughly as $db$ [1], we let $C$ be the proportionality constant. This computes $\nabla\_\Theta \ell(\Theta, x)$, with a cost $C\cdot d\cdot b$. For the SoE, we need to compute the gradients w.r.t. the experts $E\_i$ and shared parameters $S$. Let $\Theta = S + \sum \alpha\_i E\_i$ the merged model. Then, the chain rule gives the gradients $\nabla_S L = \nabla\_\Theta \ell(\Theta, x)$ and $\nabla\_{E\_i}L = \alpha_i \nabla\_\Theta \ell(\Theta, x)$. In other words, the gradients w.r.t. the experts are simply obtained by rescaling the gradient wrt the merged parameters, which gives $n\cdot b$ floating point operations. Hence, the cost of A) is $C\cdot d\cdot b + n\cdot b$ for the SoE. Crucially, the cost of backprop through the SoE is only midly affected by n, as the gradients of the experts are obtained trivially from those of the merged model; the cost is not $C\cdot b\cdot n\cdot d$, which would be the same cost as that of training a much larger $nd$ full model.

For B, the generic pretraining needs to update the two adam EMA parameters, and then update the parameters. In total, the flops count is 14d. For the SoE, we need to use Adam for all the parameters, so the cost of B) is 14$n\cdot d$.

Overall, one training iteration for generic pretraining costs $C_{gen} = Cdb + 14d$, while it is $C_{soe} = Cdb + 15nd$ for the SoE. Let us discard the cost of Adam for the generic pretraining, the relative increase is $C_{soe} / C_{gen} = 1 + \frac{15n}{Cb}$. Hence we see that it all depends on whether $n \gg Cb / 15$; the practical value of the constant $C$ depends on the architecture and other factors. In a large batch size $b$ setting, the added cost is negligible.

We insist that in fig 4 and in https://anonymous.4open.science/r/rebuttal_soup-E8F9/reb_soe_time.pdf we report time as x-axis, where the additionnal (small) overhead is already factored in.

• "Furthermore, it is unclear which “standard pre-training” is being referenced for comparison"

We will clarify that standard pre-training trains one single model of size $d$, which in section 2.5. is compared to the cost of training one soup of experts with $n$ experts of the same size $d$.

-"the experimental results indicate that selecting only two domains (s = 2) and sampling uniformly from them yields the best performance. This seems counterintuitive[...]"

We want to clarify a possible misunderstanding: as explained in 2.4, we train using only $s$ active domains at a time, but these domains are sampled at random at each iteration. For instance, if we had only 3 domains and use $s=2$, sampling from $\pi$ might give these 4 samples: [0, 0.2, 0.8], [0, 0.6, 0.4], [0.1, 0.9, 0], [0.7, 0, 0.3], where there are only 2 non-zero domain at each time, but their index change. Hence, this strategy allows us to explore the space of all pairs of domains, not just one fixed pair of domain. In practice, we use 256000 iterations, so the model has seen data coming from each domain multiple times. We hope that this clarifies the "counterintuitive" observation that the SoE accelerates training. We will add this explanation in the text to make it clearer.

• "the model includes a shared branch [...]. Given the sparse domain sampling strategy, it is unclear whether this setup might lead to conflicts across different training iterations."

Indeed, as you mention, there is a tension between domains, and the shared branch has to accomodate between all domains, just like a standard LLM has to share its capacity across different domains. However, we did not observe any particular instabilities during training, the training hyperparameters (gradient clipping, adam parameters) are standard, and the curves in fig.4 are very smooth.

We hope that our answers have addressed your concerns, and thank you again for your helpful review !

[1]: Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).

Dear reviewer,

We thank you for your review and for your insightful comments, which will help us improve the paper. We are happy to hear that "The paper is overall well-written" and that "The task setup that is being studied is very relevant and important."

We now address your concerns.

• "the models trained are very small for generative LLMs (110M parameters). I suggest running an experiment on a larger model."

We have now trained larger 350M models, using 32 experts for the Soup-of-Experts. The training curves are available here: https://anonymous.4open.science/r/rebuttal_soup-E8F9/reb_soe_time.pdf. The behavior is essentially similar to that in fig.4 in the paper. Due to lack of resources, we could not train larger models.

• "the major drawback of the paper is that is does not show evaluation on downstream data sets and only shows the loss."

This is a good point. We have score the 350M models that we have trained on three datasets: MMLU, Arc-easy and Arc-challenge. We report the corresponding average accuracy as a function of time in this figure: https://anonymous.4open.science/r/rebuttal_soup-E8F9/accs.pdf

We observe that the SoE consistently outperforms generic pre-training. We will add these results to the paper. We believe that this result indeed makes the paper more convincing.

• "Additionaly, perplexity would be another interesting metric to see."

In all our experiments, we train and report with the next-token prediction loss; the perplexity can therefore be obtained by taking the exp of our curves in fig. 4, 5, 6, 7. It will not change the ordering of methods.

• "Another important drawback of this approach is that in this experiment, we need to keep n=128 experts trained, which will require a lot of space. Further, in order to efficiently do the model updates, all these experts need to be stored in memory, which is prohibitive for large models."

Indeed, we will clarify in 2.5 that the method requires storing all the expert's weights, which might be problematic for large model sizes. However, in the case where memory is not an issue, Soup-of-Experts achieves significantly better performances than standard pretraining at a fixed computational budget, as reported in Fig. 4 and https://anonymous.4open.science/r/rebuttal_soup-E8F9/reb_soe_time.pdf.

• "Claims are not fully supported by the experiments as the downstream fine-tuning metrics and perplexity are lacking."

See above, this is now fully fixed.

• "The major weakness in my view of this paper is that is lacks evaluation on downstream fine-tuned models with tasks in the specalist domain distribution. Even though the loss looks better, further validation on fine-tuning is critical to prove the effectiveness of the approach."

As seen above, we have validated that these gains in loss also transfer to downstream tasks.

• "Another model trained on this setup will add additional strength and show more generality."

As explained above, we have now trained larger 350M models. We think that this indeed improves the generality of our paper. We acknowledge that we focus solely on the decoder-only LLM architecture, which is one of the most widely used architecture nowadays.

• "The Pile data sets may not be very different to the data that is present in the RedPajama2 set, as that contains common crawl snapshots, and the pile data sets are most likely present inside these snapshots (e.g. wikipedia)."

Indeed, this is a judicious and important remark. We used the Pile as it has a variety of downstream domains that overlap more (wikipedia) or less (dm_mathematics) with the pretraining domain, in order to cover a variety of practical cases. We will clarify this in the text.

• "Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models - Li et al 2022."

Thanks for the relevant reference that we missed, we will discuss it in the paper.

• "Page 5: sentence-bert (Devlin, 2018) is the wrong citation for sentence-bert. It would be good to know which specific model version was actually used for similarity."

Thanks for spotting this ! Indeed, the model we use is a fine-tuned version of MPNET [1], obtained at https://huggingface.co/sentence-transformers/all-mpnet-base-v2. We will properly acknowledge this in the paper.

We thank you again for your work reviewing our paper and for your comments which help us improve the paper. We hope that your concerns have been alleviated.

[1] Song, Kaitao, et al. "Mpnet: Masked and permuted pre-training for language understanding." Advances in neural information processing systems 33 (2020): 16857-16867.

Dear reviewer, Thanks a lot for your review and for your questions and remarks; they will help us improve the paper. We are please to see that you found that the main claim is well supported, that "the proposed method makes sense", and that the experiments are well designed.

We now answer your questions and the points you have raised:

• "There are a number of similar approaches in the domain adaptation realm in the fine-tuning setting which are not discussed, such as rewarded soups [1] and personalized soups [2]. It would be useful for the authors to provide comparisons to these approaches,[...]"

Thanks for raising this point. While we already mentioned [1], we will discuss [2] and [3] in the paper. Thanks to your remark, we have implemented a comparison to standard model merging. To do so, we train both a SoE and a generic model during 64000 iterations, and for the SoE we use k=64 domains. Then, for each k domain, we fine-tune the generic model for T steps on that domain, yielding models $\theta_1,\dots \theta_k$ which are specialists for each domain.

We then take $d$ domains among the $k$ at random, and either a) instantiate the SoE with weights uniform over the $d$ domains, 0 otherwise, or b) merge the specialized models $\theta_i$ where i covers the $d$ domains. We use a linear combination of models to merge. We then report the average loss of the corresponding model a) or b) on the $d$ chosen domains. For the model merging, we always pick the number of fine-tuning steps T in [1000, 2000, ..., 10000] that yields the smallest loss for each individual experiment. The results are here: https://anonymous.4open.science/r/rebuttal_soup-E8F9/soe_vs_merge.pdf

When there is only $1$ domain, fine tuning is the best method, since it has trained a specialist. As soon as we need to merge >= 4 specialists, Soup-of-Experts become advantageous. Importantly, the SoE did not have to be trained on each domain individually: in total, it has done 64000 iterations, while the model merging required 64000 + 64 * 10000 steps, so about 10 times more compute ! We will add this insightful experiment to the paper.

• "This approach seems very computationally expensive, since we first need a large number of domain experts which then need further pre-training to be adapted"

We will clarify the text, the computational cost of the method is not high as long as we have enough memory to train the models. For instance, the x-axis in fig4 is the number of hours to train the models, the fact that the SoE is better means that it reaches better losses in less time; it is therefore not more expensive.

• "The authors do not compare to or discuss approaches for souping fine-tuned models for domain adaptation."

See above.

• "Experiments are limited to small models (33M - 110M parameters)"

We have now trained a 350M model with 32 experts for 33B tokens,, see https://anonymous.4open.science/r/rebuttal_soup-E8F9/reb_soe_time.pdf . The method scales: it still is about twice as fast as generic pretraining to reach a certain level of specialized loss. We will add this to the paper.

• "When is this method advantageous over using a single pretrained model of the same total number of parameters as is contained in the experts? What is the total training time comparison between this larger model when compared to training the experts and Soup-of-Experts?"

This is an interesting question. To give a concrete order of magnitude, in the experiments on the paper we use 110M models with 128 experts, yielding in total 13B parameters. There are many differences between the SoE and a 13B model. First, the SoE will most likely have a much higher perplexity, since it instantiates a model 100 times smaller. However, as a trade-off, it will have a much higher throughput since it is a 100 times smaller model, and because of that the training costs are also orders of magnitude different. For instance, training a 13B model may take weeks with 128 GPUs (see table 2 in [3]), while our training takes ~1 day with 8 gpus. Overall, these models are very different. We will clarify this in the text.

• "How does this method compare to existing approaches combining fine-tuned models from multiple domains?"

See response above. Indeed, Soup-of-Experts is a pretraining method, it should be seen as complementary to usual model merging.

• "Is this approach feasible for larger models? The large number of experts seems like a computational challenge if the models are scaled."

As said above, we could train larger 350M models on one node of 8 Gpus. We could train larger models with more computational ressources, but this is beyond our constraints.

We hope that we have convincingly answered your questions, and we thank you again for your review !

[3]: Touvron, Hugo, et al. "Llama 2: Open foundation and fine-tuned chat models." arXiv preprint arXiv:2307.09288 (2023).

Dear reviewer,

First, we thank you for your work, and for your remarks that will help us improve our work. We are happy to hear that "the claims to be supported by sound evidence", that "the methods and evaluation criteria are well aligned with the stated goal", and that "Experiments are methodologically consistent".

We now address your concerns.

• "Low-rank experts did not yield expected improvements; there is a need for further investigation."

Indeed, we could not make this idea work, and honestly report it in the paper. We think making it work would enable another avenue to scale soup-of-experts, even though we could scale them without Lora (see https://anonymous.4open.science/r/rebuttal_soup-E8F9/reb_soe_time.pdf where we train 350M models, with 32 experts for the SoE).

• "Potential scalability challenges when significantly increasing the number of experts beyond tested scenarios."

We need to select a tradeoff between the number of experts and model size, at a fixed memory size. We have been able to train 350M models with 32 experts, see https://anonymous.4open.science/r/rebuttal_soup-E8F9/reb_soe_time.pdf

• "Limited theoretical analysis or no formal guarantees regarding stability."

Indeed, we do not conduct theoretical analysis, as we are unsure what type of results we would want to prove; note that papers in a similar flavor like [1,2,3] also do conduct any theoretical analysis

• "It would have been helpful to explicitly clarify how sensitive the results are to different embedding methods used for domain clustering." and "Also, I am curious to know how sensitive is the method's performance to the choice of embedding method used during the clustering of pretraining domains?"

Thanks for the suggestion. We will clarify in the text that we use this type of embeddings since it is the one which offers the best performance are reported in [4], fig.5. Due to lack of time in the rebuttal period, we could not obtain training runs using different embeddings, but we will make sure to add this ablation to the final version.

• "Consider providing additional runtime benchmarks or memory usage comparisons explicitly against baselines."

We note that fig.4 gives a time comparison between methods, just like the newly added larger scale experiment https://anonymous.4open.science/r/rebuttal_soup-E8F9/reb_soe_time.pdf . The baseline and Soup-of-Experts training use the same hardware, that is one node of 8 A100 GPUs, and both have a GPU utilization of over 95%. We wil include these datapoints in the final paper.

• "It would be helpful if the authors clarified the effect of the MLP capacity on final performance. Does a larger routing network yield more fine-grained domain combinations?"

We have run an ablation where we change the width of the routing network. The MLP has input dimension 4096, output dimension 128, and we either take a hidden dimension of 4 * input_dimension or 4 * output_dimension. We could not see any difference between those two choices; the training curves looked similar. We will make sure to add a more thorough ablation in the final version of the paper.

We hope that our response clarifies the questions you have raised, and we thank you again for your thorough review.

[1]: Ha, David, Andrew Dai, and Quoc V. Le. "Hypernetworks." arXiv preprint arXiv:1609.09106 (2016).

[2]: Krajewski, Jakub, et al. "Scaling laws for fine-grained mixture of experts." arXiv preprint arXiv:2402.07871 (2024).

[3]: Wortsman, Mitchell, et al. "Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time." International conference on machine learning. PMLR, 2022.

[4]:Grangier, David, et al. "Task-adaptive pretrained language models via clustered-importance sampling." arXiv preprint arXiv:2410.03735 (2024).