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

## Abstract

Abstract On-device LLMs have gained increasing attention for their ability to enhance privacy and provide a personalized user experience.
To facilitate learning with private and scarce local data, federated learning has become a standard approach, though it introduces challenges related to system and data heterogeneity among end users. As a solution, we propose a novel Co llaborative learning approach with a Mi xture of G eneralists and S pecialists (CoMiGS), being the first to effectively address both. Our approach distinguishes generalists and specialists by aggregating certain experts across end users while keeping others localized to specialize in user-specific datasets. A key innovation of our method is the bi-level optimization formulation of the Mixture-of-Experts learning objective, where the router is updated using a separate validation set that represents the target distribution. CoMiGS effectively balances collaboration and personalization, as demonstrated by its superior performance in scenarios with high data heterogeneity across multiple datasets. By design, our approach accommodates users’ varying computational resources through different numbers of specialists. By decoupling resource abundance from data quantity, CoMiGS remains robust against overfitting—due to the generalists’ regularizing effect—while adapting to local data through specialist expertise.

## 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, ), Claude (Anthropic, ), Gemini (DeepMind, ) and etc. However, cloud-based inference introduces significant delays for end users, and it often fails to meet their personalized needs (Ding et al., ; Iyengar & Adusumilli, ). Recently, there has been growing interest in deploying LLMs on edge devices, which offer benefits like lower latency, data localization, and more personalized user experiences (Xu et al., ). For instance, Apple () recently launched on-device foundation models as part of its personal intelligence system.

On-device LLMs present challenges such as limited and variable computational resources, scarce and heterogeneous local data, and privacy concerns related to data sharing (Peng et al., ; Wagner et al., ). Fine-tuning is typically performed on-device to quickly adapt to users’ individual needs. While data sharing is a common solution to address local data scarcity, on-device data is often privacy-sensitive and must remain on the device. To overcome this, federated learning has been proposed as a method for enabling collaborative learning while preserving user privacy, allowing end users to collaborate by sharing model parameters (Chen et al., ; Zhang et al., ).

Federated fine-tuning of LLMs is predominately done through Low-Rank Adaptation (LoRA, Hu et al. ()) 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 aggregating LoRA modules of different ranks (Cho et al., ; Bai et al., ). 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.

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 a result, personalized solutions are necessary. Wagner et al. () explored three personalized collaborator selection protocols, allowing each end user to choose their collaborators. Although these protocols effectively address data heterogeneity, they depend on model aggregation, which can only occur when users share the same model architecture.

Figure: Figure 1: Diagram of our proposed method CoMiGS illustrated with a 4-client setup, with a dark blue frame denoting a block. Within each block, generalist experts ($\{\theta^{G}_{1},{\mathbf{\theta}}^{G}_{2},{\mathbf{\theta}}^{G}_{3},{\mathbf{ \theta}}^{G}_{4}\}$) are aggregated across users, and specialist experts ($\{{\mathbf{\theta}}^{Si}_{1},{\mathbf{\theta}}^{Si}_{2},{\mathbf{\theta}}^{Si} _{3},{\mathbf{\theta}}^{Si}_{4}\}$) are kept local. Router and expert parameters are updated iteratively using local validation and training datasets.
Refer to caption: x1.png

There has not yet been a solution to deal with both system heterogeneity and data heterogeneity. Towards this goal, we propose a novel Collaborative learning approach via a Mixture of Generalists and Specialists (CoMiGS). Our approach allows users to share part of the knowledge while keeping some knowledge user-specific, thus providing personalized solutions. We name the shared part *generalists* and 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 1.

We integrate the expertise of generalists and specialists using a learned router that determines aggregation weights, following the Mixture-of-Experts (MoE) architecture (Fedus et al., ). As in typical MoE designs for language modeling (Jiang et al., ; Fan et al., ), we also use tokens as the routing unit. Although users may have different topic preferences or linguistic styles, they still share common tokens in their vocabularies. Our goal is to route these shared tokens to the generalists so they can be jointly learned across users.

The closest work to ours is pFedMoE from Yi et al. () from the vision domain, where each user has a shared homogeneous small feature extractor, a localized heterogeneous feature extractor, and a localized routing network, with routing unit being a semantic unit – an image. The three components are simultaneously updated. Compared to pFedMoE, our method CoMiGS introduces three key updates: 1) we reformulate the learning objective into a bi-level
optimization framework, following the inherent hierarchical order of router and expert learning; 2) we refine the routing mechanism by using the smallest unit like a pixel in an image, which is a token; 3) unlike pFedMoE’s fixed two-expert limit per user, we support varying numbers of expert modules across users.

In summary, our contributions are as follows:

