Optimized Local Updates in Federated Learning via Reinforcement Learning

2025-05-31

AI Summary

To address model aggregation degradation and inefficient local data utilization caused by non-IID data in federated learning, this paper proposes a deep reinforcement learning (PPO)-based method for adaptive optimization of local training data volume per client. The core contribution is twofold: (i) the first application of deep reinforcement learning to dynamically learn class-level data sampling weights, enabling fine-grained, privacy-preserving local data selection without sharing raw data; and (ii) incorporation of full local fine-tuning after global aggregation to mitigate non-IID bias. The method requires no additional server-side supervision or data exchange, relying solely on sparse rewards derived from local loss variations. Extensive experiments on CIFAR-10, CIFAR-100, and FEMNIST demonstrate an average accuracy improvement of 2.3–4.1%, faster convergence, and reduced communication and computational overhead.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Data and user privacy-enhancing technologies for the Web

Abstract

Federated Learning (FL) is a distributed framework for collaborative model training over large-scale distributed data, enabling higher performance while maintaining client data privacy. However, the nature of model aggregation at the centralized server can result in a performance drop in the presence of non-IID data across different clients. We remark that training a client locally on more data than necessary does not benefit the overall performance of all clients. In this paper, we devise a novel framework that leverages a Deep Reinforcement Learning (DRL) agent to select an optimized amount of data necessary to train a client model without oversharing information with the server. Starting without awareness of the client's performance, the DRL agent utilizes the change in training loss as a reward signal and learns to optimize the amount of training data necessary for improving the client's performance. Specifically, after each aggregation round, the DRL algorithm considers the local performance as the current state and outputs the optimized weights for each class, in the training data, to be used during the next round of local training. In doing so, the agent learns a policy that creates an optimized partition of the local training dataset during the FL rounds. After FL, the client utilizes the entire local training dataset to further enhance its performance on its own data distribution, mitigating the non-IID effects of aggregation. Through extensive experiments, we demonstrate that training FL clients through our algorithm results in superior performance on multiple benchmark datasets and FL frameworks. Our code is available at https://github.com/amuraddd/optimized_client_training.git.

Problem

• Optimizing data selection in Federated Learning for non-IID data
• Reducing unnecessary data sharing in client model training
• Improving client performance via reinforcement learning-based local updates

Innovation

• Uses DRL to optimize local data amount
• Adapts training weights per class dynamically
• Enhances performance post-FL with full data