Federated Learning with Packet Losses

This repository is the official implementation of Federated Learning with Packet Losses.

This paper tackles the problem of training Federated Learning (FL) algorithms over real-world wireless networks with packet losses. Lossy communication channels between the orchestrating server and the clients affect the convergence of FL training as well as the quality of the learned model. Although many previous works investigated how to mitigate the adverse effects of packet losses, this paper demonstrates that FL algorithms over asymmetric lossy channels can still learn the optimal model, the same model that would have been trained in a lossless scenario by classic FL algorithms like FedAvg. Convergence to the optimum only requires slight changes to FedAvg: i) while FedAvg computes a new global model by averaging the received clients' models, our algorithm, UPGA-PL, updates the global model by a pseudo-gradient step; ii) UPGA-PL accounts for the potentially heterogeneous packet losses experienced by the clients to unbias the pseudo-gradient step. Still, UPGA-PL maintains the same computational and communication complexity as FedAvg. In our experiments, UPGA-PL not only outperforms existing state-of-the-art solutions for lossy channels (by more than 5 percentage points on test accuracy) but also matches FedAvg's performance in lossless scenarios after less than 150 communication rounds.

Requirements

To install requirements:

pip install -r requirements.txt

Overview

Federated Training

We provide code to simulate federated training of machine learning models. The core objects are Client, ClientsSampler, and Aggregator: different federated learning algorithms can be simulated by implementing the local update scheme Client.step(), the clients' sampling strategy defined in ClientsSampler.sample(), and/or the aggregation protocol defined in Aggregator.mix() and Aggregator.update_client().

We model client activity, considering packet loss as a form of client unavailability. A client experiencing packet loss in downlink or uplink will be marked as non-available for the current federated learning round.

Similarly to CA-Fed, our code includes a Markovian clients' activity simulator, implemented in ActivitySimulator. Each client's activity follows a Markovian model with two states: 'active' and 'inactive'. The ClientsSampler observes the clients' availability dynamics generated by the ActivitySimulator using ClientsSampler.get_active_clients(). Configuration parameters for clients' activity dynamics can be found in data\constants.py and data\main.py. For advanced customization, refer to the documentation in data\README.md.

Once the ActivitySimulator is instantiated, the aggregator observes the client activity pattern during ClientsSampler.sample(). Subsequently, the aggregator can define a suitable aggregation strategy using Aggregator.mix().

Algorithms

In addition to the baseline Ideal FedAvg scenario, this paper implements three aggregation strategies, each defining a federated learning algorithm:

1. Direct Model Aggregation with Packet Loss (DMA-PL)

   • Directly averages received models.
   • Reference: Chen et al. 2021

2. Unbiased Direct Model Aggregation with Packet Loss (UDMA-PL)

   • Extends DMA-PL by considering clients' different loss probabilities in the aggregation.

3. Unbiased Pseudo-Gradient Aggregation with Packet Loss (UPGA-PL, Ours)

   • Aggregates pseudo-gradients rather than model parameters.
   • Accounts for potentially heterogeneous packet losses to unbias the pseudo-gradient step.

To execute these algorithms programmatically:

• Ideal FedAvg: Use IdealClientsSampler along with either ParamsAggregator or GradientsAggregator.
• DMA-PL: Use BiasedClientsSampler along with ParamsAggregator.
• UDMA-PL: Use UnbiasedClientsSampler along with ParamsAggregator.
• UPGA-PL: Use UnbiasedClientsSampler along with GradientsAggregator.

Alternatively, execute the run_experiment.py Python file and pass values of {ideal, biased, unbiased} as the client_sampler argument and {params, gradients} as the aggregator argument.

For practical examples, refer to the paper_experiments/ section.

Datasets

We provide three federated benchmarks spanning different machine learning tasks: multi-class classification (LEAF synthetic dataset), handwritten character recognition (MNIST), and image classification (CIFAR10).

The following table summarizes the available datasets and models:

Dataset	Task	Model
Synthetic LEAF	Multi-class classification	Linear model
MNIST	Handwritten character recognition	2-layer CNN + 2-layer FFN
CIFAR10	Image classification	MobileNet v2

See data/README.md for instructions on generating data and simulation configuration file.

Paper Experiments

Scripts to reproduce the paper experiments are provided in paper_experiments/. Specify the name of the dataset (experiment), the used algorithm (client sampler and aggregator), and configure all other hyper-parameters (please refer to their values as specified in the paper). An example on one dataset (MNIST), with our proposed algorithm (UPGA-PL), is as follows:

echo "==> Run experiment with 'unbiased' sampler, 'gradients' aggregator"

sampler="unbiased"
aggregator="gradients"

python run_experiment.py \
--experiment "mnist" \
--cfg_file_path data/mnist/cfg.json \
--objective_type weighted \
--aggregator_type "${aggregator}" \
--clients_sampler "${sampler}" \
--n_rounds 200 \
--local_steps 5 \
--local_optimizer sgd \
--local_lr 0.01 \
--server_optimizer sgd \
--server_lr 1.0 \
--train_bz 128 \
--test_bz 1024 \
--device "cuda" \
--log_freq 5 \
--verbose 1 \
--logs_dir "logs_tuning/mnist/activity_${sampler}_${aggregator}" \
--history_path "history_tuning/mnist/activity_${sampler}_${aggregator}.json" \
--seed 123

To reproduce the paper experiments, please refer to paper_experiments/.

Results

In this section, we present the performance evaluation of each aggregation strategy, including Ideal FedAvg, DMA-PL, UDMA-PL, and UPGA-PL. The evaluation is conducted on the local test dataset, which was unseen during training.

The figure below illustrates the aggregated results for each algorithm, showing the averaged train loss and test accuracy. The metrics are computed over 10 clients and averaged across 10 independent runs. The evaluation spans across two distinct datasets: Synthetic Clustered and MNIST.

Citation

If you use our code or wish to refer to our results, please use the following BibTex entry:

@article{rodio2023losses,
  title={{Federated Learning with Packet Losses}},
  author={Rodio, Angelo and Neglia, Giovanni and Busacca, Fabio and Mangione, Stefano and Palazzo, Sergio and Restuccia, Francesco and Tinnirello, Ilenia},
  journal={26th International Symposium on Wireless Personal Multimedia Communications (WPMC)},
  year={2023}
}