A Fast and Flat Federated Learning Method via Weighted Momentum and Sharpness-Aware Minimization

AI summary

Overview

Research area: Federated learning (FL), specifically optimization under non-IID (heterogeneous) client data, with a focus on convergence speed and model generalization.

Technical level: Intermediate. The paper assumes working familiarity with FedAvg, server/client momentum methods (FedCM, MIME), Sharpness-Aware Minimization (SAM), and variance-reduction techniques (SCAFFOLD).

Scope: The paper diagnoses two structural failure modes that appear when momentum and SAM are combined naively in non-IID federated learning, and proposes a single algorithm—FedWMSAM—that corrects both while keeping per-round cost close to FedAvg.

What This Paper Is About

Federated learning needs models that converge in few communication rounds and still generalize across clients whose data distributions differ sharply. Momentum helps with the first goal and SAM (which seeks flat minima) helps with the second, but simply stacking them together under non-IID data does not work well. The authors identify why—local SAM perturbation directions do not reflect the global loss surface, and accumulated momentum causes late-stage oscillation—and design a method that fixes both.

Key Contributions

1. Formalized two failure modes. The paper names and characterizes local–global curvature misalignment (each client's SAM perturbation is computed from local data and may point away from the global loss geometry) and momentum-echo oscillation (persistent momentum causes instability and overfitting in late training rounds).

2. Momentum-guided global perturbation with single-backprop SAM. Instead of a separate backward pass to compute the SAM perturbation, the method derives the perturbation direction from the difference between the current local model and the server-aggregated momentum, which acts as a proxy for where the global model is heading. This aligns local SAM directions with global geometry at no extra backward-pass cost.

3. Cosine-similarity adaptive coupling. A dynamic weight α_r, updated from the cosine similarity between the global momentum and each client's personalized momentum, produces an early-momentum / late-SAM schedule: momentum dominates early for speed, then yields to SAM for a flatter, more stable finish.

4. Non-IID convergence bound. The analysis explicitly models perturbation-induced variance as σ_ρ² = σ² + (Lρ)² and gives a rate depending on (S, K, R, N), extending the SCAFFOLD-M framework to cover adaptive personalized momentum plus SAM.

Main Findings

• Largest gains under strong heterogeneity: On CIFAR-10 at β = 0.1, FedWMSAM reaches 76.64% accuracy, +6.59 points over FedAvg; on CIFAR-100 it reaches 46.46%, +8.31 points. Under pathological splits (γ = 3 for CIFAR-10, γ = 10 for CIFAR-100) it reaches 74.46% and 43.83%, improving on FedAvg by +10.2 and +7.52 points.

• Broad competitiveness: Across three datasets and twelve heterogeneity settings, FedWMSAM is best or second-best in most cases. On OfficeHome it leads on 3 of 4 domains (Art, Clipart, Product) and has the best average, trailing SCAFFOLD only on Real World.

• Faster to target accuracy at near-FedAvg cost: It hits 0.70 accuracy in 97 rounds (matching the fastest baselines) and then pulls ahead—114, 153, 241, and 356 rounds for 0.72, 0.74, 0.76, and 0.78—whereas FedAvg, FedSAM, and FedLESAM never reach the higher targets. Per-round client time is 15.03s, close to FedAvg (14.57s) and far below double-backprop SAM variants (26.9–29.8s).

• All three modules are necessary: Ablation against a FedCM baseline (0.7229 accuracy) shows the full method at 0.7664 (+4.35%). Removing the cosine-adaptive weight drops to 0.7556; removing SAM drops to 0.7326; removing momentum drops to 0.7265. The adaptive gate only helps when SAM is present.

• Robust to the perturbation radius ρ: Accuracy peaks at ρ = 0.01 and degrades gracefully (0.7664 → 0.7563 → 0.7244 → 0.5905 as ρ moves to 0.05, 0.1, 0.5), while MoFedSAM collapses (0.7102 → 0.5562 → 0.1000) under the same conditions. The adaptive gate appears to reduce the effective perturbation during noisy rounds.

• Stable as local work grows: Unlike most baselines, FedWMSAM shows no clear accuracy drop as local epochs increase from 5 to 10 to 20, and it leads consistently once the client population reaches 75 or more. It also outperforms across most client sampling rates, with gains concentrated when participation is sparse.

• Better-separated representations: t-SNE visualizations of global model embeddings on CIFAR-10 show tighter clusters and cleaner class separation than FedAvg, FedSAM, or MoFedSAM, consistent with flatter minima.

Methodology in Plain English

The approach rests on one key observation: momentum already carries information about the direction the global model is moving. If you know that direction, you can estimate where the global model would be after a step, and the gap between the current local model and that estimated global position tells you which way the loss is rising—a natural perturbation direction for SAM.

Concretely, the server keeps a global momentum and, each round, builds a personalized momentum for every selected client by adding a SCAFFOLD-style correction term scaled by α_r/(1−α_r). This scaling is chosen so the correction affects the gradient exactly as it would in the local update, which lets the server transmit a single combined vector rather than separate momentum and correction vectors, saving bandwidth.

Clients then run local steps using the perturbed gradient x + ρ·δ/‖δ‖, where δ is the difference between the inferred global position and the current local model. This avoids the second backward pass that standard SAM requires.

After each round, the server measures how similar the global momentum is to each client's personalized momentum via cosine similarity, averages these, and updates α_r with a moving-average rule clamped to [0.1, 0.9]. As similarity rises over training, α_r rises, shifting weight from momentum toward SAM. The clamp bounds matter: the upper bound keeps momentum strong enough that SAM still has a usable direction, and the lower bound prevents momentum's influence from vanishing too early.

The convergence proof adapts the SCAFFOLD-M analysis, adding the SAM perturbation term and the adaptive personalized momentum, and shows the induced variance grows with the square of the perturbation radius times the Lipschitz constant.

Why This Matters

Impact on research. The paper reframes the momentum-plus-SAM combination as a structural problem rather than a tuning problem. Naming local–global curvature misalignment and momentum-echo oscillation gives the FL community concrete, testable failure modes, and the σ_ρ² = σ² + (Lρ)² term links perturbation size to convergence variance in a way that other SAM-based FL papers do not. The result that FedWMSAM degrades gracefully where MoFedSAM collapses at large ρ suggests a general robustness advantage from coupling the two mechanisms rather than stacking them.

Real-world applications:

• Mobile keyboard and speech models, where hundreds of millions of devices hold non-IID text and audio and communication budgets are tight.
• Healthcare analytics across hospitals, where patient populations differ by institution and generalization to held-out cohorts matters more than raw training accuracy.
• Industrial IoT and predictive maintenance, where sensors on different machines produce drifted distributions and connectivity is intermittent, making partial participation the norm.
• Cross-organization vision systems such as the OfficeHome setup tested here, where each domain (Art, Clipart, Product, Real World) is effectively a separate client.

Industry relevance. The near-FedAvg per-round cost is the headline for deployment: it means the generalization benefit of SAM can be obtained without the doubled compute that has kept SAM-family methods out of production FL pipelines. Reduced communication rounds translate directly into lower bandwidth costs and longer battery life on edge devices, and the bandwidth savings from combining momentum and correction into one transmitted vector are relevant to the privacy-conscious aggregation setting.

Future Directions

• Extending beyond the Dirichlet and pathological splits tested. The evaluation uses synthetic heterogeneity; whether the same alignment mechanism holds under real, time-varying, and label-noise-heavy client distributions is untested.

• Combining with secure aggregation and differential privacy. The method currently relies on the server inspecting client gradient differences to compute correction terms—an operation that would conflict with cryptographic privacy guarantees and needs a privacy-preserving reformulation.

• Tuning or eliminating the hyperparameters λ, the α_r clamp bounds, and ρ. The bounds [0.1, 0.9] are justified partly by prior work on FedCM rather than derived, and ρ peaks sharply at 0.01, so an adaptive radius rule could remove a sensitive knob.

• Theoretical sharpening of the bound. The N^(2/3)/S term and the assumptions inherited from SCAFFOLD-M may be tighten-able, and it is unclear how much of the empirical gain is explained by the bound versus other effects.

Target Audience

This paper is most useful to federated learning researchers and graduate students working on optimization under heterogeneity, particularly those already familiar with momentum-based FL or SAM. Practitioners building production FL systems will find the per-round cost comparison and the robustness to large ρ directly actionable. Readers seeking an introduction to FL should start elsewhere, since the paper assumes fluency in FedAvg, FedCM, SCAFFOLD, and SAM notation from the first page.

Authors’ abstract

In federated learning (FL), models must \emph{converge quickly} under tight communication budgets while \emph{generalizing} across non-IID client distributions. These twin requirements have naturally led to two widely used techniques: client/server \emph{momentum} to accelerate progress, and \emph{sharpness-aware minimization} (SAM) to prefer flat solutions. However, simply combining momentum and SAM leaves two structural issues unresolved in non-IID FL. We identify and formalize two failure modes: \emph{local-global curvature misalignment} (local SAM directions need not reflect the global loss geometry) and \emph{momentum-echo oscillation} (late-stage instability caused by accumulated momentum). To our knowledge, these failure modes have not been jointly articulated and addressed in the FL literature. We propose \textbf{FedWMSAM} to address both failure modes. First, we construct a momentum-guided global perturbation from server-aggregated momentum to align clients' SAM directions with the global descent geometry, enabling a \emph{single-backprop} SAM approximation that preserves efficiency. Second, we couple momentum and SAM via a cosine-similarity adaptive rule, yielding an early-momentum, late-SAM two-phase training schedule. We provide a non-IID convergence bound that \emph{explicitly models the perturbation-induced variance} $σ_ρ^2=σ^2+(Lρ)^2$ and its dependence on $(S, K, R, N)$ on the theory side. We conduct extensive experiments on multiple datasets and model architectures, and the results validate the effectiveness, adaptability, and robustness of our method, demonstrating its superiority in addressing the optimization challenges of Federated Learning. Our code is available at https://github.com/Huang-Yongzhi/NeurlPS_FedWMSAM.