SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sara Pieri* Jose Renato Restom* Samuel Horvath Hisham Cholakkal
Mohamed Bin Zayed University of Artificial Intelligence (MBZUAI)
Federated Learning (FL) is a promising research paradigm that enables the collaborative training of machine learning models among various parties without the need for sensitive information exchange. Nonetheless, retaining data in individual clients introduces fundamental challenges to achieving performance on par with centrally trained models. Our study provides an extensive review of federated learning applied to visual recognition. It underscores the critical role of thoughtful architectural design choices in achieving optimal performance, a factor often neglected in the FL literature. Many existing FL solutions are tested on shallow or simple networks, which may not accurately reflect real-world applications. This practice restricts the transferability of research findings to large-scale visual recognition models. Through an in-depth analysis of diverse cutting-edge architectures such as convolutional neural networks, transformers, and MLP-mixers, we experimentally demonstrate that architectural choices can substantially enhance FL systems' performance, particularly when handling heterogeneous data. We study 19 visual recognition models from five different architectural families on four challenging FL datasets. We also re-investigate the inferior performance of convolution-based architectures in the FL setting and analyze the influence of normalization layers on the FL performance. Our findings emphasize the importance of architectural design for computer vision tasks in practical scenarios, effectively narrowing the performance gap between federated and centralized learning. Our source code is available at https://github.com/sarapieri/fed_het.git.
The growing focus on data privacy and protection [37, 61] has sparked significant research interest in federated learning (FL) as it provides an opportunity for collaborative machine learning in many domains, such as healthcare [16, 14, 42, 55], mobile devices [47, 19], internet of things (IoTs) [53, 25] and autonomous driving [40, 56]. In FL, data is kept separate by individual clients to learn a global model on the server side in a decentralized way.
Despite the potential benefits of federated learning, the FL environments are highly non-trivial, as the local datasets of the individual clients may not accurately represent the global data distribution. As each client is a unique end user, the heterogeneity caused by uneven distributions of features, labels, or an unequal amount of data across clients poses a serious obstacle [27]. Therefore, the setting is inherently challenging since FL methods must accommodate both statistical diverseness and exploit similarities in client data. Typically, the non-independent and identically distributed (non-IID) nature of the clients drastically affects the model's performance, resulting in lower accuracy and slower convergence than the counterparts trained with all the data collected in a single location.
Figure 1: Our study incorporates a wide range of datasets, each presenting varying data heterogeneity levels, including Fed-ISIC2019 and GLD-23K, which are inherently federated datasets due to their image acquisition and sensor mechanisms. Differently from prior literature, we extensively evaluate models from each visual architectural family. Additionally, we assess the effectiveness of four distinct optimization methods orthogonally to the architectural choice. Our findings empirically underscore the significant influence of architectural design choices in narrowing the disparity with centralized configurations and effectively handling data heterogeneity.
To this end, the problem of client heterogeneity received significant attention from the optimization community [23, 38, 58, 29, 36, 1, 75, 50, 15].
In parallel to these developments in FL, visual recognition tasks have witnessed remarkable progress in recent years, primarily due to the advancements in deep learning models. These models have achieved state-of-the-art performance when trained on centralized datasets. However, when deployed in an FL setting, these image recognition architectures often exhibit performance degradation. Recent works have investigated the utilization of pretrained models to address the non-iid issue [4, 54]. Qu et al. [57] examined the robustness of four neural architectures across heterogeneous data, introducing the idea of tackling the problem from an architectural standpoint and encouraging the development of "federated-friendly" architectures. The study suggests that self-attention-based architectures are robust to distribution shifts compared to convolutional architectures, making them more adept in FL tasks. In this work, we perform a comprehensive study across 19 visual recognition models from five different architectural families on four challenging FL datasets. We also re-investigate the inferior performance of convolution-based architectures and analyze the influence of normalization layers on the model performance in non-IID settings.
Contributions: This work strives to offer architectural design insights that enable immediate performance improvements without the need for additional data, complex training strategies, or extensive hyperparameter tuning. The key contributions of this work are the following:
- We perform an exhaustive experimental analysis comparing 19 different state-of-the-art (SOTA) models from the five leading computer vision architectural families, including Convolutional Neural Networks (CNNs), Transformers, MLP-Mixers, Hybrids, and Metaformer-like architectures. Some of the architectural families we studied were never previously introduced in FL, demonstrating the feasibility of addressing data heterogeneity in FL from an architectural perspective. These architectures were tested across four prevalent CV federated datasets under the most common non-identicalness sources, such as feature distribution skew, label distribution skew, and domain shift in highly heterogeneous data. Our study, comprehensively illustrated in Figure 1, embodies a step toward more efficient, practical, and robust FL systems.
- Our experiments in highly heterogeneous settings reveal that convolution operations are not inherently sub-optimal in FL applications and Metaformer-like architectures generally exhibit greater robustness in non-IID settings.
- We illustrate that Batch Normalization (BN) adversely impacts performance in the heterogeneous federated learning framework. To counteract this, we empirically demonstrate that
replacing BN with Layer Normalization is an effective solution to mitigate the performance drop. This study case underlines how architectural design choices can significantly influence FL performance.
- We conduct a study on complex networks with four different optimizers, establishing that the application of optimization methods does not yield substantial performance improvements in the context of complex architectures. We argue that altering the architecture, in practical scenarios, offers a more effective and simpler-to-implement choice.
Federated learning (FL) aims to train models in a decentralized fashion, harnessing edge-device computational capabilities to safeguard the privacy of client data. In this setup, clients retain their data on the device, coordinating with a server to develop a global model. A distinct challenge in FL is data heterogeneity, as client data is often non-IID. The issue of addressing this heterogeneity has garnered significant focus from the optimization domain. To enhance the established aggregation method, FedAvg [49], various strategies have emerged. These encompass optimization improvements for federated averaging such as FedAVGM [23], FedProx [38], SCAFFOLD [29], and FedExP [26] and methods to manage data heterogeneity and imbalance like FedShuffle [21], FedLC [34], and MooN [36].
A complementary angle to explore is how the choice of models can elevate FL training outcomes. It has been observed that pretrained models can significantly counteract non-IID challenges [4], [54]. Some previous studies have highlighted the impact of Batch Normalization (BN) on performance drops under heterogeneous settings [39], and few works [22], [24] suggested replacing BN with Group Normalization. Our work is related to Qu et al.'s examination [57] of neural architecture in visual recognition robustness across heterogeneous data splits. Our experiments compare 19 cutting-edge models from diverse computer vision architectural families against demanding heterogeneous datasets and different optimizers.
3 Optimization Methods for Federated Learning
In this section, we present a brief overview of the optimization techniques utilized in our study to evaluate their effectiveness in handling data heterogeneity when applied in conjunction with complex computer vision architectures.
Federated Averaging (FedAVG) [49]: FedAVG was proposed by McMahan et al. (2017) as the standard aggregation method for FL. It offers a straightforward and effective approach where local models are trained on individual devices, and their parameters are averaged to update the global model. However, FedAVG faces challenges in scenarios involving non-IID and unbalanced data, which are prevalent in FL settings.
FedAVG with Momentum (FedAVGM) [23] To address some of the limitations of FedAVG, the addition of a momentum term has been explored. This technique enhances the standard FedAVG method by updating the global model at the server using a gradient-based server optimizer that operates on the mean of the client model updates. The inclusion of momentum facilitates efficient navigation through the complex optimization landscape and accelerates learning, particularly in the presence of sparse or noisy gradients.
FedProx [38]: To handle the heterogeneity in data and device availability in FL, Li et al. proposed FedProx (Federated Proximal). This method introduces a proximal term into the local loss function, preventing local models from deviating significantly from the global model. The proximal term acts as a regularizer, penalizing large deviations from the current global model. This regularization approach improves overall performance and enhances robustness in FL scenarios.
SCAFFOLD [29]: SCAFFOLD is a stochastic algorithm designed to enhance the performance of Federated Averaging (FedAVG) by introducing a controlled stochastic averaging step. This step reduces the variance of gradient estimates by correcting the local updates. SCAFFOLD employs control variates (variance reduction) to address the "client-drift" issue in local updates. Notably, SCAFFOLD requires fewer communication rounds, making it less susceptible to data heterogeneity
(a) CNNs
(b) Transformers
(c) Hybrid
Figure 2: Model Accuracy Analysis across CIFAR-10 splits grouped by the architecture type. The presented results depict the accuracy performance of various models on distinct partitions of the CIFAR-10 dataset. Each split represents a unique partition, progressing from an IID scenario (Split-1) to a non-IID scenario (Split-3). The models under comparison exhibit similar parameters, FLOPS, and approximately matching ImageNet and centralized accuracies. Notably, Metaformer-like architectures, including Metaformers (e), Transformers (b), and MLP-Mixers (d), exhibit a comparatively smaller performance degradation on Split-2 and Split-3. The best-performing models are CA-Former and ConvFormer, which are Metaformers with advanced token mixers.
(d) MLP-Mixer
(e) MetaFormers
or client sampling. Furthermore, the algorithm leverages data similarity among clients, leading to faster convergence.
4 Architectures for Computer Vision
In this section, we provide an overview of the architecture families studied in this work.
Convolutional-based models: We include Resnet [19] and EfficientNet [62] as the study of Qu et al. [57], and extend it by adding a recently introduced architecture, ConvNeXt [45]. EfficientNet and Resnet consist of a stack of convolutional layers, pooling layers, and batch normalization layers. ConvNext was conceived by modernizing ResNet-50 to pair with the performance of vision transformers. The key changes introduced at the component-level are: replacing the typical batch normalization with layer normalization, using inverted residual blocks and larger kernel sizes, and substituting the more common ReLU [51] with the GeLU [20] nonlinearity.
Vision Transformers: ViTs were introduced to the computer vision community by Dosovitskiy et al. [13] after becoming the de-facto standard for natural language processing tasks. The Transformer architecture reconsiders the need for convolutions and replaces them with a self-attention mechanism. ViTs work by splitting the image into non-overlapping patches treated as tokens and re-weighting them based on a similarity measure between each token pair. The main differences with CNNs are the lack of inductive bias of the vanilla self-attention and the global receptive field capable of capturing long-distance dependencies as opposed to the local receptive field of convolutions. As representatives of this family of architectures, we include ViT [60], Swin [44], Swin V2 [43], and DeiT [66].
MLP-Mixers: MLP-Mixers were first proposed by Tolstikhin et al. [64]. It was the first architecture to show that although convolutions and self-attention are sufficient to achieve top performance, they are not necessary. The MLP-Mixers take inspiration from ViT's use of patches for input representation but fully use multi-layer perceptrons instead of the typical self-attention mechanism of Transformers. For our experiments, we make use of the original MLP-Mixer [64], the ResMLP variation [65], and the gMLP [41] which uses a special gating unit to enable cross-token interactions.
Hybrids: Many efforts have been made to incorporate features from CNNs into Transformers and vice versa. We group this set of mixed models under the same 'Hybrids' category. Namely ConvMixer [67], CoATNet [8], and MaxViT [69]. ConvMixer combines a transformer-like patch design with convolution-based layers to mix patch embeddings' spatial and channel locations. CoATNet vertically stacks convolution layers and attention layers. Finally, MaxViT hierarchically stacks repeated blocks composed of multi-axis self-attention and convolutions.
MetaFormers: Although the success of transformer-based architectures was longly attributed to self-attention, recent works [73, 74] have started to identify the overall structure of the transformer block as the critical component in achieving competitive performances. Based on this hypothesis,
Metaformer is constructed from Transformers without specifying the design of the token mixers. If the token mixer is self-attention or spatial MLP, the architecture becomes a Transformer or an MLP-Mixer.
Although most of the Transformers and MLP-Mixers formally fall into the 'Metaformers' structure, we group in this family only the works published by the authors under this name. Within the networks selected, we have some working with simple operators in the token mixers, published with the explicit intent of pushing the limit of the Metaformer structure. Respectively, RandFormer, IdentityFormer [74], and PoolFormer [73] have a random matrix, an identity mapping, and a pooling operator as the token mixer. Moreover, ConvFormer [74] applies depthwise separable convolutions as the token mixer, and CAFormer [74] further combines it with vanilla self-attention in the deep stages. We also add to this category RIFomer [70], an architecture that optimizes the latency/accuracy trade-off by using re-parametrization [12] to remove the token mixer without significantly compromising the performance.
5 Experimental Setup
Models: This study aims to enrich the FL landscape by introducing a variety of cutting-edge architectures for computer vision. Within the FL framework, we conduct a comparative analysis of models that have similar parameters and FLOPS and roughly equivalent accuracy in centralized settings to the benchmark model, ResNet-50. Ensuring fairness, all the models under consideration fall within a similar parameter range (21-31M) and floating point operations per second (FLOPS) range (4-6G) $^{1}$ . Each model exhibits a comparable centralized test accuracy on ImageNet-1K [11].
Unless explicitly stated, all models are pretrained on ImageNet-1K and subsequently finetuned on the specified dataset.
Federated learning: In an effort to ensure that the insights derived from this study remain independent of the training heuristics, we have standardized and simplified the training strategy across all architectures. Our approach involves performing FedAVG for a specific number of communication rounds to reach convergence: 100 rounds for both the CIFAR-10 and Fed-ISIC2019 datasets, 30 rounds for the CelebA dataset, and 200 rounds for the GLD-23K dataset. With the exception of the GLD-23K dataset, where we conduct five local steps, each round of local training encompasses one local epoch. The training regimen employs the SGD optimizer with an initial learning rate of 0.03, cosine decay, and a warm-up phase of 100 steps. We set the local training batch size at 32 and apply gradient clipping with a unitary norm to stabilize the training process.
The numerical values presented in this paper represent the average results derived from three repetitions of the same experiment, each with different initializations, to ensure the robustness and reproducibility of our findings. Regarding optimization, we conduct hyperparameter tuning for each model based on the set of values suggested by the respective authors in their original publications. The results reported in this paper reflect the optimal outcomes derived from the hyperparameter search. A more in-depth discussion is given in the supplementary material.
Datasets: We conduct experiments on four different datasets, CIFAR-10 [31], CelebA [46], FedISIC2019 [68, 5, 7] and Google Landmarks Dataset v2 (GLD-23K) [72]. Across the experiments, all the images are resized to a standard dimension of  $224 \times 224$ . The data heterogeneity of the different datasets used in our experiments is characterized in terms of Label Distribution Skew (LDS), Data Distribution Skew (DDS), and Feature Distribution Skew (FDS). The LDS is characterized by a heterogeneous number of samples per class for each client. In the DDS, the amount of data per client varies. Lastly, the FDS is defined by varying image features per class across clients due to changes in acquisition sensors, image domains, user preferences, geographical locations, etc. (see supplementary materials for more details).
The CIFAR-10 dataset comprises 60,000 32x32 color images in 10 distinct classes representing common objects. Since CIFAR-10 is not inherently a federated dataset, we artificially simulate three federated data partitions across clients following the approach outlined in [57]. The original test set is retained as the global test set, and 5,000 images are set aside from the training set for validation
<sup>1</sup>The only exception in terms of model size is MLPMixer-S/16. However, we have opted to include this model in its smallest configuration provided by the authors, as it constitutes the inaugural exemplar of MLP-Mixer architectures.
CIFAR-10	FED-ISIC2019
ARCHITECTURE	RESNET-50	CONVNext-T	CAFORMER-S18	RESNET-50	CONVNext-T	CAFORMER-S18
FEDAVG [49]	59.9 ± 1.87	93.5 ± 0.36	97.1 ± 0.35	67.6 ± 1.59	78.6 ± 0.29	83.1 ± 0.49
FEDAVGM [58]	80.2 ± 0.49	93.4 ± 0.48	97.0 ± 0.14	66.2 ± 2.16	79.0 ± 0.52	83.8 ± 0.37
FEDPROX [38]	79.0 ± 0.41	93.4 ± 0.16	97.1 ± 0.25	68.0 ± 1.93	79.4 ± 1.13	83.0 ± 0.79
SCAFFOLD [29]	78.1 ± 0.06	94.3 ± 0.30	96.9 ± 0.21	62.2 ± 0.26	76.0 ± 1.24	77.1 ± 0.61
Table 1: Accuracy of ResNet-50, ConvNext-T, and CAFormer on CIFAR-10 and Fed-ISIC2019 across four different optimization methods. The average test accuracy of the model is displayed along with the standard deviation of three repetitions of the experiments.
purposes, resulting in a revised training dataset of 45,000 images. We employ the Kolmogorov-Smirnov statistic (KS) to simulate one Independent and Identically Distributed (IID) data partition, namely Split-1 (KS=0), ensuring balanced labels per client, and two non-IID partitions, Split-2 (KS=0.5) and Split-3 (KS=1), with label distribution skew. In Split-2, each client has access to four classes and does not receive samples from the remaining six classes. In Split-3, each client strictly sees samples from two classes only. Therefore, this is the most challenging partition, featuring a pronounced class imbalance. Each data partition has five clients.
The CelebA dataset is a large-scale face attributes dataset featuring over 200,000 high-resolution celebrity images. Each image is annotated with 40 binary labels that denote the presence or absence of distinct facial attributes. In this study, we employ the federated version of the CelebA dataset proposed by the LEAF benchmark [3]. In line with the methodology outlined in [57], we test the models on a binary classification task, determining whether a celebrity is smiling. The dataset is partitioned across 227 clients, each representing a specific celebrity and collectively accounting for 1,213 images. The dataset presents significant challenges due to the large-scale setting involving numerous clients. Each client has access to only a handful of samples, five per client on average, and in some instances, data from only one class is available.
The Fed-ISIC2019 dataset contains skin lesion images collected from four hospitals. We adopt the partitioning strategy proposed by [63], which is based on the imaging acquisition systems utilized. Given that one hospital employed three distinct imaging technologies over time, the federated dataset is distributed across six clients, encompassing 23,247 images. The task involves classifying the images into eight melanoma classes with high label imbalance across clients. Additionally, the dataset exhibits quantity skew with disparity in the number of images per client. The largest client possesses more than half the data, whereas the smallest has only 439 samples.
The Google Landmark-V2 dataset is a vast collection with over 5 million images of global landmarks. We utilize the federated GLD-23K partition, involving 233 smartphone users as clients contributing 23,080 images across 203 classes. These datasets exhibit a geographically influenced, long-tailed distribution, with some landmarks having many representations while others have very few, mirroring real-world disparities. Also, photographers' images present unique distributional characteristics tied to their habits and locations, resulting in large diversity and imbalance across clients.
6.1 Optimizers and Complex Networks in Federated Learning
An effective learning system typically requires both an optimized architecture and an effective optimization technique to shape the learning process and resultant model. The federated learning community has focused considerably on optimization techniques to address challenges related to non-IID data and the need to balance statistical diversity and leverage similarities in client data for central model training. However, most recent studies [48] rarely present results incorporating the latest architectures developed by the computer vision community. Many of these publications on vision tasks present their work on simple and shallow networks ranging from a few fully connected layers to the ResNet-50 model. This discrepancy highlights a significant chasm between the research focus of federated learning and state-of-the-art computer vision architectures.