FedConv: Enhancing Convolutional Neural Networks for Handling Data Heterogeneity in Federated Learning

TL;DR

This paper investigates how architectural modifications to CNNs can improve their robustness in federated learning with heterogeneous data, showing that CNNs can match or surpass Vision Transformers with proper design.

Contribution

It provides the first systematic analysis of architectural elements influencing CNN performance in heterogeneous federated learning and offers design principles to enhance robustness.

Findings

Modified CNN architectures can outperform standard CNNs in heterogeneous FL.

CNNs with strategic design can match or exceed ViT performance in FL.

The proposed approach is compatible with existing FL techniques and achieves state-of-the-art results.

Abstract

Federated learning (FL) is an emerging paradigm in machine learning, where a shared model is collaboratively learned using data from multiple devices to mitigate the risk of data leakage. While recent studies posit that Vision Transformer (ViT) outperforms Convolutional Neural Networks (CNNs) in addressing data heterogeneity in FL, the specific architectural components that underpin this advantage have yet to be elucidated. In this paper, we systematically investigate the impact of different architectural elements, such as activation functions and normalization layers, on the performance within heterogeneous FL. Through rigorous empirical analyses, we are able to offer the first-of-its-kind general guidance on micro-architecture design principles for heterogeneous FL. Intriguingly, our findings indicate that with strategic architectural modifications, pure CNNs can achieve a level of robustness that either matches or even exceeds that of ViTs when handling heterogeneous data clients in FL. Additionally, our approach is compatible with existing FL techniques and delivers state-of-the-art solutions across a broad spectrum of FL benchmarks. The code is publicly available at https://github.com/UCSC-VLAA/FedConv