Federated Learning with Non-IID Data

Yue ZhaoMeng LiLiangzhen LaiNaveen SudaDamon CivinVikas Chandra

article2018arXiv3,469 citations

Quantifies the severe accuracy drop in non-IID federated learning through earth mover's distance-based weight divergence and provides a data-sharing strategy that substantially recovers model performance with minimal shared data.

Listen

Federated learning allows edge devices such as phones and IoT sensors to train a shared model without sending raw data to a central server, thereby preserving privacy and reducing transmission costs. In practice, however, the data held on each device often follows a very different distribution from the overall population, and this statistical mismatch sharply reduces model accuracy. The paper examines how severe that reduction becomes, identifies its mathematical cause, and tests a practical remedy.

The authors set out to measure the accuracy loss of the standard FederatedAveraging algorithm under non-independent-and-identically-distributed data, to explain the loss through weight divergence, and to demonstrate whether a modest amount of globally shared data can restore performance. They trained representative convolutional networks on MNIST, CIFAR-10, and a keyword-spotting speech dataset, comparing fully IID partitions against two extreme non-IID partitions in which each client held examples from only one or two classes. They also derived an analytic bound linking weight divergence to the earth mover’s distance between each client’s class distribution and the global distribution, then validated the bound experimentally.

Accuracy fell by as much as 55 percent on the most skewed one-class partitions; even the milder two-class partitions produced double-digit losses on CIFAR-10 and keyword spotting. The magnitude of the loss tracked the earth mover’s distance closely, and the divergence appeared in the first communication round and grew with further local updates. Introducing a globally shared subset equal to only 5 percent of the total training data, together with a warm-start model trained on that subset, raised CIFAR-10 accuracy by roughly 30 percentage points while limiting the data each device actually received to half of the shared pool.

These results show that the statistical heterogeneity long acknowledged in federated-learning literature can be mitigated without abandoning the decentralized training paradigm. The trade-off is modest centralization of a small reference set whose size and distribution fraction can be tuned to the application. Because the shared data need be distributed only once at initialization, communication overhead remains negligible and privacy exposure is limited to non-sensitive examples.

The analysis rests on three image and speech classification tasks and on the assumption that synchronization occurs after a fixed number of local steps; performance on other model families or under highly variable client participation was not measured. Within those bounds, the experimental trends and the supporting proof are consistent across data sets and hyper-parameter choices, giving decision makers a clear quantitative basis for adopting the data-sharing approach when non-IID conditions are expected.

Arm cell phone economics federated learning Federated Learning Non-IID Data FedAvg

arXiv: 1806.00582

Paper: Federated Learning: Strategies for Improving Communication Efficiency, Jakub Konečný et al. (2016). This foundational work on reducing communication overhead in federated learning establishes the basic parameter-sharing and update-compression techniques that the source paper later builds upon to address statistical heterogeneity.
Paper: On the Convergence of FedAvg on Non-IID Data, Xiang Li et al. (2019). Analyzing the convergence properties of Federated Averaging under non-IID data distributions, this paper provides critical theoretical context for understanding the accuracy drops and weight divergence phenomena investigated in the source paper.

Paper: Federated Optimization in Heterogeneous Networks, Tian Li et al. (2018). Directly extending the source paper's exploration of non-IID challenges in federated learning, this work introduces FedProx to handle variable device constraints and non-convex objectives through a proximal regularizer.
Paper: SCAFFOLD: Stochastic Controlled Averaging for Federated Learning, Sai Praneeth Karimireddy et al. (2019). Building upon the non-IID federated optimization problem defined in the source, this paper proposes SCAFFOLD to correct client drift using control variates and eliminate the convergence slowdowns caused by data heterogeneity.
Paper: Federated Learning: Challenges, Methods, and Future Directions, Tian Li et al. (2019). Expanding upon the specific statistical challenges identified in the source paper, this comprehensive survey maps out the broader landscape of federated learning challenges, methods, and future directions.

Cover for Federated Learning with Non-IID Data

Abstract

Federated learning enables resource-constrained edge compute devices, such as mobile phones and IoT devices, to learn a shared model for prediction, while keeping the training data local. This decentralized approach to train models provides privacy, security, regulatory and economic benefits. In this work, we focus on the statistical challenge of federated learning when local data is non-IID. We first show that the accuracy of federated learning reduces significantly, by up to 55% for neural networks trained for highly skewed non-IID data, where each client device trains only on a single class of data. We further show that this accuracy reduction can be explained by the weight divergence, which can be quantified by the earth mover's distance (EMD) between the distribution over classes on each device and the population distribution. As a solution, we propose a strategy to improve training on non-IID data by creating a small subset of data which is globally shared between all the edge devices. Experiments show that accuracy can be increased by 30% for the CIFAR-10 dataset with only 5% globally shared data.

