Understanding the Role of Layer Normalization in Label-Skewed Federated Learning
Authors: Guojun Zhang
Mahdi Beitollahi
Alex Bie
Xi Chen
Venue: Accepted by TMLR ()

Abstract: Layer normalization (LN) is a widely adopted deep learning technique especially in the era of foundation models. Recently, LN has been shown to be surprisingly effective in federated learning (FL) with non-i.i.d. data. However, exactly why and how it works remains mysterious. In this work, we reveal the profound connection between layer normalization and the label shift problem in federated learning. To understand layer normalization better in FL, we identify the key contributing mechanism of normalization methods in FL, called feature normalization (FN), which applies normalization to the latent feature representation before the classifier head. Although LN and FN do not improve expressive power, they control feature collapse and local overfitting to heavily skewed datasets, and thus accelerates global training. Empirically, we show that normalization leads to drastic improvements on standard benchmarks under extreme label shift. Moreover, we conduct extensive ablation studies to understand the critical factors of layer normalization in FL. Our results verify that FN is an essential ingredient inside LN to significantly improve the convergence of FL while remaining robust to learning rate choices, especially under extreme label shift where each client has access to few classes.

Published in Transactions on Machine Learning Research (01/2024)
Understanding the Role of Layer Normalization in Label- Skewed Federated Learning
Guojun Zhang
muoju..zhang@haaueeee
Huowei Noah's 8 Ark Lab
Mahdi Beitollahi
mahdi, haaddbeettolla@@@@uuuussc.o
Huawei Noah's ' Ark Lab
Alex Bie
alexbie
Huawei Noah's Ark
Xi Chen
ai. ai.chen, Chwwwei..com com
Huawei Noah's Ark
Reviewed on OpenReview: https: openreview. net/ forum? id- 6BDHUkSPHa
Abstract
Layer normalization (LN) is widely adopted deep learning technique especially in the era of foundation models. Recently, LN has been shown to be surprisingly effective in federated learning (FL) with non-i.idd data. However, exactly why and how it works remains mysterious. In this work, we reveal the profound connection between layer normalization and the label shift problem in federated learning. To understand layer normalization better FL. we identify the key contributing mechanism of normalization methods in FT., called feature normalization (FN), which applies normalization to the latent feature representation before the classifier head. Although LN and FN do not improve expressive power, they control feature collapse and local overfitting heavily skewed datasets, and thus accelerates global training. Empirically, we show that normalization leads to drastic improvements on standard benchmarks under extreme label shift. Moreover, we conduct extensive ablation studies to understand the critical factors layer normalization in FL. Our results verify that FN is an essential ingredient inside LN to significantly improve the convergence of FL while remaining robust to learning rate choices, especially under extreme label shift where each client has access to few classes. Our code is available at https:/ http://gtttbb.tt com/ htt.a///eeeeeeeteaeeaaaeeereeeee
Maraaa/LLyyenNommllaaaiiiooo Normalization
1 Introduction
Figure 1: Visualization of label shift.
Federated learning (FL, McMahan et al. 2017) is a privacy-enhancing distributed machine learning approach that allows clients to collaboratively train a model with- out exchanging raw data. key challenge in FL handle data heterogeneily, where clients may have differ- ing training data.
Clent
Cilent 2
Label shift is crucial data heterogencity scenario where individual clients have varied label distributions. For in- stance: client a cat lover and only has cat images, while client has lot of dog images and few cat images. Such difference in client label distributions can cause sub- stantial disagreement between the local optima of clients
Ting
CCV
(ww
CMM
1rg
Client 3
CH
Dag
3eg
Oeer
CMl
Dea
Oeer
Hamogeneaus data
Heterogensous data
(one cloes label shift)
Hotorogenoous

---

Published in Transactions on Machine Learning Research (01/2024)
Table 1: Comparing normalization methods with state-of-the-art federated learning approaches. n class(es) means that each client has only access data from n class es). Dir(3) denotes partitioning with symmetric Dirichlet distribution (Hsu et al., 2019). Our setup includes 10 clients for the CIFAR-10 dataset, 50 and 20 clients for the CIFAR-100 datasct with class(es) and Dir(0.1), respectively, and 200 clients for TinyIma- geNet. Normalization (FedLN & FedFN) beats alternatives in wide variety of settings; see Section 3 for further details.
CIFAR-10
CIFAR-100
TinyImageNe
Methods
class
2 classes
Dir(0.1)
2 classes
5 classes
Dir(0.1)
Dir(0.01)
Dir(0.02)
Dir(0.05)
FedAvg
55.0
65.6
71.7
32.3
36.2
38.1
21.2
20.9
21.9
FedProx
54.7
66.7
71.4
32.4
36.6
38.8
21.0
20.8
22.1
SCAFFOLD
55.3
67.1
71.1
32.3
36.6
38.7
20.7
20.7
21.7
FedLC
10.0
57.9
65.0
6.0
18.2
22.3
0.8
0.4
0.5
FedDecorr
43.9
70.2
69.9
29.2
33.0
34.2
22.2
22.4
1.5
FedRS
10.0
56.7
66.1
10.7
24.9
22.1
14.8
18.4
19.0
FodYogi
80.9
80.0
79.2
44.6
44.1
45.2
22.5
24.3
25.7
FedLN
87.2
88.0
89.1
46.6
46.2
47.3
38.2
39.1
38.1
FedFN
80.8
82.3
84.2
40.3
45.4
45.4
34.0
34.1
34.1
and the desired global optimum, which negatively impacts FL performance. Taking label shift to the limit, client may only have access to only one class data, as illustrated in Figure 1. This is the case for problems such as speaker identification, where each user has only their own voice samples. Our study focuses on such eatreme label shift, where baseline FL approaches see drastic drops in utility.
Various methods have been proposed to mitigate the data heterogencity problem in FL. For example, Fed- Prox (Li et al., 2020a) adds l regularization to avoid client. model divergence. SCAFFOLD (Karimireddy et al., 2020) applics variance reduction to control the variance of client modcls. These frameworks focus regularizing client drift in the FL process. FedLC (Zhang et al., 2022) and FedRS (Li & Zhan, 2021) attempt to address label shift by adjusting clients' model logits. We find that these techniques do not suffice to obtain good performance in the extreme label shift setting (see Table 1).
An orthogonal approach to the aforement ioned methods is through uormalizing the neural networks. In Hsieh et al. (2020), the authors point out the advantage of group normalization (Wu & He, 2018) in label-shifted FL problems over batch normalization (Loffe & Szegedy, 2015). Du et al. (2022) attributes the shorteomings of batch normalization to external covariate shift. There has also been studies such Li et al. (2020b) and Wang et al. (2023) that modily batch normalization for FT. Most closely related to our work is the very recent study of Casella et al. (2023) that experimentally compares various normalization methods in foderated learning on MNIST and CIFAR-10.
Our contributions. Previous work lacks an in-depth analysis of layer normalization, and does not explain when and why layer normalization works in FL. In this work, we dive deeply into the theoretical and experimental analysis of layer normalization for federated learning. Our main finding is that as client label distributions become more skewed, layer normalization helps more. This is because under heavy label shift, each client easily overfits its local dataset without normalization, and LN effectively controls local overfitting by mitigating feature collapse. For example, in the extreme one-class setting (Figure 1), LN can improve over FedAvg by 32% in terms of absolute test accuracy (see Table 1).
To further understand the effect of LN for FL, ask the following question: is there much simplified mechanism that works equally well as LNY The answer is positive. From our analysis of local overfitting and feature collapse, we discover that the key contributing mechanism of normalization methods in FL is feature normalization (FN), which applies normalization to the latent feature representation before the classifier head. Feature normalization simplifies LN while retaining similar performance (Table 1).
Another in-depth analysis we make is the comprehensive experiments and ablation studies. We compare LN and FN hot variety of popular datasets including CIFAR-10/100, TinyImageNei and PACS, with
2

---

