An Architecture Built for Federated Learning: Addressing Data Heterogeneity through Adaptive Normalization-Free Feature Recalibration
Authors: Vasilis Siomos
Jonathan Passerat-Palmbach
Giacomo Tarroni
Venue: Accepted by TMLR ()

Abstract: Federated learning is a decentralized collaborative training paradigm preserving stakeholders’ data ownership while improving performance and generalization. However, statistical heterogeneity among client datasets degrades system performance. To address this issue, we propose Adaptive Normalization-free Feature Recalibration (ANFR), a model architecture-level approach that combines weight standardization and channel attention to combat heterogeneous data in FL. ANFR leverages weight standardization to avoid mismatched client statistics and inconsistent averaging, ensuring robustness under heterogeneity, and channel attention to produce learnable scaling factors for feature maps, suppressing inconsistencies across clients due to heterogeneity. We demonstrate that combining these techniques boosts model performance beyond their individual contributions, by improving class selectivity and channel attention weight distribution. ANFR works with any aggregation method, supports both global and personalized FL, and adds minimal overhead. Furthermore, when training with differential privacy, ANFR achieves an appealing balance between privacy and utility, enabling strong privacy guarantees without sacrificing performance. By integrating weight standardization and channel attention in the backbone model, ANFR offers a novel and versatile approach to the challenge of statistical heterogeneity. Extensive experiments show ANFR consistently outperforms established baselines across various aggregation methods, datasets, and heterogeneity conditions. Code is provided at https://github.com/siomvas/ANFR.

