PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

Liangqiong Qu, Yuyin Zhou, Paul Pu Liang, Yingda Xia, Feifei Wang, Ehsan Adeli, Li Fei-Fei, Daniel Rubin

arXiv:2106.06047v2cs.LGcs.CV

TL;DR

Federated learning must train across private, heterogeneous devices despite convergence, weight-divergence, and catastrophic-forgetting challenges. The paper benchmarks architectures across federated algorithms and heterogeneous evaluations, finding that replacing CNNs with Transformers improves federated learning, especially under strong heterogeneity.

  • Problem

    Federated learning faces non-guaranteed convergence, model weight divergence, and catastrophic forgetting when training across non-IID devices.

  • Method

    The paper empirically benchmarks neural architectures across federated algorithms, real-world benchmarks, and heterogeneous data splits.

  • Results

    Transformers reduce catastrophic forgetting, accelerate convergence, and reach better global models, with VIT-FL outperforming competing methods on non-IID partitions.

  • Takeaways & Limitations

    Architecture choice offers a complementary route to improving federated learning alongside optimization-based methods.

Abstract

from arXiv · show

Federated learning is an emerging research paradigm enabling collaborative training of machine learning models among different organizations while keeping data private at each institution. Despite recent progress, there remain fundamental challenges such as the lack of convergence and the potential for catastrophic forgetting across real-world heterogeneous devices. In this paper, we demonstrate that self-attention-based architectures (e.g., Transformers) are more robust to distribution shifts and hence improve federated learning over heterogeneous data. Concretely, we conduct the first rigorous empirical investigation of different neural architectures across a range of federated algorithms, real-world benchmarks, and heterogeneous data splits. Our experiments show that simply replacing convolutional networks with Transformers can greatly reduce catastrophic forgetting of previous devices, accelerate convergence, and reach a better global model, especially when dealing with heterogeneous data. We release our code and pretrained models at https://github.com/Liangqiong/ViT-FL-main to encourage future exploration in robust architectures as an alternative to current research efforts on the optimization front.

1 Stanford University,2 UC Santa Cruz, 3 Carnegie Mellon University, 4 Johns Hopkins University

The supplied passage lists Stanford University, UC Santa Cruz, Carnegie Mellon University, and Johns Hopkins University affiliations.

  • The affiliation list includes Stanford University.
  • The affiliation list includes UC Santa Cruz and Carnegie Mellon University.
  • The affiliation list includes Johns Hopkins University.

1. Introduction

Federated learning trains models across private, heterogeneous devices, but non-IID data creates convergence, weight-divergence, and catastrophic-forgetting challenges. This paper rethinks federated-model architecture and finds that Transformers improve robustness and federated performance under heterogeneity.

  • Federated learning trains models across private data distributed over heterogeneous devices while exchanging parameters rather than raw data.
  • Non-IID devices create non-guaranteed convergence, model weight divergence in parallel FL, and catastrophic forgetting in serial FL.
  • The paper benchmarks neural architectures across federated algorithms, real-world benchmarks, and heterogeneous data splits.
  • VIT-FL performs especially well on highly heterogeneous splits, with its advantage over ResNet-based FL increasing as heterogeneity rises.
  • The authors attribute improvement to Transformer robustness, which reduces forgetting while enabling faster convergence and a better global model.

2. Related Work

Related work frames federated learning around parallel and serial training, whose non-IID data challenges motivate optimization remedies and architectural alternatives. Transformer robustness to shifts and heterogeneous modalities motivates testing Vision Transformers in federated learning.

  • Parallel FL trains clients concurrently, whereas serial FL trains clients cyclically.
  • Data heterogeneity causes non-guaranteed convergence and weight divergence in parallel FL, alongside catastrophic forgetting in serial FL.
  • Existing approaches address heterogeneity through momentum, shared data, proximal objectives, replay, contrastive learning, feature matching, or distillation.
  • Continual-learning methods mitigate forgetting by constraining important weights, replaying prior data, or reweighting objectives.
  • ViTs show robustness to occlusions, perturbations, domain shifts, and adversarial examples, motivating their evaluation for federated heterogeneity.

3. Transformers in Federated Learning

The paper compares convolutional architectures with Vision and Swin Transformers in federated learning, using FedAVG for parallel training and CWT for serial training. These methods provide the architectural basis for evaluating Transformers under heterogeneous data.

  • CNN baselines include ResNet and EfficientNet families built from convolutional, activation, pooling, and normalization layers.
  • Transformer comparisons use ViT(S), ViT(T), ViT(B), and Swin Transformer model families.
  • FedAVG combines local SGD with iterative server-side model averaging across sampled clients.
  • CWT trains local clients serially and cyclically, contrasting with FedAVG's synchronous parallel training.

4. Experiments

