Junkang Liu, Fanhua Shang, Yuxuan Tian, Hongying Liu, Yuanyuan Liu
8 min
Abstract
In federated learning (FL), multi-step local updates and data heterogeneity usually lead to sharper global minima, which degrades the performance of the global model. Popular FL algorithms integrate sharpness-aware minimization (SAM) into local training to address this issue. However, in the high data heterogeneity setting, the flatness in local training does not imply the flatness of the global model. Therefore, minimizing the sharpness of the local loss surfaces on the client data does not enable the effectiveness of SAM in FL to improve the generalization ability of the global model. We define the \textbf{flatness distance} to explain this phenomenon. By rethinking the SAM in FL and theoretically analyzing the \textbf{flatness distance}, we propose a novel \textbf{FedNSAM} algorithm that accelerates the SAM algorithm by introducing global Nesterov momentum into the local update to harmonize the consistency of global and local flatness. \textbf{FedNSAM} uses the global Nesterov momentum as the direction of local estimation of client global perturbations and extrapolation. Theoretically, we prove a tighter convergence bound than FedSAM by Nesterov extrapolation. Empirically, we conduct comprehensive experiments on CNN and Transformer models to verify the superior performance and efficiency of \textbf{FedNSAM}. The code is available at https://github.com/junkangLiu0/FedNSAM.
Sam: In the new method, called FedNSAM, the server first computes a global momentum—think of it as a smoothed average of the recent direction all selected hospitals have been heading, blending their changes over time like a team's shared pace memory. Each hospital then takes a lookahead step forward using a fraction of that global momentum, anticipating where the group is going. Next, it perturbs slightly in the opposite direction of that momentum to scan for flatness, just like SAM, but now guided by the group's rhythm. This lookahead—known as Nesterov acceleration—helps local searches align toward the same broad flat area, shrinking the mismatch.
Alex: So the momentum acts like a coxswain's call, getting everyone to peek ahead together before scanning locally?
Sam: Precisely. Without sharing private data or exact gradients, this approximation reduces noise in the direction signal. On a test like CIFAR-100 with high heterogeneity, it reaches about 12 percentage points higher accuracy than FedSAM. The paper shows this leads to a clear improvement in pulling the global model into flatter terrain under tough data differences.
Alex: That addresses the sync issue without extra communication costs. Worth noting for real-world setups like hospitals. How does this global momentum estimate hold up against other fixes, like just using the raw change from last round?
Sam: Other methods, say one called FedLESAM, just use the simple difference in model updates from the previous round as the global direction signal. But when only a few hospitals join each time, that raw change misses the bigger picture and gets noisy. FedNSAM's momentum instead builds an exponential moving average—like a running total that weights recent team updates more heavily while remembering past ones, smoothing out the noise for a steadier group direction.
Alex: So it's like a smoothed team velocity rather than a jerky snapshot—keeps the coxswain's call consistent even if rowers skip turns.
Sam: Yes. The paper backs this with theory under practical rules: losses are smooth, meaning small parameter tweaks don't swing wildly; local gradients vary but stay bounded by data differences. Their math shows FedNSAM converges to good spots about as fast as top methods for tricky non-convex landscapes, but tighter overall—handling partial participation better than FedSAM by reducing terms tied to heterogeneity. Experiments on image sets like CIFAR-100 with heavy data skew confirm this pattern.
Alex: That ties the theory straight to practice. Do the results hold up on something more complex, like those transformer architectures?
Sam: The paper tests this on larger vision models—think of transformers as systems that break images into small patches and weigh important connections between them, much like figuring out key links in a social network to understand a photo. They use Swin-Small, Swin-Base, and ViT-Base on Tiny ImageNet, a tougher image set with many classes. FedNSAM reaches the highest accuracies—around 70 to 71 percent—while needing far fewer training rounds than others.
Alex: Fewer rounds on heavy-hitters like that? What about when fewer hospitals join each round, or data skew gets extreme?
Sam: Under low participation, say 2 or 5 percent of clients active, or high data skew like Dirichlet-0.1, FedNSAM holds steady. On CIFAR-100, it drops less than baselines. The paper's loss landscape plots—2D maps coloring bumpy parameter areas red and flat ones blue—show FedNSAM's global surface much smoother than FedSAM variants, especially on test data. This visual ties back to why generalization improves: flatter testing losses mean better handling of unseen variety.
Alex: So the visuals confirm the flatter global spots endure real messiness. Any tweaks they tested to fine-tune this?
Sam: They ablate components, like pairing the Nesterov tweak with other bases such as Scaffold or FedDyn, gaining around 9.5 percentage points on CIFAR-100. The evidence points to robust gains under practical constraints.
Alex: That paints a picture of reliability across scales and stresses. Grounds the method nicely for uneven real-world data. But every method has tuning needs—what about the practical limits here, like picking the right settings or edge cases?
Sam: Fair point. The momentum factor λ needs tuning—around 0.85 works well in their runs, but straying can slow things or add wobble. Their math bounds rely on smooth losses and bounded data differences; extreme non-IID skews beyond tested levels remain unprobed. Still, under realistic hospital-like variety, it holds up without extra rounds or messages.
Alex: So assumptions fit common setups, but extremes call for more checks. For places like global hospitals with diverse patient groups, this means faster paths to shared models that generalize across demographics.
Sam: Precisely. It's a measured advance: deployable now for privacy-safe AI training, with room to build on those bounds. The evidence points to better tools for real mismatched data.
Alex: Makes sense—bridges the local-global gap without overpromising. That's our look at consistency in federated flatness. Thanks for listening to ResearchPod.