Published in Transactions on Machine Learning Research (09/2025)
An Architecture Built for Federated Learning: Addressing Data Heterogeneity through Adaptive Normalization-Free
Feature Recalibration
Vasilis Siomos
musilassiomosscclyygeeeeee.....
Department of Computer Science, City St. George's, ss Universily of London
Jonathan Passorat-Palmban
Department of Computer Science, City St. George's, University of London BioMedIA, Department of Computing, Imperial College London
Giacomo Tarroni
Deportment of Computer Science, City St. George s, University of London BioMedIA, Department of Computing, Imperial College London
Reviewed on OpenReview: https: openreviee. net/ forum? id=GtdYFLsblb
Abstract
Federated learning is decentralized collaborative training paradigm preserving stakehold- ers' data ownership while improving performance and generalization. However, statistical heterogeneity among client datasets degrades system performance. To address this issue, we propose Adaptive Normalization-free Feature Recalibration (ANFR), a model architecture-level approach that combines weight standardization and channel attention to combat. heterogeneous data in FL. ANFR leverages weight standardization to avoid mis- matched client statistics and inconsistent averaging, ensuring robustness under heterogene- ity, and channel attention to produce learnable scaling factors for fcature maps, suppressing inconsistencies across clients due to heterogeneity. We demonstrate that combining these techniques boosts model performance beyond their individual contributions, by improving class selectivity and channel attention weight distribution. ANFR works with any aggre- gation method, supports both global and personalized FL, and adds minimal overhead. Furthermore, when training with differential privacy, ANFR achieves an appealing balance between privacy and utility, enabling strong privacy guarantees without sacrificing perfor- mance. By integrating weight standardization and channel attention in the backbone model, ANFR offers novel and versatile approach to the challenge of statistical heterogeneity. Extensive experiments show ANFR consistently outperforms established basclines across various aggregation methods, datasets, and heterogeneity conditions. Code is provided at https://github. htt///soovas/AANEP
1 Introduction
Federated learning (FL) (McMahan et al. 2017 is decentralized training paradigm enabling clients to jointly develop model without sharing private data. By preserving data privacy and ownership. FL holds promise for applications in healthcare, finance, and mobile devices. A fundamental challenge in FL sta- tistically heterogeneous, i.e. non-independent and identically distributed (non-lID) client datasets, as they can degrade the performance of the global model and hinder convergence (Li ct al. 2020b Hsu et al., 2019). Addressing this is critical for FL's success in real-world scenarios. Most prior research focuses on aggrega- tion methods to compensate for this issue, overlooking how model architecture affects performance under

---

Published in Transactions on Machine Learning Research (09/2025)
heterogeneity. More specifically, Batch Normalization (BN) (loffe & Szegedy 2015) hinders performance in heterogeneous FT, due to mismatched client-specific statistics and inconsistent parameter averaging (Wang) et al. 2023 Guerraoui et al.) 2024). In response, using other feature normalization methods like Group
2016) has been frequent in Normalization (GN) (Wu& & He 2018) and Layer Normalization (LN) (Ba et al.)
FL research (Hsieh et al. 2020: Reddi et al.] 2021 Wang al.; 2021 Du et al. 2022]. These alternatives
slow convergence and reduce performance compared to BN (Chen & Chao 2021
Tenison et al. 2023 Zhong
eet al., 2024). Previous work on designing models specifically tailored to combat heterogeneity is limited to only
Xu et al. 2024), leaving a research gap.
We address this gap in the image domain by proposing Adaptive Normalization-free Feature Recalibration (ANFR), an architecture-leeel approach designed to enhance robustness in FL under data heterogeneity. ANFR combines weight standardization (Qiao et al. 2020) with channel attention (Hu et al. 2018 to di- rectly tackle the challenges posed by non-IID data. Weight standardization normalizes convolutional layer weights instead of activations, avoiding reliance on mini-batch statisties, which problematic in FL. This reduces susceptibility to mismatched statistics and inconsistent averaging. Channel attention generates learnable scaling factors for feature maps, suppressing features that are inconsistent across clients due to heterogeneity, and emphasizing consistent ones. By integrating channel attention with sright-standardizel models, ANFR enhances the model's ability to focus on shared, informative features across clients. This synergy boosts performance beyond the individual contributions of these components, enhancing class selec- tivity, and optimizing channel attention weight distribution. key advantage of ANFR is its versatility; it can be seamlessly integrated with any aggregation method and is highly effective in both global and person- alized FL scenarios, while adding minimal computational overhead. Furthermore, when privacy is concern, our achicves a noteworthy balance between the competing demands of privacy and utility. This enables the implementation of strong privacy guarantees through differential privacy, without substantial sacrifice in the model's overall performance.
We validate the effectiveness of ANFR through extonsive experiments on a diverse set of datascts and tasks, including medical imaging and natural image classification, multi-class classification, and cross-device sce- narios, under various types of data heterogeneity. The results show that ANFR consistently outperforms established baselines across different aggregation methods, datasets, and heterogencity conditions. By focus- ing on architectural components, our approach complements advances in aggregation strategies and addresses a crucial gap in FI. research. The proposed model offers robust and flexible solution to the challenge of statistical heterogeneity, contributing to the advancement of federated learning by improving performance, stability, and grivacy-preserving capabilitics.
2 Related Work
Since McMahan el al.] (2017) introduced FL, most. research has focused on developing aggregation algorithms to address challenges like data heterogeneity. In global FL (GFL), methods such as proximal regularization [Li et al.] 2020a) and cross-client variance reduction (Karimireddy et al. 2020) aim to reduce client drift. Techniques like discouraging dimensional collapse through correlation matrix norm regularization (Shi et all, 2023), adopting relaxed adversarial training (Zhu et al. 2023 and performing amplitude normalization in frequency space (Jiang et al., (022) have also been proposed. Other recent ideas are constructing global pseudo-data to de-bias local classifiers and features (Guo et al. 2023), introducing concept drift-aware adaptive optimization (Panchal et al. 2023), and hyperbolic graph manifold regularizers (An et al. 2023). Tn personalized FL (pFL), personalizing layers of the model can mitigate heterogeneity. The simplest approach shares all model parameters except the classification head (Arivazhagan al. 2019). More advanced hods replace lower layers and mix higher ones (Zhang et al. 2023) or adjust mixing ratios based convergence rate approximations (Jiang et al. 2024h. While these algorithmic approaches have advanced both GFL and LFL, they often overlook the impact of the underlying architecture on performance.
We address this gap by exploring how model components can enhance FL performance. This is orthogonal to algorithmic advancements, representing a crucially underdeveloped area. Previously, Qu et al. (202) found using vision transformers instead of convolutional networks increased performance. Studies by Pieri et al. 2023) and Siomos et al.] (2024) evaluated different architectures and aggregation methods, showing that
2

---

Published in Transactions on Machine Learning Research (09/2025)
changing the architecture, rather than the aggregation method, can be more beneficial. These works did not design models specifically tailored to combat heterogeneity. More recently, Xu et al. (2024) investigated the impact of CNN micro-design choices, such as activation functions and block structure, in FL,, and proposed FedConv, an architectural design based on ResNet informed by these investigations. Our method follows in this line of work, namely integrating architectural components that enhance robustness across diverse client. distributions into the model, directly addressing data heterogeneity.
The normalization layer has been focal point of component examination as Batch Normalization (BN) [offf & Szegedy] 2015) has been shown both theoretically (Li et al., 2021 Wang et al. 2023 and em- pirically (Hsich et al. 2020 Du et al., 2022 Guerraoui et al.) 2024) to nogatively impact performance in heterogeneous FL. Mismatched local distributions lead to averaged batch statistics and parameters that fail to accurately represent any source distribution. The primary approaches addressing this issue are modifying the aggregation rule for the BN layer or replacing entirely. Some methods keep BN parameters local (Li ct al., 2021 Andreux et al. 2020) or stop sharing them after cortain round (Zhong et al. 2024). Others replace batch-specific statistics with shared running statistics when normalizing batch inputs match local statistical parameters (Guerraoui et al.] 2024) or leverage layer-wise aggregation to also match associated gradients (Wang et al. 2023). These mcthods rely on decently sized batches to accurately approximate statistics and are incompatible with differential privacy. To replace BN, Group Normalization (GN) (Wu & He 2018) has been frequently used (Hsich et al., 2020; Reddi et al.] 2021 Wang et al.; 2021 since it does not rely on mini-batch statistics. However, tuning the number of groups in GN is required to maximize effectiveness and Du et al. (2022) showed that Layer Normalization (LN) (Ba et al., 2016) performs better than GN in some settings. Separate studies have shown both GN and LN offer inconsistent benefits over
BN, depending the characteristics and heterogencity of the datasct Tenison al. 2021
2023 Chen & Chao,
Zhong et al. 2024).
We circumvent these issues by applying weight standardization (Qiao et al. 2020 to normalize the weights of the model instoad of the activations. Inspired by Brock et al. (2021a), who showed that such Normalization- Free (NF) models can train stably and perform on par with BN in centralized learning, we explore this concept in FL. Previously, Zhuang & Lyu (2024) proposed an aggregation method specific NF models for multi-domain FL with small batch sizes. Similarly, SSomos et al.) (2024) showed that NF-ResNets improve upon vanilla ResNets under different initialization schemes and aggregation methods, while [Kang et al. 2024) proposed personalized aggregation scheme that replaces each BN layer with weight normalization Salimans & Kingma 2016) followed by a learnable combination of BN and GN. Additionally, our method adaptively recalibrates the resulting feature maps using channel attention modules, such as the Squecze-and- Excitation block (Hu et al., 2018). By doing so, the model can focus more on relevant features across clients, effectively addressing data heterogeneity. [Zheng et al. (2022) previously explored channel attention for pFL, proposing modified channel attention block that is kept personal each client. Unlike previous methods limited to specific aggregation strategies settings, our approach can complement any heterogeneity- focused aggregation method, is effective even with large batch sizes, and supports various attention modules. By integrating weight standardization with channel attention, ANFR provides a robust and flexible solution to data heterogeneily in FL, overcoming limitations of activation normalization techniques and complementing aggregation methods. Table presents a tabular comparison of ANFR with related work to make it easier to understand our contributions.
3 Adaptive Normalization-Free Feature Recalibration
3.1 Background and Notation
We consider FL setting with C' clients, each owning dataset of image-label pairs D; {xx yb) and optimizing a local objective (,) = B(x.))-D, y; f)), where is a loss function and @ the model parameters. Heterogeneity among D; can degrade the global model performance and slow convergence (Kairouz el al., 2021). In this study, we modify the backbone model to address this, As they are the most widely used family, and they perform better or on par with others (Picri et al.) 2023 Siomos et al., 2024), we focus specifically on convolutional neural networks (CNNs). Let X E B mxCmmxxxx represent a batch of B image samples with Cm channels and dimensions HxW. For convolutional layer with weights W and a kernel size of 1, the
3

---

Published in Transactions on Machine Learning Research (09/2023)
Client Client 1
CR
CA
CNR Blue Detector
CA
Edge Detector
Scr
X
Illustratin how Channel Attention (CA) can boost CR and suppress CNR. Left: the two clients have beterogeneous datasets. Middle: an edge detector robust to this feature shift; the activations are consistent for both clients. Right: blue detector is robust and its activations cause conflicting gradients.
Figure 1:
Table 1: Comparison of desirable attributes between study and related work. O condition is not met, inconsistently met, and fully met, respectively. FedConv and ANFR fill a gap in the literature as they simultaneously work in GFL, pFL, are compatible with any aggregation method and offer a robust increase in performance. ANFR additionally is very adept at handling private FT. scenarios (compatible with DP), as we will demonstrate.
symbolize
Sccnario
Aggregation Agnostic
Compatible with DP
Performance
Method
GFL
pfL
Increase
FodBN
Li ct al.
2021
O
FixBN
Zhong et al.
2024)
O
O
FBN Guerraoui et al.
2024
O
O
ChannelFed Zhenn et al.
2022
FedWon (Zhuang & Lyu
2024
GN LN Ww & He.
2018
Ba al. 2016
FedConv Xu al.
2021
ANFR (Ours)
outputs are given by:
A X W >
Ctn W. X ;ei
(1)
with the dimensions of A being [B, Couts H, W] and those of W [Cout, Cinl In typical CNNs, the activations are then normalized:
À 7 (A, p2) + B, with: μs
Σ Ak
of
Σ(A, (Ax -1,)2 kE
(2)
|S
S
kES:
indexing vector and S, is the set of pixels over which pi o, computed. BN computes statistics along the (B, W) axes, LN along HC
where B,y Cout learnable parameters, i (in ic, (ir,,w)
HCH.WW W), and GN along (C, H, W) separately for each of g groups of channels. Channel Attention (CA) mechanisms, like the Squeeze-and-Excite modeling inter-channel relationships. The channel descriptor Z E BBBCout is obtained via Global Average Pooling (GAP):
(SE) block (Hu et al. 2018 recalibrate feature responses by
Z = (HW) - h,.0 H,W A...h.w
(3)
This descriptor is then non-lincarly transformed to capture dependencies between channels; in SE blocks this is done via the learnable weights W1 H
Cout xCout and W2 E Com Cout where is dimensionality reduction ratio: S (W25 (W1Z)) where S E BBCaat is the sigmoid function and the ReLU function, yielding per-channel scaling factors S which applied to the normalized activations À S A.
4

---

Published in Transactions on Machine Learning Research (09/2025)
3.2 Effect of Normalization on Channel Attention
In the presence of data heterogeneity, CA can suppress fealures sensilive to client-specifie variations and emphasize consistent ones. In earlier layers, A consists of responses to filters detecting low-level features like colors and edges, while in later layers it contains class-specific features (Zeiler & Fergus 2014). For the sake of explaining how CA impacts heterogeneous FL, we virtually partition filters into two distinct groups: those eliciting consistent features (Cn) and inconsistent ones (Cw Figure] 1 Tillstrates an example. Both clients have images of airplanes and cars; Client I's images have predominantly blue backgrounds, while Client 2's images have different backgrounds. Under this feature shift, edge-detecting filters produce consistent responses across both clients, thus belonging to CR, whereas filters sensitive to specific colors like blue activate differently across clients, forming CNR- While both activation types are informative locally, inconsistent activations from CNR cause conflicting gradients during FL training. This motivates our use of CA in this context: during training, CA can assign higher weights to Acn and lower weights to ACwR without prior knowledge of which features belong to each set. The resulting adaptive recalibration aligns feature representations across clients, reducing gradient divergence and improving global model performance.
While CA mitigates the locality of convolution by accessing the entire input via pooling (Hu et al. 2018): if the normalization of A is ill-suited to heterogeneous FL, the input to 3 becomes distorted, leading to sub-optimal channel weights:
H.W Cmn W. X e.h,,
Z O.HW ΣΣ N.W
p4y
a + +
(4)
Activation normalization techniques suffer from this issue. BN is known to be problematic in heterogeneous settings for two reasons: mismatched client-specific statistical parameters lead to gradient divergence separate from that caused by heterogeneity- between global and local models (Wang et al. 2023): and biased rumning statistics are used at inference (Guerraoui et al.] 2024). Both contribute to well-established performance degradation (Li et al., 202 Du et al. 2022). Since μ4; and σ; depend batch-specific statistics, AA varies across clients due to local distribution differences, leading to inconsistent channel descriptors, which in turn results in non-ideal channel weights. Aside from data heterogeneity, BN needs sufficient batch sizes to estimate statistics accurately, and is incompatible with differential privacy; these are limiting factors in resource-constrained and private FL scenarios. GN and LN also have drawbacks: GN normalizes within fixed channel groups, which may not align with the natural grouping of features, limiting its effectiveness under heterogeneity. LN assumes similar contributions from all channels (Ba et al., (0166), which is generally untrue for CNNs, and clashes with our goal of reducing the influence of Aevn- Crucially, both normal- ize across channels to produce Pi, o, This introduces additional channel inter-dependencces in 4 thus interfering with extracting representative channel descriptors.
3.3 Adaptive Normalization-Free Feature Recalibration
To address these problems, we propose applying CA after normalizing the convolutional weights instead of the activations using Scaled Weight Standardization (SWS) from NF models (Brock et al. 2021a), which adds learnable affine parameters to weight standardization (Qiao et al. 2020):
Cin Σ
Cin
W Couts Yoff.rout (W eoossfim (roaa) Meam
W
Σ ( vomt. Meon )2
Cin
Cout
Oin
Here, Yefl = 7gr7/V/Ca. incorporates a learnable scale parameter g and fixed scalar depending on the networks' non-linearity. We replace the normalized activation À with A' = X W 3. From (5) we observe that SWS does not introduce a mean shift (E[A') = E = 0), and preserves variance (Vr(AA)= Var(A)) for the appropriate choice of y, allowing stable training. By replacing normal convolutions with the ones described by 5 and following the signal propagation steps described in Brock et al. (2021a), we can train
stable CNNs without activation normalization. We term this combination of weight standardization and channel attention Adaptive Normalization-free Feature Recalibration (ANFR). The input to (3) when using
5

---

Published in Transactions on Machine Learning Research (09/2025)
Before FL Training
After FL Training
4.5
4.5
3,5
SE-ReoNet BN-ResNet
NF-ResNeet
3.5
SE-ReaNet BN-ReeNer
NF-RoNa
3.0
ANFR
3.0
2.5-
2.5
2.0
2.0
1.5
1.0
1.0
0,5
0.5
0.3
0.6 Selectivity Index
0.7
1.0
0.1
0.7
0.8
1.0
Selectivity Index
Figure 2: Left: CSI distributions before FL. training, queried after the last CA module. Both normalizations (BN and SWS) show similar behavior, and CA has a minor impact. Right: after FL training, CA increases class selectivity, especially in conjunction with SWS in ANFR.
Before FL Training
A fler FL Training MMA yorui bf
Be-ResNet
Choraal Weight
0.8
Naglace
Bini
Dest
Froa
Slip
0.8
Asplasz
Bini
Dear
Frog
Slip
Autemobiie
Track
0.7
Automobile
Horeo
Tnack
500
750
1250
1500
2000
250
750
1000
12.50
1500
1750
Charne Waggh
Airplane
Ship
Arplans
Deor
Pngg
ANER
0.9
Hune
0.9
Autonobile
0.8
0.8
0.7-
0.7
250
500
750
1000
1250
1500
1750
2000
250
750
1000
1500
1750
2000
Channel Index
Channel Index
Figure 3: Top: weights of the last CA module for SE-ResNet-50. Bottom: same for ANFR-50. Left: before FL training, CA provides diverse signal varying across classes and channel indices to both models. Right: after FL training, the CA module in SE-ResNet degenerates to identity, In ANFR, CA shows increased variability as it works to combat heterogencity.
ANFR is:
H.W Cin
H.W Cin soh + B
z ANFR
Yeff
ΣΣ W W.o..eee keh
Σ 6220
(6)
HWW
HWW
h,W0
Comparing 4 and 6 we note several advantages of ANFR. First, Or and are computed from convolutional weights, not the activations. Since weights are initialized identically and synchronized during FL, these weight-derived statistics are consistent across clients. Morcover, the second term of (6 now captures statistics of the input before convolution, providing an additional calibration point for CA and bypassing the effect of CNR- By applying CA after SWS, we ensure channel descriptors are not distorted by batch-dependent statistics cross-channel dependencies introduced by activation normalization. This allows CA to adjust channel responses effectively, improving the model's capacity to learn stable feature representations that are consistent across clients with diverse data distributions. Therefore, the combination of SWS and CA overcomes the drawbacks of traditional normalization methods in federated learning, providing a novel and effective solution for improving model performance in the presence data variability. Lastly, we note ANFR operates at the model level and inherits the theoretical convergence guarantees of the aggregation method it is used with.
3.4 Mechanistic Interpretability Analysis
Next, we conduct mechanistic interpretability analysis comparing the effects of BN and SWS on class selectivity and attention weight variability to further substantiate the effectiveness of integrating CA with

---

Published in Transactions on Machine Learning Research (09/2025)
SWS. We examine how well the ANFR model discriminates between classes beford and after training on the heterogeneous 'split-3' partitioning of CIFAR-10 from Qu et al. (002)). This evaluation helps understand how our method improves class discriminability under data heterogeneity. We isolate the effect of different components by comparing ANFR (using SWS with CA), BN-ResNet (using BN), NF-ResNet (using SWS without CA), and Net. (using CA with BN). Class selectivity is quantified by the class selectivity where
index (CST) (Morcos et al., 2018), defined for each neuron as CSI (Fmax /fwmm////////////////mmmm Pmax is the class-conditional activation that elicits the highest response and y-max is the mean activation for all other classes. A right-skewed CSI distribution indicates higher class selectivity, crucial for effective classification under heterogeneous data. Lastly, we examine the distribution of attention weights, like done in Wang et al. (2020), for models using CA, to understand its contribution to class discrimination.
Figure 2 shows CSI distributions for the last layer before the classifier, where class specificity is maximized in CNNs. Before FI, training, incorporating CA in SE.-ReNeet slightly increases class selectivity compared to BN-RosNet. Combining CA with SWS in ANFR shows negligible change in class selectivity compared to Ne-ResNet, indicating CA's 's minimal impact al this stage. However, afler training on heterogeneous data, we observe a notable shift: BN reduces class selectivity (compared to before training), evidenced by left-skewed distributions for BN-ResNet and SE-ResNet. Adding CA incrcases class selectivity for both normalization methods, but due to receiving inconsistently normalized inputs 4 cannot fully mitigate BN's i negative effect. The ANFR model, however, shows a significant increase in class selectivity compared to NF-ResNet, with strong class selectivity (CSI>0.75) units nearly doubling from 111% to ~21%. This improvement manifests only after FL training, indicating that combining CA and SWS in ANFR enhances the model's ability to specialize and discriminate classes under data heterogeneity.
In Figure 3 we use the variability of attention weights across channels and classes indicator of adapta- tion: high variability suggests CA is actively re-weighing features to adapt to different class characteristics. Before FL training (left panel), both SE and ANFR models display high variability, as, when heterogeneity is not a factor, CA provides a diverse and informative signal for both activation and weight normalization. After FL training (right panel), the altention mechanism of Se-ResNet turns into an identity operator, with attention weights converging to across all channels and classes, meaning SE-ResNet fails to preserve the discriminative power of CA under hetcrogeneity. In contrast, ANFR maintains high variability in CA weights across channels and classes. This sustained variability implics that CA remains active and continues to provide clas--iiscininntive signals when combined with weight standardization.
These insights support our design choices. BN's adverse effects in heterogeneous FL are highlighted by dimin- ished class selectivity and inactive CA in SE-ResNet, while ANFR maintains and improves class selectivity, demonstrating that integrating CA with weight standardization effectively counters data heterogeneity. The enhanced class selectivity in ANFR correlates with improved downstream performance in heterogeneous FL settings, as we show in Section 4 Additional details and extended CSI and attention weight results from other layers are presented in Appendix D
4 Experiments
4.1
Experimental Settings
Datasets.
We evaluate our approach on five classification datasets, including Fed-ISIC2019 (Ogier du Terrail 2022 containing dermoscopy images from 6 centers with 8 classes where label distribution skew and heavy quantity skew is present; FedChest, a novel chest X-Ray multi-label dataset with 4 clients and 8 labels with label distribution skew and covariate shift; highly non-IID partitioning of CIFAR-10 (Krizhevsky
et al.
et al.
2009) which simulates heavy label distribution skew across 5 clients using the Kolnogorov-Smire
KS) 'split-2' as presented in Qu et al. 2022);; CelebA (Liu et al. 2015) from the LEAF suite (Caldas et al. 2018), a binary classification task in a cross- setting with a large number of clients, covariate shift and high quantity skew; and FedPathology, a colorectal cancer pathology slide dataset with 9 classes derived from Kather et al. 2 featuring challenging concept drift as the images, which we do not color-normalize, were produced using two different staining protocols. FedChest contains images from PadChest (Bustos et al.
All networks are pre-trained on ImageNet.

---

Published in Transactions on Machine Learning Research (09/2025)
2020 CXR-14 (Wang et al. 2017 and CheXpert Irvin et al. (0019 which present or more of 8 common d disease labels. For FedPathology, used for DP training in Section 4.4 Dirichlet distribution sampling (Hsu ett a 2019) with a=0.5 is used to simulate a moderate label distribution skew and partition the data to 3 clients. Each task covers a different aspect of the multi-faccted problem of data heterogeneity in FL, including different domains and sources of heterogeneity, to provide a robust. test bed. More details are presented Appendix A.1 including instructions to replicate FedChest C.1.
Compared models. In Sections 4.2 4.7
we compare ANFR with typical ResNet (utilizing BN), ResNet where BN is replaced by GN,, SE-ResNet (Hu et al., 2018), and a NE-ResNet. This selection isolates the effects of our architectural changes compared to using BN, using its popular substitution GN, and using weight standardization and CA separately. We choose a depth of 50 layers for all models to balance performance with computational expense. In Section 4.8] we broaden our investigation to include models from different familics. All models used Section 4 are pre-trained on ImageNet (Russakovsky et al. 2015) using timm [Wightman 2019), but additional experiments with randomly initialized models are presented in Appendix B.2 ANFR follows the structure of NF-ResNet, with the addition of CA blocks in the same
position as NeeRe Except for Section 4.5 employ Squeeze-and -xudddittttiinn (Hu et al., 2018) as the attention mechanism. Additional model and computational overhead details are provided in Appendix A.3
Evaluated methods. We use 4 global FL (GFL.) and 3 personalized FL (pFL.) aggregation methods as axes of comparison for the models, each representing a different approach to model aggregation: the seminal FedAvg (McMahan et al. (0117) algorithm, FedProx (Li et al. 2020a), which adds a proximal loss term to mitigate drift between local and global weights, SCAFFOLD (Karimiredyy et al.) 2020), which corrects client drift by using control variates to steer local updates towards the global model, FedAdam (Reddi et al. 2021 , which decouples server-side and client-side optimization and employs the Adam optimizer [Kingma & Ba, 2017) at the server for model aggregation, FedBN (Li et al., 2021) which accommodates data heterogeneity by allowing clients to maintain their personal batch statistics, and by construction is only applicable to models with BN layers, FixBN (Zhong et al.) 2024), which fixes the batch statistics halfway through training in order to combat. some of BN's issues, and FedPer (Arivazhagan et al., 2019) which personalizes the FL process by keeping the weights of the classifier head private to each client. We note our proposal is an architectural one which is aggregation method-agnostic, thus we sclected these widely known aggregation methods to represent a spectrum of strategies, from standard averaging to methods addressing client drift and personalization. This provides robust comparison concentrated on the model architectures.
Evaluation metrics. For Fed-ISIC2019, we report the average balanced accuracy due to heavy class- imbalance as in (Ogier du Terrail et al. 2022). For FedChest, a multi-label classification task with imbalanced classes, we report the mean AUROC on the held-out test. in this section and more metries in Appendix C.2 We report the average accuracy for the other datasets. In pFL settings, the objective providing good in-federation models so we report the average metries of the best local models, as suggested in (Zhang et al. 2023
Implementation details. We select parpee-psreestrss for each dataset by tuning the RN-ResNet (using the ranges detailed in Appendix A.2) and then use the same parameters for all models. This means the results in Section 4.2 are a conservative floor of the improvements that can be achieved, and tuning for ANFR can further increase improvements, as we will show in the next section. In Fed-ISIC2019 clients use Adam with a learning rate of 5e-4 and a batch size of 64 to train for 80 rounds of 200 steps. This setup is distinct from the one used in Ogier du Terrail et al. (2022) resulting in performance improvements for all models. In Appendix B.1 we provide additional results using the original settings. In FedChest clients use Adam with a lcarning rate of 5e-4 and a batch size of 128 to train for 20 rounds of 200 steps. For DP-training in FedPathology, we set the probability of ion leakage 8 to Dal, as is common, the noise multiplier to 1.1, the gradient max norm to 1.0, and train for 25 rounds, which is the point where the models have expended a privacy budget of s=1. For CelebA and CIFAR-10 we follow the settings of Qu et al. (2022); ; Pieri et al. (2023) which were tuned by the authors. All experiments are run in a simulated FL environment with NVFLARE (Roth et al., 22022) and Pyrch (Paszke et al. 2019), using 2 NVIDIA A100 GPUs for training. We report the mean and standard deviation across 3 seeds.
8

---

Published in Transactions on Machine Learning Research (09/2025)
Table 2: Performance comparison across all architectures under different global FL aggregation methods and different datascts. Best in bold, second best underlined. ANFR consistently outperforms the basclines, often by wide margin.
Architecture
Dataset
Method
BN-ResNet
GN-ResNet
SE-ResNet
NF-ResNet
ANFR (Ous)
FedAvg
66.01
65.09:10.42
65.29+1..2
72.49-10.60
7..78
FedProx
66.4910.41
66.51#1.21
66.2910.63
71.28412.1
75.6110.71
Fed-ISIC2019
FedAdam
65.88:0.67
64.6
65.1811.90
69.96:10.14
73.02:0.98
SCAFFOLD
65.41=0.722
68.84+0.44
68.99:0.18
7..30
76.5210.60
FedAvg
82.80-+0.13
83.40++0.25
82. 18
83.40++0.1
83.49-10.14
FodProx
82.14:1.110
82.04+0.008
01.50:0.26
0..26
82.1410.100
FedChest
FedAdam
83.02::..11
82.1110.110
82.72:10.16
83.10
83.3310.007
SCAFFOLD
83.52:00114
83.95=00.05
83.50
84.06
84.26-10.10
FedAvg
91.71=00774
96.60=t0.11
94.07=0.00
96.7210.05
97.4211.011
FedProx
95.03=0.04
96.05=0.04
94.60=0.07
96.8210.04
96.3310.09
CIFAR-10
FedAdam
91.23:10.29
95.80:10.24
94.0991.117
95.54
96.93:10006
SCAFFOLD
92.5110.99
96.78:10.01
94.3011.003
96.84
97.3810.03
Table 3: pFL aggregation method comparison on Fed-ISTC2019 and FedChest. FedBN and FixBN are only applicable models using BN layers. ANFR remains the top performer.
Architecture
Datasct
Method
BN-ResNet
GN-ResNet
SE-ResNet
Ne-ResNet
ANFR (Ours)
FedPer
82.36:0.80
80.6610.47 N/A
81.22::0.77
84.220.43
84.9410.46
Fed-ISIC2019
FedBN
82.82+0006
81.84+0.28
N/A
N/A
FixBN
82.79=0.05
N/A
83.27+0.38
N/A
N/A
FedPer
83.39=0.10
$3.73::0.10
83.36#0.14
83.70 03.104
83.8+0.14
FedChest
FedBN 8..3
N/A
83.33+0.114
N/A
N/A
FixBN
83.30+10.09
N/A
83.29::0.10
N/A
N/A
4.2 Performance Analysis and Comparison
GFL scenario. Average results for all datasels, models, and GFL aggregation methods are presented in Table [ First, we observe that GN does not consistently outperform the vanilla ResNet, supporting pursuit of aggregation methods on Fed-ISIC2019 and FedChest. Second, the sub-optimality of CA operating normalized features is evident, as the SE model frequently performs worse than BN-ResNet, notably across all aggregation methods on FedChest. NF-ResNet shows strong performance across all tasks and methods, confirming the potential of replacing activation normalization with weight standardization in FL. However, our proposed ANFR model consistently outperforms NF-ResNet, often by a considerable margin. For ex- ample, on Fed-ISIC2019 with SCAFFOLD, ANFR surpasses NF-RosNetts mcan balanced accuracy by more than 3%. For the FedChest dataset, we employ large batch size of 128 to maximize the probability that all classes represented in each batch, following best practices for multi-label, bass-imbalanced datasets. This is further analyzed in batch size ablation in Appendix C.3 ANFR emerges as the top-performing model across aggregation methods and our results indicate that integrating CA with SWS networks provides significant, performance gains, suggesting that channel attention is a crucial component in designing effective FL models.
reliable alternative. For instance, GN is outperformed by BN in half of the tested
BN-
PFL scenario. Table presents the results for pFL scenarios on Fed-ISIC2019 and FedChest. FodChest, where images grayscale and large batch size, FedBN and FedPer virtually equal: BN-ResNet

---

Published in Transactions on Machine Learning Research (09/2025)
achieves AUROC of 83.38% with FedBN and 83.39% with FedPer, indicating that the estimated BN statis- ties closely match the true ones. GN-ResNet attains 83. 73% with FedPer, slightly outperforming BN-ResNet,, but ANFR with FedPer is the most performant option across both aggregation methods, yielding a mean AUROC of 83.8%. Conversely, under the severe label and quantity skew on Fed-ISIC2019, employing FedBN improves performance over FedPer for models employing BN. ANFR achieves the highest balanced accuracy of 84.94% nonetheless. Notably, GN performs worse than BN on Fed-ISTC2019, and the ineffectiveness of combining BN and CA further evidenced, as SE-ResNet is outperformed by BN-ResNet in all scenarios. These findings demonstrate that adopting ANFR enhances performance across both datasets, leading to the best overall models. Unlike the trade-offs observed with BN-FedBN and GN-FedPer combinations, ANFR consistently outperforms other architectures across varying levels of data heterogeneity.
Tuning in favor of ANFR in Fed-ISIC2019. As noted in Section 4.1 and further detailed in Appendix A.2 our hyper-paramettrs are chosen after tuning the baseline BN-ResNet and not ANFR, meaning the reported improvement the results tables is a conservative floor of the improvement that can be achieved. To illustrate the real impact of our approach, we double the number of local steps in Fed-ISIC2019, keeping all other settings constant. As seen in Table 4 the performance of ANFR increases by 1.56% compared to Table 2 while its improvement over the best bascline becomes twice as big. While this experimental setting favors ANFR, the performance of BN-ResNeet is now lower, so this is not the setting we report in Table 2 The same methodology has been applied for all experimental settings. Despite optimizing for the baselines, ANFR still remains the best option, which greatly bolsters how exciting our results are.
Table 4: Results on Fod-ISIC2019 when doubling the local steps (tuning in favor of ANFR as opposed BN- ON-ResNet). ANFR performs better than the results in Table 2 setting used in the main paper.
but BN-ResNet worse, so this not the
BN-ResNet
GN-ResNee
SE-ResNet
NF-ResNet
ANFR
FedAvg 64.52
66.16
67.55
71.76
76.34
4.3 Cross-device experiments on CelebA
Table 5 presents the results of our models on the cross-device setting of CelebA, which contains 200,288 samples across 9,343 clients. While the binary classification task is relatively straightforward for individual clients, it poses challenges at the server level due to the vast number of clients and significant quantity and class skews-soe some clients have only a few samples or labels from a single class. We observe that ANFR outperforms the baseline models, demonstrating its adaptability across diverse FL scenarios.
Table 5: Performance comparison in cross-device setting, training with FedAvg on CelebA. The training setup follows Pieri et al. (2023), where 10 clients participate each round until all clients have trained for 30 rounds. ANFR outperforms the basclines.
Architecture
BN-ResNet
Ne-ResNet
SE-ResNet
NF-ResNet
ANFR (Ours)
Average Accuracy 82.2+1.21
85.41+0.68
85.55+0.84
88.1740.3
88.91=0.28
4.4 Sample-level Differentially Private Training
In privacy-preserinng scenarios involving differential privacy (DP), BN cannot be used as calculating mini- batch statistics violates privacy-preservalion so is customarily replaced by GN. We demonstrate the utility of ANER in such settings using the FedPathology setup described in Section 4.1 We train using DP-SGD with strict sample-level privacy guarantces: following good practices, the probability of information Icakage set to D, the noise multiplier is to 1.1 and the gradient max norm to 1. We employ a privacy budget of 1=1, followed by training without privacy constraints (e=oo), to illustrate the privacy/utee /utility
10

---

Published in Transactions on Machine Learning Research (09/2025)
trade-off of each model. From the results presented in Table 6 observe that with unrestricted privacy budget, GN and ANFR perform comparably. However, when strict budget is en GN suffers sharp performance decrease of 17%, as expected following previous research (Klause et al. 2022), whereas ANFR's average accuracy is reduced by only 3%. ANFR's robustncss under DP may be attributed to its reliance on weight standardization, which has been shown to benefit from additional regularization (Brock et al. 2021b; Zhuang & Lyu 2024) such as that provided by DP-SGD's gradient clipping and gradient noising. Our experiments show DP training induces regularization effect that disproportionatel benefits NF models like ANFR, observation also reported by De et al. (2022). These findings make ANFR a promising candidate for furthering development and adoption of DP training in FL, thereby enhancing the privacy of source data contributors, such as patients.
Table 6: Accuracy on the validation set of FodPathology when training with and without DP. Performance degrades severely for GN, while ANFR retains good performance.
Privacy Budget
8= 00
GN-ResNet
84.79=2.72 84.4713.08
67.27
ANFR (Ours)
81.1110.33
4.5 Attention Mechanism Comparison
Next, we investigate the impact of different altention mechanisms performance. We compare the SE module used in previous sections with ECA (Wang et al.) 2020), and CBAM (Woo et al, 2018). ECA replaces SE's fully-connected layers with a more efficient 1-D convolution to capture local cross-channel interactions. CBAM combines channel and spatial attention and utilizes both max and average pooling to extract channel representations. From Table we observe that even the towest-performing module on each dataset outperforms all baseline models from Tables [ and [ proving the robustness of our approach. No single mechanism consistently performs best, making further exploration of attention modules interesting avenue for future work.
Table 7: Comparing different channel attention modules after FL training with FedAvg. No module is consistently the best, but even the worst outperforms the best baseline (NF-ResNet).
CIFAR-10
Fed-ISIC2019
FedChest
CelebA
Ne-ResNet
96.72 + 0.05
72.49 + 0.60
83.40
0.11
88.17
0.30
ANFR (SE)
97.42 + 0.01
74.78 + 0.16
83.49 + 0.14
88.91 + 0.28
ANER (ECA)
97.13 + 0.11
75.07 = 0.48
83.62 + 0.10
89.07 + 0.43
ANER (CBAM)
97.05 + 0.08
74.19 + 0.68
83.47 0.15
89.31 0.41
4.6 Computational Overhead
Additionally, to gauge the computational overhead of ANFR, and by extension its applicability in low- resource environments, compare training times for BN-ResNet-28 26 with those for ANFR-26 using ECA as the attention mechanism. The batch size is set to 32, and we mcasure the average time per itcration of forward + backward pass across 100 ilerations using Porch'ss profiler. We do this for two distinct scenarios: devices without a CUDA-enabled GPU (e.g., smartphones). and devices with CUDA-enabled GPUs (c.g., edge devices such as Nvidia Jetson). The results in Table 8 show ANFR introduces marginal overhead (~10% without CUDA, ~00 with CUDA) while providing a significant performance advantage, showcasing its practicality in resource-cosstrinned settings.
11

---

Published in Transactions on Machine Learning Research (09/2025)
Table Computational demand comparison in a simulated low-resource setting.
Scenario
Without CUDA
With CUDA
Metric
Forward
Backward
Total
CPU time
GPU time
BN-ResNet-26
297ms
672ms
969ms
12ms
22ms
ANFR-26 (ECA)
353ms
717ms
is 70ms
9ms
26ms
ResNet-50
GN EN-ResNet-50
SF- DE-ResNet-550
MF-ResNet-500
ANFR-ResNet-S0
Atelectasis
Pieummni
Figure 4: Comparison of the saliency maps generated by Crad-CAM++ from different architectures for Pneumonia and Atelectasis image, overlaid with ground-truth bounding boxes. We note ANFR improves localization and reduces activations outside the area of interest.
4.7 Qualitative Localization Performance
We assess the localization capability of cach architecture after FL training with the best aggregation method FedChest, SCAFFOLD. We compare the bounding box annotations provided by Wang el. al. (017 with C [Chattopadhay et al. 2018 heatmaps generated for samples labeled Atelectasis or Pneumonia from the FedChest test set. Figure shows that ANFR's heatmaps more closely align with the annotated bounding boxes. This improved localization aids model interpretability, which is crucial in areas like medical imaging.
4.8 Comparing Different Architectural Families
Finally, we relax the constraint of comparing between models of similar computational overhead and struc- ture by broadening our ovaluation to include models from different familics. We compare ResNet-50 and ANFR-50 with Efficient.eet2-SSi
(Tan & Le 2021), ViT-Small (Dosovitskiy et al., 2020 Wightman 2019), Sol-SeXt-Small (Liu al. 2022) and FedConv (Xu et al. (024). We compare, using FedAvg as the aggregation method, on the bases of accuracy, theoretical overhead, and practical overhead to provide more informative insights. The results, presented in Table g show that care is needed when choosing network in order to balance performance and computational demands. This is especially true in cross-device FL. scenarios, where straggler handling and limited compute concerns might render prohibitive the cost of otherwise performant networks like ConvNeXt (2 x memory) or FedConv (2.64 x latency). ANFR is shown to more mefory-efficient than a vanilla ResNet (due to removing BN), and occurring minimal overhead terms of latency while being a top competitor terms of accuracy.
5 Conclusion
We introduce ANFR, a novel design-level approach address the challenges of data heterogeneity design level in FL. ANFR fills research gap by simultancously working in GFL, pFL, and private FL scenarios while being compatible with any aggregation method and offering robust increase in performance. Extensive
12

---

Published in Transactions on Machine Learning Research (09/2025)
Table 9: Comparison across different model architectures terms of accuracy on datasets, parameter count, theoretical FLOPs, forward latency and GPU memory cost. Actual overhead measured using the DeepSpeed profiler (Raseey et al. 2020) for images of size (3.224,224) with a batch size of 64 and automatic mixod-precision training enabled Nvidia A100 GPU.
Accuracy/UUROC()
Theoretical Overhead
Measured Overhead
Architecture
FedChest
Fed-1SIC2019
CIFAR-10
Params
GFLOPs
Memory in GB
Latency in ins
ResNet-50
82.8
66.01
91.71
23.52
4.09
4.01 (1x)
32.38 (1x)
EeficientNetV25
79.53
69.24
93.93
20.19
2.85
5.91 (1.47x)
79.98 (2.47x)
CouvNeXII-Smal
82.28
77.71
97.73
49.46
8.68
8.16 (2.03x)
53.32 (1.65x)
ViT-Small
83.53
73.21
97.48
21.68
4.96
3.36 (0.84x)
23.26 (0.71x)
BedConv-LnvertUl
82.68
78.32
97.26
25.62
4.61
5.35 (1.33x)
85.46 (2.64x)
ANFR-50 w/ RCA
83.62
75.07
97.13
23.52
4.09
3.72 (0.93x)
40.16 (1.24x)
ANFR-50 w/ SE
83.19
74.78
97.42
26.06
4.09
3.86
(0.96x)
44.67 (1.38x)
experiments demonstrate the superior adaptability and performance of ANFR, as it consistently surpasses the performance of baseline architectures, regardless of the aggregation method employed. Our results position ANER as a compelling backbone model suitable for both global and personalized FL scenarios where statistical heterogeneity and privacy guarantees are important concerns. Our findings highlight the need to look beyond aggregation methods as the core component foderated performance and the critical role of architectural innovations in reaching the next frontier in private and collaborative settings.
References
Xuming An, Li Shen, Han Hu, and Yong Luo. Federated learning with manifold regularization and normalized update reaggregation. Advonces in Neural Information Processing Systems, 11455097-55109 2023.
Mathieu Andreux, Jean Ogier du Terrail, Constance Beguier, and Eric W Tramel. Siloed federated learning for multi-centric histopathology datasets. In Domain Adaptation and Representation Transfer, and Dis- tribuled and Collaborative Learning: Second MICCAI Workshop, DART 2020, and First MICCAI Work- shop, DCL 2020, Held in Conjunction with MICCAI 2020, Lima, Peru, October 4-8, 2020, Proceedings 2. pp. 129-139. Springer, 2020.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary, lcarning with personalization layers, 2019. URL htttps///rxiv.. org/ abs//9////1111/ 1001101111001118 00818
Federated
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL
https:
7/axxiv wwwpiiv..rgg/bb///1oeo 06450
Andrew Brock, Soham De, and Samuel L. Smith. Characterizing signal propagation to close the performance gap in unnormalized resnets, 2021a. URL http::///aaaiii.o//aaa//////
Andy Brock, Soham De, Samuel Smith, and Karen Simonyan. Hiigh-erffrmance large-scale image recog- nition without normalization. In International Conference on Mochine Learning, pp. 1059 1071. PMLR, 2021b.
Aurelia Bustos, Antonio Pertusa, Josa-Maria Salinas, and Maria De La Aatesia-Vaya Padchest: A large chest x-ray image dataset with multi-label annotated reports. Medical image analysis, 66:101797, 2020.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097 2018.
Aditya Chattopadhay, Anirban Sarkar, Prantik Howlader, and Vineeth N Balasubramanian. Grad-cam++:: Generalized gradient- gradient-based visual explanations for deep convolutional networks. In 2018 IEFE winter conference applications of computer vision (WACV), pp. 839 847. IEEE, 2018.
13

---

Published in Transactions on Machine Learning Research (09/2025)
Hong- You Chen and Wei-Lun Chao. Fedbe: Making bayesian model ensemble applicable to federated wwtp://rr.aab///....11o
learning, 2021. URL
Soham De, Leonard Berrada, Jamie Hayes, Samuel L. Smith, and Borja Balle. Unlocking high-accuracy diffcrentially private image classification through scale, 2022. URL htttt::/t/xxxi/i/ http:s///arrtl..og//b/////////// 13650
Terrance DeVries. Improved regularization of convolutional neural networks with cutout. orXiv preprint 708.:11008.0452, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alesander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition at scale. arXiv preprinl trxiv:2010. 1 11929, 2020.
Zhixu Du, Jingwei Sun, Ang Li, Pin- Yu Chen, Jianyi Zhang, Hai' Helen' Li, and Yiran Chen. Rethink- ing normalization methods in federated learning. In Proceedings of the 3rd International Workshop on Distributed Machine Learning, pp. 16 22, 2022.
Rachid Guerraoui, Rafael Pinot, Geovani Rizk, John Stephan, and François Taiani. Overcoming the chal- lenges of batch normalization in federated learning. arXiv preprint or Xiv:2400 14670, 2024.
Yongxin Guo, Xiaoying Tang, and Tao Lin. Fedbr: Improving federated learning on heterogeneous data via local learning bias reduction. In International Conference on Machine Learning, pp. 12034-12054 PMLR, 2023.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip Gibbons. The non-iid data quagmire of de- centralized machine lcarning. In International Conference on Machine Learning, pp. 4387-4398. PMLR, 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non non-identical data distribution for federated visual classification, 2019. URL httpss://arxvv.orrg org/ abs 1909. 06335
Jie Hu, Li Shen, and Gang Sun. Squeeze-and-exxittaiiin networks. In Proceedings of the IEEE conference on computer vision and pottern recognition, pp. 7132-7141, 2018.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pp. 448 456. pmir, 2015.
Jeremy Irvin, Pranav Rajpurkar, Michael Ko, Yifan Yu, Silviana Ciurea-Ileus, Chris Chute, Henrik Mark- lund, Behzad Haghgoo, Robyn Ball, Katie Shpanskaya, et al. Chexpert: large chest radiograph dataset with uncertainty labels and expert comparison. In Proceedings of the AAAI conference on artificial intel- ligence, volume 33, pp. 590-597, 2019.
Meirui Jiang, Zirui Wang, and Qi Dou. Harmoff: Harmonizing local and global drifts in federated learning heterogeneous medical images. In Proceedings of the AAAI Conference Artificial Intelligence, volume 36, pp. 1087-1095, 2022.
Mcirui Jiang, Anjie Le, Xiaoxiao Li, and Qi Dou. Heterogencous personalized federated learning by local- global updates mixing via convergence rate. In The Twelfth International Conference on Learning Repre- sentations, 2024.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurólien Bellet, Mchdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummiins, et al. Advances and open problems in federated learning. Foundations and Trends0 in Machine Learning, 14(1 (((1-2):1-100)) 2021.
Mycongkyun Kang, Soopil Kim, Kyong Hwan Jin, Ehsan Adeli, Kilian M Pohl, and Sang Hyun Park. Fednn: Federated learning on concept drift data using weight and adaptive group normalizations. Pattern Recognition, 149:1 10230, 2024.
14

---

Published in Transactions on Machine Learning Research (09/2025)
Sai Pranceth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning, pp. 5132 5143. PMLR, 2020.
Jakob Nikolas Kather, Johannes Krisam, Pormpimol Charoentong. Tom Luedde, Esther Herpel, Cleo-Aron Weis, Timo Gaiser, Alexander Marx, Nektarios A Valous, Dyke Ferber, et al. Predicting survival from colorectal cancer histology slides using deep learning: A retrospective multicenter study. PLoS medicine, 16(1)=e0002730, 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL https:
7/arxiv. /a////////////////////////////// 1t////112..989
Helena Klause, Alexander Ziller, Daniel Rueckert, Kerstin Hammernik, and Georgios Kaissis. Differentially private training of residual networks with scale normalisation. arXiv preprint arXiv:2203.00 2022.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers features from tiny images. University of Toronto, 2009.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogencous networks. Proceedings of Machine learning and systems, 2:429-450, 2020a.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang, On the convergence of fedavg on non-iid data, 2020b. URL https: http:/:////xviooo////a// aaaa///s//999 110...199...00000
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization, 2021. URL htttps://axxv..ortt wwwpg//////./..//...../ 07623
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976-11986, 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pp. 3730-3738, 2015.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Commmication--ffiiicttt learning of deep networks from decentralized data. In Artificial intelligence and slatistics, pp. 1273 1282. PMLR, 2017.
Ily Mironov. Rényi differential privacy. In 2017 IEEE 30th computer security foundations symposium (CSF), pp. 263 275. IEEE, 2017.
Ari S Morcos, David GT Barrett, Neil C RRbinooitz, and Matthew Botvinick. On the importance of single directions for generalization. Xixiv preprinl arvi 1r.Xi:1803.0000 180.00699, 2018.
Jean Ogier du Terrail, Samy-Safwar Ayed, Edwige Cyffers, Felix Grimberg, Chaoyang He, Regis Loeb, Paul Mangold, Tanguy Marchand, Othmane Marfoq, Frum Mushtag, et al. Flamby: Datasets and benchmarks for cross-silo federated learning in realistic healthcare settings. Advances in Neural Information Processing Systems, 35:5315-523 5334, 2022.
Kunjal Panchal, Sunav Choudhary, Subrata Mitra, Koyel Mukherjee, Somdeb Sarkhel, Saayan Mitra, and Hui Guan. Flash: concept drift adaptation in federated learning. In International Conference Machine Learning, pp. 11931226962. PMLR, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, "Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Sara Pieri, Jose Restom, Samuel Horvath, and Hisham Cholakkal. Handling data heterogeneity via archi- tectural design for federated visual recognition. Advances in Neural Information Processing Systems, 36: 4115-4136 2023.
15

---

Published in Transactions on Machine Learning Research (09/2025)
Siyuan Qiao, Huiyu Wang, Chenxi Liu, Wei Shen, and Alan Yuille. Micco-batch training with batch-channel normalization and weight standardization, 2020. URL
wttp://ww.....r/////.////0 10520
Liangqiong Qu, Yuyin Zhou, Paul Pu Liang, Yingda Xia, Feifei Wang, Ehsan Adeli, Li Fei-Fei, and Daniel Rubin. Rethinking architecture design for tackling data heterogeneity in foderated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10061 10071, 2022.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong Hc. Decpspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery 8 Data Mining, pp- 3505-3506, 2020.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H. Brendan McMahan. Adaptive federated optimization, 2021. URL https://arxiv.orgg
abs/2003.000 002955
Holger R Roth, Yan Cheng, Yuhong Wen, Isaac Yang, Ziyue Xu, Yuan-Ting Hsieh, Kristopher Kersten, Ahmed Harouni, Can Zhao, Kevin Lu, et al. Nvidia flare: Federated learning from simulation to real- world, 2022. URL https: haapss://arxivvv/ wwwexxiv..og/aee. 22100. 13291
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-2522 2015.
Tim Salimans and Durk P Kingma. Weight normalivation A simple reparameterivation to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016.
Yujun Shi, Jian Liang, Wenqing Zhang, Chuhui Xue, Vincent YF Tan, and Song Bai. Understanding and mitigating dimensional collapse in federated learning, IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
Vasilis Siomos, Sergio Naval-Marimont Jonathan Panebaat,a and Giacomo Tarroni. Aria: On the interaction between architectures, initialization and aggregation methods for federated visual classification. In 2024 IEEE International Symposium on Biomedical Imaging (ISBI), pp. 1 5. IEEE, 2024.
Leslic N Smith. Cyclical learning rates for training ncural networks. In 2017 IEEE winter conference on applicalions of computer vision (WACV), pp. 464 472. IEEE, 2017.
Mingxing Tan and Quoe Le. Effiiiettnet22: Smaller models and faster training. In International conference on machine learning, pp. 10096- 10106. PMLR, 2021.
Trene Tenison, Sai Aravind Sreeramadas, Vaikkunth Mugunthan, Edouard Oyallon, Trina Rish, and Fugene Belilovsky. Gradient masked averaging for federated learning, 2023. URL http::httt///tt org/oogg//aoo.ooo hbs//bb//200 /2011 11986
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H. Brendan McMahan, Blaise Aguera y Arcas, Maruan Al Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, Suhas Diggavi, Hubert Eichner, Advait Gadhikar, Zachary Garrett, Antonious M. Girgis, Filip Hanzely, Andrew Hard, Chaoyang He, Samuel Horvath, Zhouyuan Huo, Alex Ingerman, Martin Jaggi, Tara Javidi, Peter Kairouz, Salyen Kale, Sai Praneeth Karimireddy, Jakub Koneeny, Sanni Koyejo, Tian Li, Luyang Liu, Mehryar Mohri, Hang Qi, Sashank Reddi, Peter Richtarik, Karan Singhal, Virginia. Smith, Mahdi Soltanolkotabi, Weikang Song, Ananda Theertha Suresh, Sebastian U. Stich, Ameet Talwalkar, Hongyi Wang, Blake Woodworth, Shanshan Wu, Felix X. Yu, Honglin Yuan, Manzil Zaheer, Mi Zhang, Tong Zhang, Chunxiang Zheng, Chen Zhu, and Wennan Zhu. A field guide to federated optimization, 2021. URL htttp:///arx/ org/abs/2107.00
Qilong Wang, Banggu Wu, Pengfei Zhu, Peihua Li, Wangmeng Zuo, and Qinghua Hu. Eca-net: Efficient channel attention for deep convolutional neural networks. In Proceedings of the IEEE/CVF conference on computer vision and patlern recognilion, pp. 1 11534-11542, 2020.
16

---

Published in Transactions on Machine Learning Research (09/2025)
Xiaosong Wang, Yifan Peng, Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald M Summers. Chestx- ray8: Hospital-soale chest x-ray database and benchmarks. on weakly-supervised classification and local- ization of common thorax diseases. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2097-2106, 2017.
Yanmeng Wang, Qingliang Shi, and Tsung-Hui Chang. Why batch normalization damage federated learning on non-iid data? IEEE Transaclions on Newral Networks and Learning Systems, 2023.
Ross Wightman. Pytorch image models. httgg://ggthuu..co////ggatttaatt 1mmeee---mmmmmmeelll 2019.
Ross Wightman, Hugo Touvron, and Hervé Jégou. Resnet strikes back: An improved training procedure in timm. Xiv preprinl or xiiv:2110000477 2021.
Sanghyun Woo, Jongchan Park, Joon- Young Lee, and In So Kweon. Cbam: Convolutional block attention module. In Proceedings of the Europeon conference compuler vision (ECCV), pp. 3 3-19, 2018.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on compuler vision (ECCV), pp. 3 3-19, 2018.
Peiran Xu, Zeyu Wang, Jieru Mei, Liangqiong Qu, Alan Yuille, Cihang Xie, and Yuyin Zhou. Fedconv: En- hancing convolutional neural networks for handling data heterogeneity federated learning. Transactions on Machine Learning Research, 2024.
William J Youden. Index for rating diagnostic tests. Cancer, 3(1):32-35, 1950.
Ashkan Yousefpour, Igor Shilov, Alexandre Sablayrolles, Davide Testuggine, Karthik Prasad, Mani Malek, John Nguyen, Sayan Ghosh, Akash Bharadwaj, Jessica Zhao, Graham Cormode, and Ilya Mironov. Opa- cus: lssr-friendly differential privacy library in pytorch, 2022. URL https://arxiv http://arxiv.oor///////0000
12298
Matthew D Zeiler and Rob Fergus. Visualizing and understanding convolutional networks. In Computer Iision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part I 13, pp. 818-833, Springer, 2014.
Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. Fedala: Adaptive local aggregation for personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 11237- 2023.
Kaiyu Zheng, Xuefeng Liu, Guogang Zhu, Xinghao Wu, and Jianwei Niu. Channelfed: Enabling personalized federated learning via localized channel attention. In GLOBECOM 2022-2022 IEEE Global Communica- tions Conference, pp. 2987-2992. IEEE, 2022.
Jike Zhong, Hong- You Chen, and Wei Chao. Making batch normalization great in federated deep learning, 2024. URL https://arxiv..ortt orgp:/s/2333333333 (abbb//333006300
Jianing Zhu, Jiangchao Yao, Tongliang Liu, Quanming Yao, Jianliang Xu, and Bo Han. Combating exacer- bated heterogeneily for robust models in federated learning. arXiv preprint ar Xrx:22330 00250,, 2023.
Weiming Zhuang and Lingjuan Lyu. Fedwon: Triumphing multi-domain federated learning without normal- ization. In The Twelfth International Conference on Learning Representations, 2024.
17

---

Published in Transactions on Machine Learning Research (09/2025)
A Additional Implementation Details
A.1
Datasets
Skin Lesion Classification on Fed-ISIC2019. Fed-ISIC2019 (Ogier du Terrail et al., 2022) contains 23,247 dermoscopy images from 6 centers across 8 classes and is subset of the TSTC 2019 challenge dataset. We follow the original pre-processing, augmentation, loss, and evaluation metric of (Ogier du Terrail et al, 2022). This means the loss function is focal loss weighted by the local class percentages at each client, and the reported metric is balanced accuracy, as meaturess against class imbalance. The augmentations used include random scaling, rotation, brightness changes, horizontal flips, shearing, random cropping to 200 x 200 and Cutout (DeVries) 2017). We train for 80 rounds of 200 local steps with a batch size of 64. The clients locally use Adam (Kingmma & Ba 2017), learning rate of 5e-4, and a cyclical learning rate scheduler (Smith 2017). In terms of heterogeneity, Fed-ISIC2019 represents a difficult task due to class imbalance and heavy dataset size imbalance, with the biggest client owning more than 50% of the data and the smallest client 3%.
CIFAR-10. Krizhevsky et al. (2009) consists of 50,000 training and 10,000 testing 32 x 32 images from 10 classes. We follow the sctup of Pieri et al. (2023) specifically the 'split-2' partitioning where each client has access to four classes and does not receive samples from the remaining six classes. This means we train for 100 rounds of local epoch with batch size of 32. Clients use SGD with learning rate of 0.03 and a cosine decay scheduler, in addition to gradient clipping to 1.0. During training the images are randomly cropped with the crop size ranging from 5% to 100% and are then resized to 224 224.
CclebA from LEAF. A partitioning of the original CelebA (Liu et al. 2015) dataset by the colebrity in the picture, this dataset contains 200,288 samples across 9,343 clients. The task is binary classification (smiling vs not smiling). We follow the setup presented in Pieri et al. (2023). training with 10 clients each round until all clients have trained for at least 30 rounds. The other settings are the same as those for CIFAR-10.
FedPathology Slide Classification Dataset. A colorectal cancer pathology slide dataset (Kather et al. 2019), consisting of 100k training images of Whole Slide Image (WSI) patches with labels split among 9 classes, is used to simulate a federation of 3 clients. We mimic one of the most important challenges in the WSI field by not color-normalizing the images, which come from two different labs with differences in staining protocols. The original 7k color-normalized validation set from Kather et al. (2019) is kept as a common validation set. We follow common practice (Hsu et al., 2019) to simulate label skew data heterogeneity by using a Dirichlet distribution with a = 0.5 to partition the data, Since this artificial partitioning is random, we make sure to use the same seeds across architectures and privacy settings compare on exactly the same partitioning instances. Our pipeline is built using Opacus (Yousefpour et al. 2022 and (a, R)-Renyi Differential Privacy (RDP) (Mironov 2017). Following good practices, the probability of information leakage of
8 is set to 0.1/|Dal where D,] represents each client's dataset size. The Dp-specific the noise multiplier and gradient max norm are set to 1.1 and 1, respectively. Data augmentation includes random horizontal and vertical flips, random color jittering, and random pixel crasing. Clients use Adam with a learning rate of 5e-5, training for 500 local steps with a batch size of 64. Federated training is stopped after 25 rounds, which is the point where both architectures have expended, on average, a privacy budget of = 1. Finally, we train without using DP under the same settings form clearer picture of the privacy/ulility trade-off of each model.
Chest X-Ray Multi-Label Classification on FedChest. Please refer to Appendix C.1
A.2
Tuning
Hyper-parameters were optimized for the BN-ResNet and then the same parameters were used for all net- works. The ranges were as follows:
Local Steps: {100, 200, 500}
Rounds: (20, 50, 75, 100)
18

---

Published in Transactions on Machine Learning Research (09/2025)
Batch size: {32, 64, 128}
Gradient Clipping: (None, Norm Clipping to 1, Adaptive Gradient Clipping (Brock et al. 2021b
Learning rate: (5c-5 - 1o-2)
Optimizer: (Adam, AdamW, SGD with momentum]
Scheduler: (None, OneCycleLR, Cosine Annealing, Cosine Annealing with Warm-up)
FedProx μc: (1e-3, le-2, le, 2)
Addm Server learning rate: (5e-4, 1e-3, le-2, le-13
Discussion. We found both FL aggregation methods that introduce aaper-parameters difficult to tune: FedProx Li al. (2020a) made negligible difference for small j values and decreased performance we increased it; the server learning rate in FedOpt has to be chosen carefully, as large (1c-2, lc-1) learning rates led to non-convergence and small ones to disappointing performance. Gradient clipping helped ANFR but was detrimental to the vanilla ResNet. We found the use of scheduler to be very beneficial for performance, as well as making the optimizer and initial learning rate choice less impactful. We store the intermediato learning rate at each client between rounds and resumee the scheduler, and also follow this for the momentum buffers of the adaptive optimizers.
A.3 Model Details and Computational Overhead
Table 10 presents pre-training details, parameter counts, multiply-accumuulte counts (GMACs)) and floating point operation counts (FLOPs) and ImageNet. (Russakovsky et al.) 2015)) validation set top-1 performance for all models. For models which are pre-trained by us, links the pre-trained weights will be made public after acceptance.
Comparison of model details. Profiling results obtained using DeepSpeed's (Rasicy al. 2020 model profiler, for batch size of 1 and an image size of 3x224x224. Training recipe refers to the recipes presented in Wightman al. (2021). ImageNet-Ile 1K eval performance obtained from timm (Wightman 2019) results and our own training. performance evaluated on 256x256 size.
Table 10:
Model
Parameters
GFLOPs
GMACs
IN-1K performance Training Recipe
BN-RosNet-50
25.56 M
4.09
8.21
78.81
B
GN-RosNet-50
25.56 M
4.09
8.24
80.06
A1
SE-ResNe1-50
28.09 M
4.09
8.22
80.26
B
NF-ResNet-50
25.56 M
4.09
8.32
80.22*
B
ANFR-50 (SE)
28.09 M
4.09
8.32
80.4
B
ANFR-50 (ECA)
25.56 M
4.09
8.32
80.61
B
ANFR-50 (CBAM)
28.07 M
4.1
8.33
80.37
B
19

---

Published in Transactions on Machine Learning Research (09/2025)
B Additional Results
B.1 Results on Fed-ISIC2019 using FLamby hyper-parametert
The experimental setup we use for Fed-ISTC2019 in the main paper is an improved version of the example benchmark presented in Section 4.1 of Ogier du Terrail et al. (2022), so one might wonder how the compared models perform under the original seltings. To answer this we repeat Centralized, FvdAvg and SCAFFOLD training on Fed-ISIC2019 after aligning our ayper-parameters with Ogier du Terrail et al. (2022). meaning we reduce local steps to 100 without a schoduler, perform 9 federated rounds, and use pre-computed class weights in the focal loss. Results are presented in Table 11 showing ANFR comprehensively beats com- peting baselines, with
wider performance gap compared to original setting. The overall level of performance, including the gap between centralized and FL training, aligns with the results presented in Ogier du Terrail et al. (2022). as we expect. Additionally, SE-ResNet performs better than ANFR in central- ized training, but the opposite is true in FT. training, further validating our core claims in Section 3 that CA nceds Weight Standardization to optimally adjust channel responses in hetcrogencous FL. Although these results further support our claims, we believe the optimized version of Fed-ISIC2019 training we provide in the main paper more of use the community.
Table 11: Results on Fed-ISIC2019 using the original ANFR and the baselines even wider.
from FLamby.
The gap between
NN-ResNet
GN-ResNeet
SE-ResNet
NF-ResNet
ANFR (Ours)
FedAvg SCAFFOLD
59.5
25.26:2.96 57.62+2.95
61.92:11.58
60.76110.75
65.3411.29
57.61
0..34+0.42
57.35+0.73
71.07#1.27
Central
61.26+2.92
57.09=1.85
73.00#1.09
61.28+1.53
72.03=1.55
B.2 Results Using Randomly Initialized Models
Given the ubiquity and demonstrated utility of ImageNet pre-trained models in FL (Qu et al. 2022 Pieri et al. 2023 Siomos et. al., 2024), we use pre-trained models in the main paper. Nevertheless, we conduct uutt additional experiments with FedAvg on CIFAR-10, FedChest and Fed-ISIC2019, using randomly initialized models. Although the results below bolster our claims, we avoided this setting initially as random weight initialization is not representative of the current standard settings adopted by FL practitioners. The only changes made to accommodate the absence of pre-training are to change the optimizer to AdamW and the learning rate to 0.001 for CIFAR-10, and to double the number of local steps for Fed-ISIC2019. Our results in Table 12 show the same trend, of a gap existing between FL and centralized training but being smaller when using pre-trained models. In this setting, too, ANFR is the best performer.
Table 12: Results using randomly initialized models CIFAR-10, Fed-IS1C2019 and FedChest.
Dataset
CIFAR-10
Fed-ISTC2019
FedChest
Model
FedAvg
Central
FodAvg
FedAvg
Central
BN-ResNet
80.89
89.05
54.02
78.44
82.58
GN-ResNee
78.52
86.69
54.92
73.68
80.82
SE-ResNet
81.19
88.65
53.20
78.79
82.16
Ne-ResNet
81.66
88.96
56.75
79.06
83.55
ANFR (Ours)
83.20
89.58
57.71
79.41
83.67
20

---

Published in Transactions on Machine Learning Research (09/2025)
B.3 CIFAR-10 experiment without early-stopping
The results presented in Section 4.2 follow the experimental set-up of Pieri et al. (2023), where the validation set is used a form of early stopping in the following way: at every round the performance on the test set is only evaluated if the accuracy on the validation set has increased. While this methodologically valid set-up, is also interesting to see how the models perform when no early-stopping is used. To compensate for this and avoid overfitting we disable gradient clipping and increase the batch size to 64. The results are presented Table 13 showing how ANFR continues to beat the baselines.
Table 13: Alternative CIFAR-10 setting where we do not use validation-busee early stopping, but instead report final round test accuracy. NaNs indicate training instability.
Model
BN-ResNet
SE-ResNet
GN-ResNet
NF-ResNet
ANFR
FodAvg
67.39
74.75
96.73
96.62
97.45
FedProx
86.3
94.23
95.98
NaN
96.63
FedAdam
57.43
88.93
95.32
NaN
96.96
SCAFFOLD
61.37
78.99
96.57
96.84
97.49
21

---

Published in Transactions on Machine Learning Research (09/2025)
B.4 Performance Plots
To gauge convergence, it can be helpful to examine performance plots showing how accuracy progresses throughout federated training. Below we provide four such plots, comparing all models when training from scratch on CIFAR-10 using FedAvg and SCAFFOLD, comparing all models for the experiment in Table and a Fed-ISIC run from the top performing model Table 2. ANER with SCAFFOLD.
Figure 5: Top Left: training from scratch on CIFAR-10 using FedAvg. Top Right: training from scratch on CIFAR-10 using SCAFFOLD. Bottom Left: training from scratch on FedISIC using FodAvg. Bottom Right: top performing model run, ANER with SCAFFOLD on FedISIC.
GN GN-ResNet vs NF-ResNet vs ANFR
GN-ResNet vs NF-ResNet vs ANFR
0.8
0.8
0.7
0.6
Test Accuracy
Test Accurscy
0.5
0.5
ResNet
0.3
NF ResNet
0.3
NF ResNett
GN ResNet
GN ResNet
NF-ResNet
NF-ResNet
0.1
ANFR
0.1
ANFR
80
100
Communication
Communication Round
FedAvg and 400 local steps on Fed-ISIC
FedAvg and 400 local steps Fed-ISIC
0.8
0.7
Average Tost Accuracy
0.6-
0.5
0.5
0.4
0.4
Model resnet5o supervised
0.3
anfr50 pretrained
sesnel_50 50_pretrained Limm
resnet50 on supervised
Model
seresnet 30, superibddd
antr. 50, 30_pretra ned, swa.told.cxamppcccoooooooo jun, seed, 2024
20
30
50
60
70
80
20
20
50
70
90
22

---

Published in Transactions on Machine Learning Research (09/2025)
C FedChest construction and additional results
C.1 Construction and hyper-parameters
To create FedChest we use three large-scale chest X-Ray multi-label datasets: CXR14 (Wang et al., 2017) PadChest [Bustos et al. 2020) and CheXpert (Irvin et al., 2019). To derive a common dataset format for all three, we need to take several pre-processing steps:
We remove lateral views where present, keeping only AP/PA views.
2. We discard samples which do not contain at least one of the common diseases, which are: Atelectasis, Cardiomegaly, Consolidation, Edema, Effusion, No Finding, Pneumonia, and Pneumothorax.
3. We remove "duplicates" which, in this context, means samples from the same patient that have the same common labels but different non-common labels.
4. We remove 5% from the edge of each image to avoid blown-out borders and artifacts.
We resize the images to 224x224 pixels.
We apply contrast-limited histogram equalization (CLAHE) the images.
In addition to these common steps, some dataset-specific additional pre-processing steps are necessary, namely setting NaN and 'uncertain' labels of CheXpert to (not present), removing corrupted NA rows from CXR14, and removing corrupted images from PadChest.
After pre-processing, CheXpert has twice many samples the other datasets, so we further split it into two clients, exp_young and cxp_ xpp_od using the median age of the patient population (63 years). Icading to a lotal of 4 clients with train/val/test splits of (given in thousands): 23.7/15/10 for CXR14, 26/ 15/10 for PadChest, 29.7/15/7.5 for cxp_old and 31/15/7.5 for cxp_young. The task is multi-label classification across the common classos.
After luning, clients perform 20 rounds of 200 local steps with batch size of 128, the loss function is weighted Binary Cross- Eotss-Entropy (BCE), and the optimizer Adam with learning rate of 5e-4, annealed over training. Data augmentation includes random shifts along both axes, random scaling and rotation, Cutout, and random cropping.
C.2 Additional FedChest Metrics
Further to the results presented in the main texl, since some of the diseases have unbalanced label distribution, and to also gauge model performance in deployment, we use the validation Receiver Operating Curve (ROC) to find the optimal class thresholds for each client using Youden's Index (Youden, 1950). Having lixed the thresholds to these values, at test-time we measure the average accuracy and F'l score of each model and present the results in Table 14
C.3 Batch Size Ablation Study
The absence of a performance gap between NN-ResNet and ANFR on the FedChest dataset when using FedProx (Table 2 motivates us to perform study ablating the batch size to examine how inconsistent averaging, which expected to happen for small batch sizes, affects results. We compare BN-ResNet and ANFR, varying the batch size while keeping all other experimental settings unchanged.
Table 15 shows BN-Ros.ectss performance degrades more than that of ANFR for small batch sizes (16 and 32). ANFR offers significant advantages compared to BN-ResNet for small batch sizes due to the absence of BN. In the main paper
are tuned based on BN-ResNet's performance; as the best BN- ResNet result achieved with a batch size of 128, this the one used. While using a large enough batch size can mitigate intra-client variance to degree, we see that increasing the batch size 256 reduces BN- ResNet's 's performance, indicating diminishing returns. This reinforces that increasing batch size is ultimately
23

---

Published in Transactions on Machine Learning Research (09/2025)
Table 14: Classification results on the held-out test set FedChest obtained by finding the optimal decision threshold on the validation and using binarize predictions. Top part refers GFL while bottom refers to pFL.
Model
BN-ResNet-50
GN-ResNet-50
SF-ResNet-50
NF-ResNet-50
ANFR
Metric
Accuracy
F1
Accuracy
F1
Accuracy
F1
Accuracy
F1
Accuracy
F1
FedAvg
74.92
42.83
75.78
43.37
75.62
42.85
75.76
43.28
75.80
43.50
FedProx
74.72
42.28
73.41
41.76
74.14
41.60
74.11
41.47
74.16
41.85
FedAdam
74.57
42.60
74.00
41.90
74.57
42.2
74.92
42.84
75.28
43.18
SCAFFOLD
75.55
43.34
76.38
43.85
76.18
43.65
76.27
44.02
76.41
44.07
FedPer 75.23
43.11
75.54
43.59
75.40
43.18
75.66
43.60
75.91
43.75
FedBN
75.62
43.22
N/A
N/A
75.43
43.12
N/A
N/A
N/A
N/A
Table 15: Batch size ablation study on FedChest using FedProx. Smaller batch sizes more strongly affect BN-ResNet due inconsistent mini-batch statistics.
Batch Size 16
32
64
128
256
BN-ResNet
78.67+0003
80.02+0.18
81.79+0.18
82.14+0.1
81.33+0.07
ANFR 79.20+0.009
80.57+0.03
81.71+0.16
82.14+0.1
82.19+0.07
not a viable solution for addressing BN's limitations in non-IID FL, and new methods, such as ANFR, are necessary to effectively combat statistical heterogeneity.
24

---

Published in Transactions on Machine Learning Research (09/2025)
D Extended CSI and Attention Weight Analysis
D.1 Setup details and performance
FL training is performed on the extremely heterogencous 'split-3' partitioning of CIFAR-10 from Qu et al. (2022), which consists of clients who each have samples only from 2 classes. The training paramelers are the same as in Qu et al. (2022) and Section 4.1 All the compared models are pre-trained on ImageNet and have a depth of 50 layers, which results in 16 attention blocks for cach model that uses channel attention. To calculate the channel attention weights and class selectivity index distributions, we use the entire lest. set of CIFAR-10, passing each class separately through the models to extract class-conditional activations; this is done both before and after FL training.
For channel attention weights, this allows us to store the distributions of weights of each model for each class and channel index. For the CST, we query the nearest ReI.U-activated feature maps before and after each channel attention block or the equivalent points for the models that do not use such blocks. timm (Wightman 2019) terminology, we are referring to the output of act2 as before, and act3 as after, Comparing before and after distributions for the same network, allows us isolate the effect of CA in the of SE-ResNet and ANFR, and observe the bascline effect of moving through the convolutional block the CSI distribution in BN-ResNet and Ne-ResNet. Finally, the histogram of CSI values for each layer is used to draw an approximation of the contimuous probability density function for the layer.
D.2 CSI plots of all layers
From Figure D which shows the CSI plots for every layer in the models, make several observations regarding the class selectivity of cach model.
SE-ResNet. Before FL. training, CA reduces selectivity in all but the last block, in which it normalizes it. This is how CA was designed to function in the centralized setting, aiding feature learning in the first layers and balancing specificity and gencralizability in the last layer (Hu al.) 2018). After FL training, the CSI distribution much more left-skewed in the final block, showeasing how BN, under FL data heterogeneity, prohibits the network's last layers from specializing compared to centralized training.
Ne-ResNet. Before FL training we see that selectivity generally increases as we move towards the last layers. The CSI distribution of each layer after FL training is very similar to the one before it, indicating that replacing BN with SWS removes the limitation of the last layers to specialize.
The distributions are generally similar to those NF-ResNet except for some where CA reduces selectivity, adding to the evidence that part of the role of CA in centralized training is aiding general future learning. After heterogencous FL training, ANFR inherits Ne-ResNets robustness against heterogeneity, and by comparing the last layer of NE-ResNet and ANFR, we note that ANFR overall becomes more specialized.
ANFR.
D.3 Channel attention plots of all layers
Finally, Figure shows the CA weights for every CA module of SE-ResNet and ANFR.
25

---

Published in Transactions on Machine Learning Research (09/2025)
Bo JiAcA CA
AUT al
TEE
Selenitty Iado
Sdictiity Indox
Seberidty Iado
Selesiiity Indet
Seleciviiy laaks
Selecivit lakes
Sextiity Indos
Lid
375
A50
6.0
0.15
cle
6.0
115
6/5
300
aid
als
58-Bahh belac
AME
201 Salectory
Selectirry
Seaotiayy Indee
Seleestry hodes
C4
Sdaridty Taden
Saboliiiyy Indr
Sevaliiyy Indee
Selersty Iide
C4
Sl
prcbolis Seiees AMF bd
ANT behee
0.5
64
Gid
275
400
615
345
Cl
6/6
R15
945
030
6/5
Cia
300
Al5
Si-Batti Ni- dalaaaa CA
waciii Ci
k75
615
M65
sio
als
Sedectsity Iades
Saletiiiy Index
Seketray Indes
Sakensey
1.2
Jerie alto
Selectiity Indix
Sileshvity Irdex
Sakctviry Sodir
Salectity 20dee
Figure 6: Full CSI results before and after FL training for cach layer, moving first across cach column then to the next row. In earlier layers CA reduces selectivity, helping the model learn robust features, while in the later ones selectivity is increased to adapt to heterogeneity.
26

---

Published in Transactions on Machine Learning Research (09/2025)
kyyel 0
laserl
Clardd Inlos
Chooad Dakc
Chomclabee
Chande 200
kyyel
layer2.
00
Leyor
layer2
layer5.
Leyer]
layer5.2
Chassilaai
Layor?
kyys1_5
layer
Layork
laverd
Cleraal Endes
Chonact
Chomcllaii
Figure 7: Channel Attention weights for every CA module of Ne- and ANFR (top and bottom row of each layer plot, respectively), before and after FL training (left and right). Note the increased variability for ANFR, particularly in the last layer.
27


Reviews:
[{"review_id": "3BpwmpE6z6", "reviewer": "Reviewer_8K4R", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths - \n1. The paper is well written with a good explanation.\n2. The paper performs thorough evaluations across methods and architectures.\n\nWeaknesses-\n1. The proposed method is a combination of existing works applied to FL. Although the method looks effective, it is not an entirely new method or contribution. \n2. The paper lacks experiments in high non-iid settings (Dirichlet distribution alpha=0.1) which is a very common practice in FL community. Since the paper is focussed on addressing the statistical heterogeneity issue, it is important to demonstrate the effectiveness of the method in highly heterogeneous settings which are commonly used now in FL.\n3. Some relevant baseline methods should be added in the comparison (FedBN).\n4. Several FL works [1] also studies feature-shift settings beyond statistical heterogeneity. Can the authors discuss how the proposed architectural changes will affect in those settings (for example, in DomainNet)?\n\n[1] Xiaoxiao Li, Meirui JIANG, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. In International Conference on Learning Representations, 2021.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper focus on tackling statistical heterogeneity among client dataset\ts in federated learning and propose adaptive normalization-free feature recalibration. The authors propose to use weight standardization and channel attention to explicitly address the issues arising from non-iid data. The proposed method is validated on several federated learning benchmarks with consistent improvements in performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would request the authors to address the four points in weaknesses mentioned above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rPNFsffDkb", "reviewer": "Reviewer_3Xf6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The proposed method shows potential in improving traditional normalization techniques, which is crucial in Federated Learning, where data heterogeneity is a significant challenge.\n2. The method is demonstrated to perform well across various datasets and aggregation methods, which suggests its adaptability in practical scenarios.\n\nWeaknesses:\n1. The proposed approach largely builds upon existing work in weight standardization and channel attention, offering limited novelty. The paper does not sufficiently highlight its contributions in a distinct manner compared to prior research.\n\n2. The paper overstates its contribution, particularly in the abstract, where the method is described as \"the first architecture-level approach.\" It is unclear what is meant by \"architecture-level approach\"？\n\n3. There are numerous methods addressing heterogeneity in Federated Learning, such as regularization-based methods (e.g., FedProx, Ditto), personalized aggregation methods (e.g., FedALA, APPLE), model partitioning approaches (e.g., FedAS, FedRoD), and feature matching methods (e.g., FedProto), among others. The paper focuses solely on normalization, but it does not explain the advantages of this approach over the alternatives. \n\n4. The literature review is not comprehensive, and the discussion of the method’s advantages is insufficiently clear.\n\n5. While the authors list several similar approaches to their proposed method in Table 1, they do not include experimental comparisons with the latest methods, such as FixBN, FBN, and FedWon.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a method that integrates weight standardization and channel attention to improve traditional normalization techniques, such as Batch Normalization, within heterogeneous Federated Learning (FL) settings. The proposed approach aims to tackle challenges in FL by enhancing the adaptability of model training across different devices with varying data distributions. Through experiments, the method demonstrates effectiveness across different datasets and FL aggregation strategies.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The authors should emphasize how their method builds upon existing work in weight standardization and channel attention and what distinct contributions it makes. \n2. The claim that the method is \"the first architecture-level approach\" should be explained more clearly. \n3. The literature review should be more comprehensive, particularly those that address heterogeneity such as, FedProx, Ditto, FedALA.\n4. The paper should include experimental results comparing the proposed method with the latest techniques such as FixBN, FBN, and FedWon.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8xjblK9Nud", "reviewer": "Reviewer_i4b7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n1. Unlike prior works that mostly focus on aggregation strategies, this work targets with model architecture directly to solve heterogeneity. \n2. ANFR achieves great results across all tested datasets and aggregation strategies, when comparing with ResNet variants.\n\n**Weaknesses**\n1. The authors evaluate ANFR exclusively against ResNet-based variants. However, they omit analysis or benchmarking against other modern architectures such as EfficientNet, ConvNeXt, or Vision Transformers, which have demonstrated advantages in federated heterogeneous settings.\n2. While the empirical results are comprehensive, the paper would benefit from a theoretical explanation or analysis to support why SWS+CA complement each other effectively in mitigating data heterogeneity.\n3. Authors overclaim the novelty of their approach. Prior works[1] have already explored architectural design tailored to heterogeneous federated learning, but these are neither cited nor compared in the related work.\n\n*[1] Xu, P., Wang, Z., Mei, J., Qu, L., Yuille, A., Xie, C. and Zhou, Y., 2023. Fedconv: Enhancing convolutional neural networks for handling data heterogeneity in federated learning. arXiv preprint arXiv:2310.04412.*", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces ANFR, an architectural level method tailored for heterogeneous FL. It integrates scaled weight standardization with channel attention to enhance model robustness to data distribution shifts across clients.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weakness part.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)