On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists
Authors: Dongyang Fan
Bettina Messmer
Nikita Doikov
Martin Jaggi
Venue: ICML 2025 poster (2025)
Keywords: Federated Learning
Collaborative Learning
On-device LLMs
Mixture of Experts
Alternating Minimization
Primary area: deep_learning->algorithms

Abstract: On-device LLMs have gained increasing attention for their ability to enhance privacy and provide a personalized user experience. To facilitate private learning with scarce data, Federated Learning has become a standard approach. However, it faces challenges such as computational resource heterogeneity and data heterogeneity among end users. We propose CoMiGS ($\textbf{Co}$llaborative learning  with a $\textbf{Mi}$xture of $\textbf{G}$eneralists and $\textbf{S}$pecialists), the first approach to address both challenges. A key innovation of our method is the bi-level optimization formulation of the Mixture-of-Experts learning objective, where the router is optimized using a separate validation set to ensure alignment with the target distribution. We solve our objective with alternating minimization, for which we provide a theoretical analysis. Our method shares generalist experts across users while localizing a varying number of specialist experts, thereby adapting to users’ computational resources and preserving privacy. Through extensive experiments, we show CoMiGS effectively balances general and personalized knowledge for each token generation. We demonstrate that CoMiGS remains robust against overfitting—due to the generalists' regularizing effect—while adapting to local data through specialist expertise. We open source our codebase for collaborative LLMs.

On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists
Dongyang Fan 1
Bettina Messmer
Nikita Doikov
Martin Jaggi
Abstract
Personalized word
Sug suggestions
In my free time, like to skiing
On-device LLMs have gained increasing attention for their ability enhance privacy and provide a personalized user experience. To facilitate private learning with scarce data, Federated Learning has become standard approach. However, it faces challenges such as computational resource hetero- geneity and data heterogeneity among end users. We propose CoMiGS (Collaborative learning with Mixture of Generalists and Specialists), the first approach address both challenges. A key inno- vation of our method is the bi-level optimization formulation of the Mixture-of -f-Expets learning objective, where the router is optimized using a separate validation set to ensure alignment with the target distribution. We solve our objective with alternating minimization, for which we pro- vide theoretical analysis. Our method shares generalist experts across users while localizing varying number of specialist experts, thereby adapting to users' computational resources and preserving privacy. Through extensive experi- ments, we show CoMiGS effectively balances general and personalized knowledge for each to- ken generation. We demonstrate that CoMiGS re- mains robust against overfitting-de eralists' regularizing effect-whii adapting to local data through specialist expertise. We open source our codebase for collaborative LLMs.
In my free time, like to read
universal word suggestions
Have you followed the presidential election? What is the result?
yes, the 47th president Donald Trump.
Chat box between two users with different charac- teristics. Next word prediction for smart keyboards should tailored to users' topic preferences for personalization. However, ensure factual accuracy and linguistic consis- tency, the results of next word prediction should maintain universality.
Figure I:
Routdr
Export Modules
Figure 2: Diagram of our proposed method CoMiGS il- lustrated with simplified 2-heterogenous-moddee setup (corresponding to the two users in Figure 1). Generalist ex-
the gen-
perts
0 0, are aggregated across users, and specialist
experts S013 J==1
(0))
and Routers (
h2 are
kept local.
ference introduces significant delays for end users, and it often fails to meet their personalized needs (Ding al., 2024; Iyengar & Adusumilli, 2024). Recently, there has been growing interest in deploying L.MMs on edge devices, which offer benefits like lower latency, data localization, and more personalized user experiences (Xu al., 2024). For instance, Apple (2024) recently launched on-device founda- tion models part of its personal intelligence system. Meta (2024), Qwen (2024) newly released lightweight models with less than 3B parameters targeting edge AI.
1. Introduction
Large Language Models (LLMs) have been showing great success serving as foundation models, evidenced by their capability to understand a wide range of tasks, such as ChatGPT (OpenAI, 2023), Claude (Anthropic, 2023), Gem- ini (DeepMind, 2023) and etc. However, cloud-based in-
"Equal contribution 'EPFL, Switzerland. Correspondence to: Dongyang Fan <dongyang fan @cpt
On-device LLMs present challenges such as limited and vari- able computational resources, searce and heterogeneous lo- cal data, and privacy concerns related data sharing (Peng
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

CoMiGS: On-Device Collaborative Language Modeling
al., 2024; Wagner et al., 2024). Fine-tuning is typically performed on-device to quickly adapt to users' individual needs. While data sharing is common solution to address local data scarcity, on-device data is often privacy-sensitiy and must remain on the device. To overcome this, Federated Learning has been proposed as a method for enabling col- laborative learning while preserving user privacy, allowing end users to collaborate by sharing model parameters (Chen et al., 2023; Zhang et al., 2023).
ent topic preferences or linguistic styles, they should still share common phrases, for example, when talking about factual knowledge (the result of the US presidential election should be a universal fact in Figure 1). Our goal is to route these shared tokens to the generalists so they can be jointly learned across users.
We further notice a hierarchical structure between the router and the experts: the router dynamically assigns tokens based emerging expert specializations, while the experts refine their roles to optimize token processing under the router's s guidance. Towards addressing this, we formulate our learn- ing objective as a bi-level optimization problem and propose a new first-order algorithm based on alternating minimiza- tion as a solution. Our method enjoys convergence guaran- tees and is resource-efficient for deployment.
Federated fine-tuning of LLMs is predominately done through Low-Rank Adaptation (LoRA, Hu et al. (2021)) due to its lightweight nature so that the communication costs can be largely mitigated. Yet end devices may have different capacities, resulting in different LoRA ranks or different numbers of LoRA modules allowed on devices. Previous works have proposed various techniques for aggre- gating LoRA modules of different ranks (Cho et al., 2023; Bai et al., 2024). However, in both works, the devices are only equipped with shared knowledge, which makes the methods unsuitable when there is data heterogeneity across users. In such cases, a more personalized solution is needed.
In summary, our contributions
follows:
We propose a novel approach (CoMiGS) for on-device personalized collaborative fine-tuning of LLMs. Key parts of our approach are: 1) an innovative bi-level formulation of the MoE learning objective (Section 3.2); 2) a new al- gorithm based on alternating minimization (Alg. I); 3) a theoretical analysis with a proof showing linear conver- gence rate under suitable assumptions (Section 3.4).
End users' local data distributions can exhibit significant statistical heterogeneity. For instance, mobile device users may have distinct linguistic habits, topic preferences, or language usage patterns, leading to widely varying word distributions, as illustrated in different next word predictions to the same prompt "In my free time, I like to" in Fig. 1, As a result, personalized solutions are necessary. Wagner al. (2024) explored three personalized collaborator selection protocols, allowing each end user to choose their collabo- rators. Although these protocols effectively address data heterogeneity, they depend on model aggregation, which can only occur when users share the same model architecture.
Our collaborative framework effectively addresses both data heterogeneity (Section 4.2), concerning diverse local data distributions across users, and compulational resource heterogeneity (Section 4.3), with respect to varying local model architectures, making it the first model to accom- plish both.
Our framework separates model heterogeneity from data quantity (Section 4.4). Users with larger local datasets ben- efit from a bigger model, while users with more powerful models but smaller datasets are less prone to overfitting.
There has not yet been a solution to deal with both model heterogeneity and data heterogeneity. Towards this goal, we propose a novel Collaborative learning approach via a Mixture of Generalists and Specialists (CoMiGS). Our ap- proach allows users to share part of the knowledge while keeping some knowledge user-specific, thus providing per- sonalized solutions. We name the shared part generalists and
CoMiGS is resource-efficient: it adds marginal (+1.25%) computational overhead and memory requirement com- pared to FedAvg, while reducing communication costs by 50% (Section 4.5).
We release a codebase' for collaborative LLMs that allows users to easily implement various collaboration strategies, facilitating and advancing future research efforts in this field.
the user-specific part specialists. Like all previous works, the generalists and specialists are simply LoRA modules. At the same time, as long as the shared part can be aggregated, the user-specific part can be of different sizes, which can be adapted to various device capacities, as illustrated by different numbers of specialists across users in Figure 2.
2. Related Work
Collaborative Fine-Tuning for LLMs. Recently, re- searchers have been investigating the application of Fed- erated Learning in language tasks. Due to the substantial number of model parameters in LLMs, the research has largely targeted the stages following pre-training, often uti-
We integrate the expertise of generalists and specialists us- ing a learned router that determines aggregation weights, following the Mixture-of- -xperts (MoE) architecture (Fe- dus et al., 2022b). As in typical MoE designs for language modeling (Jiang et al., 2024; Fan et al., 2024), we also use tokens as the routing unit. Although users may have differ-
Our code base is available https: / htt.ttub.ccco/////// epfml/CoMiss /CMMiGS
2

---

CoMiGS: On-Device Collaborative Language Modeling
lizing parameter-efficient techniques such as adapters. Mo- htashami et al. (2023) explored teacher-student social learning framework to aggregate private-sensitive instruc- tions. Zhang et al. (2023) directly applied FedAvg (McMa- han et al., 2017) to aggregate LoRA parameters during in- struction tuning, and reported increased performance in downstream tasks. Following that, there are various works focusing on addressing resource heterogeneity where users are equipped with different LoRA ranks. HetLoRA (Cho et al., 2023) and FlexLoRA (Bai et al., 2024) provide dif- ferent ways of aggregating and distributing LoR A modules of heterogenous ranks. However, these approaches are not designed to cope with heterogeneous data on device. In contrast, Sun et al. (2024) demonstrated that freezing LoRA A matrices initialization leads to improved performance on heterogeneous data. Building on this, Guo al. (2024) showed that consistently aggregating LoRA A matrices can yield even greater performance gains. Meanwhile, Wagner et al. (2024) introduced personalized approaches to effec- tively address data heterogeneity by employing three dis- tinct collaborator selection mechanisms. However, these approaches require all users to utilize the same model archi- tecture. Unlike previous works, our framework deals with both model heterogeneity and data heterogeneity. Moreover, our method offers personalized solutions at token level, as opposed to the client-level approach in Wagner et al. (2024).
3. Method
We aim to improve personalized performance for each end user, through fine- ane-grained mixture of general and person- alized knowledge. Building on the hierarchical insights of MoE learning, we formulate our learning objective into bi-level optimization problem, where expert parameters are learned using the relatively large-sized training sets, while routing parameters are updated using the small-sized vali- dation sets. We further let experts diversify into generalists and specialists via parameter aggregation or localization. As the problem solver, we provide a multi-round gradient- based algorithm, of which the pseudo codes are presented in Appendix A.
3.1. Notions and Problem Setup
Each user has a training set X a small validation set X and test set X and the task is next token pre- diction. The validation set X and the test set X! are sampled from the same distribution Plauger (note this is a fuzzy concept in the language domain, by the same distribu- tion we mean from the same topic/category). The training set, X Xtrain can be sampled from a different distribution than plarget This is to address scenarios where distribution shifts may occur over time, such as changes in topics reflected in the typing data of mobile phone users. As illustrated in Figure 2, there are two sets of model parameters within each user: expert parameters, denoted as θ = 0 U ( where 0c is shared across the users and (0.))
Mixture of Global and Local Experts. Gaspar & Seddon (2022) introduced a fusion of global and local experts for activity prediction based molecular structures. Each local expert is tailored to specific chemical series of interest us- ing loss masking, while global expert is trained across all series. Simultancously, a routing network learns to assign soft merging scores. This approach yielded superior empir- ical results compared to single experts. Dai et al. (2024) developed DeepSeekMoF) by deterministically assigning ev- ery token "shared" experts, whereas "routed" experts are assigned tokens based on a learnable router. DeepSeekMoE is able to approach the upper bound performance for MoE models. For both works, the notion of shared/global is with respect to input samples, i.e. a shared/global expert should see all input samples. In a collaborative setup, FDLoRA (Qi et al., 2024) learns dual LoRA modules on each client to capture personalized and global knowledge respectively. Client-wise fusion weights are learned towards the end of training to combine the two sets of knowledge. In com- parison, PFedMoE (Yi et al., 2024) jointly learns a shared homogeneous small feature extractor, a localized heteroge- ncous feature extractor, and a localized routing network in an end-to-end fashion, demonstrating strong performance in the vision domain. Our work is closely related to these two works, while introducing key innovations that adapt to model heterogeneity and allow for fine-grained adaptation to target distributions.
} are user-specific specialist parameters; and routing parameters, denoted as & { i E f1, N. the user index. We use linear models as our routers. Thus, Φ, E R" produces gating values which are used to combine the n; experts at each layer, as in (I), where x is input token, y is the corresponding layer output and E, is the jth expert.
(exp Φh P5(()) Exxexp D.x); V Σ (x)
(L)
Our experts are simply LoRA modules, which approximate model updates AW E Rmmx with a multiplication of two low-rank matrices A E Rmmr and B E Bg with rank
< m,n. n. @ and 0 are disjoint sets of LoRA A and B matrices.
3.2. A Bi-Level Formulation
Instead of learning routing and expert parameters simultane- ously like the conventional way in LLMs (Zoph et al., 2022; Fedus et al., 2022a), we update the two sets of parameters in an alternating fashion. We observe a natural hierarchy between the experts and the router: the assignment of tokens to experts depends on the router's 's outputs, while the experts' parameters are updated based on the assigned tokens. In this way, the experts' development follows the router's 's decisions,
3

---

CoMiGS: On-Device Collaborative Language Modeling
establishing an inherent leader-follower structure. Follow- ing Von Stackelberg (2010), we formulate the hierarchical problem as a bi-level optimization objective as follows:
A discussion over the load balancing term is included in Appendix C.4. It is observed that a load-balancing term can improve test performance compared to not having one. However, directing more tokens to the generalists has no noticeable effect.
min Σ C( ((x: (*)(),
(upper)
Given that the data is distributed among clients, when solv- ing optimization problem from equation (3), we first obtain the solutions 0, and 0,¹ to local problems, for each client i. A parameter aggregation is then performed the user- specific 0e via a trusted server to establish a shared Oc across all users.
s.t. (T) E arg min Σ C( train 0c, 00,00 0.) (lower)
where L is the language modeling loss. The routing param- eters + = } are updated based on the validation loss, which reflects the target distribution (upper optimization), while the expert parameters O = 0c U {0 } are updated using the training loss (lower optimization). This formu- lation further brings in the following benefits: 1) routing parameters are smaller in size, making them casier to overfit. By separating the two losses, the routing parameters can be updated less frequently using the smaller validation set (visual evidence of less frequent router update leading to improved performance is provided in Figure 8 in the Ap- pendix); 2) this approach handles situations where target distributions differ from training distributions more effec- tively, as the router outputs (i.e., how the experts should be weighted) can be tailored to specific tasks.
(0 @ocht @ N
= arg min fuaia (O,
& $x+1),
S i=1
(4)
+ k+1 = N
k+1 (09
In the next round, each user replaces their of with the global Og. while their 0. remains local.
3.4. Convergence Results
We study the convergence properties of our alternating min- imization process. First, we establish a linear rate of con- vergence under general assumptions on our objectives, that always hold locally, when the parameters are close to the training solution (assuming the pretrained model is not far from the fine-tuned models). Then, we show that the case of linear experts, the same optimization procedure possesses global linear convergence.
3.3. Our Algorithm
To solve our bi-level problem, we use alternating updates of the two sets of parameters. The pseudo-code of our proposed algorithm is detailed in Alg. I in the Appendix.
We denote partial minimization operators from (2), (3) by
Alternating Update of and 4. Alternating up- date of two sets of parameters is a standard way to solve bi-level optimization problems (Chen et al., 2021). The alternating updates of expert and rout- ing parameters are performed using local training and validation sets separately. To simplify notations, we denote fvalia ( b) = Σ, L(X yalid 0 0 Φi) and ( (0. +) := E, C(X 0 0., <b.). Note that in con-
U1 -
arg min fvatial (O, Φ),
U2 + arg min frain ( ()
and their compositions by T == u2 . in and P == 2 U2.
Note that both T and P act on the corresponding spaces of + and o: T : |l|| + @² and P g + R $1
trast to (upper) bi-level formulation, we allow parameter O to be free in fvatid, which makes it easier to optimize. We can write the alternating update steps follows.
Assumption (Shared Optima). There exist O' and T
such that
- T(@*) and x'=P($))
(5)
T k++1 arg min fvatia Oc2, )) +
(2)
Eq. (5) means that fvalid and frain share the same global optima, which is reasonable when the train and validation data are similar, X X;alid, and, hence, fvalid R frain- It also holds for overparametrized models, such as LLMs.
= arg min Jiii (x+1..
(3)
Since the updates of O and & are disentangled, they do not need to be updated at the same frequency. The rout- ing parameters are smaller in size and thus can be updated less frequently. When updating model parameters, we in- clude an additional baldnnannn term as in Fedus et al. (2022a), which is standard in MoE implementation and en- courages even distribution of token assignments to experts.
Assumption 2 (Contraction Property). Let un and u2 be Lipschitz with some constants X1, A2 > 0, for any θ O and $ क;
| ( - un (+)
< A -011
(6)
||uz (+) - uzz (2(
<
dal Φ - 511.
4

---

CoMiGS: On-Device Collaborative Language Modeling
Theorem 3.1 (Convergence under Contraction). If As- sumptions 1, 2 hold, and λ1 X2 < <, then the weights (O&, 2x) generated by alternating updates (2), (3) con- verge to (0*, Φ*) with a linear rate.
single category, but their validation and test sets consist of a uniform mixture of all categories. This approach accounts for potential shifts in topics within users.
4.1.2. EXPERIMENTAL DETAILS
The proof is provided in Appendix F.2. We show that the contraction property holds when the objectives are convex quadratics. As a consequence, it also holds locally for any sufficiently smooth models, using their Taylor expansions around a local minimum. Large models are usually smooth, and when initialized with a well-pre- pre--traieeed model, it suf- fices to guarantee local convergence.
We choose the following base model architectures: GPT2 (124M, English only) and Llama 3.2 Multilingual)". We incorporate LoRA modules into every linear layer, including MLP and Self-Attention Layers, following the recommen- dations of Fomenko et al. (2024). A routing mechanism is exclusively implemented atop MLP layers. The number of LoRA experts in MLP blocks depends on the local resource abundance. For more experimental details, we refer readers to Appendix B.
Our alternating minimization process can guarantee global convergence as well, when the experts are linear models, see Appendix F.4. This indicates wide applicability of solving MoE objectives via alternating minimization.
4.2. Data-Driven Selection: Generalist vs. Specialist
Theorem 3.2 (Global Convergence for Linear Experts). If fvaiid = f and all the expert modules are linear models, we have a global linear convergence rate for a practical instance of our method.
We start by equipping users with the same model architec- ture locally, to illustrate the effectiveness of our hierarchical learning of routing and expert parameters. We compare our one generalist one specialist GS-1155) method to the following baselines. In order to match the trainable pa- rameter count of our method, we use 2 times LoRA modules within each user.
4. Experiments
4.1. Setup
Upper and lower bounds: 1) Pretrained: Pretrained checkpoints. 2) Centralizest A single model trained using data from all users. (Note this method is an unre- alistic baseline as data cannot leave the devices due to privacy concerns.)
1. DATASETS
We selected the following datasets to demonstrate the effi- cacy of our proposed algorithm: 1) Multilingual Wikipedia: Wikipedia articles in four languages: German, French and Italian from
and Dutch from Guo et al. (2020): 2) SlimPajama: We pick the following four categories- StackExchange, Github Codes, ArXiv, Book from Soboleva et al. (2023); 3) AG News: News from cat- egories of World, Sports, Business, and Sci/Tech (Zhang et al., 2016). 4) Common Corpus (pleias, 2024): specifi- cally the following three categories- - CeeeTuub-Common Public Domain Books, and EU Tenders collections, and the Harvard US Patent dataset from Suzgun et al. (2022).
Baselines: 1) Local: Training individually using only local data. 2) FedAvg: Aggregating LoRA parameters across users using uniform weights, which is equivalent to applying FedAvg (McMahan et al., 2017). 3) PCL: Aggregating LoRA parameters using client-level col- laboration graph. The graph is updated using validation performances. (Strategy 2 in Wagner et al. (2024)). 4) pFedMoE: We directly apply the method from Yi et al. (2024) in the language domain where we update routing and expert parameters at the same time and choose tokens as a routing unit. 5) FDLORA: Global and local parame- ters are learned, and a client-wise fusion weight is learned to combine global and local parameters. (Qi et al., 2024)
A distinct category is assigned to user, as it simulates the most challenging scenario for collaboration. Given our emphasis on next token prediction, we anticipate shared pre- dictions among users while maintaining category- -ppccffc distinctions. We further create the following two scenarios to showcase the wide applicability of our method:
Ablations: 1) CoMi -5-25: Both experts are specialists, meaning their weights are neither shared nor aggregated. The routing parameters are updated using a separate vali- dation set like in CoMi GS-1G15 s. 2) CoMi GoMiiGs-200 -2:::: Both experts are generalists, meaning their weights are always shared and aggregated. The routing parameters are up- dated like in CoMiGS-
In-Distribution Tasks. For each user, we construct vali- dation and test datasets that follow the same distribution as the training data. We address two scenarios in this context: (i) variation in language usage across users (Multilingual Wikipedia), and (ii) variation in topic coverage across users (SlimPajama, Common-Corpus).
'We adopt the codes from https://githin com/ karpptthy/nanooP and https://githib: com/ hoop://mnielggiitterrrmneooooooo / //github. wttphhhu....oo/ppaaaaaapaaaaaa // http://ppp.bbbssss//p//mmmmmmee
Dittr--Dittibbttion Tasks. For each user, we create val- idation and test datasets from a distribution different from the training data. During training, each user assigned a
https:
5

