Astraea: Self-balancing Federated Learning for Improving Classification Accuracy of Mobile Deep Learning Applications

Moming Duan; Duo Liu; Xianzhang Chen; Yujuan Tan; Jinting Ren; Lei Qiao; Liang Liang

Source-linked AI summary

Astraea: Self-balancing Federated Learning for Improving Classification Accuracy of Mobile Deep Learning Applications

Moming Duan, Duo Liu, Xianzhang Chen, Yujuan Tan, Jinting Ren, Lei Qiao, Liang Liang

arXiv:1907.01132v2Jul 2, 2019cs.LGcs.DCstat.ML

TL;DR

Federated learning must train shared models from private, decentralized device data despite imbalanced distributions that can reduce accuracy. The paper proposes Astraea, combining data augmentation with KLD-guided mediator rescheduling, and reports higher imbalanced-dataset accuracy than FedAvg alongside lower communication traffic.

Problem
Imbalanced distributed training data can bias federated-learning models and decrease application accuracy, while mobile and IoT data remain decentralized and private.
Method
Astraea uses global-distribution-based data augmentation and mediator-based multi-client rescheduling according to KLD between client distributions and the uniform distribution.
Results
Astraea improves top-1 accuracy by 5.59% on imbalanced EMNIST and 5.89% on imbalanced CINIC-10 over FedAvg.
Takeaways & Limitations
Astraea addresses global and local imbalance while reducing communication traffic relative to FedAvg.

Abstract

from arXiv · show

Federated learning (FL) is a distributed deep learning method which enables multiple participants, such as mobile phones and IoT devices, to contribute a neural network model while their private training data remains in local devices. This distributed approach is promising in the edge computing system where have a large corpus of decentralized data and require high privacy. However, unlike the common training dataset, the data distribution of the edge computing system is imbalanced which will introduce biases in the model training and cause a decrease in accuracy of federated learning applications. In this paper, we demonstrate that the imbalanced distributed training data will cause accuracy degradation in FL. To counter this problem, we build a self-balancing federated learning framework call Astraea, which alleviates the imbalances by 1) Global data distribution based data augmentation, and 2) Mediator based multi-client rescheduling. The proposed framework relieves global imbalance by runtime data augmentation, and for averaging the local imbalance, it creates the mediator to reschedule the training of clients based on Kullback-Leibler divergence (KLD) of their data distribution. Compared with FedAvg, the state-of-the-art FL algorithm, Astraea shows +5.59% and +5.89% improvement of top-1 accuracy on the imbalanced EMNIST and imbalanced CINIC-10 datasets, respectively. Meanwhile, the communication traffic of Astraea can be 82% lower than that of FedAvg.

I. INTRODUCTION

Federated learning trains shared models without moving private device data, but imbalanced distributions can degrade accuracy. Astraea addresses global and local imbalance through augmentation and mediator-based client rescheduling, improving accuracy while reducing communication traffic.

Federated learning context: Federated learning lets clients train a shared model locally while a server aggregates updates from a random client subset.This approach preserves local training data and supports collaborative model adaptation.
Problem: 7.92% accuracy loss occurs for FedAvg when distributed training data has global class imbalance.The paper distinguishes this global imbalance from local non-IID imbalance and notes that balanced global distributions are often assumed despite being uncommon in practice.
Approach: Astraea combines global-distribution-based data augmentation with mediator-based client rescheduling using KLD relative to the uniform distribution.The mediators combine skewed-client training to seek a partial equilibrium.
Results: 5.59% top-1 accuracy improvement on imbalanced EMNIST and 5.89% on imbalanced CINIC-10 are reported over FedAvg.The framework also reduces mean mediator-to-uniform-distribution KLD below 0.2.
Results: Astraea can reduce communication traffic by 92% versus FedAvg when achieving 75% accuracy on imbalanced EMNIST.The contribution list also reports recovery of 70.5% of accuracy loss on imbalanced EMNIST and 47.83% on imbalanced CINIC-10.

A. Background

The paper frames imbalance as a practical challenge for federated learning on heterogeneous mobile and IoT data. Experiments using distributed EMNIST show that global class imbalance, rather than every form of local imbalance, can substantially reduce accuracy and bias predictions toward majority classes.

