Source-linked AI summary
Federated learning with class imbalance reduction
Miao Yang, Akitanoshou Wong, Hongbin Zhu, Haifeng Wang, Hua Qian
TL;DR
Federated learning must train across private, heterogeneous device data while selecting only a subset of clients under communication constraints. The paper estimates class distributions from updated gradients without raw data, then uses a CMAB-based selection algorithm to reduce class imbalance. Numerical results show that the proposed algorithm selects properly balanced client sets and improves global-model convergence performance.
Problem
Non-IID client data and constrained participation create class imbalance in federated learning, while raw device data cannot be shared with the server.
Method
The paper estimates class distributions from updated gradients and uses a CMAB-based client-selection algorithm to reduce class imbalance without requiring client data information.
Results
The proposed algorithm picks properly balanced client sets and improves the convergence performance of the global model.
Takeaways & Limitations
Selecting clients toward minimal class imbalance can improve federated-learning convergence under heterogeneous data distributions.
Takeaways & Limitations
The approach assumes a balanced auxiliary dataset at the global server and operates under a spectrum setting with fewer available channels than client devices.
Abstract
from arXiv · showhide
Federated learning (FL) is a promising technique that enables a large amount of edge computing devices to collaboratively train a global learning model. Due to privacy concerns, the raw data on devices could not be available for centralized server. Constrained by the spectrum limitation and computation capacity, only a subset of devices can be engaged to train and transmit the trained model to centralized server for aggregation. Since the local data distribution varies among all devices, class imbalance problem arises along with the unfavorable client selection, resulting in a slow converge rate of the global model. In this paper, an estimation scheme is designed to reveal the class distribution without the awareness of raw data. Based on the scheme, a device selection algorithm towards minimal class imbalance is proposed, thus can improve the convergence performance of the global model. Simulation results demonstrate the effectiveness of the proposed algorithm.
1. INTRODUCTION
The paper addresses class imbalance and client-selection challenges in federated learning with heterogeneous private data. It proposes estimating class distributions without raw data and selecting balanced client combinations through reinforcement learning.
- Federated learning trains models on private device data while transmitting model updates rather than raw data to a global server.
- Wireless connectivity, device availability, and non-IID client datasets restrict participation to a subset of devices in each training round.
- Existing client-selection methods often assume IID client data, limiting their applicability to practical heterogeneous settings.
- Non-IID device data can degrade FedAvg accuracy and increase the communication rounds required for convergence.
- The proposed approach estimates class-imbalance severity without raw client data and uses reinforcement learning to select a minimally imbalanced client combination.
2. SYSTEM MODEL
The system models multiclass DNN training in federated learning, where selected clients perform local SGD and the server aggregates their updates with FedAvg. Non-IID data creates divergence between local objectives and the global objective, potentially destabilizing or even diverging training.
- The system consists of multiple clients with local datasets, a global aggregation server, and communication over a multiple access channel.
- Limited spectrum makes the number of available communication channels smaller than the number of client devices.
- The multiclass DNN maps input features to probabilities over C classes and uses cross-entropy loss to evaluate classification training.
- In each communication round, the server selects a client subset, distributes the global weights, and selected clients perform local SGD updates.
- The server aggregates participating-client updates with FedAvg to update the global model.
- With non-IID data, local optimization objectives differ from the global objective, causing accumulated weight divergence that can degrade learning.
3. ONLINE LEARNING ALGORITHM FOR CLIENT SELECTION
The paper estimates clients’ class distributions from updated gradients without accessing raw data, then combines CUCB-based online learning with class-balanced client selection.
- 3.1. Class Estimation Scheme: The scheme uses a balanced auxiliary dataset and updated-model gradients to estimate each client’s class distribution without obtaining raw client data.The output layer’s neurons correspond to classes, and auxiliary gradients are linked to those class-specific neurons.
- 3.1. Class Estimation Scheme: Gradient-square expectations provide the relation between class frequencies and gradients, producing a composition vector R for each client.A tunable β controls normalization between classes, while KL divergence evaluates client class imbalance.
- 3.2. Online Learning Framework for Client Selection: Client selection is formulated as a combinatorial multi-armed bandit problem in which individual clients are arms and client sets are super arms.The selected set receives a nonlinear reward, motivating use of the CUCB algorithm.
- 3.2. Online Learning Framework for Client Selection: CUCB maintains client-selection counts and perturbed reward estimates, using exploration to avoid repeatedly selecting only frequently chosen clients.The perturbation artificially increases expected rewards for clients selected less often.
- 3.2. Online Learning Framework for Client Selection: The class-balancing algorithm initializes with the highest estimated reward client and iteratively adds clients to form a combination with minimal class imbalance.Together, the estimation and selection algorithms identify a suitable balanced client set for FL updating.
4. NUMERICAL RESULTS
Numerical experiments on CIFAR10 evaluate the proposed client-selection algorithm against alternative schemes and parameter settings. The results show improved convergence and accuracy, with diminishing gains from selecting more clients and a trade-off governed by exploration.
- Experimental setup: The experiments use CIFAR10 with 50,000 training examples, 10,000 test examples, 10 classes, and 100 clients.Clients receive random numbers of classes and data samples to model imbalanced distributions.
- Experimental setup: The evaluation compares the proposed algorithm with greedy and random selection schemes, and also considers an IID setting.The proposed method uses α = 0.2, ρ = 0.99, and β = 1 in the reported simulation.
- Selection schemes: The proposed algorithm converges faster and reaches higher test accuracy than the greedy and uniform schemes.The reported difference from random selection is attributed to reducing class imbalance through balanced client selection.
- Client quantity: Increasing the number of selected clients improves training performance, but the improvement becomes smaller as the client count grows.The paper notes that selecting too many clients is not essential when communication consumption is considered.
- Exploration factor: The exploration factor α controls the exploitation–exploration trade-off, and a suitable value is necessary for convergence performance.Small α limits exploration, whereas excessive exploration can deteriorate performance because exploratory clients may be unsuitable.
5. CONCLUSIONS
The paper addresses class-imbalanced client selection in federated learning without requiring client data information. It combines class-distribution estimation with CMAB-based selection, and numerical results show improved global-model convergence.
- Conclusion: The proposed scheme reveals client class distributions from updated gradients without requiring clients’ data information.This estimation supports selection of client sets with more balanced class composition.
- Conclusion: The CMAB-based client-selection algorithm reduces the effect of class imbalance in federated learning.The algorithm is designed to select properly balanced client sets.
- Conclusion: Numerical results confirm that selecting properly balanced client sets improves the convergence performance of the global model.The conclusion reports this as the main demonstrated outcome of the proposed approach.