---

CoMiGS: On-Device Collaborative Language Modeling
Table I: Mean (std) test perplexity over the users with ho- mogeneous models, averaged across 3 seeds (the lower the better). Light grey denotes in-distribution tasks and dark grey denotes of--iiiittiton tasks.
possible to determine the language structure in advance. Even drastically different users still share many of the same tokens. A data-dependent combination of generalists and specialists is required.
Bose Model
GPT2-124M
LLAMA3.2-1B
Danant
Madittinuu
StinPauams
AC Nexe
Con-Curqus
AG Nenss
4.2.2. ROUTING ANALYSIS
Priteaiond
156.12
$7.19
00 65
30:40
20.37
Contratioad
55.41
(0.14)
(0.52)
17.97
16.12 (0.05)
Tiken-wise Analysis. Figure 3 visualizes token-level rout- ing results for models fine-tuned on the SlimPajama dataset. In the first layer, function words (e.g., "and," "a," "on," "the") are predominantly routed to generalists. In contrast, in the last layer, content words are more frequently assigned to generalists. This pattern is particularly evident for the first two users, trained on math and programming texts, where domain-specitee terms are primarily routed to specialists. These findings suggest that experts in later layers develop more distinct role specializations. Importantly, only the top choice is highlighted here. The abundance of blue does not imply that generalist experts play no role in predict- ing the next token. As compared to when only specialists are present (CoMi GS-2S), our CoMiGS- -1s- s gives more consistent results. More detailed token-wise routing result visualization including -ftt--i-iitritttion tasks can be seen in Appendix E. When dealing with ditto--distrbbttion texts, there is an increasing tendency seek generalists, as shown in the off-diagonal entries in Figure 14-19.
Land
$4.38(0.332
26.05
(0.06)
20.19
19:56 (0.01)
Podlag
(0.34)
23:27 (0.05)
(0.02)
15:36 (0.05)
PCL
54.53.00.10
26.99 19)
12)
(0.65)
16. (0.05
pfedMoo
(0.17)
25.40/0.09
38.72(00.21)
20:41 (0.05)
17.34 (0.05)
FDLARA
57.45
22.71 (0:40)
13.6 (0.07)
22:11 (0.05)
16.64 (0.02)
CAMMES 25
46.36.00.16
22.51 (0.08)
35 (0.13)
(8.15(0.113
18.03 (0.11)
CAMASS 2G
(6.11(0.11))
11.36 (0.01)
31.18 (0.05)
20.15:00.09
15.41 (0.05)
CAMIGS 1G18
47.19
21.79 (0.04)
33.53
18.37
16.31 (0.05)
4.2. 1. RESULT ANALYSIS
The comparison between our method and the baseline meth- ods is summarized in Table 1.
Effectiveness of Our Routing Mechanism. Depend- ing on the dataset, either CoMiGS-2G or CoMiCS-25 achieves the highest performance. The key distinction com- pared to Local 1 or FedAvg is the existence of layer-wise token-level router, which learns to combine the two general- ists or specialists. This emphasizes while knowledge might be present, the way it's combined is the key. Moreover, pFedMoE, despite having a learned router as well, under- performs our method, even in the in-distribution scenario. The reason is that the routing parameters are updated simul- taneously with the expert parameters using the training set, and thus cannot effectively adapt to the target distribution. When a separate validation set is not available, CoMi GS can alternatively sample a new training batch to update the routers and still offer competitive results for in-distribution tasks (see Table 5).
Layer-wise Analysis. Figure depicts the evolution of averaged layer-wise router outputs for the generalist and specialist experts on the out-of-distribution task, compar- ing NoMiGS-1G1s S and pF As training progresses, CoMiGS 1G1s undergoes a phase transition: the layer- wise routers initially favor generalists but gradually shift towards specialists. This shift not observed in pFedMoF, highlighting the critical role of our routing mechanism in handling
Token-level Collaborative Decisions Outperform Client- Level. Compared to the state-of-the-art bascline PCL and FDLORA, our method demonstrates a clear performance improvement. While both methods require separate valida- tion set as in our method to determine collaboration weights, PCL determines the weights to combine each client's mod- els iteratively while FDLORA determines the weights for the global and local model at the end of training. Our method, in contrast, decides the collaboration pattern based on each input token, allowing the router weights to co-adapt with the expert parameters throughout training. This enables a more flexible and fine-grained collaboration.
tasks. Additionally, we notice different layers converge to different expert score distri- bution. When applying our CoMiGS-1G1 s, for each user, there are always certain layers where the routers consistently prefer generalists, which aligns with the fact that our tar- get distribution is a union of all local training distributions. For in in-distribution tasks (Figure 10), during early stage of training, some layers favor generalists. When close to con- vergence, all layers favor specialists. We attribute this to the fact that generalists are updated with more tokens and are thus knowledgable from an early stage, while it takes longer for specialists to refine their knowledge with small local training data.
The Necessity of the Co-existence of Generalists and Specialists. The performances of CoMiGS-22 2G and CoMiGS-2S are not consistent across the different sce- narios, while our CoMiGs S can always closely track the best-performing model, which is clearly visualized in Figure 9. Even for in-distribution tasks, it is unclear whether CoMiGS-28 or CoMiGS 2S-2S will outperform, suggesting both generalists and specialists are necessary as it is im-
4.3. Adaptation to Computational Resource Heterogeneity
4.3. BASELINE COMPARISON
In this section, our focus is to deal with computational resource heterogeneity, where users can have different num-
6

---

CoMiGS: On-Device Collaborative Language Modeling
L1 and L. roguan lption indreicuss
Ricessela lill For congei1 15 Iedain 18b
We popjett rovel aipeieera en lneal sanis prioe, Our antiraatonn and anoditute anpiliis lan
lhe quit villgo Ediota reated
modees 1 mec Aen bekon sssoo unit
Ludd young alried slays site
coaliicants an pareily tann to the lost
soker kann spen, alwang yeur oing io Dadta day. whth ectidor no Frosain puldan light the bennir hee,
modes wilh bee sooit coert. colccedd beeen (dddee ladds soueeed
poinl flbon wo0k 100
dessed courte carts meo domunaaiue qe complettyy lights wthre under
pamill can prowert lange coofloordss
IV'l divendion
stood
lave.
conet einutttonn coul in the iolked Sar
feivorm, thans
epprosctedd sii ronatid spsninci
StackEx User
Codes User
Arxiv
Book User
to penern martfaee in readains bouring mong
FB 10, 11
the apricoet curdy ingenes the wil dincy sceroadi loctisi quariun phane pe tiee of roportiaa speed owe dessnnerue complntty 1Ol kaull soxer
batween nilino hills and anddere beeadl
the ubroiiale vikee magelaade of
No aspond (1) 16
posanied a cursao ri and
funcion: Hhl con souuf in scoroo pettoring bear arfictioty
rec lus beword faviko telfe Onn while
prie iiboaaciiiinoom 105
by
sesets, dee decttwee hakken aa puddon ioo Sraret
(Ttkge lakks
pcregi
denn. lims doodant that sparso moodls
spocks cord bois Muneeiia
coikkiinng towneds Fenvers Elara appeonctee))
baie
Stac KEX User
Codes User
Arxiv User
Book User
Figure 3: Visualization of in-distribution token-level routing results for CoMiGS -eriiss-1Ge 1G1S trained on SlimPajama. Tokens are colored with the Top1 expert choice at the first layer (top) and last layer (bottom). Orange denotes the generalist and blue denotes the specialist. Texts are generated by ChatGPT. Further colored text plots are provided in Appendix E.
Sports
count of CoMiGs-Icxs with (4, 2, 2, LoRA experts (rank 8), LoRA modules of ranks (32, 16, 16, 16) would be required. With Top2 routing, to match the active parameter count, each user would need LoRA modules of rank 16.
35
3:5
150
130
150
200
330
200
20
150
0.50
353
Our results, presented in Table 2, are based on allocating varying numbers of experts to users, with computational resource availability decoupled from local task complexity. Our method outperforms the baseline methods for all in- distribution tasks, regardless of matching the full parameter count or the active parameter count. This advantage stems from the fact that both LetLoRA and F elexLAAA average model parameters across users without allocating param- eters for local adaptations, focusing on building strong generalist model. In contrast, our approach adaptively inte- grates both generalist and specialist knowledge, excelling in scenarios where specialized knowledge is crucial.
101
2:01 1.0
391
393
206
100
Figure 4: Expert Scores for the generalist expert and the specialist expert, averaged across all tokens and multiple batches for the out-of-distribution task (AG News). X- axis: number of iterations. Top: CoMiGs-1G1s, Bottom: pFedMoE. Darker colors indicate deeper layers.
bers of experts n4. We denote different experimental setup by specifying the list of nıs. We still keep one generalist expert per device, but the number of specialists can vary across the users (the variation is called One- Generalist-X. Specialists, in short, CoMi GSX) It's important to note that the richness of computational resources doesn't 't always correlate with the complexity of local data. For instance, some users may have ample computational resources but local data in small quantities. In such cases, crucial objec- tive is to prevent overfitting due redundant model-fitting abilities.
Table 2: models, averaged across 3 seeds. Light / dark grey denote in-distribution and out-of-distribution tasks respectively.
Mean test ppl (std) over users with heterogeneous
HETLORA ACTIVE
FLEXLORA
DOMIGS-IGXS
ACTIVE
POLL
GF22 GGP2
MULTELENGUAL (222444)
(6.46)
51. (6.76.0)) 10)
10.6010.20)
32.71 (0.15)
73.65 (0.
(44.22)
47.24
53.76 17.76.0.11 10)
59.14 (0.04)
72.74 (0.15)
75.64 (0.19)
SLIMPARAMA (244.2)
22.10 (8.17)
23.33 10)
23.15 (0.09)
22.95 (0.10)
23.00 (0.07))
(422224)
22.25(0.09)
23.33 (0.10)
23.15 (0.09)
22.95 (0.10)
23.03 (0.05)
We compare our approach to two state-of-the-art baselines: LoRA from Cho et al. (2023) and FleXLoRA from Bai et al. (2024), both of which adapt LoRA ranks based on the resource capacity of each user. HetLoRA aggregates LoRA matrices A and B by zero-padding to the maximum rank and then distributes them back using rank truncation. In contrast, F FlexLoRA first reconstructs model updates AW and redistributes the aggregated updates using SVD. We compare our method to these baselines by matching the number of tunable parameters, measured as both active and full parameters. For example, to match the full parameter
AC NEWS (4.2222)
(0.07)
31.58 (8.14)
31.05 (0.13)
36.41 (0.18)
36.62 (0.11)
(2444)
54.23 (0.09)
31.58 (8.14)
$2.53 (0.19)
36.41 (0.18)
36.4610.04)
Z1asa2.2.18
COMMON- CORPUS
(24422)
(8.74
21.41 (8.13)
21.74 (0.09)
24.63 (0.12)
25 (0.05)
(42.444)
21:41 (0.13)
31.61 (0.10)
24.63 (0.12)
(0.05)
NEWS
(0.04)
(0.04)
12.70
4.4. User-specific Analysis
In this section, we investigate how each user can benefit from our CoMiGS- -xx.. It is observed that our approach is more robust to overfitting due to the regularizing effect
7

---

CoMiGS: On-Device Collaborative Language Modeling
Buirk User 20v0
Rrechh ther (14R)
of the generalist, while at the same time better fitting local data through the incorporation of specialist knowledge.
We conduct experiments using the Multilingual Wikipedia dataset, where there are enough tokens to allocate different data quantities to users. In practice, users may not know their local data complexity, leading to a potential mismatch in resource allocation relative to data quantity. To simulate such scenarios, we allocate model aapabilities-meastee by n, (the number of LoRA modules per user) either posi- tively or negatively correlated with their local data size. It is important to note that one generalist is always assigned. Top2 routing is always performed is n; 2 2.
250
54
1.0
180
350
Figure 7: Test Perplexity vs. the number of iterations. Low and high denote data quantity. Legend denotes ni-
Similar conclusions can be drawn from our empirical results. However, there is a limit to how much generalists can help prevent overfitting when the local tasks are easy.
4.5. Computational and Communication Overhead
More Specialists Help with Higher Data Quantity. High data quantity users (French and Italian) consistently benefit from having more specialists locally, their test perplexi- ties decrease when the number of specialists increases from 1 to 3 to 7. This suggests that when sufficient local training data is available, adding more specialists leads to improved performance.
Our approach offers a significant advantage for on-device deployment due to its minimal computational and communi- cation overhead. We compare the resource consumption of our CoMi CS- to FedAvg in Table 3, matching the parameter count for LoRA modules.
The communication costs are halved compared to standard FedAvg, as only the weights of generalist experts are ex- changed. Our framework employs a first-order algorithm, ensuring that computation and memory requirements re- main on par with those of standard FedAvg algorithms. The additional memory and computational overhead primar- ily stem from the inclusion of the router, which is minimal (1.25% increase) since the router consists of a single-layer MLP.
Cernen User doed
Dutch User lhal
Pranch User 31y0
talknn Uierl 24gh1
100
112
130
1.5
160
305
12
Figure 5: Test Perplexity vs. the number of iterations. Low and high denote data quantity. Legend denotes nt;
Generalists Help to Prevent Redundant Specialists from Over-Fitting. For users with low data quantities, local model training with just two LoRA modules already results in overfitting (a trend observed in Figure 9). Our method suc- ceeds to suppress overfitting, even when fine-tuning twice or four times many expert parameters. We attribute this to the existence of the generalists.
Table 3: Extra resource consumption (per device) CoMi DOMiGS-1G1 8 compared to standard FedAvg, assuming base model is GPT-124M with 6 training.
COMP. OVERHEAD
MEMORY
COMM. COSTS / /ROUND
FORWARD PASS
5 MFLOPS
+0.035 MB (+1.25%)
-1.41 MB (-50%)
(+1.25%)
Germen User dow!
Quich Usar (bwl
Frnnc User 3igh
salkin Jaar 24gh
5. Conclusions
We propose a novel framework for on-device personalized collaborative fine-tuning of L.MMs, grounded in an innova- tive bi-level formulation of the Mixture-of-Expert learning objective. Our fine-grained integration of generalist and specialist expert knowledge achieves superior performance in balancing personalization and collaboration within Feder- ated LLMs.
1005 2223
150 4032
LBC 24.22
Figure 6: Test Perplexity vs. the number of iterations. Low and high denote data quantity. Legend denotes n.,
Specialists Can Benefit Generalists. What happens if users can only support a maximum of one expert? In our setup, such users must rely on the generalist expert when participating in collaboration. Interestingly, even when their collaborators are allocated more specialists, low-resourced users with only one generalist still benefit from the refined role diversification between generalists and specialists. As result, the generalists become more powerful, as demon- strated in Figure?? 7.
Furthermore, our framework is the first to address both model and data heterogeneity in collaborative LLM train- ing. It further decouples local data quantity from resource availability, allowing high-resourced users to leverage larger datasets for improved performance while remaining resilient against overfitting in low-data scenarios. CoMiCS is both theoretically sound and resource efffrreetf for practical de- ployment.
We provide an additional example of the impact of local data quantities in Appendix using SlimPajama dataset.
8

---

CoMiGS: On-Device Collaborative Language Modeling
Impact Statement
Federated Learning in the Age of Foundation Models in Conjunction with NeurIPS 2023, 2023. URL https:
We offer a collaboration framework for edge devices, aim- ing to enable smaller devices to leverage large language models (LLMs) despite limited resources and data availabil- ity. Our approach enhances fairness and mitigates privacy concerns by ensuring data remains on end devices. The pri- vacy aspects can further be enhanced by differential private aggregation of generalist weights, which we do not pursue here.
/ //oppneeevew... newt//ee/////..
Dai, D., Deng, C., Zhao, C., Xu, X., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., Xie, Z., Li, K., Huang, P., Luo, F., Ruan, C., Sui, Z., and Liang, W. Deepseekmoe: Towards ultimate expert specialization in mixture-of-expert language models, 2024.
DeepMind, G.
Gemini ai model, 2023. URL https: / //www. wwwpmiieeniimmmmmmmm nd. www./eeearcc/gemtiiiiiceeeeeeeee gemmiii. Ac- cessed: 2024-09-24.
The robustness towards attackers is beyond the scope of our work. Our collaboration framework has no guarantee of resilience towards adversarial attackers through the aggrega- tion of the generalist weights, which could potentially lead to misuse by certain parties. Further research is required on top of our framework to guarantee its safe deployment.
Ding, Y., Niu, C.,, Wu, F., Tang, S., Lyu, C., and Chen, Enhancing on-device Ilm inference with histori- cal cloud-based Ilm interactions. In Proceedings of the 30th ACM SIGKDD Conference Knowledge Discovery and Data Mining, KDD '24, pp. 597-608, New York, NY, USA, 2024. Association for Comput- ing Machinery. ISBN 9798400704901. doi: 10.1145/ 3637528.3671679. URL setps://don 00. 3671679.
Acknowledgements
This work supported by the Swiss State Secretariat for Education, Research and Innovation (SERI) under contract number 22.00133 and received funding from the European Union's Horizon 2020 research and innovation programme under grant agreement No 101017915 (DIGIPREDICT))
Fan, D., Messmer, B., and Jaggi, M. Towards an empir- ical understanding of moe design choices, 2024. URL /ttpaab///22222.. 13089.
References
Fedus, W. Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022a.
Almansoori, J., Horváth, S., and Takáč, M. Collaborative and efficient personalization with mixtures adaptors. arXiv preprint arXiv:2410.03490 2024.
Anthropic. Claude ai model, 2023. URL https:://// //tp:///wwwwwww Accessed: 2024-
Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity, 2022b. URL https: httpp:ss//xxxvv.ogg../o abs/21011 03961.
anthropic, anthropic.comm/ndeee 09-24.
Apple. Apple intelligence foundation language mod- els, 2024. URL https: ftttts::/arxiv..oo wbw.o/////b////////// 21075.
Fomenko, V., Yu, H., Lee, J., Hsieh, S., and Chen, W. A note on lora, 2024.
Bai, J., Chen, D., Qian, B., Yao, L., and Li, Y. Feder- ated fine-tuning of large language models under heteroge- neous language tasks and client resources. arXiv preprint arXiv:2402. 1.:402111500 2024.
Gaspar, H. and Seddon, M. Glolloc: Mixture of global and local experts for molecular activity prediction. In ICIR22022 Machine Learning for Drug Discovery, 2022.
Guo, M., Dai. Z., Vrandecic, D., and Al- Wiki-40b: Multilingual language model dataset. In LREC 2020, 2020. URL http: //t:://wwwww.///. www.- cetecccon.ooggo www.:/oggg/poooee.ggg..... 0w122220ppd//2000
Chen, C., Feng, X., Zhou, J., Yin, J., and Zheng, X. Feder- ated large language model: A position paper, 2023. URL 2207.... 08925.
Rfou,
ht https: //tp:i/aaiiv..//aa wwtpxivvorg//aee
1.trecl1.266. 29..296.p... pdf.
Chen, T., Sun, Y., and Yin, W. Closing the gap: Tighter anal- ysis of alternating stochastic gradient methods for bilevel problems. Advances in Neural Information Processing Systems, 34:25294-25307, 2021.
Guo, P., Zeng, S., Wang, Y., Fan, H., Wang. F., and Qu, Selective aggregation for low-rank adaptation in federated learning, 2024. URL https: http///axvvvv.//// orwp/.iioo.oooo.ooooooooo abs/ 01463.
Cho, Y.J., Liu, L.. Xu, Z., Fahrezi, A., Barnes, M., and Joshi, G. Heterogeneous loRA for federated fine-tuning of on- device foundation models. In International Workshop on
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L.,, and Chen, W. Lora: Low-rank adaptation of

---

CoMiGS: On-Device Collaborative Language Modeling
large language models. arXiv preprint arXiv:2/06.09685, 2021.
in Natural Language Processing, pp. 91-96, Bangkok, Thailand, August 2024. Association for Computational Linguistics. URL http:::/aclaatthhaooo http///aaaaatttooooy.oooo org/
Iyengar, A. and Adusumilli, P. Bigger isn't always better: How hybrid pattern enables smaller lan- guage models, 2024. URL https://wwl wim.
2024. www.sriiiitee....1....
-imm.
pleias.
Common corpus, 2024.
URL https:
wttpc//tttt////eee......... //ogggligeeecc//tttttt///aaaa common_corpus
/ //huggingface.
Qi, J., Luan, Z., Huang, S., Fung, C., Yang. H., and Qian, Personalized federated learning of large language model via dual lora tuning. arXiv preprint WrXiv:2406.07925, 2024.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L.,, Lachaux, M.-A., Stock, P., Subramanian, S., Yang, S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of experts, 2024.
D. Fdlora:
Qwen. Qwen2.5; A party of foundation models!. September 2024. htthttps: http:///...nmn..bbbbbbbb io /b/llo// qwen2. 5-11m/.
Jordan, M. I. and Jacobs, R. A. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2): 181-214, 1994.
Raschka, S. Practical tips for finetuning Ilms using lora (low-rank adaptation), 2023. URL https: / //wggarine..ebattta..e. practical-tin for for-ffietuunnnnnnn teetuninng-lmms
com/p/
Kalajdzievski, D. A rank stabilization scaling factor for fine-tuning with lora, 2023.
Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, J. R., Hestness, J., and Dey, N. SlimPajama: A 627B token cleaned and deduplicated version of RedPa- jama. https: http://wwwwerbb...t.....eeee
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. Communication--ffitii learning of deep networks from decentralized data. In Artificial intelli- gence and statisticy, pp. 1273-1282. PMLR, 2017.
s slimpajama-ae a oalaammm-----ttttttt---eeeeeeeea
- cecupppiiate---errriioo--------a
2023.
URL https::/huggingggtttt co/
Meta. Llama 3.2: Revolutionizing edge ai and vi- sion with open, customizable models, Septem- ber 2024. URL https://ai.mete. com/ blog//Lammes22--cnnett-22----1 typhenchar\fond (char) metegee--mbil--deeeieeee
627B.
Sun, Y., Li, Z., Li, Y., and Ding, B. Improving lora in privacy-preserving federated learning. ICLR, 2024.
Suzgun, M., Melas-Kyriazi, L., Sarkar, S. K., Kominers, S. D., and Shieber, S. M. The harvard uspto patent dataset: A large-scale, well-structured, and multi-purpose corpus of patent applications. 2022. URL httttts::/aaaxi o org/ abs / 2202 04043.
Accessed: 25.11.2024.
Mohtashami, A., Hartmann, F., Gooding, S., Zilka, L., Shar- ifi, M., et al. Social learning: Towards collaborative learning with large language models. arXiv preprint arXiv:2312. 1441, 2023.
Von Stackelberg, H. Market structure and equilibrium. Springer Science & Business Media, 2010.
Nesterov, Y. Lectures on convex optimization, volume 137. Springer, 2018.
Wagner, N., Fan, D., and Jaggi, M. Personalized collabo- rative fine-tuning for on-device large language models. Conference on Language Modeling, 2024.
Nesterov, Soft clustering by convex electoral model. Soft Computing, 24(23):17609 17620, 2020.
Wikimedia-FFoundation. Wikimedia downloads. URL https: //tppp://dmmppss wikimedia.
OpenAl. Chatgpt (september 26 version), 2023. URL https: htops::/chat... chat. http://hhha...eii.....o opw.ei..cccoooommmmeee..oo com/. Large language model.
Xu, J., Li, 7.,, Chen, W., Wang, Q. Gao, X., Cai, Q., and Ling, Z. On-device language models: A comprehen- sive review, 2024. URL https: wawpiv.iiivv////ii www.aiiwwwa.aaaaaaaaaaoaoooeeeee 2409.00088. 00088.
Peng, D., Fu, Z., and Wang, J. Pocket.LM: Enabling on- device fine-tuning for personalized LLMs. In Habernal, I., Ghanavati, S., Ravichander, A., Jain, V., Thaine, P., Igamberdiev, T., Mireshghallah, N., and Feyisetan, O. (eds.), Proceedings of the Fifth Workshop on Privacy
2409.
Yi, L., Yu, H., Ren, C., Zhang, H., Wang, G., Liu, X., and Li, pfedmoe: Data-level personalization with mixture of
10

---

CoMiGS: On-Device Collaborative Language Modeling
experts for model-heterogeenous personalized federated learning, 2024.
Zhang, J., Vahidian, S., Kuo, M., Li, C.. Zhang, R., Wang, G., and Chen, Y. Towards building the feder- ated gpt: Federated instruction tuning. arXiv preprint wrXiv:2305.05644 2023.
Zhang. X., Zhao, J., and LeCun, Character-level convo- lutional networks for text classification, 2016.
Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., and Fedus, W. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.0888 08906, 2022.

---

CoMIGS: On-Device Collaborative Language Modeling
A. Our Algorithm
The pseudo codes our proposed CoMi GS method are presented Alg. 1. While the scheme requires a server, it can alternatively be implemented serverless all2all fashion, which requires N times more communication overhead and we do not further pursue this here.
Algorithm Pseudo code of our proposed algorithm
Input: Expert parameters (0 (.0) ) routing parameters { Local training data and validation data ( (X!ran X {1, 2, .., J. Communication round 7 and routing update period т. Load balancing weight A.
for t 1..,, T do
Server aggregates generalist parameters: 0 N E for i [)) do
0 1,1-1
Users download aggregated generalist weights and prepure model parameters for training {0f-,, - 0s,-1 Do gradient steps on (0 1,00 i,t-1) towards minimizing (7) and get (0 0
min (( trann 0 0. 11.-1), X mmin +
@ .0
7
A L X X 0 0 22.-1))
if t% then
Do gradient steps Φr,--1 towards minimizing (8) and get Φi.t
min valid 0 0 Φb), x x
(8)
Cl (X. valid 0 0 Φ.)
end if end for
Each device i (1, .., N) sends generalist weights 0 the server
end for
Return: Expert parameters (0C i.T', 0 and routing parameters {Φur)
B. Extra Experimental Details
B.1. Training Details
Following Kalajdzievski (2023), we choose y to rank-stabilized value, technique which helps stabilize gradient norms. and the rank hyper-parameters choose from. The LoRA modules function follows:
W W +7 AB,
7
All our experiments except the centralized ones were conducted on a single A100-SXM4-40G1 GPU. The centralized learning baseline experiments were conducted on a single A100-SXM4-800 80GB GPU, as a batch size of 64*4 requires larger storage capacity.
We constant learning rate of 2 x 10 for updating routing parameters and a 2 x 10 learning rate with one-cycle cosine schedule for expert parameters during fine-tuning. The LoRA rank is set to unless otherwise specified, with LoRA alpha set to 16, following the common practice of setting alpha to twice the rank (Raschka, 2023). A load balancing weight 0.01 always applied.
GPT2 Experiments. For AG News and Multilingual Wikipedia data splits, we conduct 20 communication rounds. For SlimPajama data splits, due to greater category diversity, we conduct 50 communication rounds. Between each pair of communication rounds, there are 10 local iterations. each iteration, batch size of 64 is processed with a context length
12

---

CoMIGS: On-Device Collaborative Language Modeling
of 128. We the routing update period 30 iterations, and every time update routing parameters, do 10 gradient steps the validation loss. The choice of the hyperparamters is from sweep and provide the evidence in Figure 8.
Llama3.2 Experiments. For AG News data splits, we conduct communication rounds. For Common-corpus data splits, due to greater category diversity, we conduct communication rounds. Between each pair of communication rounds, there 10 local iterations. In each iteration, batch size of 64 is processed with a context length of 128. We the routing update period to 30 iterations, and every time update routing parameters, do 10 gradient steps the validation loss.
Test Perplexity
30
20
10
0
(1, 10)
(1, 20)
(1, 30)
(10. 10) (10. 20) (10, 30) (20. 10) (20. 20)
(20, 30)
(5. 10)
(5. 20)
(5, 30)
( RRouting Update Steps, Routing Update Period)
Figure 8: Sweep results on SlimPajama data splits using GPT2-124M base model. We ablate the impact of the update period (T) and the number of update steps (s) on model performance.
B.2. Datasets
The number of tokens for experiments within each user is shown in Table 4.
Given the extensive pre-training of Ilaaaa 3.2 models on over 15 trillion tokens from public sources, and the multilingual capabilities of Llama 3.2 - IB, fine-tuning on multilingual Wikipedia or SlimPajama resulted in negligible improvements likely due to significant overlap with the pre-training data corpus. We curated another more difficult fine-tuning dataset Common Corpus to show case the distinctions of the bascline methods.
Table 4: Number of tokens in each dataset splits
User 1
User
User
User 4
TRAINING
557' '62
407'498
556'796
451 581
Multilingual
VALIDATION
300'764
216'318
220'071
165'984
TEST
229'720
219'741
210'570
172'547
TRAINING
1'000'000
7000'000
1'000' 000
1'000'000
SlimPajama
VALIDATION
2000000
200'000
200'000
200'000
TEST
2000000
200'000
200'000
200'000
TRAINING
761'924
756'719
814'131
771'460
AG News
VALIDATION
48'800
48'730
50'398
48'249
TEST
48' 167
47'721
48'344
497377
TRAINING
1'000' 000
1'000' 000
1'000' 000
1'000' 000
Common Corpus
VALIDATION
200'000
200'000
200'000
200'000
TEST
200'000
200'000
200'000
200'000
C. More Tables and Figures
C.1. Learning Curves of Different Methods
See Figure
13

---

CoMIGS: On-Device Collaborative Language Modeling
Multilingual Wikipedia
German
Dutch
French
Italian
Perplexity
60
80
50
60
40
100
100
200
SlimPajama
StackEx
Codes
ArXiv
Book
Test Perplexity
14
11
25.0
10
13
22.5
20.0
200
400
200
400
200
400
200
400
AG News
World
Sports
Business
Sci/Tech
Test Perplexity
40
40
40
35
30
50
100
150
200
50
100
150
200
50
100
150
200
50
100
150
200
FedAvg
Local
CoMIGS-2G
CoMIGS-2S
CoMiGS-1G1S
Figure 9: Test Perplexity during training (base model: GPT2-124M): our method closely follows the best performing method
Multilingual Wikipedia
German User
Dutch User
French User
Italian User
Average Scores
1.0
0.75
0.75
0.75
0.5
0.50
0.50
0.50
0.25
0.25
0.25
50
100
150
200
50
100
150
200
50
100
150
200
50
100
150
200
SlimPajama
StackEx User
Codes User
ArXiv User
Book User
Average Scores
0.75
0.75
0.75
0.75
0.50
0.50
0.25
0.25
0.25
0
200
400
0
200
400
0
200
400
200
400
Generalist
Specialist
Figure 10: Expert Scores for the generalist expert and the specialist expert from our CoMiGS 1G18 method, averaged across all tokens and multiple batches for the in-distribution task, with x-axis being the number of iterations. Darker colors represent deeper layers.
C.2. Extended Baseline Comparison
An extended version of Table is presented in Table 5. In this extension, we incorporate two additional ablations: I) Integration of a routing mechanism, updated simultaneously with the expert networks; 2) Iterative updates alternating between routing and expert parameters, with the routing parameters updated using newly-sample training batches instead of dedicated validation set. is to address the scenario where a validation set not available.
14

---

CoMIGS: On-Device Collaborative Language Modeling
Moreover, include other baseline methods FFA-LORA from Sun al. (2024) and FedSA from Guo et al. (2024). FFA-LORA keeps the LoRA A matrices fixed at initialization, while FedSA always aggregates LoRA A matrices but leave LoRA B matrices localized.
Notably, the comparison between scenarios ii) and iii) reveals minimal disparity, underscoring the significance of having an independent validation set exclusively for routing parameter updates.
Table 5: Mean test perplexity over users with homogenous models, averaged across seeds. Mean (std) with rank locator for the mean (the lower the better). Green denotes the best performing methods and red denotes our method.
IN DISTRIBUTION
OUT OF DISTRIBUTION AG News
Multilingual
SlimPajama
1) WITHOUT ROUTING
Pretrained
156.12
37.19
90.65
Centralized
55.41 (0.12)
+
19.53 (0.14)
28.19 (0.52)
Local
54.38 (0.32)
I
26.95 (0.14)
41.46 (0.06)
FedAvg
58.80 (0.34)
23.27 (0.05)
31.84 (0.02)
FFA-LoRA
66.80 (0.20)
22.85 (0.12)
33.13 (0.09)
FedSa-LoRA
57.60 (0.14)
23.40 (0.13)
31.57 (0.10)
PCL
54.53 (0.19)
26.99 (0.19)
32.25 (0.12)
UPDATE ROUTING AND EXPERT PARAMS SIMULTANEOUSLY ON TRAINING LOSS
Local-MoE
55.27 (0.40)
27.16 (0.16)
41.49 (0.01)
FedAvg-MoE
56.77 (0.37)
+-
23.32 (0.07)
32.24 (0.08)
PFedMoE
52.27 (0.17)
22.91 (0.18)
+
38.72 (0.21)
II1 ALTERNATING UPDATE ROUTING PARAMS NEWLY SAMPLED BATCHES FROM TRAINING SET
Local-MoE- tr
53.78 (0.33)
27.78 (0.06)
41.46 (0.03)
FedAvg-MoE- ir CoMiGS
59.39 (0.13)
23.00 (0.01)
1
31.70 (0.16)
50.86 (0.14)
25.45 (0.01)
I
38.93 (0.08)
IV) ALTERNATING UPDATE ROUTING PARAMS
VALIDATION
CoMiGS 25
46.36 (0.16)
22.51 (0.08)
35.81 (0.13)
CoMiGS 2G
58.31 (0.17)
21.36 (0.01)
4
31.18 (0.05)
1
CoMiGS IGIS
47.19 (0.10) I
21.79 (0.04)
4
33.53 (0.03)
1
C.3. HetLoRA
Analogously to the baseline experiment comparison in FlexLoRA (Bai al., 2024), we use 7 0.99 as pruning strength and sweep the regularization parameter in (5 x 10 10-2,5 x 10 5 10-4).
C.4. Is the Standard Load Balancing Loss Sufticient?
The standard load balancing loss encourages equal assignment of tokens to each expert. When the number of experts gets larger, there might not be enough tokens routed to the generalists, which might lead under-developed general knowledge. We will verify if this is indeed true.
To encourage enough tokens to be routed to the generalist expert such that more general knowledge can be developed, we modify our load-balgn loss by introducing importance weighting. As we separate the 0-th expert to be the generalist expert and conduct Top-2 routing, the modified load balancing loss as follows:
CLB
(n4 - 1)2 fo Po
1
Re
n2:
S3.P.
(10)
[ng
1)2
+
15

