Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

• Federated learning trains a global model across heterogeneous devices while keeping private data local and communicating model parameters instead.The introduction notes applications including healthcare, mobile devices, smart cities, and communication networks.Source: §1 · p. 1
• Across non-IID devices, current global-model methods face non-guaranteed convergence and weight divergence in parallel FL, and severe catastrophic forgetting in serial FL.These issues limit the readiness of FL for real-world data distributions.Source: §1 · pp. 1–2
• The paper rethinks federated model architecture and benchmarks neural architectures across federated algorithms, real-world benchmarks, and heterogeneous splits, using Vision Transformers for image and medical image classification.The motivation is the hypothesized robustness of Transformers to distribution shifts.Source: §1 · p. 2
• VIT-FL performs especially well on the most heterogeneous device splits, with its advantage over ResNet-based FL increasing as heterogeneity rises.The figure also reports that ViTs and Swin Transformers significantly outperform CNNs on highly heterogeneous CIFAR-10 partitions.Source: Figure 1 · p. 1, §1 · p. 2
• Transformer robustness to heterogeneous data reduces forgetting of previous devices when training on substantially different new devices, helping models converge faster and reach a better global model.The paper says VIT-FL improves on FL methods designed for heterogeneous data without training heuristics, additional hyperparameter tuning, or extra training, and is orthogonal to optimization-based methods.Source: §1 · p. 2

Federated Learning

• FL includes parallel methods such as FedAVG [47] and serial methods such as CWT and split learning [62]; heterogeneity can cause divergence or catastrophic forgetting.Parallel approaches face non-guaranteed convergence and weight divergence, while serial approaches face catastrophic forgetting.Source: §2 · p. 2
• Existing FL responses largely modify optimization to address heterogeneity, leaving architecture design as a parallel direction for improving robustness to distribution shifts.Methods include server momentum, shared data, proximal objectives, collaborative replay, contrastive learning, feature matching, and distillation.Source: §2 · p. 2
• Serial-learning methods include EWC [30], generative replay [?, 58], and cyclically weighted objectives to reduce forgetting or performance loss.These approaches constrain updates to important weights, mimic data from previous clients or tasks, or mitigate loss under label-distribution skew.Source: §2 · p. 2

Transformers

• Transformers progressed from sequence-to-sequence NLP [61] to NLP and image/video tasks, with ViT [12] achieving ImageNet state of the art using global self-attention.Self-supervised training has been especially successful across NLP tasks [11].Source: §2 · p. 2
• Prior studies find that ViTs are highly robust to severe occlusions, perturbations, and domain shifts [3], as well as synthetic and natural adversarial examples [44] [51].Source: §2 · p. 2
• These findings motivate the hypothesis that ViTs can adapt effectively to FL data heterogeneity, which the paper tests through detailed empirical analysis.Source: §2 · p. 3

3.1. Vision Architectures

• CNN baselines use ResNet [19] and EfficientNet [59], while ViT [12] and Swin Transformer [41] extract image features through sequentialization and patch embedding rather than conventional convolutions.The evaluated families include ResNet-50/152, ResNeXt-101, EfficientNet-B1/B5/B7, ViT(S/T/B), and Swin(T/S/B).Source: §3.1 · p. 3

3.2. Federated Learning Methods

• FedAVG [47] samples clients for synchronous local SGD and server averaging, whereas CWT cyclically trains and transfers one global model across clients.Each selected client trains for E epochs; CWT repeats its client cycle until convergence or a predefined number of communication rounds.Source: §3.2 · p. 3

4.2. Results

• On highly heterogeneous CIFAR-10 Split-3 (KS-1), ViT(S)-CWT and ViT(S)-FedAVG improved test accuracy over ResNet(50) counterparts by 77.70% and 37.34%, respectively.Across IID partitions, CWT and FedAVG achieved results comparable to centralized training regardless of architecture.Source: Figure 3 · p. 4, §4.2 · pp. 4–5

