Optimized Model Training with Reinforcement Learning and Domain Preservation via Divergence Reduction in Distributed Learning

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 work, 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.

Additionally, in recent years, Large Language Models (LLMs) have shown promising results on a variety of tasks, and FL presents a gateway to train LLMs in a distributed setting. As stated above, FL aims to preserve privacy while improving performance. However, two constraints; the size of LLMs and the Non-IID nature of participating clients limit the performance gain which can be achieved through FL. We use Parameter Efficient Fine-Tuning (PEFT) and design a novel centralized server-side FL aggregation algorithm, Divergence Reduction in Federated Training (DRIFT), which leverages graph-like properties of participating clients to perform custom aggregation for each participating client. We design two variants of DRIFT and, through extensive experimentation on a range of diverse tasks, show how DRIFT preserves the distribution of client model parameters on their local domain and outperforms well-established baselines.