---

CoMiGS: On-Device Collaborative Language Modeling
Table 6: Test perplexity with different load balancing terms with (hetero) or without (homo) resource heterogeneity.
No LB
LB (uniform)
LB (generalist-favored)
News (homo)
33.69 (0.21)
33.53 (0.03)
33.53 (0.03)
AG News (hetero)
34.31 (0,05)
34.28 (0.11)
34.22 (0.09)
Multi- Wiki (homo)
47.31 (0.15)
47.19 (0.10)
47.19 (0.10)
Multi-Wiki (hetero)
46.36 (0.16)
46.15 (0.04)
46.48 (0.16)
SlimPajama (homo)
21.77 (0.02)
21.79 (0.04)
21.79 (0.04)
SlimPajama (hetero)
22.15 (0.07)
(0.10
22.10 (0.17)
where
f = T Σ }i Top2 indices of p(z)}
P3 = 7 Σ E p)(2)
(11)
j is the expert index and p(z) (,,())] is the logit output from the routing network for specific token The idea is that one of the top 2 tokens should always be routed to the generalist expert, i.e. the O-th expert. Thus, f should be equal to 17266.11 for j O. As the original load balancing loss encourages uniform distribution, this modification encourages the generalist expert to have routing probability of 0.5 expectation. Note that when n, 2, this L is the same as the original load balancing loss as proposed in Fedus et al. (2022a).
We present the results in Table 6: in both scenarios, whether users have the same or different numbers of experts, including a load-balancing term leads to a slight improvement compared omitting it. However, encouraging more tokens be routed the generalists does not make a significant difference.
D. Additional Experiments
We replicate the experiments in Section 4.3 with the SlimPajama dataset, where we assign four times as many tokens to ArXiv User and Book User as to Stack Exchange User and Codes User.
More Specialists Help with Higher Data Quantity. From Figure 11, it evident that ArXiv User and Book User, with abundant local data, benefit from having more local experts.
StackEx User (low)
Codes User (low)
ArXiv User (high)
Book User (high)
11
Perplexity
14.0
22
20
18
13.0
33
9
16
200
400
200
400
200
200
400
1GXS: 2,2,4,4
1GXS: 2,2,8,8
1GXS: 2,2,2,2
Figure 11: Test Perplexity during training for the SlimPajama setup. ArXiv User and Book User have more local data and thus benefit from having more experts. The numbers in the legend indicate the number of experts n, within each user. Top-2 routing performed.
Generalists Help to Prevent Redundant Specialists from Over-Fitting? From Figure 12, observe more prominent overfitting than in Figure 6, likely because the tasks are objectively easier, indicated by lower test perplexity from the beginning of fine-tuning, Generalists have limited power to prevent overfitting with easy tasks.
Specialists Can Benefit Generalists. Low-resourced users that can only support single expert setup still benefit from collaboration, as the generalist knowledge is refined through a more detailed distinction between specialist and generalist roles via other bigh-resourced users. This indicated by the enhanced performances for Stack Exchange and Codes Users.
16