- •
We propose a novel approach (CoMiGS) for on-device personalized collaborative fine-tuning of LLMs, introducing an innovative bi-level formulation of the Mixture-of-Experts learning objective. Our approach can effectively tackle distribution shifts in local data.
- •
Our collaborative framework effectively addresses both *system heterogeneity*, with respect to varying local model architectures, and *data heterogeneity*, concerning diverse local data distributions across users, making it the first model to accomplish both.
- •
Our framework separates resource heterogeneity from data quantity. Users with larger local datasets benefit from a bigger model, while users with more powerful models but smaller datasets are less prone to overfitting.
- •
We release a codebase(^1^11Our code base is available at [https://github.com/epfml/CoMiGS](https://github.com/epfml/CoMiGS).) for collaborative LLMs that allows users to easily define their own collaboration strategies, facilitating and advancing future research efforts in this field.

## 2 Related Work

##### Collaborative Learning for LLMs.

Recently, researchers have been investigating the application of Federated 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 utilizing parameter-efficient techniques such as adapters.
Mohtashami et al. () explored a teacher-student social learning framework to aggregate private-sensitive instructions. Zhang et al. () directly applied FedAvg (McMahan et al., ) to aggregate LoRA parameters during instruction tuning, and reported increased performance in downstream tasks. Following that, there are various works focusing on addressing system heterogeneity where users are equipped with different LoRA ranks. HetLoRA (Cho et al., ) and FlexLoRA (Bai et al., ) provide different ways of aggregating and distributing LoRA modules of heterogenous ranks. However, these approaches are not designed to cope with heterogeneous data on device.
In contrast, Sun et al. () found better performances with respect to heterogeneous data can be achieved through freezing LoRA A matrices at initialization; Wagner et al. () proposed personalized solutions that can sufficiently tackle data heterogeneity, through three different collaborator selection mechanisms. Yet for both works, the users must be equipped with the same model architecture. Unlike previous works, our framework deals with both model heterogeneity and data heterogeneity. Moreover, our method offers personalized solutions at a token level, as opposed to the client-level approach in Wagner et al. ().

##### Mixture of Generalist and Specialist Experts.

Gaspar & Seddon () introduced a fusion of global and local experts for activity prediction based on molecular structures. Each local expert is tailored to a specific chemical series of interest using loss masking, while a global expert is trained across all series. Simultaneously, a routing network learns to assign soft merging scores. This approach yielded superior empirical results compared to single experts.
Dai et al. () developed DeepSeekMoE by deterministically assigning every token to \qqshared experts, whereas \qqrouted 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, Yi et al. () proposed pFedMoE, where each user has a shared homogeneous small feature extractor, a localized heterogeneous feature extractor, and a localized routing network, with routing unit being an image. The three components are jointly updated in an end-to-end fashion, demonstrating strong performance in the vision domain. Our work builds on the foundations of pFedMoE and extends it to the language domain. Furthermore, we introduce key innovations that enable more effective handling of distribution shifts and achieve a more refined balance between personalization and collaboration.

## 3 Method

We aim to improve personalized performance for each user on their target distributions, where distribution shifts can be allowed. Building on the hierarchical insights of MoE learning, we formulate our learning objective into a 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 validation sets. We further let experts diversify into generalists and specialists via parameter aggregation or localization, to leverage both collective power and specialized knowledge. As the problem solver, we provide a multi-round gradient-based algorithm.

### 3.1 Notions and problem setup

Each user has a training set ${\bm{X}}_{i}^{\text{train}}$, a small validation set ${\bm{X}}_{i}^{\text{valid}}$ and a test set ${\bm{X}}_{i}^{\text{test}}$, and the task is next token prediction. The validation set ${\bm{X}}_{i}^{\text{valid}}$ and the test set ${\bm{X}}_{i}^{\text{test}}$ are sampled from the same distribution ${\mathcal{P}}_{i}^{\text{target}}$ (note this is a fuzzy concept in the language domain, by the same distribution we mean from the same topic/category). The training set, ${\bm{X}}_{i}^{\text{train}}$, can be sampled from a different distribution than ${\mathcal{P}}_{i}^{\text{target}}$. 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 1, there are two sets of model parameters within each user: expert parameters, denoted as ${\mathbf{\Theta}}={\mathbf{\theta}}^{G}\cup\{{\mathbf{\theta}}_{i}^{S}\}$, where ${\mathbf{\theta}}^{G}$ is shared across the users and $\{{\mathbf{\theta}}_{i}^{S}\}$ are user-specific specialist parameters; and routing parameters, denoted as ${\mathbf{\Phi}}=\{{\mathbf{\phi}}_{i}\}$. $i\in\{1,2,..,N\}$ is the user index. Our ultimate goal is to optimize the average target performance across all users.

Our experts are simply LoRA modules, which approximate model updates $\Delta{\bm{W}}\in\mathbb{R}^{m\times n}$ with a multiplication of two low-rank matrices ${\bm{A}}\in\mathbb{R}^{m\times r}$ and ${\bm{B}}\in\mathbb{R}^{r\times n}$ with rank $r\ll m,n$. ${\mathbf{\theta}}^{G}$ and ${\mathbf{\theta}}^{S}$ are disjoint sets of LoRA A and B matrices.

### 3.2 A bi-level formulation

Essentially, we adopt an MoE architecture, apart from aggregating certain expert parameters. Instead of learning routing and expert parameters simultaneously like the conventional way in LLMs (Zoph et al., ; Fedus et al., ), 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 outputs, while the experts’ parameters are updated based on the assigned tokens. In this way, the experts’ development follows the router’s decisions, establishing an inherent leader-follower structure. Following Von Stackelberg (), we formulate the hierarchical problem as a bi-level optimization objective in (1).
Notably, one of the earliest MoE works (Jordan & Jacobs, ), also demonstrates a hierarchical structure, though for a probabilistic interpretation. In contrast, we approach the hierarchical structure from an optimization perspective, formulating the learning process as two nested optimization problems.

$$ $\displaystyle\min_{{\mathbf{\Phi}}}\ \sum_{i}{\mathcal{L}}(f({\bm{X}}_{i}^{ \text{valid}};{\mathbf{\Theta}}^{\star}({\mathbf{\Phi}}),{\mathbf{\phi}}_{i}), {\bm{X}}_{i}^{\text{valid}})$ (1) $\displaystyle\qquad s.t.\quad{\mathbf{\Theta}}^{\star}({\mathbf{\Phi}})\in \operatorname*{arg\,min}_{{\mathbf{\Theta}}}\ \sum_{i}{\mathcal{L}}(f({\bm{X}} _{i}^{\text{train}};{\mathbf{\theta}}^{G},{\mathbf{\theta}}^{S}_{i},{\mathbf{ \phi}}_{i}),{\bm{X}}_{i}^{\text{train}})$ $$

where ${\mathcal{L}}$ is the language modeling loss. Note we write ${\bm{X}}_{i}$ as the label here, as this is a self-supervised task. Labels are simply shifted inputs. The routing parameters ${\mathbf{\Phi}}=\{{\mathbf{\phi}}_{i}\}$ are updated based on the validation loss, which reflects the target distribution (outer optimization), while the expert parameters ${\mathbf{\Theta}}={\mathbf{\theta}}^{G}\cup\{{\mathbf{\theta}}^{S}_{i}\}$ are updated using the training loss (inner optimization). This formulation further brings in the following benefits: 1) routing parameters are smaller in size, making them easier to overfit. By separating the two losses, the routing parameters can be updated less frequently using the smaller validation set (a visual evidence of less frequent router update leading to improved performance is provided in Figure 7 in the Appendix); 2) this approach handles situations where target distributions differ from training distributions more effectively, as the router outputs (i.e., how the experts should be weighted) can be tailored to specific tasks.

Figure: Algorithm 1 Pseudo code of our proposed algorithm

### 3.3 Our Algorithm

To solve (1), we use a multi-round gradient-based algorithm as shown in Alg.1, where only generalist parameters are shared and aggregated, and specialist and router parameters are updated locally. While the scheme requires a server, it can alternatively be implemented in a serverless all2all fashion, which requires $N$ times more communication overhead and we do not further pursue this here.

##### Alternating Update of 𝚯 𝚯 {\mathbf{\Theta}} bold_Θ and 𝚽 𝚽 {\mathbf{\Phi}} bold_Φ :

Alternating update of two sets of parameters is a standard way to solve bi-level optimization problems (Chen et al., ). In between two communication rounds, we perform alternating updates of expert and routing parameters using local training and validation sets separately. The updates optimize the objectives given in (2) and (3) respectively. Since the updates of ${\mathbf{\Theta}}$ and ${\mathbf{\Phi}}$ are disentangled, they do not need to be updated at the same frequency. The routing parameters are smaller in size and thus can be updated less frequently.

