# Compare Where It Matters: Using Layer-Wise Regularization To Improve Federated Learning on Heterogeneous Data

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 27 references that are in the map.
- First posted on arXiv 1 Dec 2021.

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

- [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)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Federated Learning via Over-the-Air Computation](https://mlsanity.com/paper/1812.11750/) (2018)
- [A Survey on Distributed Machine Learning](https://mlsanity.com/paper/1912.09789/) (2019)
- [FedMix: Approximation of Mixup under Mean Augmented Federated Learning](https://mlsanity.com/paper/2107.00233/) (2021)
- [Are All Layers Created Equal?](https://mlsanity.com/paper/1902.01996/) (2019)

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

- [FedCC: Robust Federated Learning against Model Poisoning Attacks](https://mlsanity.com/paper/2212.01976/) (2022)
- [Contrastive encoder pre-training-based clustered federated learning for heterogeneous data](https://mlsanity.com/paper/2311.16535/) (2023)
- [Federated Learning with Intermediate Representation Regularization](https://mlsanity.com/paper/2210.15827/) (2022)
- [Defending Against Backdoor Attacks via Alignment Checking in Model-Contrastive Federated Learning](https://mlsanity.com/paper/2607.26933/) (2026)

## Abstract (arXiv)

Federated Learning is a widely adopted method to train neural networks over distributed data. One main limitation is the performance degradation that occurs when data is heterogeneously distributed. While many works have attempted to address this problem, these methods under-perform because they are founded on a limited understanding of neural networks. In this work, we verify that only certain important layers in a neural network require regularization for effective training. We additionally verify that Centered Kernel Alignment (CKA) most accurately calculates similarity between layers of neural networks trained on different data. By applying CKA-based regularization to important layers during training, we significantly improve performance in heterogeneous settings. We present FedCKA: a simple framework that out-performs previous state-of-the-art methods on various deep learning tasks while also improving efficiency and scalability.