4.1. Experimental Setup

• Evaluation used Retina, CIFAR-10, and federated CelebA, with simulated IID and label-skewed non-IID partitions for Retina and CIFAR-10.Retina and CIFAR-10 used 4 and 5 clients, respectively; CelebA used 227 identity-based clients for smile classification.Source: §4.1 · p. 4

4.2. Results

• On non-IID Retina and CIFAR-10, ViT-FL outperformed FedProx [37] and FedAVG-Share [68], which both dropped in performance on highly heterogeneous partitions.FedAVG-Share reached 94.4% on CIFAR-10 Split-3 versus 97% on Split-1 despite sharing 5% of local data; Transformers are orthogonal to optimization methods, so combining them can yield stronger performance.Source: Figure 4 · p. 5, §4.2 · p. 5

4.3.1 Transformers generalize better in non-IID settings

• The authors suggest ViTs’ self-attention may aid robustness to distribution shifts because it learns global interactions and is less biased toward local patterns than CNNs.They contrast this with CNN local convolutions’ reliance on local high-frequency patterns, which may be sensitive to heterogeneous devices and image-acquisition variation.Source: §4.3.1 · pp. 5–6
• On CIFAR-10 Split-3, ResNet(50)-CWT prediction accuracy on prior Client-3 data fell from >98% to <1%, while ViT(S)-CWT maintained 98%.EWC [30], designed to alleviate forgetting, barely solved the problem on these highly heterogeneous partitions.Source: Figure 5 · p. 6, §4.3.1 · p. 6

4.3.2 Transformers converge faster to better optimum

• 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 from catastrophic forgetting.The target test accuracies are 77.5% for Retina and 91.5% for CIFAR-10; ResNet(50)-FedAVG and ResNet(50)-CWT can plateau as heterogeneity increases.Source: Table 3 · p. 7, §4.3.2 · p. 7

4.4. In Conjunction with Existing Methods

• Combining Vision Transformers with FedProx [37] or FedAVG-Share [68] further improves convergence speed and performance for heterogeneous clients.Figure 6 shows the combinations reaching target performance with fewer communication rounds.Source: Figure 6 · p. 7, §4.4 · pp. 7–8, Table 3 · p. 7

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

• On real-world federated CelebA, VIT-FL outperforms state-of-the-art FL methods on the union of clients’ test data and reduces variance.In an extreme setting with one sample per client, ViTs significantly outperform ResNets across 6,000 Retina clients and 45,000 CIFAR-10 clients, with reported results from 50% to 80% on Retina and 30% to 90% on CIFAR-10.Source: §2 · pp. 6–7

4.5. Take-aways for Practical Usage

• Larger local epoch counts accelerate ViT(B)-FedAVG on homogeneous partitions but can worsen final performance on heterogeneous partitions; ViT(B)-CWT favors frequent client transfers on non-IID data.The authors recommend large E for homogeneous data, but E ≤5 for VIT-FedAVG and E = 1 for VIT-CWT in highly heterogeneous cases.Source: §4.5 · p. 8, Figure 7 · p. 8
• Pretrained Vision Transformers are recommended for practical use, although Swin(T) trained from scratch reaches 64.50% on highly heterogeneous CIFAR-10 Split-3 versus 59.68% for pretrained ResNet(50)-FedAVG.Swin(T) performance drops when trained from scratch in both centralized and federated settings; self-supervised pretraining is suggested when large-scale pretraining data are unavailable.Source: §4.5 · p. 8, Table 4 · p. 8
• Linear warm-up, learning-rate decay, and gradient clipping can be inherited from Vision Transformer training, while small learning rates and gradient-norm clipping stabilize VIT-FL in highly heterogeneous CWT settings.Gradient-norm clipping can also reduce divergence between local updates and the current global model in heterogeneous CNN federated learning.Source: §4.5 · p. 8

