PaperlayerView exact revision on arXiv

Source-linked AI summary

FedYolo: Augmenting Federated Learning with Pretrained Transformers

Xuechen Zhang, Mingchen Li, Xiangyu Chang, Jiasi Chen, Amit K. Roy-Chowdhury, Ananda Theertha Suresh, Samet Oymak

arXiv:2307.04905v1cs.LGcs.DC

TL;DR

Federated learning must contend with scarce heterogeneous client data and diverse tasks, while large pretrained transformers raise questions about device and communication constraints. This paper studies scale and modular adaptation, finding benefits for local learning, robustness, communication, and multitask use, and proposes FedYolo with task-specific modules. The computational cost of deploying pretrained transformers remains unexplored.

  • Problem

    The paper asks how to support diverse client goals with scarce heterogeneous data in federated learning, where a unified strategy remains lacking.

    Source: §1 · pp. 1–2

  • Method

    The paper studies pretrained transformers across training schemes and uses modular updates that freeze the backbone while clients train and communicate modules and a classifier head.

    Source: §1 · p. 2

  • Results

    Scale and modularity make Local-only learning fairly competitive with federated approaches, improve heterogeneity robustness, and enable communication savings of orders of magnitude.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    FedYolo assigns isolated modules to tasks over one frozen pretrained-transformer backbone, performing on par with isolated task learning while full-update multitasking suffers catastrophic forgetting.

    Source: §1 · p. 3

    The work does not explore the computational cost of deploying pretrained transformers, a stated drawback of large models.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Federated learning faces scarce, heterogeneous client data and diverse tasks, but the paper finds no clear unified strategy for these challenges. It investigates pretrained transformers’ scale and modularity, finding benefits for communication, robustness, local learning, and multitask adaptation.

1 Introduction

  • The paper frames its problem around scarce, heterogeneous client data, diverse tasks, and the absence of a clear unified federated-learning strategy.The introduction situates heterogeneity, multitasking, and catastrophic forgetting among the challenges.Source: §1 · pp. 1–2
  • Modular updates communicate only small modules while freezing the pretrained-transformer backbone, and FedYolo assigns isolated modules to tasks.Clients also communicate a classifier head under modular updates.Source: §1 · pp. 2–3
  • Combined scale and modularity make Local-only learning fairly competitive with federated approaches.Scale shrinks accuracy gaps between FedAvg and FedAvg+Local and between FedAvg+Local and Local-only learning.Source: §1 · p. 2
  • Scale improves robustness to data heterogeneity, while modularity particularly improves robustness of smaller pretrained transformers.FedAvg+Local remains competitive with FedAvg on the global distribution after local learning.Source: §1 · p. 3
  • Modular updates save orders of magnitude in communication while maintaining, and often accelerating, convergence in communication rounds.Larger scale can also reduce communication rounds by letting clients run more local SGD epochs without sacrificing global accuracy.Source: §1 · p. 3
  • FedYolo assigns each task its own module over a frozen pretrained-transformer backbone; it performs on par with isolated task learning, unlike full-update multitasking, which suffers catastrophic forgetting.The paper presents FedYolo as a way to use one pretrained transformer across multiple disparate tasks.Source: §1 · p. 3

2 Related Work

Prior federated-learning work has studied heterogeneity and personalization, while pretrained-model studies often used full updates or did not examine scale alongside modularity. The paper distinguishes its work by studying modular updates across scales and addressing multitask learning with pretrained transformers.

Federated Learning

  • Earlier federated-learning studies of pretrained models used full updates, while related modular-update studies did not explore scale.The paper reports that some related NLP studies found modular updates worse than full updates, whereas it finds benefits in few-shot or heterogeneous settings.Source: §2 · p. 3
  • The paper identifies its multitask pretrained-transformer solution as absent from prior work it reviews.It contrasts this with prior work on heterogeneity, multitasking, personalization, pretrained models, and modular updates.Source: §2 · p. 3

Parameter-efficient tuning and pretrained transformers

  • Parameter-efficient tuning adjusts limited parameters or adds small modules, and the paper reports modular methods can help smaller models under few-shot or heterogeneous data.The passage notes prior evidence that prompt tuning approaches full fine-tuning as scale grows, but falls short for smaller transformers trained with full data.Source: §2 · p. 3

3 Preliminaries and Experimental Setup

The paper evaluates scale and modularity in few-shot settings with federated, personalized, and local-only training schemes. Experiments use image benchmarks and real-world client-partitioned datasets, varying data heterogeneity and pretrained-transformer size.

