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

Thank you for your feedback. It appears that some aspects of the mentioned weaknesses may have been misunderstood. We address these points below. If our clarifications resolve some of your concerns, we kindly request you to reconsider the score provided.

W1: We have never claimed in the write-up that our contribution is the proposal of a mixture of shared and personalized models. We refer to the first point of the public comment for a detailed clarification of the novelty.

Thanks for pointing FedCP out. We indeed missed this work and will include a discussion in our future iteration. While FedCP indeed has a global and local feature extractor and a conditional policy network (CPN) that acts as a router, our CoMiGS method differs from it in the following aspects: 1) communication units: CoMiGS only requires clients to share global experts, while FedCP requires sharing of personalized feature extractors and the CPN (router). CoMiGS is more private and communication efficient. 2) tackling out-of-distribution tasks: FedCP updates the router and feature extractors both on the same loss, while our CoMiGS updates the router via a separate validation loss, making the method more robust and able to tackle out-of-distribution tasks. 3) tackling system heterogeneity: according to our understanding, FedCP only addresses data heterogeneity, not system heterogeneity. As FedCP initializes personal feature extractors with the global parameters, this requires all personal extractors to share the same structure.

W2: Our bi-level formulation does not add much computation and memory overhead. It is worth noticing that we solve the bi-level optimization via an approximate iterative updating algorithm, which removes the calculation of Hessian and thus the router update only requires first-order gradients over the router, which is a small one-layer MLP (of size 768*2 in our setup). CoMiGS even has less computational complexity than FedCP, as FedCP updates the router each iteration with the feature extractors, while CoMiGS updates the router every 30 iterations, see Appendix B. A detailed numerical analysis regarding computational complexity can be found in the second point of the public comment.

W3: We apologize for not making it clearer in the main texts. We indeed pointed out that the validation set can alternatively be replaced by a newly sampled batch from the training set (Line 837-838). For in-distribution tasks, this achieved decent performance. However, for out-of-distribution tasks, this would not necessarily work.

W4: We would like to draw your attention to Table 1 in our manuscript, where we investigate the necessity of having both generalists and specialists through an ablation study. Specifically, CoMiGS-2G and CoMiGS-2S are variants that only allow two generalists or two specialists, respectively. The empirical evidence demonstrates that neither approach consistently performs well across different task scenarios. You are correct that when tasks are more homogeneous, generalists suffice, while for highly heterogeneous tasks, specialists excel. However, in practice, it is often impossible to predict task heterogeneity in advance to determine the most suitable method. This is where our CoMiGS method stands out, as it closely tracks the best-performing approach in each scenario, thanks to its token-dependent routing mechanism.

Dear Reviewer 3sa7,

Thanks for acknowledging the fact that we managed to address your concerns. Regarding the methodological novelty part, we would like to remind you kindly that not all published works from ICLR have to offer groundbreaking novel methods, for example, FFALoRA [1] is in the same lieu and accepted by ICLR 2024, which suggested a way to improve LoRA in the context of privacy-preserving, instead of proposing a novel method to preserve privacy.

Our work indeed offers novel contributions in other aspects: 1) we are the first to address both data and system heterogeneity in collaborative LLMs with marginal compute overhead over the standard FedAvg; 2) our bi-level formulation of the MoE learning objective is novel; 3) our method brings clear empirical gains with a theoretical convergence guarantee (see Theorem 3.1 in our modified manuscript).

Thank you for your time in advance!

Best,

Submission 3571 authors

[1] Sun et al. Improving LoRA in Privacy-preserving Federated Learning, ICLR 2024

Thank you for your critical feedback, we address your concerns and questions as follows:

W1: Yes, there is a growing body of work focusing on using LoRA MoE for multi-task learning. We will incorporate these works into the related works section in future iterations. Additionally, we reviewed [5] and [6] as referenced by you. [5] is a study in the vision domain. [6] is a concurrent work that shares similarities with our CoMiGS method, which we missed during our literature search. Thanks for pointing it out. While FDLoRA also incorporates global and local LoRA modules alongside a learnable set of weights to combine them, there are significant differences from our approach:

Training Strategy: FDLoRA trains global and local experts sequentially, where local experts are initialized from global ones, causing the local experts to remain close to the global experts. In contrast, our CoMiGS method trains these experts in parallel, enabling them to co-adapt dynamically, which is supposed to lead to a greater distinction between the global and local experts.

Weight Granularity: FDLoRA employs client-wise learnable weights, whereas our method takes a more fine-grained approach by learning token-wise weights.

