Research map / Federated learning and privacy / Communication-efficient FL

Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based Stochastic Controlled Weight Averaging

ICML 2025 · Federated learning and privacy › Communication-efficient FL · approach: Client participation analysis

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes FedSWA (SWA-style local learning-rate decay with EMA server aggregation) and FedMoSWA (momentum-based stochastic control variates) for federated learning under highly heterogeneous data, with uniform-stability generalization bounds and CIFAR/Tiny ImageNet experiments.

Problem

Generalization of federated learning under highly heterogeneous (Non-IID) data, where FedSAM-type local flat-minimum methods perform poorly relative to FedAvg.

Approach

Replace SAM-style local perturbation with SWA-style cyclical/decaying local learning rates and EMA server aggregation (FedSWA), and add momentum-based variance-reduction control variates (server m, client c_i) to align local and global models (FedMoSWA), with uniform-stability generalization bounds that depend on data heterogeneity.

Key results

• On CIFAR-100 ResNet-18 Dirichlet-0.1, FedMoSWA reaches 61.9% vs 45.8% FedAvg, 40.1% FedSAM, and 51.5% MoFedSAM.
• On CIFAR-100 ResNet-18 Dirichlet-0.6, FedMoSWA reaches 67.9% in 330 rounds vs 54.2% for FedAvg and 60.1% for MoFedSAM.
• On Tiny ImageNet ViT-Base Dirichlet-0.1, FedMoSWA reaches 73.8% vs 70.9% for FedAvg.

Measured effects reported in the paper

Surprising findings

• FedSAM performs worse than FedAvg under highly heterogeneous data (e.g., 40.1% vs 45.8% on CIFAR-100 Dirichlet-0.1).
• MoFedSAM collapses to 1.0% accuracy on CIFAR-100 VGG-11 Dirichlet-0.6 in the table, an anomalous result the paper does not explain.

Limitations

• Generalization bound retains a dependence on data heterogeneity sigma_g and does not eliminate it, as the authors acknowledge.
• Theoretical analysis relies on strong convexity or smoothness and uniform-stability assumptions, while experiments use non-convex deep networks.
• Full-participation and 10% participation ablations use different client counts and round budgets from main experiments.

Experimental setup

Datasets: CIFAR-10 CIFAR-100 Tiny ImageNet

Benchmarks: CIFAR-10 CIFAR-100 Tiny ImageNet

Metrics: test accuracy communication rounds to target accuracy loss surface sharpness (visualization)

Baselines: FedAvg FedDyn SCAFFOLD FedSAM MoFedSAM FedLESAM FedASAM FedACG FedAvgM FedGAMMA

Builds on: FedSAM SCAFFOLD Stochastic Weight Averaging LookAhead MoFedSAM

Closest papers by meaning

• FedGPS: Statistical Rectification Against Data Heterogeneity in Federated Learning (NeurIPS 2025 · similarity 0.89)
• Improving Generalization in Federated Learning with Model-Data Mutual Information Regularization: A Posterior Inference Approach (NeurIPS 2024 · similarity 0.89)
• Scaffold with Stochastic Gradients: New Analysis with Linear Speed-Up (ICML 2025 · similarity 0.88)
• Momentum Benefits Non-iid Federated Learning Simply and Provably (ICLR 2024 · similarity 0.88)
• A New Theoretical Perspective on Data Heterogeneity in Federated Optimization (ICML 2024 · similarity 0.88)
• FedWMSAM: Fast and Flat Federated Learning via Weighted Momentum and Sharpness-Aware Minimization (NeurIPS 2025 · similarity 0.88)