GRO: Rectifying Local Gradient Conflicts via Swarm Dynamics for Robust Federated Learning

This repository provides an official PyTorch implementation of GRO (Federated Geometric Reflection Optimizer), a plug-and-play optimization strategy designed to improve the stability and generalization of federated learning (FL) under heterogeneous (non-IID) data distributions.

GRO rectifies local gradient conflicts via swarm dynamics, effectively suppressing 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.

• Federated Learning Baselines: PyTorch implementations of representative FL baselines and their GRO-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)

• GRO Aggregation Module: Implementation of Federated Geometric Reflection Optimizer (GRO) 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.

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

Install all required packages by running:

pip install -r requirements.txt

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"

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

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