On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists
Dongyang Fan, Bettina Messmer, Nikita Doikov, Martin Jaggi
TL;DR
CoMiGS tackles on-device, privacy-preserving language model personalization under data and resource heterogeneity by introducing a bi-level Mixture-of-Experts framework that shares generalist knowledge across users while localizing specialist expertise. A token-level router, trained on separate validation data, dynamically assigns tokens to generalists or specialists, with the expert updates driven by on-device training data; alternating minimization solves the bi-level objective and yields linear convergence under mild conditions. Empirically, CoMiGS demonstrates strong performance across in- and out-of-distribution tasks on multilingual and domain-diverse corpora, outperforming FedAvg and other MoE baselines, especially when routing decisions are made at the token level. The framework decouples data quantity from resource availability, enabling high-resource users to leverage larger local models while protecting low-resource users from overfitting, with manageable on-device overhead and half the communication costs. The work also provides theoretical convergence guarantees and open-sources the codebase to foster collaborative, privacy-preserving LLM development on edge devices.
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.