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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Fan, Dongyang, Messmer, Bettina, Doikov, Nikita, Jaggi, Martin
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866918037872967680
author Fan, Dongyang
Messmer, Bettina
Doikov, Nikita
Jaggi, Martin
author_facet Fan, Dongyang
Messmer, Bettina
Doikov, Nikita
Jaggi, Martin
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2409_13931
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists
Fan, Dongyang
Messmer, Bettina
Doikov, Nikita
Jaggi, Martin
Machine Learning
Computation and Language
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.
title On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists
topic Machine Learning
Computation and Language
url https://arxiv.org/abs/2409.13931