5. Conclusion

• Transformers alleviate catastrophic forgetting, accelerate convergence, and reach a better optimum in both parallel and serial federated learning.The paper presents these findings as advantages for heterogeneous data.Source: §5 · p. 8
• The paper offers architecture design as a new perspective on federated learning, alongside prior efforts to improve optimization.It frames robust architectures as a parallel direction for future work.Source: §5 · p. 8
• The authors attribute Transformers’ advantages to their robustness to heterogeneous data and distribution shifts.Source: §5 · p. 8
• The conclusion identifies convergence and forgetting as continuing challenges when federated learning uses heterogeneous data.Source: §5 · p. 8
• The authors release code and models to encourage further work on robust architectures.They position this work alongside research on optimization.Source: §5 · p. 8

A. Experimental Details

• The section provides additional details on the datasets, preprocessing, and experimental methodology, with code available to reproduce the experiments.Source: A. Experimental Details · p. 13

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

• RETINA uses balanced healthy/diseased left-eye images, with 6,000 for training and 3,000 each for global validation and testing after excluding mild cases.Images are processed and cropped to 224 × 224; three four-client partitions cover IID data (KS-0) and label-skewed non-IID data (KS-0.49 and KS-0.57).Source: §A.1 · p. 13
• CIFAR-10 is split into 45,000 training, 5,000 validation, and 10,000 test images, then resized and cropped to 224 × 224 [21].Five-client partitions include IID data (KS-0), KS-0.65 skew with one client receiving two classes and four receiving four, and extreme KS-1 skew with two classes per client.Source: §A.1 · p. 13
• CelebA uses the LEAF federated partition by celebrity for smile classification, retaining 227 clients and 1,213 samples after dropping clients with more than 8 samples [5].Clients average 5.34 ± 1.11 samples, and images are resized to 256 × 256 and cropped to 224 × 224.Source: §A.1 · p. 13

Implementation Details

• All models use ImageNet ILSVRC-2012 pretraining, batch size 32, and 224 × 224 inputs; experiments run in PyTorch on a TITAN V or GeForce RTX 2080 GPU.Local training uses one epoch and FL runs 100 communication rounds unless stated otherwise; implementations use SGD with momentum 0.9 and no weight decay or AdamW [29] with weight decay 0.05.Source: §A.2 · p. 13

Training hyperparameters

• Swin-FL uses AdamW [29] and ViT-FL uses SGD; CNN optimizers are chosen by parameter search, and training uses model-specific learning-rate schedules with global-norm-1 gradient clipping.Transformer training uses 500 warmup steps followed by cosine decay to zero; CNNs use linear warmup and decay or step decay, halved every 30 FL rounds.Source: §A.2 · p. 13

Hyperparameter selection

• CNN learning-rate schedules and FedProx penalty µ are grid-searched on Split-2 of RETINA and CIFAR-10, then the same parameters are applied to remaining partitions, including the extreme large-scale edge case.Detailed model hyperparameters for RETINA and CIFAR-10 are reported in Table 5.Source: §A.2 · pp. 13–14

FL hyperparameters

• RETINA and CIFAR-10 use E=1 and 100 rounds with all clients participating; FedAVGM β is 0.5/0.3 and FedProx µ is 0.001/0.1, respectively.FedAVGM β is selected from {0.1, 0.3, 0.5, 0.7, 0.9, 0.97, 0.99, 0.997}; FedProx µ is selected from {0.001, 0.01, 0.1, 1}.Source: §A.2 · p. 14
• CelebA samples 10 clients per round; CWT runs 30 rounds and other parallel methods 1,000, aiming for roughly 30 participations per client.FedProx µ is 0.001 for CelebA, and FedAVG-Share lets clients share 5% of their data on all compared datasets.Source: §A.2 · p. 14

B.1. Take-aways for Practical Usage

