Research map / Federated learning and privacy / Personalized and heterogeneous FL
On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes CoMiGS, a bi-level mixture-of-LoRA-experts method where shared generalist LoRA experts are aggregated across users and user-specific specialists are kept local, with routers updated on validation data via alternating minimization.
Problem
On-device collaborative fine-tuning of LLMs faces both data heterogeneity across users and computational (model) heterogeneity in how many adapter modules each device can host; prior federated LoRA methods handle only one of these.
Approach
Split experts into globally aggregated generalists and locally kept specialists, combined by a token-level router; formulate MoE training as bi-level optimization where routers are fit on a validation set and experts on the training set, solved by alternating first-order updates with linear convergence analysis.
Key results
- On GPT2-124M Multilingual Wikipedia with homogeneous models, CoMiGS-2S reaches 46.36 test perplexity vs 54.38 for Local and 58.80 for FedAvg.
- With heterogeneous expert counts, CoMiGS-1GXS beats HetLoRA and FlexLoRA on in-distribution tasks, e.g. 46.48 vs 57.76 on Multilingual Wikipedia (GPT2).
- CoMiGS adds about 1.25% memory/computation overhead over FedAvg and halves communication cost since only generalist weights are exchanged.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add a router (token-level learned mixing) trained on a separate validation set (bi-level) vs. routers trained on training loss (pFedMoE) | CoMiGS-1G1S · Multilingual Wikipedia, GPT2-124M | test perplexity: 47.19 vs 52.27 | Helped |
| Use a separate validation set for router updates instead of newly sampled training batches (Local-MoE-tr / CoMiGS-tr vs CoMiGS) (ablation) | CoMiGS routing · Multilingual Wikipedia, GPT2-124M | test perplexity: 50.86 (CoMiGS-tr) vs 47.19 (validation set) | Helped |
| Replace the mix of one generalist and one specialist with two specialists only (ablation) | CoMiGS-2S · Multilingual Wikipedia vs AG News, GPT2-124M | test perplexity: 46.36 vs 47.19 on Multilingual (better); 35.81 vs 33.53 on AG News (worse) | Mixed |
| Replace the mix with two generalists only (ablation) | CoMiGS-2G · SlimPajama and AG News, GPT2-124M | test perplexity: 21.36 vs 21.79 on SlimPajama (better); 58.31 vs 47.19 on Multilingual (worse) | Mixed |
| Use CoMiGS (1 generalist + 1 specialist, token-level routing) vs. client-level collaboration (PCL) | CoMiGS-1G1S · Common-Corpus, Llama3.2-1B | test perplexity: 18.37 vs 19.65 | Helped |
| Use CoMiGS-1G1S vs FedAvg on out-of-distribution AG News with Llama3.2-1B | CoMiGS-1G1S · AG News, Llama3.2-1B | test perplexity: 16.31 vs 15.86 (FedAvg) | Hurt |
| Heterogeneous-resource setting: one generalist plus varying specialists (CoMiGS-1GXS) vs HetLoRA/FlexLoRA | CoMiGS-1GXS · In-distribution tasks (Multilingual, SlimPajama, Common-Corpus) | test perplexity: e.g. Multilingual 46.48 vs 57.76 (HetLoRA active); Common-Corpus 18.74 vs 21.41 | Helped |
| Heterogeneous-resource setting on out-of-distribution AG News: CoMiGS-1GXS vs HetLoRA | CoMiGS-1GXS · AG News (GPT2 and Llama3.2-1B) | test perplexity: 33.66 vs 31.58 (GPT2); 16.39 vs 15.89 (Llama3.2-1B) | Hurt |
| Add load-balancing loss term during training (ablation) | CoMiGS-1G1S · AG News, Multi-Wiki, SlimPajama (homo and hetero) | test perplexity: e.g. Multi-Wiki homo 47.31 -> 47.19; hetero 46.36 -> 46.15 (uniform LB) | Helped |
| Generalist-favored load-balancing loss (importance weighting) vs uniform load-balancing (ablation) | CoMiGS load balancing · AG News, Multi-Wiki, SlimPajama | test perplexity: differences within about 0.1 perplexity | No clear effect |
| Local-only training on low-data users with two LoRA modules vs. adding generalists (CoMiGS) with more specialists | CoMiGS-1GXS with generalist · Multilingual Wikipedia low-data users | overfitting (test perplexity trend over iterations) | Helped |
| Replace FedAvg with FedAvg on LoRA with CoMiGS communication: exchange only generalist weights | CoMiGS-1G1S vs FedAvg · GPT2-124M communication cost | communication cost per round: -1.41 MB per round | Helped |
Surprising findings
- Generalist-only (CoMiGS-2G) is best on out-of-distribution AG News, while specialist-only is best on in-distribution Multilingual, so no single configuration dominates.
- Routing parameters updated simultaneously with experts on training loss (pFedMoE) underperform separate validation-based router updates.
Limitations
- Robustness to adversarial attackers is not guaranteed; the paper notes aggregation of generalist weights is vulnerable.
- Differential privacy for generalist aggregation is not pursued.
- Experiments use small models (GPT2-124M, Llama-3.2-1B) and only four simulated users per setting.
Experimental setup
Datasets: Multilingual Wikipedia SlimPajama AG News Common Corpus
Benchmarks: Multilingual Wikipedia SlimPajama AG News Common Corpus
Metrics: test perplexity
Baselines: Local FedAvg PCL pFedMoE FDLoRA HetLoRA FlexLoRA FFA-LoRA FedSA-LoRA
Builds on: LoRA Mixture-of-Experts FedAvg pFedMoE FDLoRA
Closest papers by meaning
- PrE-Text: Training Language Models on Private Federated Data in the Age of LLMs (ICML 2024 · similarity 0.87)
- Federated Learning with Heterogeneous LLMs: Integrating Small Student Client Models with a Large Hungry Model (AAAI 2025 · similarity 0.86)
- Open LLMs are Necessary for Current Private Adaptations and Outperform their Closed Alternatives (NeurIPS 2024 · similarity 0.85)
- FlowerTune: A Cross-Domain Benchmark for Federated Fine-Tuning of Large Language Models (NeurIPS 2025 · similarity 0.85)
- FedBPT: Efficient Federated Black-box Prompt Tuning for Large Language Models (ICML 2024 · similarity 0.85)
- FedLLM-Bench: Realistic Benchmarks for Federated Learning of Large Language Models (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.