When updating model parameters, we include an additional load-balancing term as in Fedus et al. (), which is standard in MoE implementation and encourages even distribution of token assignments to experts. 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.

##### Update of θ G superscript 𝜃 𝐺 {\mathbf{\theta}}^{G} italic_θ start_POSTSUPERSCRIPT italic_G end_POSTSUPERSCRIPT and θ i S superscript subscript 𝜃 𝑖 𝑆 {\mathbf{\theta}}_{i}^{S} italic_θ start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_S end_POSTSUPERSCRIPT :

The update of generalist parameters ${\mathbf{\theta}}^{G}$ follows a standard FedAvg scheme, through aggregating model parameters. Specifically, we simultaneously update both ${\mathbf{\theta}}^{G}$ and ${\mathbf{\theta}}_{i}^{S}$ by optimizing equation (2), which results in ${\mathbf{\theta}}_{i}^{G}$ and ${\mathbf{\theta}}_{i}^{S}$. A parameter aggregation is then performed on the user-specific ${\mathbf{\theta}}_{i}^{G}$ via a trusted server to establish a shared ${\mathbf{\theta}}_{G}$ across all users. In the next round, each user replaces their ${\mathbf{\theta}}_{i}^{G}$ with the global ${\mathbf{\theta}}_{G}$, while their ${\mathbf{\theta}}_{i}^{S}$ remains locally updated.

## 4 Experiments

### 4.1 Setup

#### 4.1.1 Datasets

We selected three diverse distributed datasets to demonstrate the efficacy of our proposed algorithm:

- i)
*Multilingual Wikipedia*: This dataset constitutes Wikipedia articles in four languages: German, Dutch, French, and Italian. We take German, French and Italian from , and Dutch from Guo et al. ();
- ii)
*SlimPajama*: We pick the following four categories – StackExchange, Github Codes, ArXiv, Book from Soboleva et al. ();
- iii)
*AG News*: This dataset covers News from categories of World, Sports, Business, and Sci/Tech (Zhang et al., ).

Opting for the most challenging scenario, each user is assigned a unique category, as shown in Figure 8, where users can have varying data quantities.
Given our emphasis on next token prediction, we anticipate shared predictions among users while maintaining category-specific distinctions. For details of our user data splits, please refer to Appendix B. We further create the following two scenarios to showcase the wide applicability of our method:

##### In-Distribution Tasks.

For each user, we construct validation 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).

##### Out-of-Distribution Tasks.

For each user, we create validation and test datasets from a distribution different from the training data. During training, each user is assigned a single News category from AG News, 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

Our base model is the GPT2 model with 124M parameters(^2^22We adopt the codes from [https://github.com/karpathy/nanoGPT](https://github.com/karpathy/nanoGPT)and [https://github.com/danielgrittner/nanoGPT-LoRA](https://github.com/danielgrittner/nanoGPT-LoRA)), which is suitable for on-device deployment. We incorporate LoRA modules into every linear layer, including MLP and Self-Attention Layers, following the recommendations of Fomenko et al. (), specifically in the [attn.c_attn, attn.c_proj, mlp.c_fc, mlp.c_proj] layers. A routing mechanism is exclusively implemented atop MLP layers. This means that each attention layer has only one LoRA expert applied to it, which is always aggregated during synchronization. The number of LoRA experts in MLP blocks depends on the local resource abundance.

We use a constant learning rate of $2\times 10^{-3}$ for updating routing parameters and a $2\times 10^{-3}$ learning rate with a one-cycle cosine schedule for expert parameters during fine-tuning. The LoRA rank $r$ is set to 8 unless otherwise specified, with LoRA alpha $\alpha$ set to 16, following the common practice of setting alpha to twice the rank (Raschka, ). A load balancing weight $0.01$ is always applied. For more experimental details, we refer readers to Appendix B.

### 4.2 Data-Driven Selection: Generalist vs. Specialist

We start by equipping users with the same model architecture locally, to illustrate the effectiveness of our hierarchical learning of routing and expert parameters. We compare our one generalist one specialist (CoMiGS-1G1S) method to the following baselines. In order to match the trainable parameter count of our method, we use 2 times LoRA modules within each user.

- i)
Upper and lower bounds:
•
Pretrained: A pretrained GPT-2 model using weights from .
•
Centralized: A single model trained using data from all users. (Note this method is an unrealistic baseline as data cannot leave the devices due to privacy concerns.)
- ii)
Baselines:
•
Local: Training individually using only local data without collaboration.
•
FedAvg: Aggregating LoRA parameters across users using uniform weights, which is equivalent to applying FedAvg (McMahan et al., ).
•
PCL: Aggregating LoRA parameters using a client-level collaboration graph. The graph is updated using validation performances. (Strategy 2 in Wagner et al. ()).
•
pFedMoE: We directly apply the method from Yi et al. () in the language domain where we update routing and expert parameters at the same time and choose tokens as a routing unit.
- iii)
Ablations:
•
CoMiGS-2S: Both of the LoRA experts are specialists, meaning their weights are neither shared nor aggregated. The routing parameters are updated using a separate validation set like in CoMiGS-1G1S.
•
CoMiGS-2G: Both of the LoRA experts are generalists, meaning their weights are always shared and aggregated. The routing parameters are updated using a separate validation set like in CoMiGS-1G1S.

**Table 1: Mean test perplexity over the users with homogeneous models, averaged across 3 seeds. Mean (std) with a rank locator for the mean (the lower the better). Green denotes the best performing methods and red denotes our method. An extended version of the table can be found in Table 4.**
|  | In Distribution | Out of Distribution |  |
| --- | --- | --- | --- |
|  | *Multilingual* | *SlimPajama* | *AG News* |
| *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) | 26.95 (0.14) | 41.46 (0.06) |
| *FedAvg* | 58.80 (0.34) | 23.27 (0.05) | 31.84 (0.02) |
| *PCL* | 54.53 (0.19) | 26.99 (0.19) | 32.25 (0.12) |
| *pFedMoE* | 52.27 (0.17) | 25.40 (0.09) | 38.72 (0.21) |
| *CoMiGS - 2S (ours)* | 46.36 (0.16) | 22.51 (0.08) | 35.81 (0.13) |
| *CoMiGS - 2G (ours)* | 58.31 (0.17) | 21.36 (0.01) | 31.18 (0.05) |
| *CoMiGS - 1G1S (ours)* | 47.19 (0.10) | 21.79 (0.04) | 33.53 (0.03) |

#### 4.2.1 Result Analysis

The comparison between our method and the baseline methods is summarized in Table 1.

##### Effectiveness of our routing mechanism:

