Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

Area of Science:

• Machine Learning
• Artificial Intelligence
• Computer Science

Background:

• Federated learning enables collaborative model training across organizations while preserving data privacy.
• Key challenges in federated learning include slow convergence and catastrophic forgetting, especially with heterogeneous data distributions.
• Current research primarily focuses on optimization techniques to address these challenges.

Purpose of the Study:

• To investigate the efficacy of self-attention-based architectures, such as Transformers, in enhancing federated learning robustness.
• To empirically evaluate Transformer architectures against traditional models on real-world heterogeneous datasets.
• To demonstrate that architectural choices can mitigate federated learning challenges like catastrophic forgetting and improve convergence.

Main Methods:

• Conducted a rigorous empirical investigation comparing different neural network architectures (Transformers vs. Convolutional Networks).
• Evaluated performance across various federated learning algorithms and real-world benchmarks with heterogeneous data splits.
• Focused on metrics including convergence speed, catastrophic forgetting, and final global model accuracy.

Main Results:

• Self-attention-based architectures (Transformers) show superior robustness to data distribution shifts inherent in federated learning.
• Replacing convolutional networks with Transformers significantly reduced catastrophic forgetting and accelerated model convergence.
• Transformers achieved better global model performance, particularly under highly heterogeneous data conditions.

Conclusions:

• Transformer architectures offer a promising, architecture-centric approach to improving federated learning on heterogeneous data.
• This study provides empirical evidence that architectural design is a critical factor for robust federated learning, complementing optimization efforts.
• Released code and models to foster further research into robust architectures for private, collaborative machine learning.