On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists

Thank you very much for your time and expertise in reviewing our work, as well as for your encouraging and positive feedback. Below, we address all your questions and concerns. If any issues remain, please let us know, and we’ll be happy to provide further clarification.

Novelty of the algorithm

We agree that the alternating minimization methods and its various formulations, such as EM-algorithm, have a long and rich history, which we also mention in our paper. At the same time, we believe that our optimization formulation for a particular MoE problem along with a particular open-sourced implementation (Algorithm 1 on page 11) is substantially novel. We provide it with extensive experiments, showing effective balance between general and personalized knowledge. Thank you for providing an additional reference, which we are happy to add to our work.

Violation of Assumption 1

We would like to note that our theoretical analysis demonstrates that the method is expected to converge quickly, at least under favorable conditions, a property that we believe is naturally expected from a well-designed method. However, we do not claim that this theory accounts for all possible practical scenarios. The primary validation of our algorithm for language modelling is provided through extensive experiments in Section 4.

For over-parameterized models (which may include LLMs), it is possible that different $f_{valid}$ and $f_{train}$ share the same optima, which we also note in our paper (Lines 209-212, right). At the same time, as you also highlighted in your review, we see experimentally that the method performs favorably even for different distributions from the training datasets. We leave it as an interesting open question for further research: what are the general and joint conditions on $f_{valid}$ and $f_{train}$ that imply Assumption 1. Thanks.

Proof of strong convexity

Thank you for pointing this out. In the final version of our paper, we include the formal proof of strong convexity for our decoupled objective from Section F.3. Below, we provide a brief overview of our reasoning.

Our goal is to show that the following function (see eq. (26) on page 27), $F(\Theta, \Phi, \Lambda) = l( \langle \lambda, \Theta^{\top} x \rangle) - \mu \langle \lambda, \Theta^{\top}x \rangle + \frac{\alpha}{2}( \| \Theta \|_F^2 + \| \Phi \|_F^2 ) + \mu d(\lambda) + \mu s(\Phi^{\top} x )$ is strongly convex, for sufficiently large $\alpha > 0$ and $\mu > 0$.

We notice that we can separate our objective in $\Phi$ and $\Theta$, thus constructing $g_1(\Theta, \Lambda) = l(\langle \Lambda, \Theta^\top x\rangle) + \frac{\alpha}{4}\|\Theta\|_F^2 +\frac{\mu}{2}d(\Lambda)$ and $g_2(\Phi, \Lambda) = -\mu \langle \Lambda, \Phi^\top x\rangle + \frac{\alpha}{4}\|\Phi\|_F^2 +\frac{\mu}{2}d(\Lambda)$, so that $F(\Theta, \Phi, \Lambda) = g_1(\Theta, \Lambda) + g_2(\Phi, \Lambda) + [\frac{\alpha}{4} (\|\Phi\|_F^2+\|\Theta\|_F^2) + \frac{\mu}{2} d(\Lambda) + \mu s(\Phi^\top x)]$.

Since the log-sum-exp function $s(\cdot)$ is strictly convex, while the negative entropy $d(\cdot)$ is strongly convex, as well as the Frobenius norm, we have the term in square brackets being strongly convex. Therefore, it suffices to prove the convexity for $g_1$ and $g_2$. For that, we compute the Hessian directly and show that it is positive semidefinite, for sufficiently large regularization parameters.

For the formal proof, please also see Proposition 1 in the following link, which we will include in our updated appendix: https://anonymous.4open.science/r/CoMiGS/Strongly_Convex_Proof.pdf

Minor typos

Thanks for pointing out the typos and dimensions of the simplex, we will make sure these are fixed.

Thank you very much for your time and expertise in reviewing our work, as well as for your encouraging and positive feedback. Below, we address all your questions and concerns. If any issues remain, please let us know, and we’ll be happy to provide further clarification.

Regarding your question about router training, we make the following clarifications: Within each user, the router training is the same as in standard MoE training, i.e. updated using gradient methods, apart from the following changes: 1) instead of updating router and expert parameters at the same time, we update the two sets of parameters in an alternating fashion. 2) router parameters are updated less frequently than expert parameters.