Published in Transactions on Machine Learning Research (01/2024)
various data heterogencity and neural architectures such CNN and ResNet. We are the first to show that layer normalization is the best method so lar for label shift, problems, outperforming existing algorithms like FedProx, SCAFFOLD and FedLC. Moreover, our ablation studies thoroughly analyze each factor of LN, inchuding with/ without mean-shift, running mean/varian variance, and before/after activation. Our empirical analysis is much deeper and broader than any of the previous results.
We summarize our contributions as follows:
Under extreme label shift, provide comprehensive benchmark to clearly show the dramatic advantage of LN over popular FI. algorithms for data heterogeneity;
We are the first to propose and carefully analyze the suitability and properties of layer normalization in label-skewed FL problems, both theoretically and empirically;
We discover the key mechanism of LN in such problems, feature normalization, which simply nor- malizes pre-classificatio layer feature embeddings.
2 Preliminaries and Related Work
In this section, we review the necessary background for our work, including federated learning and layer normalization.
Federated learning (FL) aims to minimize the objective S(0, W) E S' (0, W), with each client loss as:
fa f = -(,,,))D, W; æ, v),.
(1)
Here m& is the number of samples of client k, and m = En my. We denote Dx as the data distribution of client & and we have my i..dd samples from to estimate fe. We use O to represent the parameter collection of the feature embedding network go and W (wi,..., wc) as the soft weights for C-class classification. The per-sample loss is thus:
exp(w g0
( 2z = - -log
(2)
Eee[c] exp(w) ge(z))
In the classical FL protocol (FedArg, McMahan et al., 2017), a central server distributes the current global model to a subset of participating clients. During each communication round, selected clients perform multiple steps of gradient update on the received model with its local data and upload the updated model to the server. Finally, the server aggregates the clients' models to refine the global model.
Class imbalance problem in FL. Since its inception, foderated learning studies the class imbalance problem, i.e., different clients have different label distributions. In the seminal FL paper (McMahan et al., 2017), the authors considered partitioning the dataset with different shards, and giving each client two shards. In this way, each client would have two class labels. Another common way to create such class imbalance is to use Dirichlet distribution (Hsu et al., 2019). FedAwS (Yu et al., 2020) considered the extreme case when we have one class per each client. FedProx (Li et al., 2020a), FedI.C (Zhang et al., 2022) and SCAFFOLD (Karimireddy et al., 2020) addressed class imbalance by modifying the objective or the model aggregation. FedOpt (Reddi et al., 2020) considered adaptive aggregation as inspired by adaptive optimization such Adam (Kingma & Ba, 2014). More recently, FedDecorr (Shi et al., 2023) pointed out the feature collapse in federated learning and proposes a new algorithm to mitigate data heterogeneity. Shen et al. (2022) proposed CLIMB which studies class imbalance FL problems using constrained optimization.
Normalization. Suppose æ E IRd is a vector, we define the mean-variance (MV) normalization of as:
(MV(x) = x - p(()) u(x)= = 7
d
1 d
Σ Tzi (x)) =
1 11>(6 (x; - -2())²²
(3)
o(x))
\
i=1
3

---

Published in Transactions on Machine Learning Research (01/2024)
Here i and are standard notions of mean and variance. Symbolically, layer normalized feed-forward neural net is function of the following type:
MMY p A o NMV o po AL-1... DMV o po A
(4)
where p is the activation and A. = U.a + b; is an affine function for a vector a.
We can generalize MV normalization to tensors by considering the mcan and variance along each dimension. Based on the discussion above, batch normalization (BN, Ioffe & Szegedy 2015) is just MV normalization along the direction of samples, and layer normalization (LNN, Ba et al. 2016) is MV along the direction of hidden neurons. Unlike BN, layer normalization can be applied to batches of any size and does not require statistical information of batches of each client. Other normalization methods include group normalization (Wu & He, 2018), weight normalization (Salimans & Kingma, 2016) and instance normalization (Ulyanov et al., 2016). Other than normalization on the neural network, Francazi et al. (2023) proposed normalizing the per-class gradients to address the class imbalance problem.
MV normalization can be rewritten as:
(=) vd x 1(x))
Vd. n'(a u(=)1),,
- u(x)
with nx(x) = ///l/>/ In other words, it is a composition of mean shift, division by its norm, and scaling operation with factor vd (see also Brody et al., 2023). If u(æ) = 0, we obtain
n(x) vd n'xx).
We call the function n as the scale normalization, which shares similarity with RMSNorm (Zhang & Sennrich, 2019). This function retains scale invariance but loses shift invariance. To improve stability, we can replace //||a=|| with æ/ max(s, [|=]|]]
We may similarly construct another normalized neural network by replacing nMv with n in eq. 6:
n nopo AL on oo po A1,
(6)
which we will call feature normalized (FN) neural net. This name will be clear with Proposition 1.
Normalization in FL. The fact that group norm (Wu & He, 2018) is better than batch norm in FL. was observed in Hsich et al. (2020). Li et al. (2020b) proposed FedBN and adapted batch normalization in FL, where each client model is personalized. Analysis of batch norm in FL can be found in Du et al. (2022) and Wang et al. (2023). In Casella et al. (2023), the authors tested different normalization methods in FL on non-i.i.d. settings. Compared the these papers, we are the first to observe and analyze the connection between layer/feature normalization and label shift.
3 Efficacy of Layer Normalization
Experimental setting. To show the efficacy of layer normalization in label-skewed FL problems, we conduct an extensive benchmark experiment to compare FedAvg + LN (FedLN) and FedAvg + FN (FedFN) with several popular FI. algorithms, including the original FedAvg (McMahan et al., 2017), in addition to:
Regularization based methods: FedProx (Li et al., 2020a), SCAFFOLD (Karimireddy et al., 2020) and FedDecorr (Shi al., 2023);
Logit calibration methods for label shift: FedLC (Zhang et. al., 2022) and FedRS (Li & Zhan, 2021):
Adaptive optimization: FedYogi (Reddi et al., 2020).
To simulate the label shift problem, we create two types of data partitioning. One is n class(es) partitioning where each client has only access to data from n class(es). Another is Dirichlet partitioning from Hsu et al. (2019). As discussed in Hsieh et al. (2020), the label skew problem is pervasive and challenging for decentralized training.
We test the comparison on several common datasets including CIFAR-10, CIFAR-100 (Krizhevsky et al., 2009) and TinyImageNet. (Le & Yang, 2015) with CNN and ResNet-18 (He et al., 2016).
4

---

