PaperlayerView exact revision on arXiv

Source-linked AI summary

A Multi-agent Reinforcement Learning Approach for Efficient Client Selection in Federated Learning

Sai Qian Zhang, Jieyu Lin, Qi Zhang

arXiv:2201.02932v1cs.LGcs.AI

TL;DR

FL faces non-IID client data, heterogeneous computing and networking resources, and communication overhead, while existing solutions do not jointly address these objectives. FedMarl uses multi-agent reinforcement learning for run-time client selection. It reports the best accuracy alongside lower processing latency and communication cost, with performance across varied FL conditions.

  • Problem

    Existing FL solutions do not jointly address non-IID convergence challenges, processing latency, and communication overhead, although optimizing one objective can worsen another.

    Source: Introduction · p. 1

  • Method

    FedMarl uses multi-agent reinforcement learning to select clients at run time, using device-performance statistics and training behavior to guide decisions toward application-designated objectives.

    Source: Introduction · p. 2

  • Results

    FedMarl achieves the best model accuracy while reducing processing latency by 1.7× and communication cost by 2.2×, and outperforms benchmark algorithms across various FL conditions.

    Source: Introduction · p. 2, Conclusion · p. 7

  • Takeaways & Limitations

    The evaluation reports that FedMarl's performance extends across different FL settings without retraining the MARL agents.

    Source: Introduction · p. 2

    The theoretical stability analysis assumes the dataset aggregated from clients' non-IID data is approximately IID.

    Source: Theoretical Analysis · p. 5

Abstract

from arXiv · show

Federated learning (FL) is a training technique that enables client devices to jointly learn a shared model by aggregating locally-computed models without exposing their raw data. While most of the existing work focuses on improving the FL model accuracy, in this paper, we focus on the improving the training efficiency, which is often a hurdle for adopting FL in real-world applications. Specifically, we design an efficient FL framework which jointly optimizes model accuracy, processing latency and communication efficiency, all of which are primary design considerations for real implementation of FL. Inspired by the recent success of Multi-Agent Reinforcement Learning (MARL) in solving complex control problems, we present \textit{FedMarl}, an MARL-based FL framework which performs efficient run-time client selection. Experiments show that FedMarl can significantly improve model accuracy with much lower processing latency and communication cost.

Introduction

Federated learning (FL) lets clients train a shared model through local updates without exposing raw data, but non-IID data and heterogeneous devices make it difficult to optimize accuracy, processing latency, and communication cost together. FedMarl uses multi-agent reinforcement learning for run-time client selection and reports gains across these objectives and across different FL settings.

Introduction

  • FL enables clients to jointly learn a shared model by aggregating local DNN weight updates without exposing raw data.The server selects clients each round, receives their updates, and aggregates them into the next global model.Source: Introduction · p. 1
  • Non-IID data can degrade convergence, while heterogeneous client resources and data sizes create stragglers; iterative model updates also incur high communication cost.Prior approaches address these objectives separately, and optimizing one can worsen another, as illustrated by Scaffold's added-gradient latency.Source: Introduction · p. 1
  • The relative importance of accuracy, processing time, and communication cost varies by application, making a hand-tuned client-selection policy costly to develop.The paper notes that such a policy for each combination of objectives requires substantial time and effort.Source: Introduction · pp. 1–2
  • FedMarl achieves the best model accuracy while reducing processing latency by 1.7× and communication cost by 2.2×, and its agents perform well across different FL settings without retraining.The agents are trained using real traces of DNN training and transmission latencies collected over multiple mobile devices.Source: Introduction · p. 2, Figure 1 · p. 2

Background and Related Work

The background describes FedAvg's round-based client training and reviews FL work on statistical heterogeneity and system efficiency. It then introduces cooperative MARL, where agents seek a shared reward and Value Decomposition Networks combine their value functions.

Federated Learning

  • FedAvg selects a fixed number of clients each round to train locally and return updates; its performance degrades with non-IID data.Prior FL methods address statistical diversity through local/global model divergence control or client-personalized models, but do not consider system performance objectives.Source: Federated Learning · p. 2
  • Latency- and communication-focused FL methods improve their respective objectives, but none of the reviewed approaches jointly optimizes accuracy, processing latency, and communication cost.Examples include prioritizing faster devices, filtering or quantizing updates, and sketch-based compression.Source: Federated Learning · p. 2

Multi-Agent Reinforcement Learning

  • Cooperative MARL agents observe states, select actions, and receive a joint reward, aiming to maximize the total expected discounted reward.The agents proceed to the next state after completing their actions.Source: Multi-Agent Reinforcement Learning · p. 2
  • A Value Decomposition Network trains each agent's DNN to infer actions and represents the joint Q-function as the sum of individual Q-functions.The agents select actions with maximum Q-values, and a replay buffer stores transition tuples for training.Source: Multi-Agent Reinforcement Learning · p. 2
  • Figure 3 plots test accuracies for three algorithms separately on LeNet and VGG6.The caption identifies the architectures and plotted metric but does not state which algorithm performs best.Source: Figure 3 · p. 3

Motivation and Problem Definition

Non-IID data complicates FL convergence, while client-side training and uploads create substantial, uneven round costs. The formulation jointly targets accuracy, processing latency, and communication cost.

