Layer-Selective Sharpness-Aware Minimization for Efficient and Accurate Federated Learning

Federated learning (FL) is a distributed learning paradigm that trains a model on distributed private data while preserving data privacy. To improve the global model's generalization, recent studies have increasingly applied Sharpness-Aware Minimization (SAM), known to guide the model toward a flat region in the parameter space, to local training in FL. However, naive application of SAM to local training doubles the computational cost without noticeably improving the validation accuracy of the global model. Previous works have focused on the discrepancy between global and local loss landscapes, however, our study finds that layer-wise sharpness plays a more important role in determining the impact of local SAM on the global model's generalization. In this paper, we propose a general FL framework, FedLSAM that perturbs only a subset of layers with a significant impact on reducing global sharpness, rather than perturbing the entire model. Experimental results show that the proposed method improves generalization performance compared with perturbing all layers, while having a negligible additional computational cost.