CoMiGS splits on-device fine-tuning into shared generalists and private specialists
Four researchers at EPFL published CoMiGS on September 20, 2024, a scheme in which many devices fine-tune a language model together, with some of the added model parts shared across everyone and the rest kept on the device that trained them. In a simulation of four users, each holding Wikipedia articles in a different language, the authors measure test perplexity at 47.19 with GPT-2 124M as the base model, against 54.38 when each user trains alone on local data and 58.80 for federated averaging, which pools every user’s update into one shared model. Perplexity measures how surprised a model is by the next token, so lower is better. The paper appeared at ICML 2025.
Nobody retrains the base model. Each device trains small LoRA adapters on top of pretrained weights that stay frozen, and CoMiGS sorts those adapters into two kinds. Generalist adapters are averaged across all devices at every communication round, specialist adapters and a small router never leave the device, and the router decides per token which adapters to use. The authors train that router against a separate validation set rather than the training data, and report that in the first layer of the network common function words such as “and”, “a” and “the” go mostly to the generalists, while for the users trained on maths and programming text the domain-specific terms in the last layer go mostly to the specialists.
Nothing here ran on a phone. The appendix states that every experiment except the centralised baseline ran on a single NVIDIA A100-SXM4-40GB, with four simulated users holding between roughly 400,000 and 1 million training tokens each, drawn from multilingual Wikipedia, SlimPajama, AG News and Common Corpus, and averaged over three seeds. The base models are GPT-2 124M and Llama 3.2 1B. Results in the paper’s table are mixed rather than uniform, with CoMiGS ahead of both baselines on multilingual Wikipedia, SlimPajama and Common Corpus, and behind federated averaging on AG News, at 33.53 against 31.84 with GPT-2 and 16.31 against 15.86 with Llama 3.2 1B, which is the one setup where the test data does not follow the training distribution.
The second set of experiments gives users different numbers of adapters, two to four per device with one of them always a generalist, which is how the method is meant to cope with devices of unequal capacity. Against HetLoRA and FlexLoRA, two methods that vary the LoRA rank per device instead, the authors report lower perplexity on every in-distribution task, for example 18.68 against 21.41 and 24.63 on Common Corpus with Llama 3.2 1B. They put the overhead of their extra router at 1.25 percent more computation per forward pass and 1.25 percent more memory, with communication per round halved against federated averaging because only the generalist weights are sent.
The paper names its own limits. The authors write that differential privacy on the aggregated generalist weights is something they “do not pursue here”, that robustness against attackers is beyond the scope of the work and that their aggregation carries no guarantee against adversarial parties, and that there is a limit to how much the generalists can prevent overfitting when the local task is easy. They also note that fine-tuning Llama 3.2 1B on multilingual Wikipedia or SlimPajama produced negligible improvements, which they attribute to overlap with its pre-training data, and that is why the Llama runs use Common Corpus instead. The code is on GitHub under the Apache 2.0 licence.
