# Understanding the Role of Layer Normalization in Label-Skewed Federated Learning

arXiv: https://arxiv.org/abs/2308.09565 · HTML: https://mlsanity.com/paper/2308.09565/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 38 references that are in the map.
- First posted on arXiv 18 Aug 2023.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://mlsanity.com/paper/1607.08022/) (2016)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Deeper, Broader and Artier Domain Generalization](https://mlsanity.com/paper/1710.03077/) (2017)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [How Does Batch Normalization Help Optimization?](https://mlsanity.com/paper/1805.11604/) (2018)
- [On the Expressive Power of Deep Neural Networks](https://mlsanity.com/paper/1606.05336/) (2016)
- [The Non-IID Data Quagmire of Decentralized Machine Learning](https://mlsanity.com/paper/1910.00189/) (2019)
- [Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction](https://mlsanity.com/paper/2206.07085/) (2022)
- [Federated Learning with Only Positive Labels](https://mlsanity.com/paper/2004.10342/) (2020)
- [Towards Understanding and Mitigating Dimensional Collapse in Heterogeneous Federated Learning](https://mlsanity.com/paper/2210.00226/) (2022)
- [On the Expressivity Role of LayerNorm in Transformers' Attention](https://mlsanity.com/paper/2305.02582/) (2023)
- [Why Batch Normalization Damage Federated Learning on Non-IID Data?](https://mlsanity.com/paper/2301.02982/) (2023)
- [Proportional Fairness in Federated Learning](https://mlsanity.com/paper/2202.01666/) (2022)
- [A Theoretical Analysis of the Learning Dynamics under Class Imbalance](https://mlsanity.com/paper/2207.00391/) (2022)
- [Rethinking Normalization Methods in Federated Learning](https://mlsanity.com/paper/2210.03277/) (2022)
- [Optimization Theory for ReLU Neural Networks Trained with Normalization Layers](https://mlsanity.com/paper/2006.06878/) (2020)
- [Experimenting with Normalization Layers in Federated Learning on non-IID scenarios](https://mlsanity.com/paper/2303.10630/) (2023)

## Built on by (most cited first, 2 of 2)

- [Parametric Feature Transfer: One-shot Federated Learning with Foundation Models](https://mlsanity.com/paper/2402.01862/) (2024)
- [Hybrid Batch Normalisation: Resolving the Dilemma of Batch Normalisation in Federated Learning](https://mlsanity.com/paper/2505.21877/) (2025)

## Abstract (arXiv)

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. Our code is available at urlhttps://github.com/huawei-noah/Federated-Learning/tree/main/LayerNormalization.