Depending on the dataset, either CoMiGS-2G or CoMiGS-2S achieves the highest performance. The key distinction compared to Local or FedAvg is the existence of a layer-wise router, which weighs the two generalists or two specialists for each token according to the validation performances, as opposed to assigning equal weights. This emphasizes that even with the same expert knowledge, the way it’s combined is crucial. Moreover, pFedMoE, despite having a learned router as well, underperforms our method, even in the in-distribution scenario. The reason is that the routing parameters are updated simultaneously with the expert parameters using the training set, and thus cannot effectively adapt to the target distribution.

##### Token-level collaborative decisions outperform Client-level:

Compared to the state-of-the-art baseline PCL, as proposed by Wagner et al. (2024), our method demonstrates a clear performance improvement. PCL assigns a pairwise collaboration weight between users by evaluating how well user $i$’s model performs on user $j$’s validation set. On the two in-distribution tasks, PCL exhibits performance similar to Local, where the learned collaboration matrices are nearly identity matrices, thereby limiting effective collaboration between users. Our method, in contrast, decides the collaboration pattern based on each input token, and thus can harness the collective power more effectively.

##### The necessity of the co-existence of generalists and specialists:

The performances of CoMiGS-2G and CoMiGS-2S are not consistent across the different scenarios, while our CoMiGS-1G1S can always closely track the best-performing model, which is clearly shown in Table 1 and visualized in Figure 9.
Depending on the task type, generalists and specialists alone may not be sufficient. A balanced combination of personalization and collaboration is required, and our approach achieves this balance effectively.

#### 4.2.2 Routing Analysis

##### Token-wise analysis:

We further present a token-level routing result visualization on models fine-tuned with SlimPajama dataset in Figure 2: The first two users are fine-tuned with very specific math and programming texts, and they tend to utilize the generalist more in the last layer. Function words (”and”, ”a”, ”on” ”the” etc) are more routed to generalists, as expected. This can be seen in the top right panel of Figure 2. It is important to note that only the top choice is highlighted here. The abundance of blue does not imply that generalist experts play no role in predicting the next token. To see this, compared to when only specialists are present (CoMiGS-2S), our CoMiGS-1G1S gives more consistent results. More detailed token-wise routing result visualization including out-of-distribution tasks can be seen in Appendix E. When dealing with out-of-distribution texts, there is an increasing tendency to seek for generalists, as shown in the off-diagonal entries in Figure 14-19.

##### Layer-wise analysis:

Figure 3 depicts the evolution of averaged layer-wise router outputs for the generalist and specialist experts on the *out-of-distribution* task, comparing CoMiGS-1G1S and pFedMoE. As training progresses, CoMiGS-1G1S undergoes a *phase transition*: the layer-wise routers initially favor generalists but gradually shift towards specialists. This shift is not observed in pFedMoE, highlighting the critical role of our routing mechanism in handling out-of-distribution tasks. Additionally, we notice different layers converge to a different expert score distribution. When applying our CoMiGS-1G1S, for each user, there are always certain layers where the routers consistently prefer generalists, which aligns with the fact that our target distribution is a union of all local training distributions. This phenomenon no longer occurs with in-distribution tasks, as shown in Figure 10.

Figure: Figure 2: Visualization of in-distribution token-level routing results for CoMiGS-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.
Refer to caption: x2.png

Figure: Figure 3: 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), with x-axis being the number of iterations. Upper row: our CoMiGS-1G1S, bottom row: pFedMoE. Darker colors represent deeper layers. Expert score plots for in-distribution tasks can be seen in Figure 10 .
Refer to caption: x4.png

### 4.3 Adaptation to Resource Heterogeneity

#### 4.3.1 Baseline Comparison

In this section, our focus is to deal with system heterogeneity, where users can have different numbers of experts. We still keep one generalist expert, but the number of specialists can vary across the users (our method is denoted as One-Generalist-X-Specialists, named CoMiGS-1GXS). It’s important to note that the richness of computational resources doesn’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, a crucial objective is to prevent overfitting due to redundant model-fitting abilities.