3 Preliminaries and Experimental Setup

  • The study asks how scale affects personalization, whether scale and modularity improve heterogeneity robustness, how they affect communication, and how modules should be assigned across tasks.These questions span the preliminary experiments and the FedYolo multitask method.Source: §3 · p. 4
  • The training comparisons include Local-only learning, FedAvg, and FedAvg+Local, each with full-update or modular-update variants.FedAvg+Local fine-tunes the average global model locally after federated training.Source: §3 · p. 4
  • Experiments use few-shot CIFAR-10, CIFAR-100, CelebA, and FEMNIST data, with CelebA and FEMNIST using 2.6% and 1.8% of their total samples, respectively.CIFAR experiments use homogeneous and heterogeneous client partitions; CelebA and FEMNIST are partitioned by celebrity and writer.Source: §3 · p. 4
  • The evaluated pretrained-transformer sizes are ViT-L, ViT-B, ViT-S, and ViT-T, and modular methods include Adapter, LoRA, and VPT.The main paper reports Adapter results, with LoRA and VPT results in supplementary material; full-update is the exception to frozen backbones.Source: §3 · p. 4
  • Figure 2 plots CIFAR-100 accuracy against samples per class and reports that larger models improve both update strategies and narrow the local-versus-federated gap, especially with modular updates.Figure 3 reports that larger models outperform smaller ones under heterogeneous data, while Adapter can outperform full-update and is more helpful for smaller models.Source: Figure 2 · p. 5, Figure 3 · p. 5

4 Experiments

Experiments show that pretrained-transformer scale and modular updates improve performance and robustness across federated settings. Modular updates reduce communication, while larger models support more local training and stronger performance under a fixed communication budget.

4.1 PTF Scale Boosts Performance

  • Larger PTFs consistently outperform smaller ones on homogeneous and heterogeneous data, while modular updates can outperform full updates in some low-sample or heterogeneous settings.For ViT-T, modular updates outperformed full updates with four or fewer samples per client; the modular-update advantage was also observed under heterogeneous data.Source: §4.1 · pp. 5–7
  • As model scale increases, Local-only learning becomes more competitive with FedAvg, and modular updates further narrow their accuracy gap.With 16 samples per client under modular updates, the Local-only versus FedAvg accuracy gap was 8.10 for ViT-L and 15.44 for ViT-T.Source: §4.1 · p. 6

4.2 Heterogeneous client data distributions

  • On CIFAR-100, the homogeneous-to-heterogeneous performance decrease was 8.5% for ViT-L with modular updates, 19.7% for ViT-T with modular updates, and 40.5% for ViT-T with full updates.Larger PTFs consistently performed better across heterogeneity levels, while modular updates could bolster smaller models’ robustness.Source: §4.2 · p. 7
  • The gap between FedAvg and FedAvg+Local shrinks as PTF scale grows, while full updates tend to widen it, especially with smaller backbones.The experiments suggest larger PTFs and modular updates could reduce the need for computationally intensive personalization.Source: §4.2 · p. 7
  • Modular updates and larger PTFs reduce forgetting, helping personalized models retain global accuracy while improving local accuracy.The forgetting ratio decreased with modular updates and as PTF scale increased.Source: §4.2 · p. 8

4.3 Reducing Federated Communication Cost with Modular Updates

  • As more tasks are added, FedYolo with task-specific modules outperforms conventional FedAvg on Task 0 in both global and locally adapted accuracy.Figure 10 reports the Task 0 CIFAR-100 comparison for multi-task federated learning.Source: Figure 10 · p. 9
  • Modular updates reduce communication cost by over 100x and typically converge faster than full updates in the federated setting.Modular updates reached target accuracy in 37.25 (±1.08) epochs on average, versus 47.25 (±1.48) for full updates.Source: §4.3 · p. 9
  • With communication cost limited to 10^7 on CIFAR-100, ViT-L with E = 20 achieved 75.04 accuracy in 3 rounds, versus 54.61 for ViT-T with E = 1 in 33 rounds.Larger PTFs can maintain or improve performance with larger local epoch counts, whereas larger E harmed performance for small PTFs.Source: §4.3 · p. 9

5 Federated Multitask Learning via FedYolo

FedYolo assigns each task its own module within a single frozen pretrained transformer, which clients load once. In multitask experiments, it outperforms FedAvg, particularly with more tasks and smaller models, while per-task module uploads raise a privacy concern.

