Focus on Where You Aggregate: Restricted SAM for Non-IID Federated Learning

Data heterogeneity and multiple local update steps in federated learning can increase client drift and make global parameter aggregation less stable. Existing Sharpness Aware Minimization (SAM) based federated learning methods usually try to ease this problem by flattening local neighborhoods around client models. However, they do not explicitly optimize on the geometry region where the global parameters are actually aggregated. As a result, the aggregated global model often cannot fully benefit from local flattening. To address this problem, we propose a federated restricted SAM method called FedRSAM, which improves the stability of federated optimization by flattening the parameter aggregation region. First, we build a cone-shaped geometry region using the global update direction from the previous round and the variance of client parameter aggregation, in order to predict where the aggregated parameters may fall in the next round. Then, we restrict the search of SAM perturbation directions to the cone-manifold and only flatten the loss surface inside this region. Extensive experiments on multiple benchmarks and on both vision and natural language backbones verify the superiority of our method in terms of convergence stability and performance under heterogeneous settings.