1 Introduction
2 FedAvg on Non-IID data
2.1 Experimental Setup
2.2 Experimental Results
3 Weight Divergence due to Non-IID Data
3.1 Mathematical demonstration
3.2 Experimental Validation
3.2.1 Experimental Setup
3.2.2 Weight Divergence vs. EMD
3.2.3 Test Accuracy vs. EMD
4 Proposed Solution
4.1 Motivation
4.2 Data-sharing Strategy
5 Conclusion
References
A Appendix
A.1 Test accuracy over communication rounds for a smaller batch size
A.2 Test accuracy of SGD and FedAvgFedAvg with IID or non-IID data
A.3 Proof of Proposition

Knowls

Knowl 1 — Theoretical Upper Bound on Weight Divergence in Federated Averaging
theoretical result
Consider a $C$ -class classification task over compact input space $\mathcal{X}$ and label space $\mathcal{Y} = [C] = \{1, \dots, C\}$ with population distribution $p(x, y)$ . Let a neural network parameterized by weights $w$ output class probabilities $f(x, w) = (f_1(x, w), \dots, f_C(x, w))^T$ on the probability simplex $\mathcal{S}$ .

Let there be $K$ distributed clients where each client $k \in [K]$ holds $n^{(k)}$ i.i.d. samples drawn from a local distribution $p^{(k)}$ . Suppose model synchronization in Federated Averaging (FedAvg) occurs every $T$ local SGD steps with learning rate $\eta$ , and for each class $i \in [C]$ , $\nabla_w \mathbb{E}_{x|y=i}[\log f_i(x, w)]$ is $\lambda_{x|y=i}$ -Lipschitz with respect to $w$ .

Let $w_{mT}^{(f)}$ denote the global aggregated weight vector after the $m$ -th synchronization round, and let $w_{mT}^{(c)}$ denote the weight vector obtained by centralized SGD after $mT$ total steps with matching initialization. Then the $\ell_2$ weight divergence satisfies:
$\|w_{mT}^{(f)} - w_{mT}^{(c)}\| \le \sum_{k=1}^K \frac{n^{(k)}}{\sum_{l=1}^K n^{(l)}} (a^{(k)})^T \|w_{(m-1)T}^{(f)} - w_{(m-1)T}^{(c)}\| + \eta \sum_{k=1}^K \frac{n^{(k)}}{\sum_{l=1}^K n^{(l)}} \sum_{i=1}^C \|p^{(k)}(y = i) - p(y = i)\| \sum_{j=0}^{T-1} (a^{(k)})^j g_{\max}(w_{mT-1-j}^{(c)})$
where:
- $a^{(k)} = 1 + \eta \sum_{i=1}^C p^{(k)}(y = i) \lambda_{x|y=i} \ge 1$
- $g_{\max}(w) = \max_{i \in [C]} \|\nabla_w \mathbb{E}_{x|y=i}[\log f_i(x, w)]\|$
This bound establishes that weight divergence accumulates from the prior round (scaled by $(a^{(k)})^T$ ) and is driven by the Earth Mover's Distance (EMD) between client distribution $p^{(k)}$ and population distribution $p$ , specifically quantified by $\sum_{i=1}^C \|p^{(k)}(y=i) - p(y=i)\|$ , modulated by the learning rate $\eta$ , the synchronization interval $T$ , and the maximum gradient norm $g_{\max}$ .
Knowl 2 — Globally Shared Data Strategy for Non-IID Federated Learning
model/method
To mitigate model weight divergence and accuracy loss caused by skewed non-IID class distributions across edge clients, a small globally shared dataset containing a uniform distribution over all target classes is distributed from the central server.

Let $D$ denote the union of all private client datasets, and let $G$ denote a globally shared dataset centralized in the cloud with equal representation across all classes $C$ . The method proceeds as follows:
1. Cloud Warm-Up Pre-training: Prior to federated training, a warm-up model is pre-trained directly on $G$ .
2. Initialization and Distribution: At initialization, the warm-up model parameters and a random subset of $G$ comprising a fraction $\alpha \in (0, 1]$ of $G$ are sent to each edge client.
3. Local Training: Each client trains its local model on the combined dataset consisting of its local private data and its assigned portion of $G$ .
4. Federated Aggregation: The central server collects local model weights and performs standard FedAvg aggregation at each synchronization step.
The framework introduces two trade-off hyper-parameters:
- $\beta = \frac{|G|}{|D|} \times 100\%$ , defining the overall relative size of the globally shared dataset relative to the total private dataset size.
- $\alpha$ , defining the fraction of $G$ allocated to each individual client.
Knowl 3 — Weight Divergence Metric for Federated Learning
definition
Weight divergence quantifies the relative deviation between model parameters trained via Federated Averaging (FedAvg) and model parameters trained via centralized Stochastic Gradient Descent (SGD) under identical weight initialization, learning rate schedule, and equivalent aggregate batch sizes.