We compare our approach to two state-of-the-art baselines: HetLoRA from Cho et al. () and FlexLoRA from Bai et al. (), both of which adapt LoRA ranks based on the resource capacity of each user. HetLoRA aggregates LoRA matrices ${\bm{A}}$ and ${\bm{B}}$ by zero-padding to the maximum rank and then distributes them back using rank truncation. In contrast, FlexLoRA first reconstructs model updates $\Delta{\bm{W}}$ and redistributes the aggregated updates using SVD. We compare our method to these baselines by matching the number of fine-tunable parameters, measured as both active and full parameters. For example, to match the full parameter count of CoMiGS-1GXS with $(4,2,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.

Our results, presented in Table 2, are based on allocating different computational resources to users, with resource availability decoupled from local task complexity. We find that our method outperforms the baseline methods most of the time, regardless of whether we match the full parameter count or the active parameter count. This advantage stems from the fact that both HetLoRA and FlexLoRA average model parameters across users without allocating parameters for local adaptations, focusing on building a strong generalist model. In contrast, our approach adaptively integrates both generalist and specialist knowledge, excelling in scenarios where specialized knowledge is crucial.

**Table 2: Mean test perplexity (std) over users with heterogeneous models, averaged across 3 seeds, with red being the top1 method. For example, $(4,2,2,2)$ means in our CoMiGS-1GXS setup users have 4, 2, 2, and 2 experts, respectively, and in the two baselines, all users have rank 16 to match active parameter count, or ranks 32, 16, 16, and 16 to match full parameter count.**
|  |  | Ours | HetLoRA | FlexLoRA |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  |  | CoMiGS-1GXS | Active | Full | Active | Full |
| *In Distribution* | Multilingual |  |  |  |  |  |
|  | *(2,2,4,4)* | 46.48 (0.16) | 57.76 (0.10) | 58.60 (0.20) | 77.65 (0.20) | 77.85 (0.26) |
|  | *(4,4,2,2)* | 47.24 (0.09) | 57.76 (0.10) | 59.14 (0.04) | 77.65 (0.20) | 76.29 (0.17) |
|  | SlimPajama |  |  |  |  |  |
|  | *(2,4,4,2)* | 22.10 (0.17) | 23.33 (0.10) | 23.15 (0.09) | 22.97 (0.11) | 22.99 (0.08) |
|  | *(4,2,2,4)* | 22.28 (0.09) | 23.33 (0.10) | 23.17 (0.09) | 22.97 (0.11) | 22.99 (0.09) |
| *Out of Distribution* | AG News |  |  |  |  |  |
|  | *(4,2,2,2)* | 33.66 (0.07) | 31.58 (0.14) | 31.95 (0.13) | 36.45 (0.06) | 36.49 (0.17) |
|  | *(2,4,4,4)* | 34.22 (0.09) | 31.58 (0.14) | 32.52 (0.19) | 36.45 (0.06) | 36.40 (0.08) |

#### 4.3.2 Analysis related to local data quantities

In this section, we further separate resource abundance from data quantity. It is observed that our approach is more robust to overfitting due to the regularizing effect of the generalist, while at the same time better fitting local data through the incorporation of specialist knowledge.

We conduct experiments using Multilingual Wikipedia dataset, where we allocate low data quantities to German and Dutch users, and high data quantities to French and Italian users, as shown in Figure 8. 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 capabilities—measured by the number of LoRA modules per user—either positively or negatively correlated with their local data size. It is important to note that one generalist is always assigned, and resource abundance is only reflected in the number of specialists.

##### More Specialists Help with Higher Data Quantity.

French and Italian users consistently benefit from having more specialists locally, as their test perplexities 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.

Figure: Figure 4: Test Perplexity vs. the number of iterations. Low and high denote the relative data quantity among users. The numbers in the legend indicate the number of experts $n_{i}$ within each user. Top-2 routing is performed.
Refer to caption: x5.png

##### 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 goal here is to prevent overfitting. Figure 5 demonstrates that our method succeeds to surpress overfitting, even when fine-tuning twice or four times as many expert parameters. We attribute this to the existence of the generalists.

Figure: Figure 5: Test Perplexity vs. the number of iterations. Low and high denote the relative data quantity among users. The numbers in the legend indicate the number of experts $n_{i}$ within each user. Top-2 routing is performed. German and Dutch Users despite having high resources locally, do not overfit on their small-sized local data.
Refer to caption: x6.png

##### 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 a result, the generalists become more powerful, as demonstrated in Figure 6.

Figure: Figure 6: Test Perplexity vs. the number of iterations. German and Dutch Users, despite having only one expert locally, still benefit from their collaborators having more experts, thereby enhancing the generalist’s performance. The numbers in the legend indicate the number of experts, $n_{i}$, within each user. Top-2 routing is applied when $n_{i}\geq 2$.
Refer to caption: x7.png

We provide an additional example of the impact of local data quantities in Appendix D using SlimPajama dataset. 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.

## 5 Conclusions and Future Directions

We propose a novel framework for on-device personalized collaborative fine-tuning of LLMs, grounded in an innovative bi-level formulation of the Mixture-of-Experts learning objective. Our fine-grained integration of generalist and specialist expert knowledge achieves superior performance in balancing personalization and collaboration within Federated LLMs.

Furthermore, our framework is the first to address both system and data heterogeneity in collaborative LLM training. It also 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.

An interesting future direction to explore is adopting our framework for collaborative instruction tuning of larger LLMs and evaluating its performance on downstream tasks. While our paper focused on a single generalist, it is possible to include multiple generalists, and their impact on performance remains to be seen. We hope our work paves the way for a new direction in on-device collaborative LLMs.

## Appendix A Limitations and Societal Impact

##### Limitations.

Compared to FedAvg type of methods, our method requires extra gradient steps on routing parameters and memory storage for the routing parameters. Since a routing network is usually a one-layer MLP, the extra cost in computation and storage is relatively small.

The robust performance of our method relies on the fact that we update routing parameters and expert parameters on two independent losses. This means we need some validation samples independent from training samples. When local data size is minimal, this can be an issue.

Our method, while generally robust, still has a tendency towards overfitting when there is a significant mismatch between local resource abundance and data complexity, similar to other methods.

##### Societal Impact.

We offer a collaboration framework for edge devices, aiming to enable smaller devices to leverage large language models (LLMs) despite limited resources and data availability. Our approach enhances fairness and mitigates privacy concerns by ensuring data remains on end devices. The privacy aspects can further be enhanced by differential private aggregation of generalist weights, which we do not pursue here.

The robustness towards attackers is beyond the scope of our work. Our collaboration framework has no guarantee of resilience towards Byzantine attackers, which could potentially lead to misuse by certain parties.

## Appendix B Extra Experimental Details

Following Kalajdzievski (), we choose $\gamma$ to be a rank-stabilized value, a technique which helps stabilize gradient norms. $\alpha$ and the rank $r$ are hyper-parameters to choose from. The LoRA modules function as follows:

$$ ${\bm{W}}={\bm{W}}^{0}+\gamma\cdot{\bm{A}}{\bm{B}},\qquad\gamma=\frac{\alpha}{ \sqrt{r}}$ (4) $$

All our experiments except the centralized ones were conducted on a single A100-SXM4-40GB GPU. The centralized learning baseline experiments were conducted on a single A100-SXM4-80GB GPU, as a batch size of 64*4 requires a larger storage capacity.

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. In each iteration, a batch size of 64 is processed with a context length of 128. We set the routing update period to 30 iterations, and every time we update routing parameters, we do 10 gradient steps on the validation loss. The choice of the hyperparamters is from a sweep run and we provide the evidence in Figure 7.

Figure: Figure 7: Sweep results on SlimPajama data splits. We ablate the impact of the update period ($\tau$) and the number of update steps ($s$) on model performance.
Refer to caption: x8.png

The dataset distribution and number of tokens within each user are shown in Figure 8 and Table 3 respectively.

Figure: Figure 8: The data splits across users, with bubble size denoting the relative size of the local dataset.
Refer to caption: x9.png

**Table 3: Number of tokens in each dataset splits**
|  |  | User 1 | User 2 | User 3 | User 4 |
| --- | --- | --- | --- | --- | --- |
| Multilingual | Training | 557’662 | 407’498 | 556’796 | 451’584 |
| Validation | 300’764 | 216’318 | 220’071 | 165’984 |  |
| Test | 229’720 | 219’741 | 210’570 | 172’547 |  |
| SlimPajama | Training | 1’000’000 | 1’000’000 | 1’000’000 | 1’000’000 |
| Validation | 200’000 | 200’000 | 200’000 | 200’000 |  |
| Test | 200’000 | 200’000 | 200’000 | 200’000 |  |
| AG News | Training | 761’924 | 756’719 | 814’131 | 771’460 |
| Validation | 48’809 | 48’730 | 50’398 | 48’249 |  |
| Test | 48’167 | 47’721 | 48’344 | 49’377 |  |

## Appendix C More Tables and Figures

### C.1 Learning curves of different methods

Figure: Figure 9: Test Perplexity during training for all the three datasets: our method closely follows the best performing method
Refer to caption: x10.png

Figure: Figure 10: Expert Scores for the *generalist* expert and the *specialist* expert from our CoMiGS-1G1S 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.
Refer to caption: x11.png

### C.2 Extended Baseline Comparison

An extended version of Table 1 is presented in Table 4. In this extension, we incorporate two additional ablations: 1) 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-sampled training batches instead of a dedicated validation set. Moreover, we include another baseline method FFA-LoRA from Sun et al. (), where the LoRA A matrices are fixed at initialization.

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 4: Mean test perplexity over users with homogenous models, averaged across 3 seeds. Mean (std) with a 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 |  |
| --- | --- | --- | --- |
|  | *Multilingual* | *SlimPajama* | *AG News* |
| i) 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) | 26.95 (0.14) | 41.46 (0.06) |
| *FedAvg* | 58.80 (0.34) | 23.27 (0.05) | 31.84 (0.02) |
| *FFA-LoRA* | 57.83 (0.13) | 23.42 (0.069) | 31.60 (0.14) |
| *PCL* | 54.53 (0.19) | 26.99 (0.19) | 32.25 (0.12) |
| ii) 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) |
| iii) Alternating update routing params on newly sampled batches from training set |  |  |  |
| *Local-MoE - tr* | 53.78 (0.33) | 27.78 (0.06) | 41.46 (0.03) |
| *FedAvg-MoE - tr* | 59.39 (0.13) | 23.00 (0.01) | 31.70 (0.16) |
| *CoMiGS - tr* | 50.86 (0.14) | 25.45 (0.01) | 38.93 (0.08) |
| iv) Alternating update routing params on a validation set |  |  |  |
| *CoMiGS - 2S* | 46.36 (0.16) | 22.51 (0.08) | 35.81 (0.13) |
| *CoMiGS - 2G* | 58.31 (0.17) | 21.36 (0.01) | 31.18 (0.05) |
| *CoMiGS - 1G1S* | 47.19 (0.10) | 21.79 (0.04) | 33.53 (0.03) |