Published in Transactions on Machine Learning Research (01/2024)
Results.
Our results are displayed in Table 1. We conclude that:
Under extreme label shift, all baseline algorithms do not show a clear edge over the vanilla FedAvg algorithm, except Fed Yogi. This demonstrates the frustrating challenge of label skewness.
YedYogi can drastically improve FedAvg in such cases (e.g. CIFAR-10), but larger datasets like TinylmageNNet, the improvement is marginal.
FedLN is the only algorithm that can dramatically improve FedAvg in all scenarios.
FedPN largely captures the performance gain of FedLN, despite being consistently inferior to FedLN.
Because of the drastic improvement of layer norm and feature norm label-skewed problems, we will carefully explore the reasons behind, both theoretically and experimentally, the following sections.
4 Why Is Layer Normalization So Helpful?
Inspired by the success of layer/feature /feature norms, we wish to obtain some primary understanding of the reason why they are so suitable under label shift. Our main findings include: (a) Layer norm and feature norm can be reduced to the last-layer normal lizatioin:: (b) the last-layer normalization helps address feature collapse which happens in the local overfitting of one-class datasets. This accelerates the training of feature embeddings. (c) the main advantages LN/FN lic in training process, rather than the expressive power.
4.1 Delegating scaling to the last layer
We first show that both feature normalized and layer normalized neural networks can be simplified to the last-layer scaling, under the assumption of scale equivariance.
2..... , L, and Assumption 1. The bias terms from the second layer onward all 0, i.c., b, for =2, the activation function is (Leaky) ReLU.
For an input æ, denote the ith activation as a; := po... po (zzz). Under Assumption 1, a vanilla neural network is scale equivariant w.r.t. the first layer activation, i.., suppose X > and h := p AL... p A is the neural network function from the second layer then h(Xai) = Nh(ai) holds. Assumption 1 is also necessary for such equivariance.
Proposition 1 (reduced feature normalization). Under Assumption 1, scale normalizing each layer is equivalent to only scale normalizing the last layer. That is, for all affine transformations A Ar the function
oon
is equal to
nopo Aropo po A
(7)
if all the intermediate hidden vectors after activation non-zero.
The formal statement of this proposition can be found in Proposition 1'. In this proposition, we used blue color to highlight the normalization components before and after the simplification. Details of Proposition can be found in the appendices. Eq. 7 can be considered scale normalizing the feature embedding with vanilla network, which is why we called it feature normalization in eq. 6. Note that Proposition 1 tells that starting from the same random initialization for A A AL, training both formats will result in the same final network.
Similarly, we can simplify a layer-normalized neural network, by defining the shift operator:
s(x: = x - p(x)
where is arbitrary finiteddimensional vector.
5

---

Published in Transactions on Machine Learning Research (01/2024)
Proposition 2 (reduced layer normalization). Under Assumption 1, MV normalizing each layer is equivalent to only MV normalizing the last layer and shifting previous loyers. That is, for all offine trans- formations A Ar the function
næv o po o Ar. o mMv o po Ar. DMV po A
is equal to if none of the intermediate hidden layer aclivations to normalize are proportional to the all-one vector The formal statement of Proposition 2 can be found in Proposition 2'.
DMV po Ar, osopo A..I..
(8)
Extension to ResNet. So far we have talked about simple feod-forward networks (including MLP and CNN). We can extend our results of Prop. and Prop. 2 to other model architectures like ResNet (He et al., 2016), as long as the original model is scale equivariant. Note that ResNet is a composition of multiple blocks, where each block is:
block(x) = p(æ A₂ o po A1(x))
(9)
where each of A2, Ag is a linear transformation. For any À > 0, we would have block(Aæ) = Ablock(x)) If we add layer normalization this block, becomes:
blockuN(x) nMv p(z + A nMv p A1(x)).
(10)
A similar conclusion as Prop. 2 would be to replace the block above as:
blockíN(2) s p(x x + A2 o nMv o p A (z)).
(11)
However, cannot replace the second nMv with mean shift, which causes larger gap between LN and only normalizing the last layer feature as we will see in Table 2.
Our conclusions do not immediately exlend to Vision Transformers (Dosovilskiy et al., 2020), since the multi-head attention is not scale equivariant due to the use of softmax.
4.2 Expressive power
The reduction of feature norm and layer norm allows us to analyze the expressive power of LN A/NN networks. Denote c, EF, & as the 0-1 classification error of vanillaFN/INN network respectively. We have:
Proposition 3 (expressive power). Given any model parameters (0, W) and any sample (æ, with go(x) * 0, vanilla and FN networks have the same error on sample (x. y), i.c., e(0, W; a; y) (0 w) For any model parumeler (0, W) of layer normalized nelwork, one can find (0', W of a vanilla model such that =(0', x.; = (0, W; y), for any æ,
Let explain the significance of this proposition. In this work aim to understand the role of layer normalization in federated learning. There could be two reasons why LN is helpful:
Layer normalization is more powerful in that it can express better functions to fit the data.
Layer normalization allows the training process to be faster.
Proposition 3 rules out the first explanation. It tells us that in principle, a vanilla network can Iearn any pattern that LN/PN networks generate. If an LN/FN network could reach 80% accuracy, so can vanilla network by adapting its parameters. Moreover, the classes of FN and vanilla networks are equivalent, but FN still outperforms vanilla networks in our settings. This rules out the explanation of model class restriction. Therefore, the benefit of FFN we saw in Table 1 lies in the training process, which we will present in the next subsection.
There still remains an important question regarding Proposition 3: since vanilla NNs can express LN net- works, is the reverse true? Or could LN networks express any "meaningful" NN functions that fit the data? This question is out of our scope and we leave it to future research.

---

Published in Transactions on Machine Learning Research (01/2024)
Table 2: Comparing the difference between LN and FN for CNNs and ResNet. We observe a smaller accuracy gap for CNNs, as predicted by our theory.
Methods
class
classes
Dir(0.1)
FedLN CNN
78.04
77.52
78.27
FedFN CNN
77.14
76.67
78.13
FedLN ResNet
86.15 80.81
88.01
89.06
FedFN- ResNet
82.30
84.19
Fre trainaandd modcl
Vanilla after S steps
EN ofter 5 stops
100
200
Cllnt's data Other
20
20
20
Figure 2: Local overfitting in the one-class setting on CIFAR-10. The client only has examples from class 0. The blue lines show the average global performance. (left) the test accuracics of the pro-trained model before local training; (middle) after steps of local training with a vanilla model; (right) after 5 steps of local training with FN. Best viewed color.
4.3 Normalization essential for label shift
Now that we have simplified normalized networks to the last-layer scaling, our next question is: what the relation between last-layer scaling and label shift? In order to further illustrate this connection, we consider the most extreme one-class setting, where each client has samples from only one class (see Figure 1). We defer the exploration of more general label shift to future work.
Label shift induces local overfitting. The scarcity of label variation on client poses a risk of local overfitting, which we define below:
In FL, local overfitting describes the situation when a client model performs extremely well on its local dataset, but fails to generalize to other clients.
We illustrate local overfitting in Figure 2, where the client dataset only has samples from class 0. After local training, the vanilla model easily reaches 100% on its local dataset, while the performance on other clients drastically drops. This is true even when it is initialized from a relatively good pre-trained model (left figure). This resembles the well-known phenomenon of catastrophic forgetting (McCloskey & Cohen, 1989). Comparably, feature normalization can mitigate local overfitting (Figure 2, right).
Let try to understand local overfitting more clearly. Given the dataset S₃ {(x;, k) for client k. feature embedding ge and class weight vectors (class embeddings) wy 's, the cross-cntropy loss of client k. eq. becomes:
(O, W) = B(a.ess log Σ exp((w. wx) (a)oo((())) RE(C)
Minimizing this loss function drives (w, wr) go + 00 for all æ; and e e k leading pathological behavior. For example, long as (wa-wo (x,) for all and e + k both training and test accuracies reach 100%, which often occurs in practice. We present necessary condition to minimize fx(0, W):
7

---

Published in Transactions on Machine Learning Research (01/2024)
weights
Feature norm samples
Fedfig Venilla FodFN
10
000 00000 d0lle
Log singul or values
04
60
Classes
Samples
Index
Index
Figure 3: Local training with only samples from one class. (left): vector norms of each class embedding. (middle left): the norms of different feature vectors. We randomly choose 20 images from the dataset. (middle right): singular values of features of a local overfitted model; (right): singular values of normalized features learned from FedAvg and FedFN.
Theorem 1 (divergent norms). In order to mánimize the one-class local loss fa(00 W), we must have at least one of the following: 1) /well + oc for all c. E [C] and c. + k; 2) [ + oo for all w;; 3) || + oc.
Theorem tells us that either some class embedding norms all feature embedding norms must diverge. Thus it points out the importance of controlling the feature/class embedding norms. If we do not add any feature/layer normalization, minimizing the label skewed local dataset could result in divergent norms.
Which norms diverge? Tn Figure we perform local training on client with samples from only one class of the standard CIFAR-10 dataset. In the left panel, calculate the norms of each w. with E [10]. It can be shown that ||wa|| || is relatively large compared to others. In the middle-left panel, we sample 20 images from client 1 and compute their feature norms. We see that lige ll's are all very large, compared to class embeddings. Combining with Theorem 1 and Figure 3, we can argue that controlling the divergence feature norms more important, which emphasizes the importance of LN/FN in label-skewed FL.
Feature collapse explains local overfitting. If we plot the singular values of the feature matrix of 20 random images from different classes, (z) (0()))) we observe huge spectral gap. While the first singular value 0 ~ 1.9 103, the second singular value is 02 6.9, which implies that the feature embeddings are approximately in one-dimensional subspace. This is known as feature collapse (c.f. Shi et al. 2023), as the feature embeddings of different classes are mapped to the same directions. This explains the local overfitting of one-class datasets.
Normalization addresses feature collapse. In fact, in the one-class setting, there no need for the feature embedding to distinguish images from different classes, especially when there is not much variation in the local dataset (such as the one-class setting). For example, the embedding can simply increase the feature norms without changing the directions.
In contrast, if the feature norms are constrained (like in LN/FN), then each client cannot overfit by increasing the feature norms, but has to learn the directional information of the feature embeddings. This accelerates the training of feature embeddings under heavy label shift.
We verify this claim on the right of Figure 3. While a vanilla. network learns degenerate features, an FN network mitigates feature collapse and the feature embedding has more directions that are effective, measured by the corresponding singular values (c.f. Shi et al. 2023).
Communication cost and privacy. Adding FN or LN does not increase the number of parameters of the network, and only cause minor additional computation at each client. Therefore, vanilla FedAvg and FedFN have the same communication cost for sending models to a server at each communication round. Further, as implied in Table 1, due to faster convergence, FedFN requires less communication rounds reach pre-defined accuracy, and therefore, FedFN has better total communication cost..
8

---

Published in Transactions on Machine Learning Research (01/2024)
FedAvg Vanilia
FedFN
FedLN
Centralized Vanilla
Accuracy
5000
10000
15000
20000
25000
30000
35000
40000
Batch
Comparing test accuracies of centralized training with FedAvg of models with different normal- izations in one-class label shift in CIFAR-10. The z-axis denotes how much data fed into an algorithm, measured by batches.
Figure 4:
Table 3: Comparing layer-wise vs. last-layer normalization. n and n-1 denote layer-wise wise and last-layer scaling respectively; they perform similarly as predicted by Prop. 1. My and nwys denote layer-wise and last-layer MV normalization respectively; Prop. predicts they perform similarly. The table shows that the assumption bias terms as required by Assumption 1 does not afTect the performance. Our setup includes 10 clients.
without bias classes
with bias classes
Methods
class
Dir(0.1)
class
Dir(0.1)
FedAvg
56.82
71.83
73.06
58.57
72.0
73.17
FedFN n.
77.05
76.60
77.27
77.70
76.29
76.50
FedFN nl L-1
77.14
76.67
78.14
74.46
76.27
76.78
FedLN TM
77.16
78.09
78.71
77.75
76.73
78.36
FedLN nMvs
77.71
77.65
79.94
76.89
76.74
77.26
Unlike BN which requires the record of the rumning statistics of data, during training, LN and FN do not record statistics of batches. Instoad, LN and FN, the normalization is done independently for each sample both training and inference. Therefore, adding FN and LN an FL system is unlikely to degrade privacy.
5
Experimental Analysis
We conduct extensive experiments to further analyze the relation between LN/FN and label-shifted federated learning.
5.1 FN is the essential mechanism of LN
In Figure 4, compare different FedAvg algorithms with centralized training in the one-class setting of CIFAR-10. We apply different normalization to CNN architecture.
If we apply vanilla unnormalized networks to FedAvg, the average lest accuracy grows very slowly compared to centralized training, although the same amount of data is passed to both algorithms. Comparably, the convergence of FedAvg with LN and FN networks much faster and closer centralized training. Morcover, the performance of FodFN FedLN are very similar each other, confirming that FN is the main mechanism of LN under reme label shift.

---

Published in Transactions on Machine Learning Research (01/2024)
85
Accuracy (%6
75
70
FedAvg Vanilla
redHN
FedLN
60
0.02
0.1
0.5
2.0
4.0
B
Figure 5: Effect of data heterogeneity on the performance of normalization. / represents the parameter in the Dirichlet distribution that. is used to sample client. label distributiins..
Table 4: Comparison among normalization methods in FL. FedAvg + BN means average all the parameters after using batch normalization. FedBN is from Li al. (2020b). We use 10 clients with CIFAR-10 dataset.
Methods
1 class
2 classes
Dir(0.1)
FedAvg
56.82
71.83
73.06
FedGN
77.15
77.34
79.15
FedAvg + BN
9.63
24.64
58.80
FodBN
10.00
19.90
41.25
FedFN
77.14
76.67
78.14
FedLN
77.71
77.65
79.94
We also verify Prop. and Prop. by rumning FN/LN before and after simplification. To satisfy Assump- tion 1, we removed the bias terms in the neural net and kept the RcLU activation. Table 3 verifics that up to statistical error, normalization methods can be simplified last-layer normalization. Even more, our simplification helps slightly as the last-layer normalization avoids division in the middle. In all cases, the performance of FN is very close to LN.
We also compare the performance gap between FN and LN with different architectures, including CNN and ResNet. From Table 2 we find that the gap with CNN is smaller than that of ResNet. This agrees with our discussion following 11: the larger gap is due to the LN inside each block that cannot be absorbed into FN. We present additional experimental results regarding this point in the appendices.
5.2 Does normalization always help?
One might argue that layer norm just advantagoous in general and it has no relation with label shift. We show that is not the case. In Figure 5, we compare FedFN with FedAvg with different levels of label skewness. Here / controls the level of label skewness. When 8 is small, the label shift is heavy and there is clear gap between FedFN/FedINN and FedAvg. However, when the clients are more i.i.d., the performance gap diminishes. This revcals the strong connection between LN/FN and label shift.
5.3
Do other normalization methods help?
We compare FN/LN with other candidates of normalization, including group norm (GN, Wu & He, 2018) and batch norm (BN, loffe Szegedy, 2015).
Table 4 shows the comparison in the one-class FL setting. For FedFN/FedLa we use Prop. and Prop. 2 to simplify the normalization layers. We compare FedFN /FedLN with FedGN (FedAvg GN.
10

---

Published in Transactions on Machine Learning Research (01/2024)
Centralized training
90
80
Accuracy (%)
70
50
Vanilla
FN
40
LN
GN
BN
20
40
6C
80
100
Epochs
Figure 6: Comparison of different normalization methods under the centralized setting.
Table 5: PACS dataset comparison. 2-2-3 classes means split each domain into 3 clients, with 2 classes, classes and 3 classes respectively.
Methods
2-2-3 classes
Dir(0.5)
Dir(1.0)
FedAvg
54.35
67.60
68.15
FedProx
55.96
57.14
56.58
Fog
64.56
71.81
72.20
FedFN
71.58
72.72
71.89
FedLN
70.46
69.39
72.72
group number-2) and BN. Note that there are two ways to apply batch normalization: one is simply adding BN as we did for WN/LN/GN, another is to avoid averaging the rumning mean/varii vaarince of the BN layers and keep them local, as suggested by Li et (2020b).
As argued Du al. (2022); Hsich al. (2020): Li et al. (2020b); Wang et al. (2023), plainly adding BN does not improve FedAvg. Our Table 4 verifies this conclusion and shows that FedAvg BN could worsen the performance, especially under heavy label shift. Moreover, each after personalizing the BN layers of each client (Li et al., 2020b), FedBN does not help mitigate label shift either. One explanation is that FedBN was designed for covariate shift problems, while our main setting is label shift.
On the other hand, FedGN bohaves similarly to FedLN and FedEN. This is expected since group norm can treated as slight generalization of layer norm: splits the hidden vector into subgroups and MV normalizes each group. If the group number is one, then GN reduces LN.
However, the situation differs in the centralized sctting,
in Figure 6. When aggregate all the data the server and perform different lypes of normalization, we find that BN/LN/GA improves centralized training similarly, as observed in recent literature (Wu & He, 2018). In contrast, the performance gain of FN disappears. This again verifies the strong connection between FN and label shift.
5.4 Label shift under covariate shift
In real scenarios, we may also face other challenges than label shift. For example, covariate shift may also a great challenge in FL which occurs when the input distributions differ. We test different normalization methods under covariate shift using a classic datasct called PACS (Li et al., 2017), in addition our main problem, label shift. The results can be seen in Table 5. From the table we conclude that FodFN/FodIn still has clear edge under covariate shift.
11

---

Published in Transactions on Machine Learning Research (01/2024)
100
FedAvg
FedFN
FedLN
Accuracy (%)
40
-3.0
-2.5
-2.0
-0.5
0.0
Learning rate (logıc)
Figure 7: Effect of learning rate on the performance of FN, LN and vanilla networks in one-class distribution setting.
Table 6: Comparison among different variations of FedLN.
Methods
1 class
2 classes
Dir(0.1)
FedLN learnable
78.04
77.52
78.27
FedLN static
77.17
78.41
78.27
FedLN before
77.04
78.08
78.27
5.5 Learning rate robustness
In Figure 7, we show the performance of FN, LN, and vanilla methods in the one-class setting of CIFAR-10. With different learning rates, the performance of the vanilla method changes severely, while the variation of FN/LN is relatively small within a large range of learning rates. This shows the case of hyperparameter tuning with normalization.
5.6 Variation of layer normalization
In the implementation of LN, we have in fact implemented learnable parameters (FedLN learnable), by adding an element-wise affine transformation to the vector:
70 1(x)1
B.
0 (x)
If we fix and B = O, then it reduces to usual MV normalization, which we call the static sctting, We compare both methods in Table 6, where we also add using LN before the activation (FedLN before). Such variation does not make a noticeable difference in our setting.
6 Conclusion
In this work, reveal the profound connection between layer normalization and the label shift problem in federated learning. We take first step towards understanding this connection, by identifying the key mechanism of LN as feature normalization. FN simply normalizes the last layer with scale, but is extremely helpful in addressing feature collapse and local overfitting. Under various heavy label shift settings, the advantage of LN/FN is significant compared other state-of-the-are algorithms. Some future directions inchude understanding this connection with more solid theory and testing our finding with a wider range of modalities.
12

---

Published in Transactions on Machine Learning Research (01/2024)
References
Jimmy Lei Ba, Jamic Ryan Kiros, and Geoffrey E Hinton. 1oXiv:607.064500 2016.
Layer normalization.
arXiv preprint
Shaked Brody, Uri Alon, and Eran Yahav. On the expressivity role of layernorm in transformers' attention. arXiv preprint arXiv:2305.022 02582, 2023.
Bruno Casella, Roberto Esposito, Antonio Sciarappa, Carlo Cavazzoni, and Marco Aldinucci. Experimenting with normalization layers in federated learning on non-iid scenarios. arXiv preprint orXiv:2303.1000 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander 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. In International Conference on Learning Rep- resendations, 2020.
Zhixu Du, Jingwei Sun, Ang Li, Pin-Yu Chen, Jianyi Zhang, Hai Li, Yiran Chen, al. Rethinking normal- ization methods in federated learning. arXiv preprint arXtv:2210.03277 2022.
Yonatan Dukler, Quanquan Gu, and Guido Montúfar. Optimization theory for ReLU neural networks trained with normalization layers. In International Conference on Machine Learning, pp. 2751 2760. PMLR, 2020.
Emanuele Francazi, Marco Baity-Jesi, and Aurelien Lucchi. A theoretical analysis of the learning dynamics under class imbalance. In International Conference on Machine Learning, pp. 10285-10322. PMLR, 2023.
Kunihiko Fukushima. Cognitron: A self-organizing multilayered neural network. Biological cybernetics, 20 (3):121-136, 1975.
Angeliki Giannou, Shashank Rajput, and Dimitris Papailiopoulos. The expressive power of tuning only the norm layers. arXiv preprint arXiv:2302.07997 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on compuler vision and pallern recognilion, pp. 770 778, 2016.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip Gibbons. The non-iid data quagmire of de- centralized machine learning. In International Conference on Machine Learning, pp. 4387-4398. PMLRR 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909, 00335,, 2019.
Sergey Tofle and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448 456. PMLR, 2015.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theortha Surcsh. SCAFFOLD: Stochastic controlled averaging for foderated learning. In In- lernational Conference on Machine Learning, pp. 5132 5143. PMLR, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2014.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009. Technical report.
Ya and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
Yann LeCun, Leon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural Networks: Tricks of the Trade, pp. 9 50. Springer, 1998.
13

---

Published in Transactions on Machine Learning Research (01/2024)
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain gen- eralization. In Proceedings of the IEEE International Conference on Computer Vision, pp. 5542-5550, 2017.
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.
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, 2020b.
Xin-Chun Li and De-Chman Zhan. FodRS: federated learning with restricted softmax for label distribution non-iid data. In Proceedings of the 27th ACM SIGKDD Conference on Knowledgee Discovery & Dula Mining, pp. 995 1005, 2021.
Kaiseng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. arXiv preprint arXiv:2206.0706 2022.
Michael McCloskey and Neal Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pp. 109-165. Elsevier, 1989.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiiee learning of deep networks from decentralized data. Artificial Intelligncce and Statisties, pp. 1273-1282. PMLR, 2017.
Maithra Raghu, Ben Poole, Jon Kleinberg, Surya Ganguli, and Jascha Sohl-Dickstein. On the expressive power of deep neural networks. In International Conference on Machine Learning, pp. 2847-2854. PML.R, 2017.
Sashank J Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Represenlations, 2020.
Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in Neural Information Processing Systems, 29, 2016.
Shibani Santurkar, Dimitris Tsipras, Andrew Nyas, and Aleksander Madry. How does batch normalization help optimization? Advances in Neural Information Processing Systems, 31, 2018.
Zebang Shen, Juan Cervino, Hamed Hassani, and Alejandro Ribeiro. An agnostic approach to federated learning with class imbalance. In International Conference on Learning Representations, 2022.
Yujun Shi, Jian Liang, Wenqing Zhang, Vincent Tan, and Song Bai. Towards understanding and mitigating dimensional collapse in heterogencous federated learning. In The Eleventh International Conference on Learning Representations, 2023. URL attps://openreet net http:://ppppprreeoooeeoooorr///e
Eero P Simoncelli and Bruno A Olshausen. Natural image statistics and neural representation. Annual review of neuroscience, 24(1):1193 -1111 2001.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint Xixiv Xi.:::60..0000 107.00002, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
Yanmeng Wang, Qingjiang Shi, and Tsung- Hui Chang. Why batch normalization damage federated learning on non-iid data? arXiv preprint 2XXi:2301.00982 2023.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 3-19, 2018.
14

---

Published in Transactions on Machine Learning Research (01/2024)
Felix Yu, Ankit Singh Rawat, Aditya Menon, and Sanjiv Kumar. Federated learning with only positive labels. In International Conference Machine Learning, pp. 10946-10956. PMTR, 2020.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Guojun Zhang, Saber Malekmohammadi, Xi Chen, and Yaoliang Yu. Proportional fairness in federated learning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https://openreviee het/fformmiideryUHgaaaaa
Jie Zhang, Zhiqi Li, Bo Li, Jianghe Xu, Shuang Wu, Shouhong Ding, and Chao Wu. Federated learning with label distribution skew via logits calibration. In International Conference on Machine Learning, pp. 26311 16129. PMLR, 2022.
A Limitations
There are several limitations of our work. First, we did not try to propose a new algorithm, but rather to analyze and understand existing normalization methods. The main contribution is to unveil the connection between normalization and label shift in FL, and to understand this connection.
Moreover, we did not provide a rigorous proof on the convergence rate, since the loss landscape of neural networks is non-convex and it is a longstanding problem in deep learning to rigorously analyze the optimiza- tion in usual settings. Rather, we take alternative approach using feature collapse and local overfitting to qualitatively explain the learning procedures. A more solid foundation for the optimization of normalized networks is an interesting future direction.
On the empirical side, we did not run our experiments several times for the reproducibility check, due to the expensive computation in FL lasks with heavy label shift. our V100 GPUs, takes few days finish one experiment with 10,000 global rounds and 10 local steps. However, the trend consistent that normalization methods are much faster in heavy label shift problems, from our experiments on multiple datasets and hyperparameter choices. Last but not least, have only tested our methods in vision datasets. Tt would also be interesting to test normalization on other data formats, such as natural language, molecule structure, etc.
B Detailed Theorems and Proofs
We first verify the necessity of Assumption here for the scale equivariance of the embedding from the second layer. The necessity of the bias terms should be straightforward. For the activation function, we observe the following lemma:
Lemma 1 (scale equivariant activation). If R- + R, then p(At) = Ap(t) for any X > 0 and t E R iff is a piecewise linear function, with the form:
at t > 0, where and b are two real numbers.
p(t)
(12)
bc
150,"
Constants a, can in fact be absorbed into the adjacent layers, allowing Leaky ReLU networks to express all neural networks with activation functions like eq. 12.
Proof. Taking = 1, we have p()) = Ap(1) for À > 0. Similarly, p(->) = p(-1)A. Therefore, a a= a=p(1) and = p(-1) are necessary. It remains to prove that p(0) = 0, which can be obtained with = 2, f 0 in p(At) = Ap(t). It also easy to verify that eq. 12 suffices.
15

---

Published in Transactions on Machine Learning Research (01/2024)
Proposition 1 (reduced feature normalization) Under Assumption 1, scale normalizing each layer is equivalent to only scale normalizing the last layer. That is, for all offine transformations A Ar the function
is equal to
nopo A Ar po A1
(7)
if all the intermediate hidden vectors after activation are non-2ero.
Before start the proof, we first provide more formal description of the proposition above. Proposition 1' (reduced feature normalization), Scale normalizing each layer is equivariant to only and (z( is computed from
scale normalizing the last layer. More formally, suppose
a: no p b.), U3 Rdixd.11 for i [L], b1 RR: b, 0, for i [2, I.. If (U0i0 + b1) # and p(U # for any 2,.... Tu, then go the same as go (2), with::
(13)
a' = go x = n PU,, + b,), a; = --,,, + bi), for i [L- 1).
(14)
b, 0 for i =2....,14 and ag = æ.
Proof. For I, = 1 it is easy to verify. If L 2 2, we show that there exists d > 0..... A > 0 such that a. X for 1 < <<< L. For i = 1, we have = ag//laal Ai with A = 1/1|a//l Suppose for 1 L 1, a, = X holds. Then
0:41 = no p(U++1a a.) = p(U))
p(Us+1A,a)) p(U++10))
A141d44
(15)
[p(U++1a;]]
(p(U;+1>aa)) [P(U)++a))]]
where dit1 11110/411ll Therefore, from a =x--ia'--- obtain:
no p( = (Urar-1)
p(Urat_1)
ar. no
[P(ULaL-1)] |p(Una'_ = no p(oop( a'
(16)
Proposition 2 (reduced layer normalization). Under Assumption 1, MV normalizing each layer is equivalent to only MV normalizing the last layer and shifting previous layers. That is, for all affine trans- formations A....., AL the function
nMV po O po Ar. DMV po A
is equal to
DMV o po
A
(8)
if none of the intermediate hidden layer activations to normalize are proportional to the all-one vector 1. We similarly provide a more formal description.
Proposition 2' (reduced layer normalization). Suppose our layer- neural network is:
a, = nMv o p(U,a;-1 + by). U, E Rdxdd for E [L] dd b, =0, for i E L]. with = = go(x), and P(U,a;-1 + bi) + p(pU b;))1 for E [Z]. Then go(x:) computed from eq. 17, is the same as (x), compuled from:
(17)
a'. (x( nMv p(Ura by), a; so p(U,a) bi), for i [L [L-1] 1]. b. 0 for i .,..., L and a æ.
(18)
Proof. The proof follows analogously from the proof of Proposition 1'.
16