The experiments compare CNNs and Transformers across federated algorithms, datasets, and heterogeneous partitions. Transformers retain accuracy, converge faster, reduce catastrophic forgetting, and outperform CNN-based and optimization-based alternatives in challenging settings.

  • Experimental scope: The study evaluates whether Transformers learn better global models, handle heterogeneous partitions, reduce communication costs, and improve optimization-based federated learning.Experiments use CWT and FedAVG across Retina, CIFAR-10, and CelebA, including large-scale and highly heterogeneous settings.
  • Overall performance: On IID data, CWT and FedAVG achieve accuracy comparable to centrally trained models regardless of architecture, but CNN accuracy drops substantially on heterogeneous partitions.The largest degradation occurs on extremely heterogeneous CIFAR-10 partitions such as Split 3 and KS-1.
  • Overall performance: 77.70% and 37.34%: ViT(S)-CWT and ViT(S)-FedAVG improve test accuracy over their ResNet(50) counterparts on highly heterogeneous CIFAR-10 Split-3, KS-1.Replacing CNNs with ViTs retains model accuracy in highly heterogeneous non-IID settings.
  • Comparison with alternatives: VIT-FL outperforms FedProx, FedAVG-Share, and ResNet-based counterparts on non-IID data, while also reducing variance and improving global models.The reported comparisons cover Retina, CIFAR-10, CelebA, and an edge setting with thousands of clients.
  • Catastrophic forgetting: ViT(S)-CWT preserves Client-3 accuracy at 98% after transfer to Client-4, whereas ResNet(50)-CWT falls from over 98% to under 1%.EWC-CWT barely resolves catastrophic forgetting under the most heterogeneous CIFAR-10 split.

2. Generalization of VIT-FL on real-world federated

Across real-world and highly heterogeneous federated settings, Vision Transformers improve generalization, convergence, and communication efficiency, while training choices such as local epochs and pretraining affect outcomes.

  • Generalization: VIT-FL outperforms state-of-the-art federated methods and reduces variance on CelebA, indicating a better global model than CNN counterparts.The extreme edge case contains 6,000 Retina clients and 45,000 CIFAR-10 clients, with ViTs outperforming ResNets.
  • Communication cost: Communication cost is evaluated using rounds required to reach 95% of centrally trained ResNet(50) accuracy, with targets of 77.5% for Retina and 91.5% for CIFAR-10.For serial CWT, one communication round is one complete training cycle across all local clients.
  • Convergence: 34 and 85 communication rounds let ViT(S)-CWT reach target performance on CIFAR-10 Split-2 and Split-3, while ResNet(50)-CWT diverges.ResNet convergence slows as heterogeneity increases, eventually reaching a plateau or failing to reach the target accuracy.
  • Optimization combinations: ViTs can combine with FedProx and FedAVG-Share to accelerate convergence and reach target performance in fewer communication rounds.This architectural change is presented as complementary to optimization-based federated learning methods.
  • Training choices: Large local epoch counts accelerate convergence on mild or homogeneous heterogeneity but may worsen final performance on severe heterogeneity.The authors recommend large E for homogeneous data and small E for highly heterogeneous cases, specifically E ≤5 for VIT-FedAVG and E = 1 for VIT-CWT.
  • Pretraining: Pretraining is important for VIT-FL, although Swin(T)-FedAVG trained from scratch reaches 64.50% on highly heterogeneous CIFAR-10 Split-3, above pretrained ResNet(50)-FedAVG at 59.68%.Self-supervised pretraining is suggested when large-scale pretraining datasets are unavailable.

5. Conclusion

The paper reframes federated learning under heterogeneous data as an architecture-design problem alongside optimization, finding that Transformers reduce forgetting, accelerate convergence, and improve the global solution.

  • 5. Conclusion: Transformers provide a new architecture-focused perspective on federated learning challenges involving convergence and catastrophic forgetting.The paper positions architecture design as complementary to existing optimization-focused research.
  • 5. Conclusion: Extensive experiments across parallel and serial federated methods demonstrate advantages in heterogeneous settings.The reported benefits include alleviating catastrophic forgetting, accelerating convergence, and reaching a better optimum.
  • 5. Conclusion: The authors release code and models to encourage further research on robust architectures for federated learning.

A. Experimental Details

The experimental-details section directs readers to dataset, preprocessing, and methodology specifications and provides code for reproducing the experiments.

  • A. Experimental Details: The paper documents its datasets, preprocessing steps, and experimental methodology in additional detail.These details are presented as part of the experimental setup.
  • A. Experimental Details: Reproduction code is available at https://github.com/Liangqiong/ViT-FL-main.

A.1. Detailed Image Pre-processing and Data Partitions