### C.3 HetLoRA

Analogously to the baseline experiment comparison in FlexLoRA (Bai et al., ), we use $\gamma=0.99$ as pruning strength and sweep the regularization parameter in $\{5\times 10^{-2},5\times 10^{-3},5\times 10^{-4}\}$.

### C.4 Is the standard load balancing loss sufficient?

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 to a 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-balancing 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 is as follows:

$$ ${\mathcal{L}}_{i}^{\text{LB}}={\color[rgb]{0,0.6015625,0.53515625}\definecolor [named]{pgfstrokecolor}{rgb}{0,0.6015625,0.53515625}\frac{1}{(n_{i}-1)^{2}+1}} \cdot f_{0}\cdot P_{0}+\sum_{j=1}^{n_{i}-1}{\color[rgb]{0,0.6015625,0.53515625 }\definecolor[named]{pgfstrokecolor}{rgb}{0,0.6015625,0.53515625}\frac{n_{i}-1 }{(n_{i}-1)^{2}+1}}\cdot f_{j}\cdot P_{j}$ (5) $$

where

$$ $f_{j}=\frac{1}{T}\sum_{x\in{\mathcal{B}}}\mathbbm{1}\{j\in\text{Top2 indices of }p(x)\}\qquad P_{j}=\frac{1}{T}\sum_{x\in{\mathcal{B}}}p_{j}(x)$ (6) $$

$j$ is the expert index and
$p(x)=[p_{j}(x)]_{j=1}^{n_{i}}$
is the logit output from the routing network for a specific token $x$. The idea is that one of the top 2 tokens should always be routed to the generalist expert, i.e. the 0-th expert. Thus,
$\frac{p_{0}}{1/2}$
should be equal to
$\frac{p_{i}}{1/2(n_{i}-1)}$
for $i\neq 0$. As the original load balancing loss encourages uniform distribution, this modification encourages the generalist expert to have a routing probability of 0.5 on expectation. Note that when $n_{i}=2$, this ${\mathcal{L}}_{i}^{\text{LB}}$ is the same as the original load balancing loss as proposed in Fedus et al. ().

We present the results in Table 5: in both scenarios, whether users have the same or different numbers of experts, including a load-balancing term leads to a slight improvement compared to omitting it. However, encouraging more tokens to be routed to the generalists does not make a significant difference.

**Table 5: Test perplexity with different load balancing terms with (hetero) or without (homo) system heterogeneity.**
|  | No LB | LB (uniform) | LB (generalist-favored) |
| --- | --- | --- | --- |
| AG 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) | 22.10 (0.11) | 22.10 (0.17) |

## Appendix 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 is evident that ArXiv User and Book User, with abundant local data, benefit from having more local experts.

Figure: 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_{i}$ within each user. Top-2 routing is performed.
Refer to caption: x12.png

Generalists Help to Prevent Redundant Specialists from Over-Fitting? From Figure 12, we observe more prominent overfitting than in Figure 5, likely because the tasks are objectively easier, as indicated by lower test perplexity from the beginning of fine-tuning. Generalists have limited power to prevent overfitting with easy tasks.

Figure: Figure 12: In this SlimPajama setup, Stack Ex User and Codes User despite having low resources locally, overfit slightly on their small-sized local data. Numbers in the legend denote the number of experts $n_{i}$ within each user. Top2 routing is performed.
Refer to caption: x13.png

Specialists Can Benefit Generalists.
Low-resourced users that can only support a 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 high-resourced users. This is indicated by the enhanced performances for Stack Exchange and Codes Users.

Figure: 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_{i}$, within each user. Top-2 routing is applied when $n_{i}$ $\geq$ 2
Refer to caption: x14.png

## Appendix E Visualization of Expert Specialization

To visualize which tokens are routed to the generalist and specialist experts for our CoMiGS-1G1S model trained on SlimPajama, we ask ChatGPT to 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 Top1 routed index. The routing results after the very first layer (0th), a middle layer (5th), and the very last layer (11th) 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 (0th), a middle layer (5th), and the very last layer (11th) are presented in Figure 17, 18 and 19.

For all the plots, diagonal entries are *in-distribution* texts and off-diagonal entries are *out-of-distribution* texts.

Figure: Figure 14: Visualization of token-level routing results for CoMiGS-1G1S trained on SlimPajama. Tokens are colored with the first expert choice at the 0th (first) 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.
Refer to caption: x15.png

Figure: Figure 15: Visualization of token-level routing results for CoMiGS-1G1S trained on 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.
Refer to caption: x16.png

Figure: Figure 16: Visualization of token-level routing results for CoMiGS-1G1S trained on SlimPajama. Tokens are colored with the first expert choice at the 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.
Refer to caption: x17.png

Figure: Figure 17: Visualization of token-level routing results for CoMiGS-1G1S trained on AG News. Tokens are colored with the first expert choice at the 0th (first) 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.
Refer to caption: x18.png

Figure: Figure 18: Visualization of token-level routing results for CoMiGS-1G1S trained on AG News. Tokens are 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.
Refer to caption: x19.png