---

Published in Transactions on Machine Learning Research (01/2024)
B.1 Expressive power of normalized networks
In this subsection, provide the proof steps for Proposition 3. First, recall:
0 W; æ, v) = x(y & argmax, Bi), with 9 softmax W (x)):
(19)
EF (0, W; x, = ( & argmax, Ut), with softmax W ge
where X is the indicator function. We prove the first part of Proposition 3 for feature normalized networks:
Proposition 3.A (expressive power of FN). For any model porameters (0, W) and any sample (x, y) with go(x) + O, we have e(0, W; 2. (0, w,
Proof. We first show:
=(0, w; = EF(O, W; ae y)
(20)
If the l.s.s. is true, then wy go 2 w go for any This gives
w go
w for any
|1
|Ioel|
From eq. 19 know Er(0, W; a. x) 0. Similarly, we can show
Er ( a; = - c(0, w; v)
(21)
This concludes the proof of FN.
Note that the assumption go + can be removed if we use max maxee) in our feature normalization. Similarly, we can show that vanilla models. From Prop. 2, a pre-activation LN network can be represented as:
layer normalization is not more expressive than
a = (x) = po nMv (U bL), a, po poos(U,,,.1 + bi), for i [L 1],
(22)
with b, 0 for i 2, ...... and du We use EpL to denote the accuracy function of pre-activation LN, and Etp to denote post-activation I.N.
Proposition 3.B (expressive power of LN). For any model parameter (0, W) of pre-actination layer nor- malization, one can find (0', W) of vanilla model such that e(O', w; Epz W; W; ;,,,,, for any æ, y. For any model parameter (0, W) of post-activation layer normalization, one can find (6', W') of vonilla model such that e0', W' x. y) = E (0, W. y), for any æ. y-
Proof. For the proof of LN networks, first observe that pre-activation layer normalization is equivalent to eq. 22. Each layer of the form (i 1,2,..., L):
a p: b.).
(23)
is equivalent to:
a, = P + P.b.),,,
(24)
where P. I- 111 is projection matrix and d, is the dimension of a, Thus, can take U/ P.U,
and b/ = P.b; so that:
p pos(U,a,-1 = p( b().
(25)
Also note that the normalization factor in the last layer does not affect the final prediction, seen from Proposition 3.A. Therefore, collecting all the U and by for = 1,2,.... ...., To we obtain the required (0', W).
17

---

Published in Transactions on Machine Learning Research (01/2024)
We can show similar conclusion for post-activation LN. Suppose (U1, U,,.,,.....UU b1, W) are the parameters of an LN network. We can first take U = U1 by b1 such that:
= (U1 bb) a1.
(26)
where a1 is the output of the first layer of vanilla network. Denote di = s(z,) and 2; p(U,a, + by) for i 2 1. Assume 2 a:-1 where a, is the outpul of the ith layer of vanilla network with parameters
(2 .0/2....U.... b, W then we can prove that for 2 2:
2, p(U,d'_1) p(Uis(z)-1)) - p(U;ai = aii
(27)
if we take U/- 1 Thercfore, by induction, prove that 2, a; for i 1, 1.2..... .... L. Specifically, we have 21. = (0U,.a',.1 br.) ar. az.
Since L. (0, w; v) = 0 iff:
w.J a', 2 w a/, for all c C [C],
(28)
and a1g = =nmv(z)) = s(z,) [z1) = P(zz//(zz)) with Pt. =I J. 111, eq. 28 is equivalent to:
w PLaL 2 w. PLaz, for all c e [C].
(29)
Hence, taking W Pt W gives the desired vanilla network.
Next, we can prove the following corollary:
Corollary 1. Suppose s", Sp and sú are the optimal prediction errors, by vanilla, feature-normalized and (pre/post--cclivation) layer-normalized neural nelworks respectively on a fixed datasel S:
$2,F,L = inf (E (0. W; O.W
(30)
Then we have: = S E. Note that ez has the same meaning
Proof. Suppose for an LN network, (0; W;)) is a sequence that can achieve prediction error E + 6, with 6, $ O. Then from Proposition 3.B, there exists (0"", W!") of vanilla network which can achieve prediction error Et, + & and 8, $ 0, and thus Ex < Et. Similarly, we can show g E from Proposition 3.A.
B.2 Proof of divergent norms
Theorem 1 (divergent norms). In order to minémize the one-class local loss fa W), we must have at least one of the following: 1) [.wlll oc for all [C] and + k; 2) (x; co for all æx; 3) ||w]|| + xc.
Proof. We first show that it is possible and necessary to have (w, wt) go 0 for all x, Sş and * k. Suppose the layer affine transformation is Ay(a) U,a bj, with U, Rd xdj and by Rd,
Since # 0, there exists To E IR such that p(za) * 0. For the possibility, by taking U, = 0, b = xn and U, - 11 bj 0 for j 2 2, we have a, = p(xa)1 1 for all j 2 1, and specifically go = az = p(zo)1. If we choose wE = p(xr)1 1 and w. = - for all c + k, then (w. - wxx) go < can be achieved. Moreover, by replacing W with tW and taking + co, the loss (0, W) can be arbitrarily close to zero
For the necessity, suppose otherwise (w. wx) go (a:)) 2 0 for some æ; E Sy and c c k. Then fx () 2 L log C. However, we have seen in the last paragraph that the infimum of fa is zero. my
T
Now suppose (w.- - wx) go (x:)) < 0 for all x; E S and e * & and all the three claims are false. Then || is bounded, |wwell bounded for some e # k. and (xx))] is bounded for some I;. So (0, W) E(O, a., k)/me is lower bounded. This is contradiction.
2
18

---

Published in Transactions on Machine Learning Research (01/2024)
C Experimental Setup
In this appendix, we provide details for our experimental setup.
C.1 Datasets and models
To evaluate the performance of our proposed method common FL benchmarks. We use four datascts: CIFAR-10/ 1000 (Krizhevsky el al., 2009), TinyImageNel (Le & Yang, 2015), and PACS (Li et al., 2017). For CIFAR-100, we use the CNN network (Fukushima, 1975), and for CIFAR-10, we test both CNN and ResNet-18 (He et al., 2016). For the remaining datascts, we use ResNet-18. Since we are investigating the effect normalization on the performance of the FL, we use ResNet-18 without any normalization (i.e., CNN + skip connection) as our base method. In ResNet-188 we add layer norm (LN), group norm (GN, Wu & He, 2018) or batch norm (BN) after each activation, and in ResNet-18 with feature norm (FN), we only apply scale normalization to the output features of the network before the last fully connected layer. Similarly, for CNN models, we add normalization methods after each activation. For all the normalization met hods including FN, we set = le 5 for stability (sec the preliminaries). For the implementation of LN/Gh /BN, we may track the running statistics and average all the parameters including the running means and running variances.
Table 7: CNN model for CIFAR-10 dataset. denote le as the kernel size, as the mumber of channels, and 8 as the size of the stride.
Models
Layer
Vanilla
LN
BN
FN
Input (shape)
(28, 28, 3) (5, 64, 1)
(28, 28, 3)
(28, 28, 3)
(28, 28, 3) (5, 64, 1)
Conv2d (k,c,
(5, 64, 1)
(5, 64, 1) batch norm (2, 2)
Normalization
layer norm (2, 2)
MaxPool2d
(2, 2)
(2, 2) (5, 64, 1)
Conv2d (k, c, s) Normalization
(5, 64, 1)
(5, 64, 1)
(5, 64, 1)
layer norm (2, 2)
batch norm (2, 2)
MaxPool2d
(2,
(2, 2)
Flatten (shape)
1600
1600
1600
1600
Dense (in e, out c)
(1600, 384)
(1600, 384)
(1600, 384)
(1600, 384)
Normalization
layer norm (384, 10/100)
batch norm (384, 10/100)
feature norm (384, 10/100)
Dense (in e, out c)
(384, 10/100)
In order preserve the scale invariance (cquivariance), we remove all the biases for the models. We summa- rize the CNN model structure with different normalization layers in Table 7. We the same architecture except for the last layer for both CIFAR-10 10 and CIFAR-100. We keep ReLU our activation function, and softmax for the last layer.
C.2 Data partitioning
Our client datascts are created through partitioning common datasets including CIFAR-10, CIFAR-100, TinylmageNet and PACS. We consider two types of data partitioning to simulate label shift in practice:
class(es): in this setting, each client has access to samples from only n classes.
Dirichlet partitioning: we partition samples from each class with a symmetric Dirichlet distribution and parameter 3. For each client, we collect one partitioned shard from each class. We denote this with the shorthand notation Dir See Zhang et. al. (2023) as an example.
Our setup includes 10 clients for the CIFAR-10 datasct, and 20 clients for the CIFAR-100 datasct, 200 clients for TinyImageNet, and 12 clients for PACS. We distribute the entire training and test datasets among clients in non-i.i.d manner. More specifically, we one class per client, two classes per client,
19

---

Published in Transactions on Machine Learning Research (01/2024)
Table 8: Dataset distribution for datasets.
Dataset
Train clients
Train examples
Test clients
Test examples
# of classes
CIFAR-10
10
50,000
10
10,000
10
CIFAR-100
50/20
50,000
50/20
10,000
100
TinyImageNet
20
25,000
20
25,000
200
PACS
12
7,907
12
1,920
7
class per client distribution
classes per client distribution
label skew (p=0.1) distribution
Q
1
snce
2
1
0
2100
2
1
20
2755
0
2-
0
9
5000
0
3
0
0
0
0
2-
0
0
2503
2500
0
0
0
3
2-
470
205
530
2225
5
4340
0
3
2300
220C
23
2675
Client 1d
5-
2500
2500
2
5-
400
0
53
5000
20
0551
90
6-
6
0
2
0
5QU0
6
0
0
2540
0
2300
3
6-
15
0
5/5
250
25
25
7
5000
7
2500
2500
7-
540
140
125
80
2500
9
2
3000
9
2500
5
2500
9
500
6112
2901
0 1
3
4
8
9
0
2
4 5
8
9
1
2
5
6
Label
Label
Label
Figure 8: Data distribution of CIFAR-10 dataset across clients with label shift. (left) one-class: each client has access to one-class; (middle) two classes: each client has access to two classes; (right) Dirichlet (0.1): Dirichlet allocation with 8 = 0.1.
and Dirichlet allocation with B = 0.1 for the CIFAR-10 datasot as shown Figure For the CIFAR-100 dataset, we use two classes per client (50 clients), classes per client (20 clients), and Dirichlet allocation with / 0.1 (20 clients). For TinyImageNet,
apply Dirichlet allocation with 0.1, 0.2, 0.5.
Finally, for the PACS dataset, to add distribution shift to the clients, we first split the dataset into 4 groups, photo (P), art painting (A), clipart (C) and sketch (S). For each group, we split the data into 3 clients, with Dirichlet partitioning Dir(0.5) and Dir(1.0), as well as partitioning into disjoint sets of classes (two clients have two classes each, and the other client has three classes of samples).
Statistics on the number of clients and examples in both the training and test splits of the datasets are given in Table 8. To illustrate the data partitioning, we include the label distributions for each client for CIFAR-10 and PACS. The label distributions of the CIFAR-10 dataset are shown in Figure 8. For the PACS dataset, the label distributions are shown in Figure 9.
C.3 Optimizers, hyperparameters, and validation metrics
We use SGD with a 0.01 learning rate (1r) and batch size of 32 for all of the experiments except for E experiments in CIFAR-100 in which we take Ir 0.1 as the learning rate and 1r 0.001 for PACS. We use SGD with a momentum of 0.9 only for our centralized training baseline. In each client, we take E steps of local training in which we itcrate E batches of data per client. We use online augmentation with random horizontal flip and random cropping with padding of 4 pixels for all of the datasets. Moreover, we test and utilize Fed (Reddi et al., 2020) as a server adaptive optimizer method in combination with FN, All the reported experiments are done with 10,000 global rounds. We report the accuracy of clients on their test data in the last iteration of our experiments with single run due computational limitation. Our hyperparameter choices are summarized in Table 9.
20

---

Published in Transactions on Machine Learning Research (01/2024)
2-2-3 classes client distribution
label skew (ß (8=0.51 distribution
label skew (B 1) distribution
153
0
54
25
29
231
1-
1:0
134
1
110
31
105
23
163
276
264
2-
70
25
182
500
3-
301
2015
3-
223
3
87
68
14
270
227
157
232
113
110
5
104
230
5-
70
96
18c
101
Ciient
3199
373
182
122
289
123
50
35
14
233
7
271
115
7-
353
222
108
7-
22'
82
130
151
28
263
224
107
164
182
9-
80M
191
9-
543
280
254
9
177
202
113
182
124
80
10
0
0.
500
300
10
532
110
590
10
435
122
202
22
422
11
0
0
655
60
123
11
22
300
120
11
0
265
201
300
105
43
32
0
1
2
3 Label
4
5 6
0
1
2
3
4
5
6
Label
Labol
Figure 9: Data distribution of the PACS dataset across clients with label shift. (left) 2-2-3 classes: cach client has access to two classes except the last client who has access to three classes; (middle) Dirichlet allocation with 8 = 0.5; (right) Dirichlet allocation with 3 = 1.
Table 9: Hyper-parameter choices.
Dataset
CIFAR-10 (CNN)
CIFAR-10 (ResNet)
CIFAR-100
PACS
TinylmageNet
Learning rate
0.01
0.01
0.1
0.001
0.01
Batch size
32
32
32
32
32
Optimizer
SGD V
SGD
SGD
SGD
SGD
Augmentation
V
C.4 Benchmarks
In this section, provide the implementation details and hyperparameters of our benchmark experiments.
FedProx. For the implementatinn ion of FedProx, stated in Li et al. (2020a), we add a weighted regularizer to the loss function of each client during local training to penalize the divergence between the server model and clients' models. To set the weight of the regularizer (or μ in the original paper), we sweep the datasets for weights 0.001, 0.010 0.1, 1 and select the best-performing value in the Dirichlet distribution with = 20 for each dataset. Our results show that 0.01 was the best-performing value for the weights for all the datasets. For the rest of the hyperparameters, we use the hyperparameters as FedAvg.
SCAFFOLD. For the implementation of SCAFFOLD, use Option II as described in Karimireddy et al. (2020). We use the same setting and hyperparameters as FedAvg. More specifically, clients perform E stops of local training while clients control variates using the difference between the server model and their local learned model.
FedDecorr. Similar FodProx, FodDecorr (Shi et al., 2023) applics a weighted regularization term during local training that encourages different dimensions of representations to be uncorrelated. We utilize their code for the implementation. In their results, the best-performing weight. for the regularization term (or a
21

---

Published in Transactions on Machine Learning Research (01/2024)
in the original paper) was 0.5. Therefore, we set the weight to 0.5 for all the datasets. For the rest of the hyperparameters, we use the same setting and hyperparameters as FedAvg.
FedLC. As described in Zhang et al. (2022), for the implementation of FedLC, during the local training, we shift the output logits of the network for each class depending on the number of available samples of each class for that client. This allows clients to add a weighted pairwise label margin according to the probability of occurrence of each class. We set the weight (or T in the original paper) to 1. As shown in Table 1, FedLC fails in the one-class setting since clients do not have two classes to ereate label margin. For the rest of the hyperparameters,
use the same setting and hyperparameters FedAvg.
FedRS. Similar to FedLC, FedRS (Li & Zhan, 2021) attempts to address the label shift in clients by limiting the update of missing classes' weights during the local training. FedRS multiplies the output logits of the missing classes by a weighting parameter (a in the original paper) and keeps the rest of the logits intact. As described in Li & Zhan (2021), 0.5 was the best-performing weight for all the datasets, and therefore, we set the weight to 0.5 for our experiments. For the rest of the hyperparameters, including the learning rale, we use the same hyperparameters as FedAvg.
Yog is one of the server-side adaptive optimization methods described in Reddi et al. (2020). Basically, it utilizes a separate learning rate for the server to calculate the global model for the clients. We utilize the best-performing values for the CIFAR-10 experiments in Reddi et al. (2020). More specifically, we set the server learning rate to 0.01 (7) in the original paper), 31, 82, and T to 0.9, 0.99, and 0.001 respectively. We
also initialize vector it to 1e-6. For the rest of the hyperparameters, including the clients' learning rates, we use the same hyperparameters as FedAvg.
C.5 Hardware
Our experiments are run on a cluster of 12 GPUs, including NVIDIA V100 and P100.
Experiment Configurations for Figures
We provide details for the experimental sctup of the figures in our main paper. Unless otherwise specified, we fix the local step number E 10 in our experiments.
Figure 2. In this experiment, first train vanilla CNN model with CIFAR-10 dataset in the one-class setting using FedAvg for 10000 epochs with 10 local steps and report the class accuracy of the global model (left) using the same setup as rest of the paper. Then, we locally train this model using examples belonging to class 0 (i.e. data of client 0) for 5 steps with FN (right) and without FN (middle) and report the class accuracy of these two models in the entire dataset. Note that FN networks and vanilla networks are the same except the constraint on the feature embedding.
Figure 3. In this one-class setup, for the three figures on the left, we only locally train client 1 with its local data which only contains class
Figure 4.
We use SGD with batch size 32 for FedAvg with 10 local steps and batch size of 320 for centralized training. We use 0.01 as the learning rate for both centralized and FedAvg. The batch size 320 is chosen to match federated and centralized learning.
Figure 6. We use batch size 32, learning rate 0.001, and SGD with 0.9 momentum for the centralized training.
E
Additional Experiments
Additional local steps and different neural architectures. We provide additional experiments with different local steps, on CIFAR-10/100 and CNN architecture in Tables 10 and 11. This table shows our results are consistent with different local steps. As label shift becomes more severe, the advantage of FN/LN
22

---

Published in Transactions on Machine Learning Research (01/2024)
Table 10: SoTA comparison on CIFAR-10 with CNN.
1 class
2 classes
Dir(O.1)
Methods
E = 1
E = 10
E = 20
E =
E = 10
E = 20
E
E = 10
E = 20
FedAvg
57.45
55.86
58.65
57.35
71.83
73.16
52.82
73.06
74.58
FedProx
57.39
54.24
58.31
57.42
71.86
72.78
53.94
72.85
74.72
SCAFFOLD
57.41
54.14
57.97
37.45
71.88
73.44
52.87
72.38
75.20
FedLC
9.72
9.72
9.72
39.35
62.01
64.49
45.08
70.25
71.98
FedDecorr
47.50
48.18
48.09
47.73
70.88
73.43
45.46
74.58
75.31
FedRS
10.00
10.00
10.00
37.96
63.26
64.06
44.82
70.60
72.63
FedYogi
77.61
73.13
74.21
76.77
78.97
77.12
76.46
79.35
77.99
FedFN
77.52
77.14
77.35
74.52
76.67
76.12
75.73
78.13
77.35
FedLN
79.25
78.04
78.12
77.39
77.52
76.93
77.10
78.27
77.48
Table 11: SoTA comparison CIFAR-100 with CNN.
classes
classes
Dir(0.01)
Methods
E
E 10
E 20
F
F 10
F
F.
F
F. 20
FedAvg
39.38
32.33
33.57
41.81
36.22
36.65
40.92
38.14
38.12
FedProx
42.47
32.39
31.70
41.65
36.59
36.95
44.35
38.85
38.67
SCAFFOLD
41.40
32.30
34.43
41.46
36.57
37.06
44.55
38.7
38.69
FedLC
4.41
6.04
5.75
17.63
18.17
18.78
21.31
22.34
22.35
FedDecorr
33.03
29.21
32.43
33.27
32.94
35.26
33.58
34.24
36.48
FedRS
6.35
10.73
10.81
22.30
24.88
25.29
21.78
22.19
24.13
Yogi
46.66
44.60
43.98
46.79
44.08
43.89
48.10
45.20
44.96
FedFN
46.51
40.32
43.70
45.51
45.41
44.10
48.15
45.36
44.52
FedLN
48.72
46.59
45.89
49.83
46.22
45.68
51.21
47.30
45.53
Table 12: Combining adaptive optimization with normalization.
Methods class
2 classes
Dir(O.1)
FedAvg 56.82
71.83
73.06
FegYogi
73.13
78.97
79.35
FedFN
77.14
76.67
78.14
FedLN
77.71
77.65
79.94
FedEN + Yogi 80.47
78.83
80.35
FedLN + Yogi
79.27
79.16
80.46
is clearer. Note that batch size 32), but due to label shift not every algorithm can converge well.
reduces to the contralized sctting with 320 batch size (since each client has
Combination with other techniques. We can further improve the performance of normalization, by combining it with other techniques such as adaptive optimization (Reddi et al., 2020). In Table 12, we show that combined with FodYogi (Reddi al., 2020). FedFN and FedLN has improved performance. We do not consider other combinations such as with FedProx SCAFFOLD they are not shown help FedAvg much in the unnormalized case.
23

---

Published in Transactions on Machine Learning Research (01/2024)
Table 13: Comparison between different modifications of ResNet.
Methods
class
2 classes
Dir(0.1)
FedLN ResNet
86.15
88.01
89.06
FedFN nMv inside ResNet
77.76
84.34
86.63
FedFN inside ResNet
82.14
84.05
85.17
FedFN ResNet
80.81
82.30
84.19
Table 14: Results with multiple runs for CIFAR-10 dataset in one class settings with 10 clients for both CNN and ResNet models.
Methods
CNN
ResNet
FedAvg
55.7640.35
55.34
FodYogi
73.7610.70
76.90
FodFN
76.8540.37
80.55
FedLN
77.21
83.40
Modifications of ResNet. Last but least, we test different modifications of ResNet in Table 13. We denote пиv inside the following block:
block(x)) = p(x + A nMv o A1))),
and by inside we use:
block(æ=) = p(x A Agonopo A1
for each block. Note that we always add the last-layer feature normalization for these two methods. We see that both nMv and help the improvement of FedFN, and they narrow the gap between FN and LN.
Multiple seeds. In Table 1, we reported the last-iteration accuracy for single run. In Table 14, we take three runs for CIFAR-10 with one class settings for both CNN and ResNet models and report the mean and standard deviation for the three runs. Our results show that the comparison is consistent across different runs. Due computation constraints we could not afford to run multiple seeds for all our experiments.
F Additional Related Work
Normalization not a new idea, and can be at least traced back to whitening. As suggested by Le- Cun al. (1998), we should shift the inputs so that the average is nearly zero and scale them to have the same covariance. Simoncelli & Olshausen (2001) also explained whitening in computer vision through principal components analysis. Normalization is also standard technique in training neural networks if one treats the inputs as the intermediate layers. Depending on which dimensions to normalize, it could be batch normalization (loffe & Szegedy, 2015), weight normalization (Salimans & Kingma, 2016), instance normalization (Ulanov et al., 2016), layer normalization (Ba et al., 2016), and group normalization (Wu & He, 2018). Specifically, layer normalization widely used modern neural architectures such as Trans- formers (Vaswani al., 2017). More recently, Zhang Sennrich (2019) proposed Root Mean Square Layer Normalization (RMSNorm) which dispenses with the mean shift in layer normalization. This shares some similarities with our feature normalization. However, their motivation is quite different from ours.
There has been recent work that aims to understand normalization methods. For example, Santurkar et al. (2018) explained that batch normalization can help optimization; Dukler el al. (2020) provided convergence results for two-layer ReLU networks with weight normalization with NTK analysis. Lately, Lyu et al. (2022) studied the convergence of gradient descent with weight docay for scalc-invariant loss functions. None of the previous work provided an analysis of the expressive power and local overfitting with softmax and cross-entropy loss.
24

---

Published in Transactions on Machine Learning Research (01/2024)
Expressive power. The expressive power of (normalized) neural networks has been studied in Raghu et al. (2017) and Giannou et al. (2023). In Raghu et al. (2017), the authors proposed new measures of expressivity and trajectory regularization that achieves similar advantages as batch norm. In Giannou et al. (2023), the authors argued that by tuning only the normalization layers of a neural network, we can recover any fully-connected neural work functions. Neither of them provides the same result as our Proposition
25


Reviews:
[{"review_id": "rO8I64gpPQ", "reviewer": "Reviewer_Zjcq", "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:\n\n- The paper is very well-motivated, picks an interesting problem, and is well-structured in the sense that I could understand what was happening without being an expert in Federated Learning.\n- The results in the paper are tested on a wide suite of tasks and benchmarks, and are also comprehensive in answering when and why layer normalization works.\n- The experimental procedure is thorough and well-documented.\n\nWeaknesses/Questions:\n\n1. The Theorem statements are a bit too informal for my liking (e.g., Propositions 1 and 2) --- for example, what is a \"proper domain\"? \n2. The writing could be improved - although the structure of the paper is nice, there are some grammatical errors and awkward sentences that make the paper a bit hard to read at times.\n3. I don't believe this will work, because of the authors' presented intuition, but it would be good to check that other methods for preventing overfitting (e.g., tuning learning rate, increasing weight decay) do not have the same effect as LN.\n4. Although I understand what the authors mean here, I would recommend removing or revising the paragraph that follows Proposition 3, since it's a bit misleading. In particular, Proposition 3 shows that any function that an LN network can represent can also be represented by a non-LN network---the authors use this to argue that the benefit must be something during training. However, since *not* every NN can be represented as an LN, this claim is not true ---the proposition doesn't rule out model class restriction as the underlying mechanism behind LN's effectiveness.", "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": "This paper studies the role of Layer Normalization in Federated Learning in the presence of extreme label shift (e.g., when each client in the federated learning setup has data from exactly one of the classes being distinguished). I understood the paper as making two main contributions: (a) establishing and investigating the effectiveness of Layer Normalization for Federated Learning under extreme label shift, and (b) hypothesizing that the mechanism behind this effectiveness is \"feature normalization.\"", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses above", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Cer1qDPx5t", "reviewer": "Reviewer_JL46", "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": "[S1] The paper is very well written, and the ideas are presented clearly. The concepts are well-connected and the paper is easy to follow. \n\n[S2] The topic of the paper is relevant given the ubiquitous use of normalization techniques in practice. \n\n[S3] The numerical examples are a good proof that the method works. As a comment, I would like to see experiments turning LN/FN on and off for each of the algorithms considered. \n\n[W1] What is the role of FL in this paper? What would be different than a centralized case? I fail to understand why the authors decided to write this paper in the context of FL. \n\n[W2] The fundamental finding of the paper -- the role of FN in LN -- seems to me as not convincing. That is to say, the difference between LN and FN is the quality is normalizing (or not) the last layer. Therefore, the role of FN is not properly explained or understood. Even more so, the theory of the paper is regarding the relationship between these two normalizations. The empirical results show that LN is better than FN, therefore why is FN relevant? \n\n[W3] In Propositions 1 and 2, the bias term =0 seems like an extremely difficult assumption to observe in practice and therefore, the relevance of these two propositions is seriously compromised. The authors should provide a better explanation of why these two propositions are relevant findings.\n\n[W4] The paper does not provide enough details on the number of agents. They are spread out in the appendix, but I would suggest adding this data to each individual table with results.\n\n[W5] The paper does not fully mention the topic of privacy which is of vital importance in FL. To compute LN/FN you need data from batch sizes which requires clients to provide local and private data. How is privacy assured in LN?\n\n[W6] The paper does not mention the problem of communication when addressing the normalizations. The central agent needs to combine the local information from all agents and therefore more needs to be communicated. How does this affect training time and communication costs? \n\n[W7] Proposition 3 and its explanation need more details. The way it is currently presented, and the importance/significance of this result is not clear. Even more so, the statement of the proposition is not self-contained and so it is difficult to assess its relevance. After carefully reading it, and looking at the proof (which I believe it is correct) I do not understand the point of the proposition, what message is it trying to convey? Also, adding to W1, the number of agents plays no role in it. \n\n[W8] What is the relevance of Theorem 1 with respect to LN? There is little explanation provided and intuitively, if the norms diverge, the role of LN seems to be irrelevant. The authors should further clarify this point. \n\n[W9] There are some important and relevant works that have not been included such as: \n\nAN AGNOSTIC APPROACH TO FEDERATED LEARNING WITH CLASS IMBALANCE Shen et al. \nA Theoretical Analysis of the Learning Dynamics under Class Imbalance Francazi et al.", "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 attempts to explain the role of Layer Normalization in federated learning. The authors identify feature normalization (normalizing pre-classification layer feature embeddings) as a key element that improves the quality of learning and connects it to layer normalization. They have enough numerical examples that validate the ability of their method. They also provide a theoretical analysis based on layer normalization that adds to their empirical results.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See previous section.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "K9eHWBEB7Q", "reviewer": "Reviewer_vPYs", "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\n\n1. The paper is well written and easy to follow and self-contained with sufficient background on normalization and federated learning.\n\n2. It presents a study of layer normalization (LN) in federated learning(FL) setting and provides insights into when and why LN helps in FL. The authors observe that when the distribution of labels is uneven among clients, layer normalization becomes more helpful. This is because layer normalization prevents the model from feature collapse, which helps control overfitting and avoids losing important features.\n\n3. They introduce the idea of feature normalization (FN), a process within LN that normalizes latent feature representations before the classifier head. While this doesn't increase the model's expressive capabilities, LN and FN effectively manage feature collapse and local overfitting in datasets with significant imbalances, which speeds up global training.\n\n\nWeaknesses:\n\n\n1. It is not quite clear what the empirical results are reporting? Is it the mean of the accuracy over multiple runs with different random seeds? Could you also plot/report the std. deviations? \n\n\n2. The comparison in most figures is between FedAvg and variations with different normalizations, however table 1 suggests FedYogi is pretty close to FedLN, FedFN. \n\n3. Minor comments: \n\n   a. Please organize the tables, figures in pages 8-11 to better utilize the space. \n\n   b. I personally do not like “All you need”  titles and the paper is full of this phrase. I would suggest using a shorter and more crisp title that matches with the paper’s findings.", "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 investigates effectiveness of layer normalization (LN) in federated learning (FL) with non-i.i.d. data. While LN has proven effective in this context, the underlying reasons for its success have not been well understood. The authors establish a connection between layer normalization and the label shift problem in federated learning. They find that, as the distribution of labels among clients becomes more imbalanced, layer normalization becomes increasingly beneficial. This happens because LN controls local overfitting by mitigating feature collapse. Further, they introduce the concept of feature normalization (FN), a mechanism within LN that applies normalization to latent feature representations before the classifier head.  LN and FN control feature collapse and local overfitting on heavily skewed datasets, leading to accelerated global training.  The empirical findings demonstrate that normalization, particularly in the form of FN, significantly improves performance on standard benchmarks, particularly under extreme label shift conditions. Their results confirm that FN is a vital component of LN, enhancing convergence in FL and maintaining robustness to learning rate variations, especially in scenarios with extreme label shift.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address the comments above in the weaknesses section.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the use of layer normalization for federated learning when there is label skew.  This is prompted by the recent empirical success of layer norm in federated learning in non iid settings.  The reviewers all voted to accept the paper (all "leaning accept").  They found the paper well motivated, relevant and found the empirical results convincing.  However, all the reviewers seemed to find issues with the clarity of the paper.  Therefore, I recommend accepting the paper, but I would like the authors to make some minor revisions for clarity, addressing the reviewer comments.  Also please address the comments regarding the motivation for and clarity of the theoretical contribution of the paper.
Decision: Accept (poster)