FedVSSAM: Mitigating Flatness Incompatibility in Sharpness-Aware Federated Learning

TL;DR

FedVSSAM introduces a novel approach to address flatness incompatibility in federated learning, improving global model generalization under data heterogeneity by stabilizing local updates.

Contribution

The paper proposes FedVSSAM, a variance-suppressed sharpness-aware method that aligns local and global directions to mitigate flatness incompatibility in federated learning.

Findings

FedVSSAM effectively mitigates flatness incompatibility.

It outperforms baseline methods in diverse FL settings.

Theoretical convergence guarantees are established.

Abstract

Sharpness-aware minimization (SAM) is an effective method for improving the generalization of federated learning (FL) by steering local training toward flat minima. Under data heterogeneity, however, device-side SAM searches for locally flat basins that are incompatible with the flat region preferred by the global objective. We identify this structural failure mode as flatness incompatibility, which explains why improving local flatness alone may provide limited training and generalization improvement for the global model. We reveal that flatness incompatibility arises from data heterogeneity and the friendly adversary phenomenon, and is further amplified by local updates and partial device participation. To mitigate this issue, we propose Federated Learning with variance-suppressed sharpness-aware minimization (FedVSSAM), which constructs a variance-suppressed adjusted direction and uses it consistently in local flatness search, local descent, and global update. FedVSSAM anchors both perturbation and update directions to a more stable global direction, instead of correcting only an isolated local perturbation. We establish non-convex convergence guarantees of FedVSSAM and prove that the mean-square deviation between the adjusted direction and the global gradient is effectively controlled. Experiments demonstrate that FedVSSAM mitigates flatness incompatibility and outperforms the baselines across diverse FL settings.