---

CoMIGS: On-Device Collaborative Language Modeling
StackEx User (low)
Codes User (low)
ArXiv (high)
User (high)
Perplexity
14.5
22
14.0
10
20
13.5
18
33
13.0
200
400
200
400
200
400
200
400
1GXS: 4,4,2,2
1GXS: 1GX5:8,8,2,2
1GXS: 2,2,2,2
Figure 12: In this SlimPajama setup, Stack Ex User and Codes User despite having low resources locally, overfit slightly on their mall-sized local data. Numbers the legend denote the number experts n, within each user. Top2 routing is performed.
StackEx User (low)
Codes User (low)
ArXiv User (high)
Book User (high)
Test Perplexity
14
22.5
20.0
34
17.5
200
400
200
1GXS: 1.L11
1GXS: 1,1,4,4
1GXS: 1.1,8,8
Figure 13: In this SlimPajama setup, Stack Ex User and Codes User, despite having only one expert locally, still benefit from other users having more experts, thereby enhancing the generalist's performance. The numbers in the legend indicate the number of experts, n, within each user. Top-2 routing is applied when n, 2 2
E. Visualization of Expert Specialization
To visualize which tokens are routed to the generalist and specialist experts for our CoMiGS-1G15 model trained on SlimPajama, ask ChatGPT generate texts in the style of StackExchange, Python Codes, ArXiv Paper and Books. We then feed those texts to the user-specific models and color the token with the Topl routed index. The routing results after the very first layer (Oth), middle layer (5th), and the very last layer 11h are presented in Figure 14, 15 and 16.
We perform the same experiments on AG News, asking ChatGPT to generate News text on the topics World, Sports, Business, and Sci/Tech. The routing results after the very first layer (Oth), a middle layer (5th), and the very last layer (1 are presented Figure 17, 18 and 19.
For all the plots, diagonal entries in-distribution texts and off-diagonal entries out-of-distributiin texts.
17

---

CoMiGS: On-Device Collaborative Language Modeling
and 2 requloriaation and tochniquas 10 prevent over fitling in machine learning models, L1 requllrr ication ( Lessol adits the abeclute value of the mignituue of
111 and 1 requleriznte inationtionnnn are sechniqurs 10 prevent over fitting in mechine learning modells. L1 regulerr ication ( Lessol ados the absclute value of the magnitude of confiicients as a penally term to the loss function. This can result sparee models with lew coeffizientss eflecaively performing feature selection. L22 regular nation (A idge! adds the squared magnitudo of coeflictents as penaltt term. Tiis doesn't it lead to sparse models but can prevent large coefiicients,
I11 and 1 reqularization are sochniquaná to prevent over fitting in mechine learning models. L1 reguular icalion (Lsso adds the absclute value of the magnitude of coefiicients as a penally term to the loss function. This can result in sparse models with few coefiizientts ellectiell performing feature selection. L2 regular ization 1R idgel adds the squared magnitude of coofiicentss as penalty term. This doesn't it lead to sparse modelss but can prevent large coofiicinnss,
and I requiar perrinnion are techniquass to prevent over fitting in machine learning models. L1 reguulaiication ication Lasso) adule the absclute value of the magnitude of coeflicients as a perillly term to the loss function. This can result in speree models with few coefiizeenss ellectively performing feature selection. L2 regular ization (Rddee adds the squared
StackEx
coefficients as a perillly jorm to the loss function This can result in sperse models with few coeffizientss effectively performing feature soc splectionn. L2 regular nntion 8 idge! addss the squared magritudio of coefficients as penalty sorm. This doesn't it lead lo sparse models but con prevent large coefiiientss
magnitude of coofiictents as penalty term. This doesn't it. lead to sparse modells but can prevent large coofiiients.
def lib cncci In?
die lib onaii (m2 Rib [0,
del on (n2 Rib [0,
del obcoaacal Rib (0, 1]
lib [0,
for in mange nng
for in range m|
for range 22
for rangel ng fi.. aepandd [1b Ab(/11
Rb. (Hpan Ab(/211
Rb. abb.apaanll (11] Rb(/21)
Rb. app.nddpaadd (1b Ab(/211
Codes
nntum Rbb
nnturn rib
return rbbb
return tib
print (iliofinoooooooo acsil RD])
pfintilibon acdi [D]]
if accil [1D]]
irintiiiB-onncsiiii (10))
We present novel quanttum aigorithes that sigillicatll improve the efficiency of
We present novel quantum alporithms that significantly improve the efficiency of solving lineat systems of equalions. Oui approech lever ages quantum phase estimation and amplitude ampllication to achieve an exponential speed up over classical counterperts. We demonstrate that out algonthmm can salve Ax bil with complexity 1,0 log NN 11, where if
We present novel quantum alporithms that signilicantly improve the efficiency of solving linest systems of equations, Oui approsch lever ages quantum phase estimation and amplitude ampllication to achieve an exponential speed up over classical counterperts. We demonstrate that our algonthms can solve Ay bid with complexity V.0 loo (N])1 whore ki
We present novel quantuum algorithms that significantly improve the efficiency of solving lineet systems of equalions. Our approsch lever ages cuantum phase estimation and ampiltdde ampilication to achreve an exponental speed up over classical counterpents. ts We demonssrase that our algonthms can solve Al bü with complexity V. loog (N)1) whore V/ N is the dimension of the system, under specific conditions. Numer ical simulations confimi the robustinsss of our
soking lineal systems of equalions. Out
approach lever ages queritum phase estimatinn and amplitude ampilication to achieve an exponential speed up over classical counterparts. We demonstrate that out algonthim can solue Ax bü with complexity V.O.t Moog [N) where If NVI is the dimension of the systom, under speciiic conditions, Numeriia inal simulations confim the robusiness of our
ArXiv
NV) is the dimonsion of the system, under specific conditions. Numer ical simulations confim the robustress ness of our
N is the dimonsion of the system, under specific conditions, Numer ical simulations confirm the robustness ness of cur
quiet village Eidortt, between rolling and anzient tornsts, livrr a young (irl named Flara
quiet village betweeen rolling and ancient loxests) livar young (irl named Flara possessed a curlous mind and an adventurous sprit, always year ning to explore beyond the familir paths. One faseful day, while wandering through the woods, she discovered hidden glen
restled
quiet village between rolling and ancient loxests) livd young named Flana possessed a curoou mind and an adventurous spiit,, always year ning 10 exploce beyond the familia paths. One fsteful day. while wanderiin through the woods, she discovered hidden glen bat hed in golden light. In the center stood an ancient oak tree; its branches
iestlee
quret village Eldorts, and ancient torests, livad young quill named Flara
between
possessed 8 curicus mind and an
possessed a cutoous mind and an adventurous spiill, always year ning t0 exploce beyond the familiar paths. One foteful day, while wandering through the woods, the discovered hidden gl
adventurcus spiit, ahways year ning t0 explore beyond the familir paths. One fsneful day, while wandering through the woods, she discovered hiddeen glen bat hed in golden light. in the center stood an ancient cak: tree: its branches
Book
bat hed in golden light. in the center
bat hed in golden light. In the center stood an ancient oak tree; its branches
stood an ancient oak tree, its branches
stresching sowards the heavans. As Elara approached, she notioad a shimmoring ing
stretching sowards the heavens. As Elara approached, she noticad a shimmaring ing
stretching towards the heavens. As Elara
stretching towards the heavens. As Elara approached. she nalleed a shimmaring
approached, she noticnd a shimmaring ing
StackEx User
Codes User
ArXiv User
Book User
Figure 14: Visualization of token-level routing results for CoMiGS- 1G1s trained on SlimPajama. Tokens are colored with the first expert choice at the Oth (first) layer. Orange denotes the generalist and blue denotes the specialist. Diagonal entries are in-distribution texts and off-diagonal entries are out-of-distribttion texts. Texts are generated by ChatGPT.
18

---

CoMiGS: On-Device Collaborative Language Modeling
and 2 requloriaation and tochniquas 10 prevent over fitling in machine learning models, L1 regullrr ication ( Lessol adits the abeclute value of the mignituue of
1 and 1 requleriznte inationtinnnnn are sechniqurs 10 prevent over fitting in mechine learning meddels 11 reguller ication ( Lessol ados the absclute value of the magnitude of coefiicients as a penally term to the loss function. This can result in sparee models with lew coeffizients eflecaively performing feature selection. L22 regular nation (A idge! I addss the squared magnitudo of coeflictents as penaltt term. This doesn't it lead to sparse models but can prevent large confiicients
I11 and 1 reqularization are sochniquaná to prevent over fitting in mechine learning modele. L1 reguular icalion (Laso adds the absclute volue of the magnitude of coefficients as a penally torm to the loss function. This can result in sparse models with lew coefiizientss ellectiell performing feature selection. L2 regular ization (R adgel adds the squared magnitudo of cooflicennss as penalty term. This doesn't it lead to sparse modelss but can prevent large coofiicintss
and I requiar ietion are techniquass to prevent over fitting in machine learning models. L1 resuulai ication Lasso adule the absclute value of the magnitude of coeflicients as a perillly term to the loss function. This can result in speree models with few coefiiceenss, ellectively performing feature selectionn L2 regular ization (Rddee adds the squared
StackEx
coefficients as a perillly jorm to the loss function This can result in sperse models with few coeffizients, effectively performing feature so selectionn. L2 regular nntion 8 idge! addss the squared magritudio of coefficients as penalty sorm. This doesn't it lead lo sparse models but con prevent large confiiients,
magnitude of cooffictents as penalty term. This doesn't it. lead to sparse modells but can prevent large cooficients.
def lib cncci In?
die anai (n)
del on (n2 Rib [0,
del obcoaacal Rbb (0, 1]
lib [0,
Rib [0,
for in mange nng
for in range m|
for range n/
for rangel ng fi.. aepandd [1b Ab(/11
Rb. (Hbbad Ab1/211
Rb. abbpapa [Hb [11] (b(221)
Rb. app.nddpaadd (1b Ab(/211
Codes
nntum Rbb
ineurn rib
return rbbb
return tib
print (ilioiiooonooooo acdil RD])
pfinttilbon acdi [D]]
pf accil [10]]
printifiB (10))
We present novel quanttum aigorithes that sigillicatll improve the efficiency of soking lineal systems of equalions. Out
We present novel quantum alporithms that significantly improve the efficiency of solving lineat systems of equalions. Oui approech lever ages quantum phase estimation and amplitude ampllication to achieve an exponential speed up over classical counterperts. We demonstrate that out algonihmms can selve Ax bil with complexity 1,0 log (NI 11 where if
We present novel quantum alporithms that signilicantly in improve the efficiency of solving linest systems of equations, Oui approsch lever ages quantum phase estimation and amplitude ampllication to achieve an exponential speed up over classical counterperts. We demonstrate that our algonthms can solve Ay bid with complexity 1,0 lolg (N])) whore ki
We present novel quantuum algorithms that significantly improve the efficiency of solving lineet systems of equalions, Our approsch lever ages cuantum phase estimation and ampiltdde ampilication to achreve an exponertal speed up over classical counterperts. ts..s.....ss.. We demonssrase that our algonthms can solve Al bü with complexity V. loo (N]1) whore V/ N is the dimension of the system, under specific conditions, Numer ical simulations confimm the robustiasss of our
approach lever ages queritum phase estimatinn and amplitude ampilication to achieve an exponential speed up over classical counterparts. We demonstrate that out algonthim can solue Ax bü with complexity V Mog (N))) where If NVI is the dimension of the systom, under speciiic conditions, Numeriin inal simulations confim the robusiness of nur
ArXiv
NY is the dimonsion of the system, under specitic conditions. Numer ical simulations confem the robust ness of our
N is the dimonsion of the system, under specific conditions, Numerian ical simulations confir the robustness ness of cur
quiet village Eidoort, between rolling and anzient torasts, livrr a young (irl named Flara
quiet village betweeen rolling and ancient torests) livar young (irl named Flara
iestled
quiet village between rolling and ancient loxests) lvad young named Flana possessed a curoou mind and an adventurous spirlt, aawaas year ning 10 explore beyond the familia paths. One fsteful day, while wanderring through the woods, the discovered hidden glen bat hed in golden light. In the center stood an ancient oak tree; its branches
iestlee
quret village Eldorts, and ancient torests, livad young quill named Flara
between
possessed 8 curicus mind and an
possessed a curlous mind and an
possessed a cutoous mind and an adventurous spiill, always year ning 10 exploce beyond the familiar paths. One foteful day, while wandering through the woods, the discovered hidden gl en bat hed in golden light. In the center stood an ancient oak tree; its branches stretching towards the heavens. As Elara approached. she notleed a shimmaring
adventurcus spirl, always year ning t0 explore beyond the familir paths. One fsneful day, while wandering through the woods, she discovered hiddeen glen bat hed in golden light. In the center stood an ancient cak: tree: its branches
adventurous sprit, always year 1o explore beyond the familar paths. One faseful day. while wandering through the woods, she discovered hidden glen
Book
bat hed golden light. in the center
stood an ancient oak: tree, its branches
stresching sowards the heavans. As Elara approached, she notioad a shimmoring ing
stresching sowards the heavens. As Elara approsched) she noticad a shimmaring ing
stretching towards the heavens. As Elara
approsched, she noticnd a shimmaring ing
StackEx User
Codes User
ArXiv User
Book User
Figure 15: Visualization of token-level routing results for CoMiGS- trained SlimPajama. Tokens are colored with the first expert choice at the 5th layer. Orange denotes the generalist and blue denotes the specialist. Diagonal entries are in-distribution texts and off-diagonal entries are out-of-distribution texts. Texts are generated by ChatGPT.
19

---

CoMiGS: On-Device Collaborative Language Modeling
and 1 regularization and tochniquas t0 prevent overfitiig in machine jearning modelss L1 reguller ication ( [Lsss) adile the abeclute value of the mignitude of
11 and 1 requlerization are sochniques 10 prevent over littnng in mechine learning modelss L1 regullercca ication [ (Lassl) ados the abaclute value of the magnitude of confiicients as a penalty term to the loss functionn This can result sparee models with lew coeffizientss effecaively performing feature selection. L2 regular nation [A idge! adds the squared magnitudo of coeflictntss as penalty term. This doesn't iI lead to sparse models but can prevent large coofficients,
I11 and 1 reqularization are sochniquaná to prevent over fitting in mechine learning models. L1 reguuler icalion (Lsso adds the absclute value of the magnitude of coefiicients as a penally term to the loss function. This can result in sparse models with few coefiizientts ellectieey performing feature selection. L2 regular ization 1R idgel adds the squared magnitude of coofiicents as penalty term. This doesn't it lead to sparse models but can prevent large coofiicinnss,
and I requiar istion are techniquass to prevent over fitting in machine learning models. L1 resuulai ication Lasso adule the absclute value of the magnitude of coeflicients as a perillly term to the loss function. This can result in sperze models with few coefiiceenss, ellectively performing feature selectionn L2 regular ization Ridge adds the squared
StackEx
coefficients as a perillly jorm to the loss function. This can result in sperse models with few coefficients, eflectively performing feature sel retection.. L2 regular nation a idge! adds the squared magritudoo of coefficients as penalty serm. This doesn't it lead to sparse madels but con prevent large conficients,
magnitude of cooffictents as penalty term. This doesn't it lead to sparse modells but can prevent lorge cooficients.
def lib cncci In?
die lib anaal (nn
del on (n2 Rib [0,
del obcoaacal Rbb (0, 1]
lib [0,
Rib [0,
for in mange nng
for in range m|
for range 22
for rangel ng fi.. aepandd [1b Ab(/11
Rb. (Hpbn Ab(/211
Rb. abb.apa [Hb (11] Ab(
Rb. app.nddpaadd (1b Ab(/211
Codes
nntum Rbb
return rib
return rbbb
return tib
print (iliofioooonoooo acdil RD])
print inntttiiiooooo [C]]]
pfint accil [1D]]
printifiB (10))
We present novel quanttum atgonithess
We present novel quantum alporithms that signllicanlly improve the efficiency of solinng lineat systems of equalions. Our approech lever ages quantum phase estimation and amplitude ampllication to achieve an exponential speed up over classical counterpetts. We demonstiate that our algonthmms can salve AA bil with complexity 1.0 Mog (N) 1t where V NV) is the dimonsion of the system, under spectlie conditiins, Numer ical simulations confem the robustness of our
We present novel quantum alporithms that signilicantly improve the efficiency of solving lineat systems of equations, Oui approsch lever ages quantum phase estimation and amplitude ampllication to achieve an exponential speed up over classical counterperts. We demonstrate that our algonthms can solve Ay bid with complexity V.0 log (N])1 whore ki
We present novel quantuum algorithms that significantly improve the efficiency of solving lineet systems of equalions, Our approsch lever ages cuantum phase estimation and ampiltdde ampilication to achreve an exponertal speed up over classical counterperts. ts..t......s.. We demonssrase that our algonthms can solve Al bu with complexity V. loo (N]1) whore V/ N is the dimension of the system, under specific conditions. Numer tcal simulations confimm the robustiasss of our
that sigiliiantly improve the efficiency of
soking lineal systeme of equalions. Out
approach lever ages quenitum phase estmation and ampitudd ampilication to achieve an exponential speed up over classical counterparts. We demonstrate that out algonthims can solue Ax bü with complexity V.O.t log (N))) where if NV is the dimonsion of the systom, under speatiic conditiinss Numerica ical simulations confim the robusiness of nur
ArXiv
N is the dimonsion of the system, under specific conditions, Numer ical simulations confirm the robustness ness of cur
quiet village Eidorts, between rolling and anzient tornsts) livar a young (ir) named Flata She
quiet village betweeen rolling and ancient toxests, llvar young (irl named Flara possessed a curlous mind and an adventurous sprit,, alwavs year ninnnn to explore beyond the familar paths. One faseful day, while wandering through the woods, she discovered hidden glen
nestled
quiet village between rolling and ancient loxests) livd young named Flana possessed a curoou mind and an adventurous spiilt, always year ninn 10 explore beyond the familir parhs. One foteful day, while wanderiin through the woods, she discovered hidden glen bat hed in golden light. In the center stood on ancient oak tree; its branches
iestlee
quret village Eldorts between rolling and ancient torests, livad young quill named
possessed a curccus mind and an
possessed a cutoous mind and an adventurous spiill, always year ning 10 exploce beyond the familiar paths. One foteful day; while wandering through the woods, the discovered hidden gl en bat hed golden light. In the center stood an ancient oak tree; its branches
adventurcus spiil,, shways year ning to
Book
explore beyond the familar poths. One faneful day, while wandering through the woods, lshe discovered hiddeen glen bat red in golden light. In the center stood an ancient oak: tree. its branches
bat hedl in golden light. in the center stood an ancient oak: tree, its branches
stresching sowards the heavans. As Elara approached, she notioad a shimmoring
stresching sowards the heavens. As Elara
stretching towards the heavens. As Elara
stretching towards the heavens. As Elara approached. she notleed a shimmaring ing
approached, she nolicad a shimmaring ing
approached, she noticnd a shimmaring ing
StackEx User
Codes User
ArXiv User
Book User
Figure 16: Visualization of token-level routing results for CoMiGS- trained on SlimPajama. Tokens are colored with the first expert choice at the 1 11th (last) layer. Orange denotes the generalist and blue denotes the specialist. Diagonal entries are in-distribution texts and off-diagonal entries are out-of-distribution texts. Texts are generated by ChatGPT.
20

---

CoMiGS: On-Device Collaborative Language Modeling
In landmark ovant, a hightic peace accord hes osen signed between long- standing rivels in the Midsle Cast The agreement, broks red by international mediators, ams to end decades of conficc and pave the way for economic cooperation and regaional stability:
In londmark avont, a historic poace accord hes been signed between long- standing rivels in the Middle Cast. The agreement, broker by international medietors, aims to end decades of confiict and pave the way for economic cooperation and regionel stabilit.. Leaders from both sides extrested hage and commiment peaceful future. emphasiang the accord's 's potensal to bring prospenity and security to the region. This unprecedentad move has
In landmark avant, a historic peace accord has been signed benween long- standing rivels in the Middlle Cast The agreement, brokerre by international medietors, ams to end decades of conflict and pave the way for economic cooperation and regiional slabiliy:
in landmark evant, a historic peass accord has been signed between long- standing rivels in the Midelle Cast The agreementt brokered by international medietors, amm to end decades of confict and pawe the way for economic cosperellion and regional slability
World
Leaders from both sides expressed hope peaceful future, emphasizng the accord's 's potentaa bring prospenity and security the iegionn. Thi unprecedented mave has
Leadors from both sides expressed hope and comm tment emphasiang the accord is potentaa to bring progpentt and security the region. This unprocedented move has
Leaders from both sides expressedd hope and co commitment to peaceful future, emphasiong the accord's is potential to bring p proqpenity and security the regi region. This unprosddented move has
and comm
peaceful future,
Young tonnis sensation Emma Garcia has takan the sports world by storm with her stunning victory at the Grand Slam tournament. At (ust 19, Garcia dafested top ranked plavers, shownasing exccpiional skill and composure on the court. Her aggressive play style and strategic ac ummn have carned her widesooea acdaiim from lans and analysss allice. This historic win propels Garcia into the global spollight, Perking thet se formdable contender in future
Young lonris sensation Emma Gardia has taken the sports word by storm with her stunning victory at the Grand Slam tournament. At (ust 19, Garcia defeated top ranked plavers, showsasing exceaaiiona skill and compoaure on the court. Her aggressive playstçll and stratogic ac ummn have earned her widescreed aodaim from fans and analysa allin This Iistoric wan propels Geroa into the glosal spollight marking her se formidaale contender in future
Young lennis sensation Emma Garcia has taken the sports word by storm with her stunning victory at the Grand Slam tournament. At just 19, Garcia defeated top markdd plavers, showsasing exceritonal skill and compoaure on the court. Hear apgreasive playstyle and stratogic ac ummn have carned her widespread acdeim from lans and analysss allike This Instoric win propels Garoa into the glocal spotlight marking her as formidacl contender in future
Young lennis sensation Emma Garca has taken the sports world by storm with her stunning victory at the Grand Slam tournament. At just 18, Garcia dafeated top rarked plavers, showcasing excapiioral al skill and composure on the court. Her angressive playstyle and stratogic ac comen have aarned her widespreed accisim from lans and analysts alike, This hasooii win propels Garcia into the global spolightt marking her a formiable contender in future
Sports
Major tach companies are sqpnlicantly increesing their investments artificial inteliqence (All startups, siming to spear
Major tech companies are ficently Major tech companies are sgrvfcannlly incressing their investments artilicisl intelligence (Al) startups, siring to speer head innovation and maintain compelitive edge. Giants like Google, Microsoft, and Amazon are leacingg funding rounds, pouring billion into Al- driven enterprises focusing on machine learning, natural lanquage processing. and automation. These investments ard expected to accelerate advancements in All applications across vanious indusirios,
Major tech companios are cantly
increesing their investments artificial intellgence (All startups, siming to speer head inmovation and maintain compeliiive edge. Giants liee Google, Microsoft, and Amazon are leading funding rounds, pouring billons into All driven enterpeses focusing on machine lcamning natural language p processing and automation. These investments are expected to accolerate advancements in Al applications across vatious industiios,
incressing their inveetments artificial intelligence (Al) startups, aiming to spper head innovation and mamttan compeliiive edge. Giants like Google Microsoft, and Amazon are leeding funding rounds, pousing billions into Al- driven enterpreses focusing on machine learning, natural lanquage processing and automation. Those investments are expected to accelenate advancements in Ill applications across various industrios,
head innovation and mainzain
Business
compeliltve edge, Giants like Google,
Microsoft, and Amazon are leecing
funding rounds, pouring billions into Al- driven enterprses focusing on machine leaming, intural language processing and automation. Thesse investments are expected to accelerate advancements in Al applications across varous industriss,
the sun risess city park comes to with phony of ning hirds Jogsers weaue
the sun rises, ctt park comes to with sym phony of lrping and rust ling through winding paths their footstaps creating a ihyth me best. Familles gother for pien os. speeading out blankets and unpacking baskets filld with treats.
As the sun rises, ctt park comes to with a sym phony of chirping birds and rustiing leaves Joggers wanve through winding paths, their footstaps creating a ihyth mic best. Families gather for pie ics. spresding out blankats and unpacking baskets filed with treats.
As the sun rises, ctty park comes to with sumphory of chirping airds and rustting leavss Jog whave through winding paths their footsteps creating a my/h mic beat. Families gather pie ics. spreading out blankats and unpacking baskets filedd with treats.
and rustiing
Joggars wanvi
ihrough winding paths their footstaps creeling a mi beat. Families gather for picnics. spreading out blankets and unpacking baskets fillei with treats.
Scl/Tcch
Chidrenalls a laughter fils the air as they chase each other across the playground. By noon, office workers sook refuge under shady trees, enjoying a peaceful lunch broak. As evaning fills, the park
Childrenlils is laughter fils the ai as they chase each other acroes the playground. By noon, office workers seek refuge under shady trees, enjoying a peaceful lunch break. As evening falls, the park
Childrenill. laughter fills the air as they chase each other across the playground.
Chidden883 i laughter fills the air as they chase each ather across the playground. By noon, office wark seek refuge under shady trees, ergoying a peaceful lunch break.. As everving falls, the park
By noon, office workers seek refuge under shady trees, enjoying a peaceful lunch break. As evering falls, the park
World User
Sports User
Business User
Sci/Tech User
Figure 17: Visualization of token-level routing results for CoMiGS -111 trained on AG News. Tokens colored with the first expert choice at the Oth (first) layer. Orange denotes the generalist and blue denotes the specialist. Diagonal entries are in-distribution texts and off-diagonal entries are out-of-distribttion texts. Texts are generated by ChatGPT.
21

---

CoMiGS: On-Device Collaborative Language Modeling
In landmark auant, a hishric peace accord has been signed bewwee long- standing rivels in the Midsle Cast The agreement, broke red by international mediators, ams to end decades of confict and pave the way for economic cooperation and regaional stabiliiy:
In londmark avont) a historic poace accord hes been signed between long- standing rivels in the Middle Cast. The agreement, broker by international medietors, aims to end decades of conflict and pave the way for economic cocperation and regionel stability. Leadens from both sides extrested hage and commement peaceful future. emphasiang the accord's 's potensal to bring prospenity and security to the region. This unprecedentad move has
in landmark avant, a historic peace accord has been signed benween long- standing rivels in the Middlle Casl The agreement, broke red by international mediators, ams to end decades of conflict and pave the way for economic cooperation and resiional stabiliyy.
in landmark avant, a historic peacs accord has been signed between long- standing rivels in the Midelle Cast The sgreementt broke red by international mediators, aim to end decades of conflict and pawe the way for economic couperaliion and regaooall slability
World
Leaders from both sides expressed hope peaceful future, emphasizng the accord's 's potentaa to bring prosperity and security the regionn Thi unprecedented mave has
Leaders from both sides expressed hape and commitmnt emphasiong the accord is potertaal to bring progpenty and security to the region. This unprocedented move has
Leaders from both sides expressedd hape and co commitment to peaceful future, emphasiong the accord is potertial to bringg proqpenity and security the regioon. This unprocodented move has
and comm
peaceful future,
Young tonris sensation Emma Garcia has taken the sports world by soorm with her stunning victory at the Grand Slam tournament. At (ust 19, Garcia dafested top ranked plavers, shownasing exccpiional skill and composure on the court. Her aggressive play style and strategic ac ummn have carned her widescoed acdaaim from lans and analysss allice This hastoric win propels Garcia into the global spollight, mar ling hher se formd contender in future
Young lonris sensation Emma Garcia has taken the sports word by storm with her stunning victory at the Grand Slam tournament. At (ust 19, Garcia defeated top ranked plavers, showsasing exceaaiiona skill and compoaure on the court. Her aggrossive play play-tççl and stratogic ac umen have carned her widescreed aodeim from fans and analyss allin This Iistoric wan prop Geroa into the glosal spollight marking her se formidaaee contender in futurr
Young lennis sensation Emma Gardia has taken the sports wordd by storm with her stunning victory at the Grand Slam tournamant. At just 19, Garcia defeated top marked plavers, showsasing excepitioal skill and composure on the court. Her apgressive play styyltyle and atratogic ac ummn have carned her widespread acdaim from lans and analyysss allike This Instoric win propels Garoaa into the glocal spotlight marking her as formidacle contender in future
Young tennis sensation Emma Garcia has taken the sports world by storm with her stunning victory at the Grand Slam tournament. At just 18, Garcia dafeated top ianved plavers, showsasing excepiional skill and composure on the court. Her angressive stratogic ac uman have earned her widespreed acciaim from lans and analysts alike This hasoiic win propels Garcis into the global spotligtt, marking her a formidable contender in future
Sports
style and
Major tach companies are stpnlicantly
Major tech companies are ficently Major tech companies are sign/icantlly increesing their investments artilicaal intelligence (Al) startups; simiig to ippear head innovation and maintain compelitive edge. Giants like Google) Microsoft, and Aazzon are leaciing funding rounds, pouring billion into All driven enterprises focusing on machine learning, natural lanquage processing. and automation. IThess investments ard expected to accelerate advancements in iNl applications across various industrios)
Major tech companies are sigrifcantly
increesing their investments artificial inteligence AAll startups, siming to spear
increesing their investments artificial intellgence (All startups, siming to speer head inmovation and maintain compeliiive edge. Giants like Google, Microsofti and Amazon are leading funding rounds) pouring billons into All driven enterpeses focusing on machine lcamning natural language p processing and automation. These investments are expected to accolerate advancements in Al applications across vatious industiios,
incressing their inveetments artificial intelligence (Al) startups, aiming to ispear head innovation and maintain compeliiive edge. Giants like Google) Microsoft, and Amazon are leading funding rounds, pousing billions into All driven enterprises focusing on machine learning, natural lanquage processing. and automation. Ihhsse investments ard expected to accelenate advancements in All applications across warious industries.
head innovation and mainzain
Business
compelittvv edge. Giants like Google)
Microsoft, and Amazon are leecingg
funding rounds, pouring billions into Al- driven enterprses focusing on machin leaming, nntural language processing and automation. Thesee investments ard expected to accelerate advancements in Al applications across varous industriss.
the sun iises, caty park comes to with symphhoy of nin hirds log gers weaus
the sun risess ctt park: comes to with sym phony of lrping and rust ling through winding paths, their footstaps creating a mh best. Familles gother for picn iss spreading out blankets and unpacking baskets filld with treats.
As the sun rises, ctt park comes to with a sym phony of ch piinping hirds and rust ling leaves Jog gars wenve through winding paths their footstaps creating a ihyth mic best Families gather for pie ics, spresding out blankets and unpacking baskets filed with treats.
As the sun rises, city park comes to with sumphory phony of ping airds and rustting leavs Jog pers weave through winding paths their footstaps creating a my/h mic beat. Families gather for pie ics, spreading out blankats and unpacking baskets liledd with treats.
and rustingg
Joggars wanvi
ihrough winding paths their footstaps creeling a mi beat. Families gather for picnics. spreading out blankets and unpaaking baskets fillei with treats.
Scl/Tech
Chidren a laughter fils the air as they chase each other across the playground. By noon, office workers soek refuge under shady trees, enjoying a peaceful lunch broak. As evaning fillss, the park
Childrentils i laughter fils the air as they chase each other acroes the playground. By noon, office workers seek refuge under shady trees, enjoying peaceful lunch break. As evening falli, the park
Children laughter fills the air as they chase each other across the playground.
ChildrenES i laughter fills the air as they chase each ather across the playground. By noon, office workers soek refuge under shady trees, engoying a peaceful lunch break. its everiing falls, tne park
By noon, office workers seek refuge under shady trees, anjoying a peaceful lunch break. As evering falls, the park
World User
Sports User
Business User
Sci/Tech User
Figure 18: Visualization of token-level routing results for CoMi -111 trained on AG News. Tokens colored with the first expert choice at the 5th (middle) layer. Orange denotes the generalist and blue denotes the specialist. Diagonal entries are in-distribution texts and off-diagonal entries are out-of-distribution texts. Texts are generated by ChatGPT.
22

---

CoMiGS: On-Device Collaborative Language Modeling
In landmark ovant, a historic peace accord has been signed betwee long- standing rivels in the Midsle Cast The agreement, broke red by international mediators, ams to end decades of confict and pave the way for economic cooperation and regaional stability!
a londmark avant) a historic poace accord hes been signed betwoen long- standing rivels in the Middle Cast. The agreoment, broke red by international medietors, aim to end decades of confiict and pave the way for economic cooperation and regionel stability. Leadens from both sides extrested hage and commiment peaceful future. emphasiang the accord's 's potensal to bring prospenity and security to the region. This unprecedentad move has
In landmark ovant) a historic peace accord has been signed benween long- standing rivels in the Middlle Cast The agreement, brokerre by international medietors, ams to and decades of conflit and pave the way for economic cooperation and resiinall slabiliy.
in landmark evant) a historic peace accord has been signed between long- standing rivels in the Midelle Cast The sgreementt brokered by international medi aim to end decades of confict and pawe the way for economic coupereltion and regional slability
World
Leaders from both sides expressed hope peaceful futurr, emphasizng the accord's 'a potentaa to bring prosperity and security the region. This unprecedented mave has
Leadors from both sides expressed hope and comm tment
Leaders from both sides expressedd hope and commitmet to peaceful future. emphasiong the accord's is potertial to bring p prospenity and security the neg region. This unprocddettdd move has
and commitent
peaceful future,
emphasiong the accord is potential to bring progpeniy and security the region. This unprocedented move has
Young tonnis sensation Emma Garcia has takon the sports world by soomm with her stunning victory at the Grand Slam tournament. At (ust 19, Garcia dafested top rankeri plavers, shownasing exccpiional skill and composure on the court. Her angressive placatyla and strategic ac uumnn have esrned her widesooea acdaaim from lans and analyssa allice This historic win propels Gercia into the global spollight mer ling her as formd contender in future
Young lonris sensation Emma Garcia has taken the sports word by storm with her stunning victory at the Grand Slam tournament: At (ust 19, Garcia defeated top ranked plavers, showsasing exceaaiiona skill and compcaure on the court. Her aggressive play style and strategic ac umen have carned het widescreed aodsim from fans and analyes allin This historic wan prop Gardia into the glocal spollight marking her se formidaale contender in futurr
Young lennis sensation Emma Gardia has taken the sports wordd by storm with her stunning victory at the Grand Slam tournament. At just 19, Garcia defeated top marked plavers, showsasing excepitioal skill and compoaure on the court. Her angressive playstyle and atratogic ac uumn have carned her widespread acdeim from lans and antalysss allice This Instoric win propels Garoia into the glocal spotlight, marking her as formidall contender in future
Young lennis sensation Emma Garci has taken the sports world by storm with her stunning victory at the Grand Slam tournament. At just 18, Garcia daleated top rarked plavers, showcasing excapiioral al skill and composure on the court. Her angressive playstyle and stratogic ac umen have aarned hor widespreed accisim from lans and analysts aliko, This hakori win propels Garcia into the global spotightt marking her a formidable contender in future
Sports
Major tach companies are signlicantly increesing their investments artificial inteligence (All starlups, siming to spear
Major tech companies are ficently Major tech companies are sign/icantlly increesing their investments artilicial intelligence (Al) startups; siring to ispear head innovation and maintain compeliiive edge. Giants like Google) Microsoft, and mazon are leaciing funding rounds, pouring billion into All driven enterprises focusing on machine learning, natural lanquage processing. and automation. ihhese investments ard expected to accelerate advancements in iNl applicationn across various indusitios,
Major tech companies are sigrifcantly
increesing their investments artificial intellgence (All startups, siming to speer head inmovation and maintain compeliiive edge. Giants like Google, Microsoft, and Amazon are leading funding rounds, pouring billons into All driven enterpeses focusing on machine lcamning natural language p processing and automation. These investments are expected to accolerate advancements in Al applications across vatious industiios,
incressing their inveetments artificial intelligence (Al) startups, aiming to ispear head innovation and maintain compeliiive edge. Giants like Google) Microsoft, and Amazon are leading funding rounds, pousing billions into driven enterprises focusing on machine learning, natural lanquage processing. and automation. These investments ard expected to accelenate advanzements in Ill applications across warious industrios)
head innovation and mainzain
Business
compeliiive edge. Giants like Google, Microsoft, and Amazon are leecingg
funding rounds, pouring billoons into Al- driven enterpriees focusing on machine leaming, inntrrl language processing and automationn These investments are expected to accelerate advancements in Al applications across varous industries,
the sun rises, cafy park comes to with symphore of nig hirds gers weaus
the sun rises, ctt park: comes to with sym phony of ltping and rust ling through winding paths, their footsteps creating a ihyth me best. Familles gother for pien cs. spreading out blankets and unpacking baskets filld with treats.
As the sun rises, ctt park comes to with a sym phony of ch phipping birds and rust lintingg through winding paths, their footstaps creating a thyth mic best Families gather for pie ics spresding out blankats and unpacking baskets filed with treats.
As the sun rises, city park comes to with symphore phony of ping airds and nustting leavss Jog pers weave through winding paths their footstaps creating a my/h mic beat. Families gather for pie ics, spreading out blankats and unp unkiining baskets liledd with treats.
and rustingg
sersers
Joggars wanvi
Joggers wanve
ihrough winding paths their footsteps creeting a mhy beat. Families gather for picnics spreading out blankets and unpacking baskets fillei with treets.
Scl/Tcch
Chiddennlls i laughter fils the air as they chase each other across the playgroundi By noon, office workers sook refuge under shady trees, enjoying a peaceful lunch broak. As evaning fillss, the park
Childrenlll is laughter fils the ar as they chase each other acroes the playground. By noon, office workers seek refuge under shady trees, enjoying a peaceful lunch break. As evening falls, the park
Childrenill: laughter fills the air as they chase each other across the playground. By noon, office workers seek refuge under shady trees, enjoying a peaceful lunch break. As evering fills, the park
ChildrenESS i laughter fills the air as they chase each ather across the playground. By noon, office workers soek refuge under shadk trees, engoying a peacefull lunch break. Ats everiing falls, the park
World User
Sports User
Business User
Sci/Tech User
Figure 19: Visualization of token-level routing results for CoMiGS G-111 trained on AG News. Tokens colored with the first expert choice at the 1 1tth (last) layer. Orange denotes the generalist and blue denotes the specialist. Diagonal entries are in-distribution texts and off-diagonal entries are out-of-distribution texts. Texts are generated by ChatGPT.
23

---

CoMIGS: On-Device Collaborative Language Modeling
F. Alternating Minimization Convergence
E.. Notation
Let us recall our notation from Sections 3.3 and 3.4. We have two differentiable functions fi Φ) ( )) and f2(@,, ②) = frain (e, 4) that constitute our problem. For the sake of generality, let us assume that the target variables O and d belong to their corresponding feasible convex sets Q and f2,
Q s R
Ω
R
Then, we consider the following alternating minimization process, starting from some initial (Oo. Po), for every 2 0:
4 s+1
arg min fi (+), 2),
(12)
arg min fz(O, @x+1).
=
If fi = f2 that would be standard alternation minimization as for minimizing one function f1. However, in our setting fi and f2 can different.
For fixed @ and 4, let us denote the corresponding arg min operators by
:= arg min (ie, ②)
$C0
and
(2( = arg min f2(O, Φ).
ecQ
Using this notation, rewrite algorithm (12) as follows:
4 Po+1 = U1 (O)),
+ k41 = Uz( (
k 0.
(13)
We further define the following operators:
T( (2(u1(0))
E
Q,
θ
Q,
P(Φ) un (uz )))
?
E
Ω
With this notation, we can rewrite the sequence {Ox h220 simply as
O+11 = T(Qx), 2 O.
(14)
We use the following main assumption on functions fi and f2: Assumption 1. There exist O* E Q and x* E Q such that
@* T(@*)
and
P(Ø*)
(15)
Remark 1. Note that if f1 fu = f. condition (15) holds for the global minimizer of our function (O* arg min f(O,4) $).
eEQ,
Clearly, this assumption should hold if functions fi and f2 are sufficiently close: fi RJ f2, or in case of overparametrized models. remains an interesting open question: what are the general and joint joint conditions on fi and f2 that imply (15).
F22 Contraction and Convergence
As we will see, is natural to assume that operators 2 and are contractions. We will provide working example of our setting in the next section, where this condition will hold. We assume to have some norms fixed on Q and Ω. that are not necessarily Euclidean. For simplicity, and when it clear from the context, we will use the same symbol for both norms, even though they can be different for spaces of and T
24

---

CoMIGS: On-Device Collaborative Language Modeling
Assumption 2. Let ui and ug be Lipschitz with some constants A1, A2
0:
(+) u1 [+)]
A1 O1,
AO, θ E Q,
(16)
u2 () U2( +
Aa +
+ f2.
Under these assumptions we can show the convergence of the sequence { k20 generated by (14). Indeed, for every 2 0, we have
(15)
|0k+1 - - [TT(@&) 1I
[T((&) T(
[ (O )|
X21|u1 (Ox) un un (+*)
S O) -
and we that , @+ with the linear rate. The same reasoning be applied the sequence { @2t221- Thus, have established the following general convergence result.
Theorem F.1 (Theorem 3.1). Let Assumptions 2 hold and x1 X2 Then, the sequence (O 2x)420 generated by
alternating process (/2) converges to (Q', Φh) linearly, for every k 2 0:
[ ; - 1-e S (A1
1100 -0'|,
(17)
Φ T11l S (A112)*11 T1
Example 1. Consider the following quadratic objective
f(0.0) ②)
(A0, + (BO 3) + (C
where A A E x and B
BT
| x are symmetric matrices, and C x We assume that
strictly convex, which means
A C C
H
0.
B
Clearly, for this objective, we have θ 0 and 4
0. Then
U1 (9) arg min f(O, ②)
=BB-c and
U2 (b) = arg min /(O, -A-ICT +.
Hence, the composition operator == U2 o un linear:
T(O) A-ICTB-CCO
(18)
and it holds
[T(@)
1- <
I -
Now, denoting by p > O and 2 μ the smallest and the largest eigenvalues of matrix correspondingly, and using the Schur complement, we conclude that
LI,
and μI = A 3 L..
(19)
from which we are able to hound the norm of our matrix as follows:
(19) L-и < 1. < 1
- B 'cl
|A-1/2 B
B-CCB--CAA-77
which proves the contraction property.
25

---

CoMiGS: On-Device Collaborative Language Modeling
Example 2. Note that for a general difserentiable function using the Taylor expansion, the operator T u2 w1 where W1 arg min f(O,, ②) and uz () arg min f(O, $), can be expressed as follows (compare with (18)):
T( - 0r - - H2 22
where
H1
(
81 (0' 06
- - (0) Φ'))dT,
H12
81 ( -T((()) 0- ( Φ T'))dT,
8084
Hz2
0
O
0- ( 4
+
D2')dTT
H21
8 3ede
(wu (0) bb))dT.
Therefore, assuming that the Hessian strictly positive definite and Lipschitz continuous in neighborhood of the solution, localizing the current point the neighborhood, @ x @* and t x $*, previous example (see, e.g. Theorem 1.2.5 in (Nesterov, 2018) for the local analysis of Newton's 's method).
obtain the contraction property. as in the
E3. Linear Modeling and Decoupling
In this section, let us study important example of linear models, applicable to both experts and the router. As will show, in this case and under very mild assumptions we can justify all conditions from the previous section and therefore obtain the global linear convergence for our alternating process.
Problem Formulation For simplicity, we consider the case of one client and assume that training and validation datasets are the same, Xtrain = x However, our observations can be generalized to a more general case of several clients, and different but statistically similar datasets train X valid, Hence, we have, fi = f2 = f. Note that in this case, our bi-level formulation also equivalent to joint minimization of f w.r.t. all variables.
We assume that client has generalist expert model, that denote by 0 Rd. and N specialist experts, that we denote by 01.... ON E RRd. We compose these models together as matrix θ (O",..., 0%). In principle, different models can have different expressivity, which we take into account in our modeling by convex set of constraints: θ E Q C (N+11)
We denote by Φ0 ..... NN E Rd the parameters of our Router, composed together as matrix + = ( Φ which can also be constrained by convex set: Ω C Rdx(N+1), For given data input x E Rdd, the Router decides which experts use with the SoftMax operation + Tx(x)) AN, where
AN = { VVERNY
N Σ ==
Σg =1
is the standard Simplex, and
(6) E expl(o". .))
T
(x)
(20)
Under these assumptions, we set the following structure of our optimization objective,
f(O, +)
1 12
(0):)-(0),20) Σ 1 (x:) (0) xz)) + (10 + +1
(21)
where EI,..., Fn are given data vectors, and 4(),15i5 n are the corresponding convex losses (e.g. the logistic loss for binary classification, or the quadratic loss for regression problem). We use a 2 0 as a regularization parameter, which can also be seen as the weight decay, and || the Frobenius norm of matrix.
26

---

CoMIGS: On-Device Collaborative Language Modeling
Decoupling Let us introduce the auxiliary variables, X E Av. Sis n, and A (At,,.. X²) E A C R
which is column-stochastic matrix. Employing the matrix notation, we rewrite our problem in the following form:
1 2 E e (1x; + (1o12 + 1111))) : X = (x;), [sisn]
1
min BEQ AEAR
(22)
i=1
Now, we apply the relaxation of constrained problem (22) by the following decouple of X from oe (2xi), with some parameter μ 2 0 and a distance function V AN AN R+ between distributions:
(, = 1 Σ l (Tx))) + i (1011- + 1 1=1
+ # Σ :T
min @EQ.EEN A
2m
(23)
A natural choice for V is the Dullback-Leibler divergence, which gives, for every 1 S i m
N =
Σ [1]] ( ( In 11 6) Σ [ ( In T (2:)) (3)
(20)
Σ {x [mn[>](6) - (0)) xzi)
-(6),
In Σ
exp((ot
((0) x
y=0
d(x') (2', Txx) + ((TTx))))
where
d(x) Σ In AEAN,
is the negative entropy, and
s(y) = In Σ exp e
Y V
is the log-sum-exi function. Note that both d(-)) and s(.) are convex functions on their domains. Moreover, it is well known that d(-) is strongly convex w.r.t. -n-ormm (see, e.g., Example 2.1.2 in (Nesterov, 2018)):
(Vd((A),,) h) Z | A
(24)
Therefore, obtain the following decoupled optimization formulation:
min BEQ AEAR
F,
(25)
l (A (Tx) y ( (N, TT:) -x,
- (10 13 +
+ 13:
It is clear that setting parameter μ := +o0, we obtain that (25) is equivalent to our original problem (22). However, for < +00 we obtain more flexible formulation with auxiliary distributions X' E AN. each for every data sample 1 < i < n, that makes it easier to treat the problem. Parameters (2t) has an interpretation of latent variables, which makes our approach similar to the classical EM-algorithm (Jordan Jacobs, 1994). We note a similar work from Almansoori et al. (2024), which proposes to train mixture of generalists on local routers, which resembles simplified version of our method.
It is clear that function (O, Φ A) is partially convex: it is convex w.r.t ( b when A is fixed, and it is also convex w.r.t. A when (O, 4b) is fixed.
In what follows, we show that under very mild conditions and choosing regularization parameter a, μ 2 0 sufficiently large, we can ensure that (-) is jointly strongly convex, regardless of non-convex cross terms: 4, (( e 2;)) and (X+, & x;).. Our theory generalizes a recent approach to soft clustering (Nesterov, 2020). With this technique, we will be able to show the global linear convergence rate for the alternating minimization approach that we discussed in the previous sections.
27

---

CoMIGS: On-Device Collaborative Language Modeling
Joint Strong Convexity Let consider the i-th of our objective (25) that correspond the data sample with index 1 S i n. Omitting extra indices, obtain the following function,
F = -(AA)TT))--AA)00)))
Φ -12,
(101 + |1813)) + μd(A) + ( x)),
(26)
where E Q. E E,, A Our goal is to ensure that (26) is strongly convex w.r.t to the standard Euclidean norm of the joint variable. Namely, we establish the following result.
Proposition 1. Let the loss function {(.) be convex and assume that its first derivative is bounded: 2 max, (t). Assume that the regularization coefficient is sufficiently large:
2 2||z0|/ max $
(27)
Then the objective in (26) is strongly convex.
Proof. Note that our objective (26) can be separated in d and O We construct the following functions I1 and 92-
(I(0))) ((11.T + O + + (d(A) and () == -u(A, +
+ + =d(A)
and F(O, = (0, + (b,) + - + (d(A) + uus ( TT.)............... Since the log-sum-exp function æt).
(d())
x.)) is strictly convex, and the negative entropy d(.), as well as the Frobenius norm, are strongly convex, it suffices to prove convexity for 91 and 92- Computing the second derivative of 91 and applying it to an arbitrary direction z = [H: h] of corresponding shapes, we get
(V 1 (0,
- H ((V2d(A),,)))
+ TT)) [T) H x)² (h, TT)) (A, H x)
+ ((), T x)) (h, H x))
e | 1 p| || [F |
|bll H F (VT-
p| 0,
where used Young's inequality in (+). The bound for ø2 follows by the same reasoning, substituting ((t) = jut and therefore setting p
For the decoupled optimization formulation (25) it is natural to organize iterations in the following sequential order, starting from an arbitrary O E Q and &o e Ω, for some > O:
As+1
arg min (Ox, &s, AEAS
Px+1
arg Py (OR, Φ AR+1), $
(28)
θx+1 =
arg min (O, P++1. A8+1).
eeQ
Note that each minimization subproblem in (28) is convex and can be implemented very efficiently by means of linear algebra and convex optimization. At the same time, due to decoupling of variables and strong convexity we are able to ensure the global convergence of this process to the solution of (25).
28

---

CoMiGS: On-Device Collaborative Language Modeling
F.4. Convergence for Functional Residual
Note that in our decoupled optimization formulation (25), variables θ and o are independent of each other, when A is fixed. Therefore, the second and third step in iteration process (28) can be done independently.
For the sake of notation, let denote X = A, concatenated variable Y = (θ, Φ) and the objective in new variables as f(x, ) = (O, D, A). By our previous analysis, we can assume that f is strongly convex. We denote by the parameter of strong convexity and by the constant of Lipschitz continuity of the gradient of f. Its global minimum is ((X*,**)
denoted by (x*, Y+ and correspondingly f = ==
Then, iteration process (28) can be rewritten simply as the following alternating iterations, for k 2 0:
Y111 =
arg f(x,, (), YEY
X
arg migiin/ fXX, Ys+1), Xex
where A' and y are the corresponding convex domains (A' = Ag and y = Ω x Q)- Then, the stationary condition for YR+1 (see, e.g., Theorem 3.1.23 in (Nesterov, 2018) gives
0f 8r (X,, Yx+1), Y 2 0, AY E E.
(29)
Choosing
7
4 s 1,
(30)
we obtain
xf(x", ((Xx,))=0 + + (x-) Yx+1)
( (1x, + (of (Xx, X* Xx) +x)+ + (Xt Mx+1), Y
Yk+1) + #1 Xxl13
- (11-))(Xx Y--1)
(29),
f(x,, Yo+1) + (ox (X,, ( -x,) + 11 - xx)11
W IV
f(X,, Yx+1) + 1ox (Xs, 1x+1), x X + Xx113] XEX
min
min Y(+1)) f(XE41, Yet1),
where in (*) we used strong convexity, and in (**) we used the Lipschitz continuity of the gradient. Thus, we get the following inequality:
f(X++1, Mx+1) f s (1-x) -x41) -
and using the same reasoning for Y. + Y441 update, we obtain
f(X++1, YR+1) f s (1 (1-7) (F(x))) (²).
which is the global linear rate. Thus, have established formally the following convergence result.
Theorem F.2. Let be strongly convex with constant μ O, and let its gradient be Lipschitz continuous with constant L> O. Then, for k 2 0 iteration of the alternating minimization process, we have
(x
< (1 1 +) 2k (1) Yo)
Note that this result is directly applicable for our linear models from previous sections, f(X,Y) = (,, as we show that objective (25) is jointly strongly convex, when the regularization parameter sufficiently large.
29


Reviews:
[{"review_id": "uHXK2iiqgD", "reviewer": "Reviewer_RnkV", "summary": "This paper introduces CoMiGS, a modular federated learning framework for adapting LLMs using a mixture of generalist and specialist LoRA experts. CoMiGS employs a bi-level optimization strategy, alternating between routing and expert parameter updates. Experimental results on GPT-125M and Llama-3.2-1B demonstrate its superior performance over both local and federated baselines, while also offering theoretical and empirical insights into the framework’s behavior.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**Methods And Evaluation Criteria**\n\n_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.\n\n**Experimental Designs Or Analyses**\n\n_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.\n\n_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.\n\n_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. \n\n**Weakness**\n\n_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.\nOne 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.\n\n\n_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.\n\n_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.\n\n\n**Questions**\n\n_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. \nSorry for not making it clear, but our routers are localized per user, instead of being federated, as illustrated in Figure 2.\n\n_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. \n\n_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.", "claims_and_evidence": "The paper makes a compelling case for the separation of global and user-specific parameters in a modular architecture that is federatedly trained. This enables specializing experts towards specific sub-distributions and efficient learning under non-IID data and heterogeneous target devices.\n\nHowever, while the authors claim enabling on-device training, which is a typical scenario of cross-device federated settings, they do not evaluate in either case. The evaluation measures cost in a device-independent way, while the number of participating clients and participation scheme does not resemble cross-device settings.", "methods_and_evaluation_criteria": "The proposed methods seem well motivated and the evaluation is reasonable. If anything, I would propose some additional areas of exploration to showcase the generality, scalability and applicability of the method. Concretely:\n\n* Although perplexity gives a feel on the quality of the output, the number is not definitive on the downstream quality of the model. Towards this setting, it might be beneficial to also have LLM-as-judge reports on the models, or incorporate tasks like QA to understand the Question-Answering capabilities of the models.\n* How would CoMiGS work with other PEFT methods or adapters (e.g. DORA, VERA)?\n* How would CoMiGS work under different aggregation methods?\n* Another interesting avenue for exploration, especially for on-device deployment, would be the interplay of the technique with quantization methods (or other compression schemes), where the router and adapters may operate on a lossy pretrained model.\n* Since the paper inherits a federated setup, an interesting question arise wrt the tradeoff of utility and privacy when training the generalists under DP.", "theoretical_claims": "I went over the theoretical alternative minimization convergence proofs in the appendix at a high-level. Looks reasonable.", "experimental_designs_or_analyses": "For the experimental evaluation, I have the following comments:\n\n* It is unclear from the paper how the authors have federated the datasets and whether the distribution is non-IID amongst clients.\n* Are the expert scores in Figure 4 the results of Eq.1?\n* In §4.3, where the authors evaluate the adaptation to system heterogeneity in clients, an alternative could be not to have experts (or the same number of experts) across layers, especially given the dynamics of expert selection over the depth of the network.", "relation_to_broader_scientific_literature": "CoMiGS builds on top of the fields of federated learning, LLM adaptation and modular networks, by providing a framework that enables the personalization of on-device experts (specialists), while sharing knowledge across users via federated experts (generalists). It does so by taking into consideration both data and system heterogeneity, typical features of federated setups, thus adapting to non-IID multi-device federated environments.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "### Strengths\n\n* The approach of having two sets of experts that are federated trained to specialize and route between global vs. local objectives is well motivated and work competitively to the baselines.\n* The mixture of generalists and specialists allows the model to simultaneously capture individual user preferences and linguistic styles while still allowing exchange of globally shared knowledge. I like the interplay between generalists and specialists, especially wrt regularization.\n* The evaluation explores various dimensions of the learning dynamics and modular behavior during inference, which provides interpretability in the operating dynamics of the model. I particularly like the fact that the authors have evaluated both in and out of distribution datasets.\n\n### Weaknesses\n\n* The federated paradigm put forward does not seem to be focusing on cross-device setup, but rather to assume small client sizes and full participation. Furthermore, the models have not been deployed \"on-device\".\n* The dependence on validation sets may cause lack of robustness in the technique under distribution shifts.\n* There seems to be much sensitivity of the results to the number of experts and data size. Although evaluated, there seems to be no proposed way to pick them, give a specific learning scenario.", "other_comments_or_suggestions": "* The figures in the evaluation are barely legible. The authors should probably use a larger font size to increase legibility.\n* In Table 2, the numbers in the parentheses should be explained in the caption.", "questions_for_authors": "* I am not sure I have fully understood the reasoning behind the size of the trainable parameters (i.e. in the router) and the ability to update less frequently. Also, is the only way to achieve this by skipping updates in the outer loop, or can be effectively achieved by adjusting the learning rate of the outer optimization?\n* Have the authors tested against also specializing the router function per user instead of federating it?\n* Does Figure 3 suggest that we might not need the same number of specialists across layers? If so, is there the potential for further optimization?\n* Following up from the results of Figure 5, what is a good way for practitioners to select how many specialists to dedicate to their learning model?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m0G6oLAt4y", "reviewer": "Reviewer_iYSM", "summary": "The authors focus on the problem of on-device collaborative fine-tuning of LLMs to address both computational resource heterogeneity and data heterogeneity among users. The authors try to develop a framework that can balance general and personalized knowledge for each token generation while being robust against overfitting. The experimental results show the proposed method improves the inference performance and computational efficiency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nRegarding your question about **router training**, we make the following clarifications:\nWithin 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. \n\nOur 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.", "claims_and_evidence": "Yes. The claims are easy to follow and the evidence is clearly supported.", "methods_and_evaluation_criteria": "Yes. The authors use typical LLM models and tasks.", "theoretical_claims": "Yes. The formulation of distributed model training is clear.", "experimental_designs_or_analyses": "Yes. The experiments are correctly configured and the insights obtained from the experiments are clearly explained.", "relation_to_broader_scientific_literature": "This paper is strongly related to the on-device LLM model design and deployment.", "essential_references_not_discussed": "No. I think the references are adequately covered.", "other_strengths_and_weaknesses": "The authors propose separating the experts into generalists and specialists, where generalists are shared across users to provide general knowledge and specialists are localized to provide personalized knowledge. This is a practical idea to guarantee the model robustness.", "other_comments_or_suggestions": "Overall, this paper is interesting and the technical depth is fine in most aspects.", "questions_for_authors": "The authors propose to use a learnable router to determine the aggregation weights of the experts based on the input tokens. Could you please give more details on how to train this router, especially when the data is unbalanced among the users?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A."}}, {"review_id": "CWTBrdCJrI", "reviewer": "Reviewer_TG8c", "summary": "The paper introduces a method for training language models collaboratively between multiple devices/clients and personalizing them to their on-device data at the same time. This is done by introducing a mixture of generalist experts and specialist experts, where the generalists are trained collaboratively (e.g., through federated averaging) and the specialists are trained locally (without aggregation). The routers are trained locally as well. Thus, only the generalist experts are aggregated and learned collaboratively across clients, which reduces communication overhead. The problem is formulated as a bi-level optimization problem with a training algorithm that optimizes the experts on a training dataset and then optimizes the routers on a validation dataset, done alternatively. This method addresses both data heterogeneity and computational resources heterogeneity, which is purportedly the first to accomplish that , and it also separates model heterogeneity from data quantity as shown in the extensive experiments section. The authors also provide convergence analysis under some standard assumptions.\n\n## update after rebuttal\n\nThe authors have addressed most of my concerns and I'm satisfied with the rebuttal. I will keep my original positive rating.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**Novelty of the algorithm**\n\nWe 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.\n\n**Violation of Assumption 1**\n\nWe 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.\n\nFor 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.\n\n\n**Proof of strong convexity**\n\nThank 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.\n\nOur goal is to show that the following function (see eq. (26) on page 27),\n$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 )$\nis strongly convex, for sufficiently large $\\alpha > 0$ and $\\mu > 0$.\n\nWe 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)]$.\n\nSince 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. \n\nFor 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\n\n\n**Minor typos**\n\nThanks for pointing out the typos and dimensions of the simplex, we will make sure these are fixed.", "claims_and_evidence": "The authors claim that the bi-level formulation of the MoE learning objective and the alternating minimization algorithm are new. The formulation itself might be, but an alternating minimization algorithm of the router and the experts is quite general, so it is difficult to claim that it is \"new\", especially since it resembles expectation maximization (which the authors themselves mention). For example, [this work](https://arxiv.org/abs/2410.03497) trains mixture of \"generalists\" on local (input-independent) routers, which resembles a simplified version of the proposed method.\n\nThe theoretical analysis is interesting and shows linear convergence. The assumptions are claimed to be \"suitable\". I agree that the claims are standard. However, I would argue that the existence of a minimizer might not be necessarily be a suitable assumption in language modeling. In general, the loss doesn't go to 0 on these problems.\n\nThe claim regarding not overfitting even with many experts and little data might only be applicable because this is a fine-tuning problem, but it is still an interesting result.\n\nThe other claims are true and are well corroborated in the experiments, though I'm not sure whether this work is, indeed, the first that addresses both data heterogeneity and computational resources heterogeneity.", "methods_and_evaluation_criteria": "The proposed method is very sound and the authors provide theoretical and experimental evidence that it works.\nThe  authors provide extensive experimentation and ablation studies to evaluate the method, all of which show positive evidence that the proposed method shows good performance, robustness, and efficiency with respect to the baselines.", "theoretical_claims": "I checked the correctness of the proofs in the Appendix. The analysis is strightforward and clear. I did not find any significant errors.", "experimental_designs_or_analyses": "The experimental design does capture interesting practical scenarios. For example, the out-of-distribution experiments uses a validation and test set that is from a different distribution from the training dataset, so that assumption 1 is violated. However, the method still performs favorably, which shows its robustness. Other experiments also provide ample evidence for the claims in the contributions list. The authors run experiments on various datasets and compare to many recent baselines. The analysis is also valid for the considered setup.", "relation_to_broader_scientific_literature": "The contributions of this paper are important to the literature. They tackle multiple issues of interest. For example, the data heterogeneity and systems heterogeneity in federated learning. the proposed method is also relevant to practical applications of LLMs in the real world.", "essential_references_not_discussed": "I do not know of any essential references related to this work that were not discussed.", "other_strengths_and_weaknesses": "This work provides a practical and theoretically well-grounded method for an important application: collaborative training of personalized language models. The solution is sound and the experiments show good performance.\n\nA great strength of this paper is that they share well-written code for reproducing the experimental results.\n\nOne weakness might be that this method mainly works for fine-tuning. It would be interesting to see whether the same benefits still hold for full training. Another could be that this procedure is only applicable to mixture of experts that routes tokens, i.e., transformer-based models and tokenized data.", "other_comments_or_suggestions": "In line 1435, the authors mention that the goal is to show that (26) is strongly convex, but the next line closes the section by saying that direct computation of the Hessian shows the joint strong convexity. Perhaps the author can write the details to make it clearer to the reader how it holds for sufficiently large $\\alpha$ and $\\mu$.\n\nMinor comment: the simplex defined in line 1361 should actually be $\\Delta^N$ because there are $N$ degrees of freedom/dimensions for $N+1$ values (the last value is completely determined by the rest).\n\nTypos: There is a little 'z' hanging on the right side of Table 2. Also, in line 1406, there is a typo: \"srongly\".", "questions_for_authors": "When you have many generalists, and one device can only use a few, how do you choose the generalists which will be assigned to those low-resource devices?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper studies a new Federated Learning setup to address both computational resource heterogeneity and data heterogeneity among users. The proposed approach uses a mixture of generalist and specialist LoRA experts. The method is justified both theoretically and empirically. All the concerns from the reviewers were addressed in the rebuttal.
Decision: Accept (poster)