A. Background: Federated learning research has emphasized communication and privacy, while relatively few studies address accuracy degradation from imbalance.Prior work mainly considers local imbalance while assuming balanced global data distributions.
A. Background: Mobile-device data distributions vary with device usage, producing size imbalance, local non-IID imbalance, and global class imbalance.The paper illustrates this distinction using cameras deployed in schools versus cameras deployed in the wild.
B. Motivation: The study constructs distributed EMNIST datasets because no large distributed image-classification dataset is available for this analysis.EMNIST contains 47 handwritten-letter and digit classes, and the paper builds five distributed variants with different imbalance settings.
B. Motivation: Global imbalance reduces EMNIST test top-1 accuracy from 81.60% on INS1 to 73.68% on LTRF1, a 7.92% reduction.LTRF2 shows a 6.20% reduction despite having twice the training data of LTRF1.
B. Motivation: Random local imbalance does not degrade accuracy, while scalar imbalance slightly improves test accuracy from 79.99% to 81.60%.The result separates the observed accuracy impact of global imbalance from the tested local and scalar imbalance settings.
B. Motivation: Confusion matrices show six lowest-frequency English-letter classes poorly classified under LTRF1, with predictions biased toward majority classes.BAL1 instead has most classifications concentrated along the diagonal.

III. DESIGN OF ASTRAEA

Astraea is designed to recover federated-learning accuracy by addressing both global imbalance and local client imbalance. Its design combines minority-class augmentation with mediator-based client rescheduling toward partial equilibrium.

III. DESIGN OF ASTRAEA: Astraea targets precision degradation from imbalanced federated training data by relieving both global imbalance and local client imbalance.The paper motivates the framework after mathematically demonstrating that imbalance can decrease federated-learning accuracy.
III. DESIGN OF ASTRAEA: The framework rebalances training through data augmentation for minority classes and mediator-based client rescheduling toward partial equilibrium.These two strategies address the global and local aspects of the imbalance problem, respectively.

A. Mathematical Demonstration

The mathematical analysis compares ideal SGD with federated optimization and concludes that imbalanced client data distributions prevent federated learning from reaching optimal test-set weights. This motivates rebalancing client training through Astraea.

Problem formulation: The paper formulates federated learning on imbalanced data as a precision-degradation problem and compares it with traditional SGD-based deep learning.The ideal SGD case is used to derive the optimal-weight update.
Assumptions: The analysis assumes balanced test-set distribution and identical initial weights for SGD-based and federated learning.The training-data distribution is denoted by ˆpdata, while ptest denotes the balanced image-classification test distribution.
Federated optimization: Each client optimizes its local weights by gradient descent on its own training-data distribution, which may be imbalanced.The local update uses learning rate η and client-specific weights w(k).
Result: The server aggregates client updates with FedAvg, but the resulting federated weights cannot reach the optimal weights when client training distributions are imbalanced.The proof compares the federated update with the ideal update and identifies the mismatch in data distributions as the source of degradation.
Implication: The paper therefore proposes augmenting minority classes and combining skewed clients through mediators to approach partial equilibrium.These two mechanisms become the core design of Astraea.

B. Astraea Framework

Astraea inserts mediators between the FL server and clients to rebalance training while preserving local data. It combines clients with complementary class distributions and uses KLD toward uniformity to guide this process.

Framework: Astraea introduces mediators between the FL server and clients to rebalance training without redistributing local data.The mediators coordinate client training before updates are aggregated by the server.
Client categories: Clients are categorized as uniform, slight, or biased according to data quantity and class distribution.Slight clients have little data, while biased clients hold particular classes and contribute to global imbalance.
Mediator coordination: Mediators reschedule clients with complementary distributions, such as clients holding labels 0–1 and 2–3, to form a partial equilibrium.Combining multiple clients can also expand the effective training set and expose the mediator to more patterns.
Balancing criterion: Astraea measures partial equilibrium with Kullback-Leibler divergence between the mediator distribution and the uniform distribution.The mediator seeks a combined distribution close to uniform.
Training and aggregation: During training, mediators coordinate assigned clients, collect their updates, and send aggregated updates to the FL server for weighted FedAvg aggregation.The server weights mediator updates by nm/n before producing the next global model.

C. Astraea Workflow

Astraea’s workflow initializes the global model, augments underrepresented classes, reschedules clients through mediators, trains sequentially within mediator groups, and aggregates the resulting updates. The design trades greater balancing against communication overhead.