The experiments use Retina, CIFAR-10, and CelebA federated partitions that range from IID or moderately skewed clients to extreme identity- or label-based heterogeneity.

  • Retina: Retina is binarized into Healthy and Diseased classes after excluding scale-1 samples, using left retinal images.The source dataset contains clinician-labeled retinal fundus images on scales 0–4.
  • Retina: Retina partitions include one IID split and two non-IID label-skew splits across four simulated clients.
  • CIFAR-10: CIFAR-10 uses 45,000 training images, 5,000 validation images, and the original 10,000-image test set, with images resized to 256 × 256.
  • CIFAR-10: CIFAR-10 Split 3 is an extreme label-skew case in which every client receives images from only two classes.Split 2 assigns one client two classes and the remaining four clients four classes.
  • CelebA: The federated CelebA benchmark partitions celebrity images by identity, tests smile classification, and retains 227 clients averaging 5.34 ± 1.11 samples.Clients with more than eight samples are dropped to increase task difficulty, leaving 1,213 samples.

A.2. Implementation Details and Hyperparameters

The experiments use standardized optimization and federated-learning settings across heterogeneous RETINA, CIFAR-10, CelebA, and OpenImage partitions. Gradient clipping is used to stabilize training, while Transformer-specific schedules and learning-rate choices are retained where supported.

  • Optimization: Transformer models use inherited training practices including AdamW or SGD, linear warm-up, cosine decay, and gradient clipping.ViT-FL uses SGD, while Swin-FL uses AdamW with weight decay 0.05; the warm-up lasts 500 steps.
  • Hyperparameter selection: CNN hyperparameters are tuned on RETINA and CIFAR-10 Split-2, then reused across the remaining partitions, including the extreme large-scale setting.The tuned parameters include learning-rate scheduling and FedProx’s proximal penalty constant.
  • Data partitions: Figures 8 and 9 visualize label-skewed non-IID partitions by reporting each class’s fraction of samples within RETINA and CIFAR-10.The rectangle values represent class-specific sample fractions.
  • Federated-learning settings: All methods use one local epoch and 100 communication rounds unless otherwise stated, with all local clients participating per round on RETINA and CIFAR-10.CelebA uses 10 randomly sampled clients per round, with 30 rounds for CWT and 1000 for other parallel methods.
  • Optimization: Gradient clipping stabilizes training and accelerates convergence for most FL methods on highly heterogeneous partitions.Figure 10 compares methods with and without clipping, and Table 5 specifies clipping at global norm 1.
  • Learning-rate behavior: On heterogeneous partitions, VIT-CWT favors a smaller learning rate, whereas reducing the CNN learning rate worsens performance.For VIT-CWT, larger learning rates are suggested for IID or mildly skewed partitions, while smaller rates stabilize heterogeneous training.

B.2. Experiments

The experiments extend evaluation to large-scale real-world federated data and compare Vision Transformer methods with ResNet-based methods. On OpenImage, VIT significantly outperforms ResNets in a highly heterogeneous setting.

  • OpenImage: OpenImage contains 1.3M images across 600 categories and 14k clients; the evaluated subset contains 81,088 images across 365 categories and 9,265 clients.The subset selects classes with between 20 and 800 samples.
  • OpenImage: VIT significantly outperforms ResNets on the heterogeneous OpenImage partition, exceeding centrally hosted reference accuracies of 60.56% for ResNet and 63.50% for VIT.The comparison is reported in Table 8 for the selected OpenImage subset.
  • Optimization-based comparisons: Figure 11 compares ViT(S)-CWT and ViT(S)-FedAVG with other state-of-the-art optimization-based FL methods using ResNet-50 as the backbone.The figure reports that the Vision Transformer-based FL methods outperform other methods in non-IID partitions.
  • Normalization ablation: The batch-normalization ablation tests whether VIT-FL’s performance can be explained solely by avoiding batch normalization.Replacing batch normalization with group normalization in ResNet50 is used for this comparison.

B.4. Comparisons to Existing FL Methods

VIT-FL is compared with optimization-based FL methods and normalization variants under non-IID conditions. It outperforms the compared methods, while several alternatives degrade on highly heterogeneous partitions.

  • Optimization-based methods: VIT-FL outperforms FedAVGM, FedProx, and FedAVG-Share in non-IID data partitions despite the compared methods receiving tuned optimization parameters.All methods use ResNet-50 as the backbone for the optimization-based comparison, with clients sharing 5% of local data for FedAVG-Share.
  • Optimization-based methods: FedProx and FedAVGM suffer severe performance drops on highly heterogeneous partitions despite carefully tuned optimization parameters.This result is reported in the comparison associated with Figure 11.
  • Optimization-based methods: FedAVG-Share drops to 94.2% on CIFAR-10 Split-3 versus 96% on Split-1 even after 5% of local data is shared among clients.The comparison illustrates that limited data sharing does not prevent performance loss on the highly heterogeneous split.
  • Normalization analysis: ResNet50 with group normalization still loses performance on highly heterogeneous partitions, indicating VIT-FL’s gains are not solely due to removing batch normalization.Table 7 reports this comparison for CWT and FedAVG on RETINA and CIFAR-10.
  • Real-world benchmark: VITs significantly outperform their ResNet counterparts on the large-scale OpenImage benchmark covering 365 categories and 9,265 clients.Table 8 reports the prediction-accuracy comparison.
Loading 2106.06047v2…