Figure: Figure 19: Visualization of token-level routing results for CoMiGS-1G1S trained on AG News. Tokens are colored with the first expert choice at the 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.
Refer to caption: x20.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 4
questions: - Is there any overhead (communication, computation) compared to classical FL approaches? Especially considering the alternating update of theta and fi.
- Is there a particular reason that a theoretical analysis is missing from this paper?
- On the privacy aspect mentioned above, do you think that privacy risk changes depending on the device capabilities? E.g., more capable devices (higher capacity) will aggregate more information, and thus be more exposed?
soundness: 3
strengths: - Well-written paper. Nicely structured, well motivated, and easy to follow. Overall, excellent presentation.
- I appreciate the effort to publish the code, together with the submission, while preserving their anonymity. This is not always the case. In addition, the code is of high quality.
- While the approach that the authors have followed is simple, it is highly practical and demonstrates an effective balance when addressing data and system heterogeneity.
- Carefully chosen baselines to compare the performance.
summary: This work introduces an approach for collaborative, on-device language modeling based on Federated Learning. It addresses the challenges of system and data heterogeneity, common in federated learning configurations. The focus is on language modeling tasks, specifically MoE-based LLMs using LoRA. The novelty lies in a bi-level optimization formulation for the MoE learning objective, which effectively balances collaboration and personalization. This approach is particularly advantageous in scenarios with high data heterogeneity across multiple datasets and use cases.
weaknesses: - My biggest concern is regarding the significance of its contribution, especially considering the heavy overlap with the listed related work "pFedMoE: Data-Level Personalization with Mixture of Experts for Model-Heterogeneous Personalized Federated Learning". I understand that the authors have clearly and fairly mentioned the differences. However, I can't deny that the other work is stealing some of its contributions.
- I would have expected a theoretical analysis section in such a paper. Instead, all results and discussions are derived from simulations.
- While I said the manuscript is very well written, I do think it is missing some background work that would make it easier to follow. For instance, a few sentences about what MoE is, what the router is, and its role, what is LoRA, etc. I am aware of the space limit, but I felt that this was missing from the paper.
- Similar to the previous comment, it is missing some (more generic) relevant related works on the topic of addressing system and data heterogeneity during FL configurations. I understand that this work is focused on language modeling tasks; however, I believe the authors should mention (at least briefly in 1-2 lines) a few examples in more classical DNN (non-LLM) scenarios. Some examples:
  * Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout (Horvath et al.)
  * Revisiting Sparsity Hunting in Federated Learning: Why does Sparsity Consensus Matter? (Babakniya et al.)