For weight vector $w^{\text{FedAvg}}$ and centralized SGD weight vector $w^{\text{SGD}}$ , the weight divergence is defined as:
$\text{weight divergence} = \frac{\|w^{\text{FedAvg}} - w^{\text{SGD}}\|}{\|w^{\text{SGD}}\|}$
where $\|\cdot\|$ denotes the Euclidean ( $\ell_2$ ) norm.
Knowl 4 — Accuracy Degradation of FedAvg Under Non-IID Class Skew
empirical result
Across benchmark image and audio classification tasks (MNIST, CIFAR-10, and Keyword Spotting [KWS]), training Convolutional Neural Networks with Federated Averaging (FedAvg) on skewed non-IID data causes significant test accuracy drops relative to centralized SGD:
- 1-Class Non-IID (each client holds data from only 1 class): Test accuracy drops by up to $11.31\%$ on MNIST, $51.31\%$ on CIFAR-10, and $54.50\%$ on KWS when using small local batch sizes ( $B=10$ ). With larger batch sizes ( $B=100$ for MNIST/CIFAR-10, $B=50$ for KWS), the drops are $6.52\%$ , $37.66\%$ , and $43.64\%$ , respectively.
- 2-Class Non-IID (each client holds data from 2 classes): Test accuracy drops by $1.77\%$ on MNIST, $15.61\%$ on CIFAR-10, and $15.07\%$ on KWS for small batch sizes.
- Local Epochs Invariance: Increasing the number of local epochs from $E=1$ to $E=5$ does not mitigate accuracy loss.
- Pre-training Failure: Initializing FedAvg with model weights pre-trained via centralized SGD does not prevent accuracy degradation when training continues on 1-class non-IID data.
Knowl 5 — Dependence of Weight Divergence and Test Accuracy on Earth Mover's Distance
empirical result
Controlled experiments evaluating FedAvg across varying degrees of class imbalance, quantified by the Earth Mover's Distance $\text{EMD} = \sum_{i=1}^C \|p^{(k)}(y=i) - p(y=i)\|$ between local client distributions and the uniform population distribution, demonstrate:
1. Monotonic Divergence Growth: As EMD increases from $0$ (IID) to $1.8$ (extreme non-IID), layer-wise weight divergence increases across all CNN layers on MNIST, CIFAR-10, and KWS. The effect is most pronounced in the first convolutional layer and the final fully connected layer.
2. Accuracy Governed by EMD: Across 5 distinct random class partitions constructed for each specific EMD value, test accuracy exhibits low variance (standard deviations $< 0.0009$ for MNIST, $< 0.010$ for CIFAR-10, $< 0.006$ for KWS), demonstrating that EMD itself—rather than the specific identity of allocated classes—governs performance.
3. Nonlinear Accuracy Drop: Test accuracy decreases gradually at low EMD values but drops precipitously beyond an EMD threshold of approximately $1.0$ .
Knowl 6 — Empirical Performance and Trade-Offs of the Data-Sharing Strategy on CIFAR-10
empirical result
On the CIFAR-10 dataset partitioned into an extreme 1-class non-IID setting across 10 clients (where standard FedAvg achieves only $\sim 44\%$ test accuracy):
- Impact of Shared Dataset Size ( $\beta$ ): Merging the entire shared dataset $G$ with client data increases FedAvg test accuracy to $74.12\%$ at $\beta = 10\%$ and $78.72\%$ at $\beta = 25\%$ .
- Impact of Per-Client Distributed Fraction ( $\alpha$ ): When initializing with a warm-up model pre-trained on $G$ and distributing only a fraction $\alpha$ of $G$ to each client, test accuracy reaches $73.12\%$ for $\beta=10\%$ and $77.08\%$ for $\beta=20\%$ at $\alpha=1.0$ . As $\alpha$ is reduced from $1.0$ down to $0.5$ , test accuracy decreases by less than $1\%$ .
- Overall Efficiency: By centralizing $\beta = 10\%$ data in the cloud and distributing a fraction $\alpha = 0.5$ (such that each client receives only $5\%$ shared data relative to total data), test accuracy increases by $\sim 30\%$ absolute over the baseline non-IID setting.

