FedYolo: Augmenting Federated Learning with Pretrained Transformers
Xuechen Zhang Mingchen Li Xiangyu Chang Jiasi Chen Amit K. Roy-Chowdhury Ananda Theertha Suresh Samet Oymak
The growth and diversity of machine learning applications motivate a
rethinking of learning with mobile and edge devices. How can we address diverse
client goals and learn with scarce heterogeneous data? While federated learning
aims to address these issues, it has challenges hindering a unified solution.
Large transformer models have been shown to work across a variety of tasks
achieving remarkable few-shot adaptation. This raises the question: Can clients
use a single general-purpose model, rather than custom models for each task,
while obeying device and network constraints? In this work, we investigate
pretrained transformers (PTF) to achieve these on-device learning goals and
thoroughly explore the roles of model size and modularity, where the latter
refers to adaptation through modules such as prompts or adapters. Focusing on
federated learning, we demonstrate that: (1) Larger scale shrinks the accuracy
gaps between alternative approaches and improves heterogeneity robustness.
Scale allows clients to run more local SGD epochs which can significantly
reduce the number of communication rounds. At the extreme, clients can achieve
respectable accuracy locally highlighting the potential of fully-local
learning. (2) Modularity, by design, enables $>$100$\times$ less communication
in bits. Surprisingly, it also boosts the generalization capability of local
adaptation methods and the robustness of smaller PTFs. Finally, it enables
clients to solve multiple unrelated tasks simultaneously using a single PTF,
whereas full updates are prone to catastrophic forgetting. These insights on
scale and modularity motivate a new federated learning approach we call "You
Only Load Once" (FedYolo): The clients load a full PTF model once and all
future updates are accomplished through communication-efficient modules with
limited catastrophic-forgetting, where each task is assigned to its own module.