Since the implementation of [6] is unavailable, we recreated their method in our codebase using the hyperparameters provided in their paper. The results of this implementation are presented in the table below: FDLoRA failed to outperform our CoMiGS-1G1S. We noticed that FDLoRA performs similarly to CoMiGS-2G method, where we assign two generalists per client. This indicates that indeed the local and global experts trained from FDLoRA are quite close.

We also highlight a key limitation of FDLoRA: it does not support different numbers of experts across clients. This is a crucial feature in our method, as it effectively addresses system heterogeneity, allowing for greater flexibility in diverse client setups.

W2: We would like to clarify that our method does not imply that devices should have multiple LoRA modules. Instead, the number of LoRA modules per device can vary depending on the edge device capacity. Those resource-limited edge devices can operate effectively using only the generalist expert. As demonstrated in Figure 6, edge devices equipped with just one generalist expert can still benefit from collaboration with more powerful devices: the first device achieves better performances when the numbers of LoRA modules across the clients change from [1,1,1,1] to [1,1,8,8]. To the best of our knowledge, our proposed method is the first to address both system heterogeneity and data heterogeneity simultaneously.

W3: Thanks for pointing this out. We understand your concern over the scaling ability of our method to larger and more recent LLMs. However, scaling up to larger models becomes less practical within an academic budget, as we even need to create multiple clients at one time, which is extensive in memory and thus hinders us from using larger models. So far, we can fit all our experiments on an A100 GPU.

We are currently working on a LLAMA 3.2 - 1B base model, however, the proper comparison between the baselines requires a bit more time. We will follow up with the results in the following days.

Q1: Please refer to the second point of our public comment for a detailed numerical comparison of the computational/communication overhead and memory requirements. Our CoMiGS method is rather resource-efficient. As we previously stated in W2, our framework can easily cope with resource-constrained devices. Furthermore, they can benefit from collaboration with resource-rich devices.

We're pleased to share the results of our LLAMA 3.2-1B base model comparison, following up on our previous response.

Given the extensive pre-training of LLAMA 3 models on over 15 trillion tokens from public sources [2], and the multilingual capabilities of LLAMA 3.2 - 1B [1], fine-tuning on multilingual Wikipedia or SlimPajama resulted in negligible improvements (see Table below) likely due to significant overlap with the pre-training data corpus.

To clarify the differences between the methods, we propose a new fine-tuning dataset, which is derived from Common Corpus (specifically, the YouTube-Commons, Latin-PD, and TEDEUTenders collections) and the Harvard USPTO dataset. Following our previous methodology, each client is assigned one of the datasets to maximize heterogeneity. We use this dataset to model the in-distribution scenario. Additionally, we reduced the number of training iterations for the Agnews experiment.

Our results on the Common Corpus-based dataset, which emphasizes domain-specific language and structure, demonstrate that our CoMiGS-1G1S method can outperform local training and FedAvg. In the out-of-distribution scenario (Agnews), CoMiGS-1G1S performance tracks the performance of CoMiGS-2G and FedAvg, similar to our observations with the GPT experiments. The complete results are presented below.

[1] https://ai.meta.com/blog/llama-3-2-connect-2024-vision-edge-mobile-devices/

[2] https://ai.meta.com/blog/meta-llama-3/

Thanks for your feedback. We address your raised concerns and questions as follows:

W1: In our write-up, we selected the most challenging scenario for collaboration, where each user is assigned a distinct, non-overlapping category. While this setup emphasizes the need for personalization, our results demonstrate that collaboration remains crucial and offers significant benefits.

In response to your suggestion, we conducted additional experiments using homogeneous data. Specifically, we assigned all four clients French Wikipedia texts containing 100,000, 500,000, 1,000,000, and 5,000,000 tokens each. In this scenario, we expected FedAvg to perform best, which was indeed observed. For local training, users with smaller datasets experienced overfitting. However, our CoMiGS method effectively leveraged the generalist expert to mitigate overfitting, demonstrating its effectiveness in homogeneous data settings as well.

W2: We agree that the combination of MoE and LoRA is not novel, and we have not claimed any originality regarding them. However, we would like to emphasize that the novelty of our work lies in: (1) the bi-level formulation of the MoE objective, which allows seamless handling of out-of-distribution tasks, (2) being the first to address both system and data heterogeneity in collaborative LLMs, and (3) effectively disentangling resource abundance from data quantity. We consider these significant contributions to advancing collaborative LLMs.

Furthermore, our work highlights an interesting phenomenon: the heterogeneity of language cannot be quantified in the same way as in the vision domain. While all clients are consistently assigned distinct categories (e.g., different topics or languages), the preferences over local training or FedAvg may vary, as shown in Table 1. This is because language, unlike visual data, inherently consists of numerous repeating semantic units, leading to unique dynamics in how heterogeneity manifests and is addressed. We will make this more explicit as well in our future iteration.