Knowl 7 — Comparison of Centralized SGD and FedAvg Test Accuracy Across Data Partitions

data/table

The following table records the final test accuracy of centralized SGD and FedAvg across MNIST, CIFAR-10, and Keyword Spotting (KWS) datasets under IID, 2-class non-IID (non-IID(2)), and 1-class non-IID (non-IID(1)) configurations with large and small batch sizes:

Training	B	E	MNIST (%)	CIFAR-10 (%)	KWS (%)
SGD	large	NA	98.69	81.51	84.46
FedAvg IID	large	1	98.69	80.83	84.82
FedAvg non-IID(2)	large	1	96.29	67.00	72.30
FedAvg non-IID(1)	large	1	92.17	43.85	40.82
FedAvg non-IID(1)	large	5	91.92	44.40	40.84
Pre-trained non-IID(1)	large	1	96.19	61.72	63.58
SGD	small	NA	99.01	84.14	86.28
FedAvg IID	small	1	99.12	82.62	86.64
FedAvg non-IID(2)	small	1	97.24	68.53	71.21
FedAvg non-IID(1)	small	1	87.70	32.83	31.78

Batch sizes designated as 'large' are $B=100$ for MNIST and CIFAR-10, and $B=50$ for KWS. 'Small' batch sizes are $B=10$ across all datasets. For centralized SGD, the effective batch size is $10 \times B$ to match the aggregate batch size of 10 clients in FedAvg.

Knowl 8 — FedAvg Test Accuracy Distribution Across Earth Mover's Distance Values

data/table

The following table details the mean and standard deviation of test accuracy across 5 distinct randomly generated class distributions for each specified Earth Mover's Distance (EMD) value:

EMD	0	0.36	0.72	1.08	1.44	1.62	1.764	1.8
MNIST mean	0.9857	0.9860	0.9852	0.9835	0.9799	0.9756	0.9620	0.9220
MNIST std (× 10^-4)	6.431	2.939	4.604	4.308	4.716	8.085	8.232	1.939
CIFAR-10 mean	0.8099	0.8090	0.8017	0.7817	0.7379	0.6905	0.5438	0.4396
CIFAR-10 std (× 10^-3)	2.060	2.694	2.645	3.622	3.383	2.048	9.655	1.068
KWS mean	0.8496	0.8461	0.8413	0.8331	0.7979	0.7565	0.5827	0.4475
KWS std (× 10^-3)	1.337	3.930	4.410	5.387	1.763	3.329	1.078	4.464

The table demonstrates that variations across random partitions with the same EMD are negligible relative to the steep drop in accuracy that occurs as EMD transitions from $0.72$ to $1.8$ .

Coverage note — None was omitted; all key theoretical results, methodologies, empirical findings, and primary data tables from the paper's contribution are captured.

References

