ECGR: Exploratory–Convergent Gradient Re-aggregation for Federated Learning

This repository provides an official PyTorch implementation of ECGR (Exploratory–Convergent Gradient Re-aggregation), a plug-and-play optimization strategy designed to improve the stability and generalization of federated learning (FL) under heterogeneous (non-IID) data distributions.

ECGR decomposes local update dynamics into exploratory and convergent gradient components and selectively re-aggregates local gradients to suppress noisy or conflicting updates while preserving effective descent directions. The proposed strategy can be seamlessly integrated into existing FL baselines without modifying local training procedures.

---

Features

• Federated Learning Baselines: PyTorch implementations of representative FL baselines and their ECGR-enhanced variants. The currently supported methods include:

  • FedAvg (H. B. McMahan et al., AISTATS 2017)
  • FedNova (J. Wang et al., NeurIPS 2020)
  • FedProx (T. Li et al., MLSys 2020)
  • SCAFFOLD (S. P. Karimireddy et al., ICML 2020)

• ECGR Aggregation Module: Implementation of Exploratory–Convergent Gradient Re-aggregation (ECGR) as a modular aggregation strategy that can be combined with existing FL baselines to improve robustness under non-IID data distributions.

• Dataset Preprocessing: Automated downloading and preprocessing of benchmark datasets, followed by partitioning into multiple clients according to federated learning settings. Non-IID data distributions are simulated via Dirichlet-based label skew. The currently supported datasets include MNIST, Fashion-MNIST, SVHN, CIFAR-10, and CIFAR-100. Other datasets (e.g., medical imaging datasets) need to be downloaded and organized manually.

• Postprocessing and Visualization: Tools for visualizing training dynamics and evaluating global model performance, including testing accuracy curves averaged over multiple random seeds.

---

Installation

Dependencies

• Python (3.8)
• PyTorch (1.8.1)
• OpenCV (4.5)
• NumPy (1.21.5)

Install Requirements

Install all required packages by running:

pip install -r requirements.txt

Running Federated Learning with ECGR

Test Run

All hyperparameters are specified in a YAML configuration file (e.g., ./config/test_config.yaml).
To run federated learning experiments with ECGR or baseline methods, execute:

python fl_main.py --config "./config/test_config.yaml"

Evaluation Procedures

You can place the results in the results/test directory, and then run the following command:

python postprocessing/eval_main.py -rr ../results/test