We also refer you to the first point of our public comment for a detailed clarification on the novelty/contribution of our method.

W3: Our discussion of resource heterogeneity primarily focuses on the device axis rather than the time axis, as detailed in Section 4.3. We demonstrated that resource heterogeneity can be effectively addressed by enabling devices to utilize varying numbers of specialists. Regarding your concern about potential fluctuations in device resources over time, such scenarios may indeed emerge with advancements in edge devices, such as mobile phones. Expanding existing pretrained experts by adding new ones is a research area in its own right. We refer to [1] for an example, of which the technique can be complementary to our method.

[1] MoExtend: Tuning New Experts for Modality and Task Extension. Zhong et al.

Dear Reviewer 6KmX

Thank you for the time and effort you have dedicated to reviewing our manuscript and for providing thoughtful feedback.

In response to your suggestions, we have added a new experiment exploring different data quantities, which demonstrates the consistent performance of our method. Additionally, we have updated the manuscript to include a convergence proof and a new base model (Llama3.2 - 1B). These additions further support the theoretical and empirical soundness of our proposed method.

Given these improvements, we kindly request a re-evaluation of our work. We understand this is a particularly busy period, and we greatly appreciate your time and consideration.

Thank you,

Submission 3571 Authors

Thank you for your encouraging feedback! Regarding your raised concerns, we address them as follows:

W1: Please refer to the first point of our public comment.

W2: We refer to the existing theoretical analysis of alternating minimization provided,e.g., in [1] (see their Theorem 1) or [2], which can be directly applied to our algorithm under the condition that $f(X^{\text{train}}, \cdot)$ and $f(X^{\text{valid}}, \cdot)$ are sufficiently close. This serves as the primary motivation for developing this algorithm.

We checked the convergence guarantee of our proposed algorithm in the past days. Assume that there exists a unique global optimum for both $L(f(X^{\text{train}}, \Phi, \Theta))$ and $L(f(X^{\text{valid}}, \Phi, \Theta))$. As long as this unique global optimum $(\Phi^\star, \Theta^\star)$ is the same for both functions and both functions behave properly, satisfying certain qualification assumptions as in [1] or [2], the convergence guarantee remains valid.

This assumption is reasonable, as it implies that for both the training and target distributions, there exists an optimal set of expert and routing parameters for language modeling. We plan to include this guarantee in the next iteration of our work to strengthen our results, and thanks for pointing this out.

W3 & W4: Thanks for pointing it out. We will make sure that a brief introduction of the building blocks as well as references to the mentioned related works are included in the next iteration.

W5: Yes, you are right about the privacy concerns. However, we would like to point out further that the only parameters we share are the global LoRA experts from the MLP layers, which is 38% of the total trainable parameters in the one-generalist-one-specialist case. With more specialists within each client, the ratio is even lower. Additionally, noise can be added to the global parameters to further guarantee differential privacy. We will add a more cautious and thorough discussion to the write-up.

Q1: Our extra memory and computational complexity only come from the router, which is a small one-layer MLP. The size is 768*2 in our experiments. We refer to the second point of our public comment for a detailed analytical comparison.

Q2: The absence of a theoretical analysis stems from our primary focus on the LLM community, where providing formal mathematical proofs is not a common convention. However, we agree that having a convergence guarantee would better justify our method and plan to include one in the future iteration of the manuscript.

Q3: No, the privacy leakage does not scale with device capacity, as we earlier pointed out in W5, no matter how many experts one device has, the only shared part is the global experts, which in our settings is the very first expert from each MLP layer. Having more capabilities does not make one device more exposed to privacy leakage.

[1] Alternating Minimizations Converge to Second-Order Optimal Solutions. Qiuwei Li, Zhihui Zhu, Gongguo Tang Proceedings of the 36th International Conference on Machine Learning, PMLR 97:3935-3943, 2019.

[2] On the rate of convergence of alternating minimization for non-smooth non-strongly convex optimization in Banach spaces. Jakub Wiktor Both, Optimization Letters 16.2 (2022)

Dear Reviewer HA1D,

Thank you for the time and effort you have dedicated to reviewing our manuscript and providing thoughtful and positive feedback.

Following your suggestions, we added a convergence result to the manuscript (see Theorem 3.1 from the modified manuscript). Moreover, we have clarified the communication and computational overhead as you pointed out in the global response, and added a new experiment to strengthen our empirical results. All of these changes can as well be found in the modified manuscript (highlighted in blue).

We hope these adequately address your concerns and kindly ask for a re-evaluation of our work. We understand it is a particularly busy period for you, and we greatly appreciate your time in advance.

Thank you,

Submission 3571 Authors