Our method seems to handle unbalanced data very well. For example in Section 4.4, to simulate high and low local data quantities, we assigned 10x tokens to data to French and Italian users as to German and Dutch users (Line 387-389). Our results show that our method is robust to the local data imbalance, no matter whether local resource abundance is positively (Figure 5) or negatively (Figure 6) correlated with local data quantities.

Thank you very much for your time and expertise in reviewing our work, as well as for your encouraging and positive feedback. Below, we address all your questions and concerns. If any issues remain, please let us know, and we’ll be happy to provide further clarification.

Methods And Evaluation Criteria

Choice of perplexity as the metric: First, we focus on next-word prediction for end users, and perplexity is a well-established metric for evaluating the quality of such predictions. Second, we believe that LLM-as-judge evaluations and existing QA tasks may not be the most appropriate metrics for assessing personalization, as the benchmarks we reviewed are primarily designed to test knowledge understanding. To properly evaluate personalization, we would need to curate custom personalized QA questions. However, since we are not experts in benchmark creation, we have chosen to rely on more standard evaluation metrics.

Experimental Designs Or Analyses

Dataset distribution across users: We briefly mentioned our user-specific dataset creation in Lines 259-260 - A distinct category is assigned to a user, as it simulates the most challenging scenario for collaboration. For example, with Multi-lingual Wikipedia dataset, we assign each category (in this case, a language) to each of the four users. The tokens per user are specified in Table 4 in the appendix. The data distribution is thus Non-IID, in the most extreme case.

Expert scores: Yes you are right. Expert scores in Figure 4 are the results of Eq.1, we will make it clear in the updated manuscript. Eq.1 gives expert scores for each single token, and what we report in Figure 4 has been averaged across all tokens and multiple batches.

Number of specialists per layer: This is a great question. Indeed we investigated the scenario where some users may not have experts, i.e. they are only equipped with generalists. Please check Figure 7 for this setting. Our experimental result shows that those low-resource users can still benefit from collaborating with high-resource users.

Weakness

Cross-device setting? Due to academic budget constraints, we were unable to conduct cross-device experiments. Our current setup involves a single user equipped with a model containing hundreds of millions to billions of parameters, and we were unable to scale testing to multiple users. Nevertheless, we believe our framework remains applicable in cross-device scenarios. One can envision a hierarchical structure in which neighboring devices form local clusters that apply our method independently. The same approach can then be extended across clusters to enable broader coordination. Importantly, our method does not rely on full user participation—while full participation can accelerate the learning of a generalist model, it is not a requirement for the method to function effectively.

Dependence on valid set. “The dependence on validation sets may cause lack of robustness in the technique under distribution shifts.” – Quite the opposite, for any OOD target distribution, as long as there is a small valid set, our router can learn a set of task-specific weights to combine generalists and specialists. Thus, CoMiGS gives great flexibility to tackle OOD target scenarios.

Sensitivity of CoMiGS. Such sensitivity is expected—it would be surprising if a more capable device performed identically to a less capable one. We would like to emphasize that our method effectively disentangles resource availability from data quantity. As demonstrated in Figures 5 and 6, we evaluate two contrasting scenarios: one where low-data-quality users are assigned more experts and high-data-quality users fewer, and another with the reverse setup. In both cases, our CoMiGS method exhibits strong robustness.

Questions

Router update and personalization: We followed standard MoE architecture, so the router is simply a one-layer MLP, that gives weights for each expert. You are right that avoiding router overfitting can as well be achieved with a smaller learning rate used for router update. In our approach, we went for less frequent updates, as it is more resource-efficient. Sorry for not making it clear, but our routers are localized per user, instead of being federated, as illustrated in Figure 2.

Futher optimization regarding #specialists per layer: Indeed, Figure 3 suggests that different layers may not require the same number of experts. However, since each user exhibits a unique pattern of expert utilization across layers, it becomes challenging to perform further optimization at the level of individual users.

Suggestions to practitioners. For practitioners, when there is no prior information on the local task complexity, we would suggest they select as many specialists as their devices allow. Our framework can effectively mitigate overfitting. In case the local task is complex, having more specialists can help with better performances, as illustrated in Figure 5.