• ViT federated training can inherit ViT practices such as linear warm-up, learning-rate decay, and gradient clipping.The passage presents these as training strategies that can be directly inherited from ViT training.Source: §B.1 · p. 14
• Gradient clipping stabilizes training and accelerates convergence for FL methods on highly heterogeneous splits.The figure compares methods using ViT(B) and ResNet-50 backbones, with and without clipping.Source: Figure 10 · p. 15, §B.1 · p. 15
• The reported experiments apply gradient clipping at global norm 1 to all models to stabilize training.The hyperparameter table covers RETINA and CIFAR-10 experiments with ResNets, EfficientNets, ViTs, and Swins.Source: Table 5 · p. 15
• The reported step-decay scheduler halves the learning rate every 30 epochs.Source: Table 5 · p. 15
• For VIT-CWT, use a large learning rate on IID or mildly skewed partitions, but a smaller rate to stabilize highly heterogeneous training.VIT-CWT favors a relatively smaller learning rate on heterogeneous partitions.Source: §B.1 · p. 15
• Using a smaller learning rate for CNN counterparts leads to worse performance.This contrasts with VIT-CWT’s preference for a smaller learning rate on heterogeneous partitions.Source: §B.1 · p. 15

B.2. Experiments on Real-World Federated Datasets

• OpenImage was collected from Flickr and originally spanned 1.3M images, 600 categories, and 14k clients; the evaluated filter contains 81,088 images, 365 categories, and 9,265 clients.Source: §B.2 · p. 15
• 63.50% for centrally hosted VIT versus 60.56% for centrally hosted ResNet; federated VIT significantly outperforms ResNets on heterogeneous OpenImage.The passage reports the centrally hosted results alongside the federated comparison but does not provide federated scores here.Source: §B.2 · p. 15
• For OpenImage, parallel methods sampled 10 clients per round with E=1; CWT ran 30 rounds and other parallel methods 27,000 rounds, targeting about 30 participations per client.Experiments on CelebA and OpenImage used SGD with momentum 0.9, no weight decay, and gradient clipping at global norm 1.Source: §B.2 · p. 15, Table 6 · p. 16
• With ResNet-50 as backbone, ViT(S)-CWT and ViT(S)-FedAVG outperform other optimization-based federated methods on non-IID partitions.Source: Figure 11 · p. 16

B.3. Investigating the Influence of Normalization Technique in VIT-FL

• 93.87% versus 56.46% on CIFAR-10 Split-2 shows group normalization improves CWT and FedAVG over batch normalization on mildly skewed partitions.Source: §B.3 · p. 16
• Group normalization still loses performance on highly skewed partitions, whereas VIT-FL performs promisingly on both mild and extreme skew, so its effectiveness is not purely due to normalization.Source: §B.3 · p. 16

B.4. Comparisons to Existing FL Methods

• The comparison uses ResNet-50 for FedAVGM, FedProx, and FedAVG-Share, with parameters tuned by grid search on Split-2 and reused for other partitions.FedAVG-Share allows clients to share 5% of their data with one another.Source: §B.4 · p. 16
• On CIFAR-10, FedAVG-Share reaches 94.2% on Split-3 versus 96% on Split-1, while VIT-FL outperforms all compared FL methods on non-IID partitions.FedProx and FedAVGM also suffer severe drops on highly heterogeneous partitions despite tuned optimization parameters.Source: §B.4 · p. 16
• Replacing ResNet-50 batch normalization with group normalization still leaves performance loss on highly heterogeneous partitions, so VIT-FL's gains are not purely due to avoiding batch normalization.Table 7 compares CWT and FedAVG on RETINA and CIFAR-10 using ResNet50 and ResNet50(GN).Source: Table 7 · p. 17
• On OpenImage, VITs significantly outperform their ResNet counterparts across 365 categories and 9,265 clients.Table 8 reports prediction accuracy on this large-scale real-world dataset.Source: Table 8 · p. 17