1.Y. Zhang, N. Suda, L. Lai, and V. Chandra, —Hello edge: Keyword spotting on microcontrollers,— arXiv preprint arXiv:1711.07128, 2017.
2.L. Lai, N. Suda, and V. Chandra, —Cmsis-nn: Efficient neural network kernels for arm cortex-m cpus,— arXiv preprint arXiv:1801.06601, 2018.
3.H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, —Communication-efficient learning of deep networks from decentralized data,— in Int. Conf. on Artificial Intelligence and Statistics, 2017.
4.J. Konecň ỳ, B. McMahan, and D. Ramage, —Federated optimization: Distributed optimization beyond the datacenter,— arXiv preprint arXiv:1511.03575, 2015.
5.H. B. McMahan and D. Ramage, —Federated learning: Collaborative machine learning without centralized training data,— in Google, 2017.
6.Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner, —Gradient-based learning applied to document recognition,— in Proceedings of the IEEE, pp. 2278–2324, 1998.
7.A. Krizhevsky and G. Hinton, —Learning multiple layers of features from tiny images,— in Technical Report, University of Toronto, 2009.
8.W. Shakespeare, —The complete works of william shakespeare,— in https://www.gutenberg.org/ebooks/100.
9.K. Bonawitz, V. Ivanov, B. Kreuter, A. Marcedone, H. B. McMahan, S. Patel, D. Ramage, A. Segal, and K. Seth, —Practical secure aggregation for privacy-preserving machine learning,— in Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security, pp. 1175–1191, ACM, 2017.
10.J. Konecň ỳ, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon, —Federated learning: Strategies for improving communication efficiency,— arXiv preprint arXiv:1610.05492, 2016.
11.Y. Lin, S. Han, H. Mao, Y. Wang, and W. J. Dally, —Deep gradient compression: Reducing the communication bandwidth for distributed training,— arXiv preprint arXiv:1712.01887, 2017.
12.J. Duchi, E. Hazan, and Y. Singer, —Adaptive subgradient methods for online learning and stochastic optimization,— Journal of Machine Learning Research, vol. 12, no. Jul, pp. 2121–2159, 2011.
13.T. Tieleman and G. Hinton, —Divide the gradient by a running average of its recent magnitude. coursera: Neural networks for machine learning,— tech. rep., Technical Report. Available online: https://zh. coursera. org/learn/neuralnetworks/lecture/YQHki/rmsprop-divide-the-gradient-by-a-running-average-ofits-recent-magnitude (accessed on 21 April 2017).
14.D. Kinga and J. B. Adam, —A method for stochastic optimization,— in International Conference on Learning Representations (ICLR), 2015.
15.L. Bottou, —Large-scale machine learning with stochastic gradient descent,— in Proceedings of COMPSTAT’2010, pp. 177–186, Springer, 2010.
16.A. Rakhlin, O. Shamir, K. Sridharan, et al., —Making gradient descent optimal for strongly convex stochastic optimization.,— in ICML, Citeseer, 2012.
17.J. Dean, G. Corrado, R. Monga, K. Chen, M. Devin, M. Mao, A. Senior, P. Tucker, K. Yang, Q. V. Le, et al., —Large scale distributed deep networks,— in Advances in neural information processing systems, pp. 1223–1231, 2012.
18.S. Ghadimi and G. Lan, —Stochastic first-and zeroth-order methods for nonconvex stochastic programming,— SIAM Journal on Optimization, vol. 23, no. 4, pp. 2341–2368, 2013.
19.V. Smith, C.-K. Chiang, M. Sanjabi, and A. S. Talwalkar, —Federated multi-task learning,— in Advances in Neural Information Processing Systems, pp. 4427–4437, 2017.
20.P. Warden, —Speech commands: A dataset for limited-vocabulary speech recognition,— arXiv preprint arXiv:1804.03209, 2018.
21.A. Krizhevsky, I. Sutskever, and G. E. Hinton, —Imagenet classification with deep convolutional neural networks,— in Advances in neural information processing systems, pp. 1097–1105, 2012.
22.B. Graham, —Fractional max-pooling,— arXiv preprint arXiv:1412.6071, 2014.
23.I. J. Goodfellow, O. Vinyals, and A. M. Saxe, —Qualitatively characterizing neural network optimization problems,— in arXiv preprint arXiv:1412.6544, 2014.

Citation

MLA

Zhao, Y., et al. “Federated Learning with Non-IID Data”. arXiv, 2018, https://doi.org/10.48550/arXiv.1806.00582.

APA

Zhao, Y., Li, M., Lai, L., Suda, N., Civin, D., & Chandra, V. (2018). Federated Learning with Non-IID Data. arXiv. https://doi.org/10.48550/arXiv.1806.00582

Chicago

Zhao, Y., M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra. 2018. “Federated Learning with Non-IID Data”. arXiv, ahead of print. https://doi.org/10.48550/arXiv.1806.00582.

Harvard

Zhao, Y. et al. (2018) “Federated Learning with Non-IID Data”, arXiv [Preprint]. Available at: https://doi.org/10.48550/arXiv.1806.00582.

Vancouver

1. Zhao Y, Li M, Lai L, Suda N, Civin D, Chandra V (2018) Federated Learning with Non-IID Data. arXiv. https://doi.org/10.48550/arXiv.1806.00582

BibTeX

@article{https://doi.org/10.48550/arxiv.1806.00582,
  doi = {10.48550/ARXIV.1806.00582},
  url = {https://arxiv.org/abs/1806.00582},
  author = {Zhao, Yue and Li, Meng and Lai, Liangzhen and Suda, Naveen and Civin, Damon and Chandra, Vikas},
  keywords = {Machine Learning (cs.LG), Machine Learning (stat.ML), FOS: Computer and information sciences, FOS: Computer and information sciences},
  title = {Federated Learning with Non-IID Data},
  publisher = {arXiv},
  year = {2018},
  copyright = {arXiv.org perpetual, non-exclusive license}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Authors

Federated Learning with Non-IID Data

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Theoretical Upper Bound on Weight Divergence in Federated Averaging

Knowl 2 — Globally Shared Data Strategy for Non-IID Federated Learning

Knowl 3 — Weight Divergence Metric for Federated Learning

Knowl 4 — Accuracy Degradation of FedAvg Under Non-IID Class Skew

Knowl 5 — Dependence of Weight Divergence and Test Accuracy on Earth Mover's Distance