Initialization: During initialization, clients share local distribution information, after which the server counts the global distribution and initializes model weights and the optimizer.The workflow begins only after participating devices have joined the training task.
Rebalancing: The server computes augmentation amounts for underrepresented classes, and clients apply transformations such as shifts, rotations, shears, and zooms in parallel.The augmentation set is determined from the global class distribution.
Mediator rescheduling: After augmentation, mediators greedily assign clients whose combined distribution minimizes KLD to the uniform distribution, with each mediator coordinating at most γ clients.The client-assignment procedure continues until every client is assigned.
Training: Each client performs mini-batch SGD and returns its updated model to the mediator, which passes the model to the next waiting client across mediator epochs.The local epoch affects training time, while mediator epochs determine repeated model updates before server aggregation.
Aggregation: The server aggregates mediator updates with FedAvg, producing a more balanced integrated model through partial equilibrium.Increasing collaborating clients can improve balancing but also raises communication overhead.

A. Experimental Setup

The evaluation implements Astraea by modifying TensorFlow Federated and tests it on imbalanced EMNIST and CINIC-10 against FedAvg. The setup uses dataset-specific models and client parameters.

Implementation: Astraea is implemented by modifying the TensorFlow Federated Framework and evaluated with its single-machine simulation runtime.The experiments use TFF as the execution framework.
Datasets and models: The evaluation uses imbalanced EMNIST with its CNN model and imbalanced CINIC-10 with the CIFAR-10 model described in Keras documentation.Imbalanced CINIC-10 is created by resampling the original dataset.
Parameters: For imbalanced EMNIST, the experiment sets K=500 and B=20; for imbalanced CINIC-10, it sets K=100 and B=50.These are the dataset-specific experimental parameters reported in the setup.
Baseline: FedAvg is selected as the state-of-the-art federated learning baseline.The baseline has also been applied to Google Keyboard for query-suggestion improvement.

B. Effect of Accuracy

Astraea’s accuracy depends on both global augmentation and mediator rescheduling, while client and mediator configurations affect convergence, equilibrium, and overfitting. Across the evaluated settings, combining the strategies provides the strongest supported improvements.

Augmentation vs. mediator: +5.89% top-1 accuracy is achieved on imbalanced CINIC-10 after applying the rescheduling strategy with α = 0.67.Data augmentation alone improves +4.12% top-1 accuracy when α = 1.00.
Augmentation vs. mediator: Combining augmentation and rescheduling is important because rescheduling achieves partial equilibrium but cannot solve global imbalance.The paper reports that combining both strategies achieves the maximum accuracy improvement.
c vs. γ: Larger client participation accelerates early convergence, but after 150 rounds large c can reduce accuracy through over-training and overfitting.For c = 100 and γ = 20, accuracy decreases from 79.03% to 77.79%; larger γ does not improve accuracy.
Mediator equilibrium: Mediator rescheduling reduces mean KLD from 0.550 to 0.125, whereas augmentation alone reduces it from 0.550 to 0.498.Larger γ reduces KLD variation, suggesting better partial equilibrium when more clients participate or are assigned to mediators.

C. Overhead

Astraea adds time, storage, and per-synchronization communication costs, but its experiments report lower total communication for reaching a target accuracy. Storage and timeout behavior constrain configurations with stronger augmentation.

Overhead types: Astraea’s additional overheads comprise data augmentation, client rescheduling, mediator training time, augmentation storage, and communication.The paper ignores computational overhead because the extra calculations require few resources and can run on the FL server.
Time overhead: Mediator training time per synchronization round is EmγE × T, and augmentation time is considered negligible because it occurs only during initialization.Rescheduling also adds time through its client-search process.
Storage overhead: 1.61% accuracy improvement is achieved on imbalanced EMNIST without additional storage, rising to 3.28% with 25.5% additional storage.The total additional augmentation storage is 90 MB, or 185 KB per client; α = 2 fails due to timeout.
Communication overhead: Astraea’s synchronization-round traffic is 2|w|(⌈c/γ⌉ + c), while each FL communication round costs 2c|w|.Here, |w| denotes the size of all model parameters.
Communication overhead: 81.7% reduction in communication cost is reported for reaching 75% top-1 accuracy: Astraea uses 215 MB versus FL’s 1176 MB.Astraea’s synchronization round has higher traffic than an FL communication round, but it requires less total communication to reach the target accuracy.

V. CONCLUSION

The conclusion identifies global imbalance as an accuracy problem for federated learning and presents Astraea as a two-part rebalancing solution. Experiments report accuracy gains on both evaluated datasets and communication efficiency.

Conclusion: Astraea combines minority-class data augmentation with mediator-based client rescheduling to rebalance federated training.The mediators reschedule clients toward partial equilibrium.
Conclusion: +5.59% top-1 accuracy is reported on imbalanced EMNIST and +5.89% on imbalanced CINIC-10 versus FedAvg.The conclusion also states that Astraea’s communication is effective.