- The authors mention privacy as the motivation of using on-device LLMs (that's fair for sure), but also when applying FL. There are numerous attacks on the FL space that render such claims unrealistic, and perhaps such claims should be phrased more cautiously.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: see weakness
soundness: 2
strengths: - The article exhibits clear writing and a well-organized structure, effectively distinguishing its contributions from existing works and facilitating reader comprehension.
- The article addresses the potential applications of CoMiGS within the realms of federated learning and large language models, highlighting the method's broad applicability and promising implications.
summary: The article introduces CoMiGS, a collaborative learning approach designed to address the challenges of system and data heterogeneity in on-device large language models (LLMs). By combining generalists and specialists, CoMiGS facilitates knowledge sharing while also personalizing solutions for individual user datasets. A key innovation of this approach is its bi-level optimization framework for the Mixture-of-Experts (MoE) objective, which allows for flexible aggregation and adaptation to the varying computational resources of users. CoMiGS demonstrates improved performance over existing methods, such as pFedMoE, by enhancing routing mechanisms and accommodating a flexible number of experts for each user. This ultimately leads to better privacy and personalization in federated learning.
weaknesses: 1. The study is based on a limited scale, involving only four clients. It would be beneficial to explore methods to further differentiate between clients, such as employing clients with the same domain data but varying quantities.

2. While the article emphasizes the relationship between large language models (LLMs) and personalized federated learning, the novelty appears insufficient. The approach seems to replace traditional personalization methods in federated learning with LoRA training for LLMs and uses the Mixture-of-Experts (MoE) framework for personalizing each LoRA module. The motivation section needs to reinforce the necessity of this approach to establish its distinct contribution.

3. The article does not adequately address how it resolves issues related to resource heterogeneity. The experimental validation in this area appears insufficient, particularly regarding the fixed nature of experts for each client, as client resources may fluctuate over time.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. To facilitate both data sharing and end-user personalization within the inherent leader-follower structure, the proposed approach seems to introduce additional computational and communication costs during training and data sharing. Could the authors clarify the impact of this optimization on computational efficiency and memory storage requirements? Additionally, how does this affect the feasibility of deploying CoMiGS on resource-constrained edge devices?
soundness: 3
strengths: 1. CoMiGS effectively addresses heterogeneity by integrating knowledge across diverse end users while tackling the challenges of data heterogeneity. Its bi-level optimization mechanism strategically utilizes the roles of generalists and specialists, facilitating both data sharing and personalized experiences for end users.

2. CoMiGS employs load-balancing techniques to ensure a balanced allocation among different experts, enhancing the model’s robustness and generalization capabilities by maintaining equilibrium across them.
summary: This paper presents CoMiGS, a collaborative learning framework for personalized fine-tuning of LLMs on edge devices, addressing data heterogeneity and privacy challenges. CoMiGS features a dual-expert structure with generalists and specialists for effective knowledge sharing while preserving local data privacy. It uses a bi-level optimization mechanism, updating router parameters based on validation loss (outer optimization) and expert parameters based on training loss (inner optimization). This approach mitigates issues seen in traditional MoE architectures. Overall, CoMiGS effectively manages system and data heterogeneity in federated LLM training, enabling efficient data sharing without compromising privacy and enhancing model robustness.
weaknesses: 1. CoMiGS resembles a fusion of LoRAMoE [1,2,3,4] and FL-LoRA [5,6]. However, it lacks analytical comparisons with related studies and does not provide experimental comparisons to support its claims.

2. CoMiGS is designed for edge deployment, but fine-tuning multiple LoRA modules (generalists vs. specialists) may incur more computational overhead.

3.  The title suggests a focus on LLMs, but evaluation is limited to GPT-2 (124M), which is too simplistic in terms of both model type and parameter count. The study omits more complex LLMs, such as Llama3-1B, Llama3-3B, and Gemma-2B, which are designed for mobile environments. With the increase in parameters, issues such as embedded knowledge may become more pronounced, potentially enhancing the robustness and adaptability for In-Distribution and Out-of-Distribution task scenarios. This prompts the question: could this technique scale effectively with future generations of LLMs? Further research is necessary to explore these possibilities.

References:

[1] When MOE Meets LLMs: Parameter Efficient Fine-tuning for Multi-task Medical Applications, SIGIR 2024.

[2] Mixture of LoRA Experts, ICLR 2024.

[3] Lorahub: Efficient cross-task generalization via dynamic lora composition, COLM 2024.

[4] Pushing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning, ICLR 2024.

[5] pFedLoRA: Model-Heterogeneous Personalized Federated Learning with LoRA Tuning, Arxiv 2023.

[6] FDLoRA: Personalized Federated Learning of Large Language Model via Dual LoRA Tuning, Arxiv 2024.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: - Could the authors clarify how CoMiGS fundamentally differs from existing frameworks that use a mixture of shared (global) and personalized (local) models to handle heterogeneity in federated learning? Specifically, what are the unique contributions of CoMiGS in comparison to these prior methods?
- CoMiGS relies on a validation set for tuning and routing adjustments, which may not be feasible in real-world on-device settings due to data limitations. How would CoMiGS function without a dedicated validation set, or are there alternative mechanisms for evaluation that do not require a separate validation split?
soundness: 2
strengths: - The token-level routing enables CoMiGS to handle data heterogeneity more efficiently than client-level methods, and the bi-level optimization approach aligns well with the personalized and collaborative goals of the framework.
- The paper provides robust experimental evaluations on three datasets, covering both in-distribution and out-of-distribution scenarios.
summary: This paper introduces a novel collaborative learning framework called CoMiGS (Collaborative learning with a Mixture of Generalists and Specialists) for on-device fine-tuning of Large Language Models (LLMs). CoMiGS addresses both system heterogeneity (differences in local computational resources) and data heterogeneity (differences in user-specific data) by differentiating between generalists, which handle shared knowledge across devices, and specialists, which adapt to user-specific data. The approach is formulated as a bi-level optimization problem, with generalists and specialists iteratively updated through a routing mechanism, which adapts based on the target distribution. Experimental results demonstrate that CoMiGS outperforms existing methods in scenarios with varied data and system settings.
weaknesses: - While the paper presents CoMiGS as a novel collaborative learning framework, the core idea of combining generalist and specialist models to handle heterogeneous data is not particularly new. Similar approaches have been proposed in previous federated learning and mixture-of-experts frameworks. Many existing works utilize a combination of shared global models (analogous to generalists) and personalized models (specialists) to handle data and system heterogeneity [1]. The authors do not provide a clear explanation of how CoMiGS fundamentally differs from or improves upon these prior methods, which weakens the claim of novelty.
- The bi-level optimization approach may introduce significant computational overhead, especially for on-device applications with limited resources. Each optimization level requires iterative updates, making the entire process computationally intensive.
- The paper does not propose alternative mechanisms for model evaluation or adaptation that avoid the use of a validation set. Other federated learning approaches often incorporate online learning or lightweight evaluation methods that do not require a dedicated validation split, which CoMiGS could consider.
- The rationale behind maintaining both generalists and specialists within the framework is not entirely convincing. If specialist models are designed to adapt to user-specific data, the role of generalists becomes ambiguous, especially in highly heterogeneous settings where each user's data diverges significantly from the global average. The coexistence of generalists and specialists adds complexity to the model without a clear demonstration of the complementary benefits of this setup.


[1]: FedCP: Separating Feature Information for Personalized Federated Learning via Conditional Policy. Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, Haibing Guan
、

Meta-review: The paper presents a new method for building on-device language modeling. The proposed approach seeks to provide each person with a personalized language model using ideas from mixtures and federated learning. The personalized model on each device is a mixture -- built by sharing tokens across shared models (generalists) and local models (specialists). The routing mechanism ('i.e., which tokens to share') is determined by solving a bi-level optimization problem that uses data from the target domain to account for distribution shifts. The paper benchmarks the performance of the proposed approach for salient tasks on three real-world datasets that exhibit heterogeneity.

**Feedback from Reviewers:** This paper received mixed feedback from four reviewers. On the one hand, reviewers recognized the broader need for work in this area and appreciated the technical machinery. On the other hand, they highlighted two overarching concerns: (W1) computation overhead resulting from the use of bi-level optimization; (W2) validity of experimental claims (e.g., what is driving the gain, will it scale beyond GPT-2). Given the mixed feedback, I spent more time on this paper than usual -- reading through the reviews, the author responses, and the final submission. With regards to the weaknesses identified by reviewers:

- W1: I do not share the reviewers' concerns. My view is that research papers that introduce new methods should develop methods that offer new functionality or that improve functionality in a meaningful way. In such cases, the bar for "compute" should be that it is viable -- i.e. we should treat it as a constraint rather than an end in-and-of itself. It does not make sense to view compute as an objective since: (1) the new functionality may be worth the compute; (2) there are more effective avenues to reduce compute (e.g., an efficient implementation at production). Given this, the burden for methods papers is to show viability and functionality so that the method can be adopted and implemented efficiently in practice. 

- W2: I do not share the reviewers concerns, but have different concerns. One of the reviewer's concerns in this case was whether the method would exhibit the same kinds of performance gains when using new generations of LLMs. During the rebuttal, this concern led the authors to run the experiments with LLAMA-3.2 and to report that the gains reported in their paper had disappeared for 2/3 datasets. In this case, I agree with the authors that this can be explained by test set contamination. However, the broader issue is that we have ended up in a position where we can no longer draw conclusions.

**Strengths and Weaknesses** Overall, I find that there is much to like about this paper. It targets a problem that is of broad interest (on-device learning), proposes an interesting method (federated learning with mixtures), and includes experiments that demonstrate viability and soundness. In my view, its main weakness is a failure to demonstrate significance. As it stands, the main "sell" for the method is that it can improve local performance by adapting to local datasets. This sell is not compelling because we expect such improvements when learning on-device models. On my end, it is also concerning since we have a limited ability to test such claims going forward.

In this case, I see a major opportunity to address this issue by drawing the focus back to mixtures. Given that their method is meant to learn a mixture, we should expect that these models will exhibit the same kinds of strengths as mixtures of LLMs (see e.g., [Wang et al.](https://arxiv.org/pdf/2406.04692)). I would use such papers as a starting point to identify "strengths" that can lead to meaningful improvements in functionality in the federated regime. These "strengths" could provide the basis for more compelling demonstrations in Section 4. As it stands, the experiments in this part are getting in the way of a compelling story. Specifically, the routing analysis in 4.2 simply highlights the behavior of the algorithm, and the gains observed under heterogeneity are somewhat expected. 

**Recommendation**: I am recommending rejection at this stage. I want to be clear that this was a borderline case. Having read the paper, I find that the work is well-written and well-executed, but has simply been positioned in a way that will limit its impact. In this case, I ultimately believe that this work will end up as a far stronger contribution to the field following another round of revisions.

Decision: Reject