5 Federated Multitask Learning via FedYolo

  • FedYolo assigns each task a unique module plugged into a single frozen pretrained transformer, which clients load once.Clients update and communicate only the task-specific modules, with minimal communication cost.Source: §5 · p. 10
  • Sending task-specific modules can reveal which tasks clients have to the server; zero-padded module vectors with secure aggregation are proposed as an alternative.The alternative uses a vector of length K times the single-module size, with nonzero entries only for tasks a client handles.Source: §5 · p. 10
  • FedYolo consistently outperforms conventional FedAvg in multitask learning, especially with more tasks and smaller pretrained transformers.The evaluation assigns each client one of four image-classification tasks: CIFAR-10, CIFAR-100, CelebA, or FEMNIST.Source: §5 · p. 10

6 Conclusions, Limitations, and Future Directions

The paper concludes that pretrained-transformer scale and modularity can address federated-learning bottlenecks, and presents FedYolo as an effective multitask scheme. Its computational cost was not explored, and the authors identify additional settings and design directions for future study.

6 Conclusions, Limitations, and Future Directions

  • The computational cost of deploying pretrained transformers was not explored and is identified as a clear drawback.Compression or distillation, early exits, and offloading are suggested as possible mitigations for separate investigation.Source: §6 · p. 10
  • The paper concludes that scale and modularity have potential to address federated-learning bottlenecks and presents FedYolo as an effective multitask scheme.Source: §6 · p. 10
  • The authors identify unlabeled client data and non-stationary distributions as settings for future study..Source: §6 · p. 11
  • Future FedYolo designs could assign shared modules across tasks or clients or search for optimal module placements within pretrained transformers.Source: §6 · p. 11

A Organization of the Appendix

The appendix reports additional experiments on communication rounds, comparisons with FedProx and head-tuning, and the generality of findings across modules and model scales. It also points to further experiment details and a representation-based explanation in later appendices.

A Organization of the Appendix

B.1 Comparisons to Existing FL Methods

The appendix compares FedYolo with optimization-based federated-learning methods and evaluates modular updates against head-tuning. It also reports results for other module types.

B.1 Comparisons to Existing FL Methods

  • FedProx shows no significant performance improvement over FedAvg, while FedYolo maintains a substantial advantage and outperforms FedProx with full updates in multitask settings.FedDyn results are described as similar; the authors also note that FedYolo can be combined with optimization-based methods.Source: §B.1 · p. 15

B.2 Comparison to Head-tuning

  • Modular updates consistently outperform head-tuning in performance and robustness across the reported settings.The comparison includes homogeneous and heterogeneous data settings.Source: §B.2 · p. 15
  • The appendix also reports experiments with VPT and LoRA modules across sample-count and mild-heterogeneity settings.Source: Figure 13 · p. 16

B.3 Results of Other Modules

Results across module types generally agree, while increasing prompt size eventually stops improving performance. The advantage of larger PTFs is not explained by parameter count alone.

B.3 Results of Other Modules

  • VPT and LoRA results indicate that module type affects performance, but the paper's main findings generalize across module types and experiments.Source: §B.3 · p. 16

B.4 The Impact of Trainable Parameter Count

  • At 299,108 prompt parameters, performance reaches a saturation point, with further prompt-size increases yielding no significant gain.Source: §B.4 · p. 17
  • Adding parameters to smaller PTFs does not necessarily improve performance, so parameter count alone does not explain larger PTFs' advantage.Source: Figure 14 · p. 17, §B.4 · p. 17

B.5 Comparison to Centralized Training

The comparisons find that larger models are more robust to data heterogeneity and benefit more from pretraining in few-shot federated learning.

B.5 Comparison to Centralized Training

  • Larger-scale models show greater robustness to heterogeneity when federated accuracies are compared with centralized accuracies.Source: §B.5 · p. 17, Figure 15 · p. 17

B.6 The Impact of Pretraining

  • Pretrained models outperform models trained from scratch, and larger models tend to benefit more when few-shot training is used.Source: §B.6 · p. 17
  • The pretraining comparison uses FedAvg on CIFAR-100 across varying levels of heterogeneity.Source: §B.6 · p. 17

C Experiment details and reproducibility

The reported experiments use a specified federated training protocol and repeated runs; the supplementary figure reports gains from pretraining, especially at larger model scales.

C.1 Data partition

The appendix identifies data partitions and module parameter counts, then examines representation changes as a plausible explanation for large PTF benefits.

C.1 Data partition

C.2 The setting of modules

D A Representation-based Explanation for the Benefits of Large Models

Loading 2307.04905v1…