FL Processing Time Breakdown

  • Training dominates and varies by device: Nexus 5 takes 8.9s versus 1.4s communication, while Columbus upload takes 1.68s versus 361ms download.For 60 samples and five epochs, Google Pixel XL takes 4.8s, 2× faster than Nexus 5; LeNet and ResNet-18 show a similar upload/download trend.Source: FL Processing Time Breakdown · p. 3

Problem Formulation

  • The formulation chooses clients across T rounds to maximize final global-test accuracy while minimizing processing latency and communication cost, with designer-set objective weights.The client-selection matrix is T × N; the objective takes expectation over stochasticity in DNN training.Source: Problem Formulation · p. 4

FedMarl System Design

FedMarl casts client selection as a MARL problem and uses per-client agents to choose participation at run time. Its design combines agent state and action choices with a reward reflecting accuracy, processing latency, and communication cost, within a workflow that gathers client statistics and updates the global model.

FedMarl System Design

Design of MARL Agent States

  • Each agent's state includes probing loss, latency history, communication cost, local dataset size, and the current training round.The state vector is normalized, and agent MLP weights are shared to reduce storage overhead and prevent lazy-agent problems.Source: Design of MARL Agent States · p. 4

Description on Agent Actions

  • Agents choose whether a client stops after probing training or continues local training and uploads its model update.The action is binary: 0 terminates the client after probing; the other action continues participation.Source: Description on Agent Actions · p. 4

Design of the Reward Function

  • The reward reflects changes in test accuracy, processing latency, and communication cost after client selection.Its accuracy utility is designed to remain moderately responsive to small improvements near the end of FL, and agents are trained using VDN.Source: Design of the Reward Function · pp. 4–5
  • FedMarl achieves the best accuracies across LeNet, VGG6, ResNet-18, and LSTM datasets.Table 2 reports accuracy performance for these models and datasets.Source: Table 2 · p. 5

FedMarl System Workflow

  • Each round gathers probing losses and latency statistics, makes client-selection decisions, then has selected clients complete training and upload updates for the global model.The server workflow also updates communication-latency statistics.Source: FedMarl System Workflow · p. 5

Theoretical Analysis

The analysis considers non-IID client data and compares FedMarl's global model weights with weights trained by conventional SGD on an approximately IID aggregation of client data. Under the stated assumptions, the global model weights are bounded, and the paper concludes that FedMarl converges.

Theoretical Analysis

  • The analysis assumes the aggregation of all clients' non-IID training data is approximately IID.It also assumes class-conditional gradients are Lipschitz and bounded, and uses the learning rate in the bound.Source: Theoretical Analysis · p. 5
  • The global model weights are bounded; the paper therefore concludes that FedMarl converges.The proof is provided in the supplementary materials.Source: Theoretical Analysis · p. 5

Evaluation

FedMarl is evaluated against benchmark selection methods for accuracy, latency, and communication cost, then tested across system conditions, model and dataset combinations, learned selection behavior, and reward settings.

Evaluation

  • Under the same training environment, MARL converges much faster than single-agent RL; convergence takes 300, 200, 300, and 200 episodes for LeNet, VGG6, ResNet-18, and LSTM, respectively.Source: Evaluation · p. 6

FedMarl Performance

  • FedMarl achieves the best combination of test accuracy, processing latency, and communication cost; FedMarl, HeteroFL, and FedNova outperform other methods on accuracy.Methods that select only some clients reduce latency and communication cost, unlike FedAvg, FedProx, and FedNova, which transmit all client updates.Source: FedMarl Performance · p. 6

Generalizability of the MARL Agents

Evaluation under Different System Conditions

  • Agents trained with K = 100 outperform the other algorithms at K = 400 and K = 800, with 1.36× and 1.32× higher average total reward, respectively.FedMarl also has the best overall performance at different local-training epoch counts and transfers to a pool augmented with six new mobile-device types.Source: Evaluation under Different System… · p. 7, Figure 7 · p. 7

Evaluation across Different DNN Architectures and Datasets

  • Agents generalize across DNN architectures and datasets: ResNet-18/Fashion MNIST-trained agents score 0.965 on LeNet/MNIST, versus 1.0 for agents trained on LeNet from scratch.Finetuning on the target DNN for 10 episodes further improves reward.Source: Table 3 · p. 6, Evaluation across Different DNN… · p. 7

Learned Strategy by the MARL Agents

  • The agents favor clients with low probing loss and latency, but sometimes select high-latency clients with low loss to balance accuracy and processing latency.They select 3 of 10 clients in the first round and 7 of 10 in the last round.Source: Figure 6 · p. 6, Learned Strategy by the MARL Agents · p. 7

Performance under Different Rewards

  • Increasing reward weights w2 from 0.2 to 0.5 and w3 from 0.1 to 0.3 lowers total latency and communication cost at the price of lower accuracy.The reward settings let application designers customize FedMarl according to their objective preferences.Source: Performance under Different Rewards · p. 7, Table 4 · p. 7

Conclusion

FedMarl performs run-time client selection in federated learning and outperforms benchmark algorithms across model accuracy, processing latency, and communication cost under varied FL conditions.

Conclusion

  • FedMarl outperforms benchmark algorithms in model accuracy, processing latency, and communication cost under various FL conditions.The framework uses MARL to make intelligent client selections at run time.Source: Conclusion · p. 7

Related papers

Loading 2201.02932v1…