PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalized Federated Learning via Sharpness Aware Minimization

Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, Zhuo Lu

arXiv:2206.02618v1cs.LG

TL;DR

Cross-device FL faces severe distribution shift from heterogeneous client data, while ERM can lead to sharp loss valleys and conflicting local objectives. The paper introduces FedSAM and MoFedSAM, analyzes their convergence, and reports competitive performance with significantly reduced deviation among local clients.

  • Problem

    Cross-device FL suffers distribution shift because client datasets vary considerably, while ERM can produce sharp loss valleys and local overfitting that harm global-model performance.

  • Method

    FedSAM uses Sharpness Aware Minimization as the local optimizer, while MoFedSAM additionally uses previous global-model updates during local SAM training to bridge local and global models.

  • Results

    The proposed algorithms achieve competitive performance and significantly decrease performance deviation among local clients; FedSAM also has a reported generalization bound.

  • Takeaways & Limitations

    SAM-based local optimization offers a generality-focused approach to distribution shift in FL without additional communication costs compared with existing FL studies.

  • Takeaways & Limitations

    The theoretical analysis assumes smooth local objectives and bounded variance of the global gradient without perturbation.

Abstract

from arXiv · show

Federated Learning (FL) is a promising framework for performing privacy-preserving, distributed learning with a set of clients. However, the data distribution among clients often exhibits non-IID, i.e., distribution shift, which makes efficient optimization difficult. To tackle this problem, many FL algorithms focus on mitigating the effects of data heterogeneity across clients by increasing the performance of the global model. However, almost all algorithms leverage Empirical Risk Minimization (ERM) to be the local optimizer, which is easy to make the global model fall into a sharp valley and increase a large deviation of parts of local clients. Therefore, in this paper, we revisit the solutions to the distribution shift problem in FL with a focus on local learning generality. To this end, we propose a general, effective algorithm, \texttt{FedSAM}, based on Sharpness Aware Minimization (SAM) local optimizer, and develop a momentum FL algorithm to bridge local and global models, \texttt{MoFedSAM}. Theoretically, we show the convergence analysis of these two algorithms and demonstrate the generalization bound of \texttt{FedSAM}. Empirically, our proposed algorithms substantially outperform existing FL studies and significantly decrease the learning deviation.

1. Introduction

The introduction frames cross-device FL distribution shift as a consequence of heterogeneous client data and argues for improving local learning generality. It proposes FedSAM and MoFedSAM, with theoretical analyses and empirical evidence of faster convergence and reduced client deviation.

  • Motivation: Cross-device FL faces distribution shift because clients operate in varied environments and their local datasets differ considerably.
  • Motivation: ERM can overfit local training data and place models in sharp loss valleys, harming global-model performance under conflicting client objectives.
  • FedSAM: FedSAM applies Sharpness Aware Minimization as the local optimizer to improve global-model generalization without requiring adversarial min-max optimization.
  • MoFedSAM: MoFedSAM additionally uses previous global-model updates during SAM-based local training to bridge smooth information between local and global models without accessing private client data.
  • Theory and results: MoFedSAM achieves speedup under full and partial client participation, while FedSAM matches the best convergence rate of existing FL studies and provides a generalization bound.
  • Scope and conclusion: The paper focuses on classic cross-device FL rather than Personalized FL and reports that its algorithms significantly reduce performance deviation among local clients.

2. Preliminaries and Proposed Algorithms

The paper formulates client-heterogeneous FL as an ERM problem and argues that sharp loss regions and distribution shifts can hurt generalization across clients. FedSAM replaces local ERM with SAM-based perturbed-loss optimization, while its two-step local procedure approximates the inner maximization efficiently.

  • Motivation: Client distributions may differ, so an ERM-trained global model can fit a mixture while failing to generalize across all clients.The paper links this failure to client heterogeneity and differing validation distributions.
  • Motivation: Existing FL approaches address distribution shift through local-epoch choices, proximal terms, or knowledge distillation, while fairness methods focus on average performance.The paper motivates a more general global model that also reduces performance deviation.
  • FedSAM: SAM seeks a low-loss neighborhood by perturbing model parameters, rather than only a single ERM solution.The perturbed objective is defined using a radius ρ around the model.
  • FedSAM: For small ρ, FedSAM uses first-order Taylor expansion to convert the inner maximization into a linear constrained optimization.The resulting perturbed model is formed from the gradient direction and perturbation radius.
  • FedSAM: FedSAM iteratively computes a perturbed model near the highest local loss and then performs gradient descent using the gradient at that perturbed point.This procedure applies SAM locally without the min-max optimization cost associated with adversarial FL.
  • FedSAM: SAM smoothness is presented as a mechanism that can improve FL generalization and convergence by reducing the effective loss-surface smoothness parameter.The paper notes that the perturbed loss can be smoother than the original loss surface.

3. Theoretical Analysis

The theoretical analysis establishes convergence results for FedSAM under non-convex FL assumptions and derives a non-asymptotic generalization bound. It also compares convergence terms with prior methods and explains how perturbation and network parameters enter the bound.

  • Assumptions: The analysis assumes smooth local losses, bounded global gradient variability, and stochastic gradients that are unbiased with bounded variance.The paper identifies the first two assumptions as standard in non-convex FL and imposes a tighter variance-related condition.
  • FedSAM convergence: FedSAM convergence is proved for full and partial client participation under learning-rate and perturbation schedules proportional to inverse square-root communication-round factors.The full-participation theorem specifies ηl and ρ schedules and states convergence of the generated iterates.
  • FedSAM convergence: FedSAM’s dominant convergence terms match the best rates reported for existing general non-convex FL studies, while uniform sampling does not fundamentally change the convergence structure.The rates contain an additional higher-order term associated with SAM’s local smoothness step.
  • FedSAM convergence: For partial participation, the dominant term reflects fewer participating clients and random-sampling heterogeneity, and the convergence rate improves as the number of clients increases.The paper attributes this improvement to SAM making the global model more generalizable and reducing distribution shift.
  • Comparison: FedRobust has a less favorable convergence rate than FedSAM and requires multiple gradient-descent steps per local epoch, increasing local computation.The comparison is presented as both theoretical and computational.
  • Generalization bounds: Theorem 3.5 gives a non-asymptotic FedSAM generalization bound for deep networks, with dependence on perturbation size, margins, sample size, and network norms.The bound uses a spectral-norm framework and assumes smooth 1-Lipschitz ReLU activations.
  • Generalization bounds: The margin-based risk uses client-specific SAM losses and empirical client distributions, while γ = 0 reduces it to an average misclassification rate under distribution shift.The notation distinguishes underlying client probabilities from empirical probabilities estimated from training samples.

4. Momentum FedSAM (MoFedSAM)

MoFedSAM extends FedSAM with momentum that carries global-model information into local training. The analysis reports speedups over FedSAM, while experiments compare testing accuracy, communication targets, and client deviation.

  • MoFedSAM method: MoFedSAM reuses aggregated global-model information to guide local training and directly smooth and generalize the global model.The method is motivated by the observation that FedSAM’s local optimizer does not directly control the global-model gradient term.
  • MoFedSAM method: The momentum term contributes to each local training epoch, retains information from previous updates under partial participation, and approximates the global perturbed gradient.It also acts as a correction to the local gradient direction.
  • Theoretical analysis: MoFedSAM convergence is analyzed under full and partial client participation, with learning-rate and perturbation choices stated in the theorem conditions.The theorem assumes a bounded momentum parameter and full participation for one stated result, followed by a partial-participation result.
  • Experiments: Experiments evaluate testing accuracy on different datasets and report average and standard-deviation training and testing accuracy plus communication rounds to target accuracies.The targets are EMNIST 80%, CIFAR-10 80%, and CIFAR-100 50%.
  • Theoretical analysis: For sufficiently large T relative to K, both participation strategies have stated O(…) convergence rates, and the paper reports speedup compared with FedSAM.The local-training-related convergence term is also stated separately.

5. Experiments

Experiments across multiple datasets, models, heterogeneity levels, and parameter settings show that FedSAM and MoFedSAM outperform same-category FL benchmarks while reducing client-performance deviation.

  • Experimental Setup: Experiments span three datasets, three learning models, five FL benchmarks, and varying parameter settings.The cross-device setup uses 100 clients with 20% participation, Dirichlet 0.6 heterogeneity, and K = 10 local epochs by default.
  • Performance with compared benchmarks: FedSAM outperforms FedAvg, SCAFFOLD, and FedRobust, while MoFedSAM outperforms MimeLite and FedCM on accuracy and convergence.Comparisons are made within the same momentum category because momentum FL transmits additional information.
  • Impact of Non-IID levels: 0.43%, 1.24%, and 1.52% higher test accuracy and 7, 40, and 59 fewer communication rounds are achieved by MoFedSAM as non-IID severity increases on CIFAR-10.The comparison is against MimeLite across IID, Dirichlet 0.6, and Dirichlet 0.3 settings.
  • Performance with compared benchmarks: FedSAM and MoFedSAM significantly decrease the deviation between best and worst local accuracy.The reported deviation indicates stronger generalization of the global model across local clients.
  • Loss surface visualization: FedSAM and MoFedSAM significantly improve loss-surface sharpness over FedAvg on ResNet-18 trained with CIFAR-10.The visualization uses two random sampled orthogonal Gaussian perturbations as the x- and y-axes.
  • Impact of other parameters: The best perturbation radius ρ differs by dataset: 0.2 for EMNIST, 0.5 for CIFAR-10, and 0.6 for CIFAR-100.Increasing batch size and participating clients improves learning performance, whereas increasing K does not substantially guarantee better accuracy.

6. Conclusion

The paper proposes FedSAM and MoFedSAM to improve global-model generality under heterogeneous cross-device FL, and supports them with convergence analysis, a FedSAM generalization bound, and experiments showing reduced client deviation.

  • 6. Conclusion: FedSAM and MoFedSAM address distribution shift from data heterogeneity by focusing on global-model generality.The algorithms do not generate more communication costs than existing FL studies.
  • 6. Conclusion: The paper derives convergence results for the proposed algorithms in general non-convex FL settings and presents a generalization bound for FedSAM.The convergence analysis includes full and partial client participation results.
  • 6. Conclusion: Extensive experiments strongly support significantly decreased performance deviation among local clients.The conclusion connects reduced client deviation with the proposed focus on global-model generality.

B.3. Convergence Analysis of Partial Client Participation FedSAM

For partial client participation, the FedSAM analysis bounds the stochastic update behavior under learning-rate and participation conditions, yielding a convergence result whose higher-order terms are characterized.

  • Update bounds: The analysis bounds the expected squared global update norm for partial client participation.The derivation includes the partial-participation scheme without replacement.
  • Theorem B.9: Theorem B.9 analyzes FedSAM under partial client participation with ηl ≤ 1 10KL, ηgηl ≤ 1 KL, and a positivity condition involving S, K, L, and ηl.The theorem applies under Assumptions 1–3.
  • Learning-rate choices: Choosing ηl = 1 √KS and perturbation amplitude ρ proportional to the learning rate supports the partial-participation convergence analysis.The analysis considers conditions relating the number of sampled clients S and local epochs K.
  • Convergence expression: When S ≥ K, the analysis omits higher-order terms to obtain a simplified convergence expression.The resulting expression contains terms involving R, K, and L.

C. Generalization Bounds

The paper derives a margin-based generalization bound for FedSAM by combining perturbation-based neural-network bounds with a union bound across clients and layers.

  • Bound formulation: The generalization analysis treats FedSAM’s SAM local optimizer through a margin-based error for general neural networks.The bound is formulated for a perturbation ˜w = w + δ satisfying a margin-based condition.
  • Proof strategy: The proof combines PAC-Bayesian perturbation bounds with neural-network output-change bounds based on layer perturbation magnitudes.The supporting lemmas respectively control data-independent perturbation effects and network output changes.
  • Bound dependencies: The resulting bound depends on neural-network parameters including depth, hidden-layer width, spectral norms, and the Frobenius norm.The theorem assumes bounded input norm, 1-Lipschitz activations, and layer-wise constraints.
  • Client-level guarantee: A union bound extends the result across all layers and the distributions of all N clients.The final statement bounds the average SAM loss of the clients with probability at least 1 − ζ.

D.2. Convergence Analysis of Full client participation MoFedSAM

This section analyzes MoFedSAM under full client participation, using descent and convergence results under stated learning-rate and smoothness assumptions.

  • Descent analysis: The analysis establishes a descent lemma for MoFedSAM with full client participation.The expectation is taken over algorithmic stochasticity.
  • Convergence guarantee: Theorem D.4 gives a convergence guarantee for MoFedSAM under full client participation.The theorem is stated for constant local and global learning rates subject to the paper’s conditions.
  • Convergence guarantee: The convergence proof selects learning rates and perturbation amplitude with inverse-square-root dependence on the round horizon.The perturbation amplitude ρ is chosen proportional to the learning rate.

D.3. Convergence Analysis of Partial client participation MoFedSAM

This section extends MoFedSAM’s convergence analysis to partial client participation through a participation-specific descent lemma and convergence theorem.

  • Descent analysis: The analysis first bounds the expected squared update norm for partial client participation.This bound is stated in Lemma D.5.
  • Descent analysis: The proof combines the descent relation with bounds from Lemmas A.2–A.4 and a participation-dependent constraint on the global learning rate.The derivation also substitutes the bound on the third term into the descent inequality.
  • Descent analysis: Lemma D.6 provides a descent relation for MoFedSAM when only a subset of clients participates.The expectation is with respect to the stochasticity of the algorithm.
  • Convergence guarantee: Theorem D.7 states convergence of MoFedSAM under partial client participation and prescribed local and global learning rates.The theorem is given under Assumptions 1–3 and the corresponding participation conditions.
  • Convergence guarantee: The proof uses learning rates and perturbation amplitude scaled with the number of communication rounds, with S ≥ K as an additional condition.Here, ρ is proportional to the local learning rate.

E. Experimental Setup

The experiments use EMNIST, CIFAR-10, and CIFAR-100 in a cross-device federated setting, pairing dataset-specific models with tuned optimization parameters.

  • Datasets and models: The experimental datasets are EMNIST, CIFAR-10, and CIFAR-100.EMNIST is a 62-class image-classification dataset, while CIFAR-10 and CIFAR-100 contain 10 and 100 labels, respectively.
  • Federated setting: Client heterogeneity is simulated with Dirichlet label-ratio sampling using parameter 0.6 by default.The same Dirichlet allocation parameter is used for CIFAR-10 and CIFAR-100.
  • Datasets and models: A CNN is used for EMNIST, whereas ResNet-18 is used for CIFAR-10 and CIFAR-100.The EMNIST CNN has two convolutional layers, max pooling, dropout, and a 128-unit dense layer.
  • Optimization settings: Learning rates and momentum coefficients are selected by grid search, with ηg = 1 and ηl = 0.1 as defaults.The momentum grid includes β values from 0.01 to 1.

F.1. Training accuracy on different datasets

Training accuracy alone obscures the generalization gap across heterogeneous clients. The proposed methods are evaluated through training and validation performance, with MoFedSAM showing lower local-model deviation in the reported CIFAR-100 example and greater robustness to heterogeneity than benchmarks.

  • Training accuracy: Training accuracy on different datasets is compared in Figure 4.The paper contrasts these results with validation accuracy to assess performance divergence.
  • Training accuracy: 85.26% training accuracy and 4.41% local-model deviation for FedCM contrast with 86.02% and 3.23% for MoFedSAM on CIFAR-100.These are the reported training-accuracy and deviation values for the two algorithms.
  • Validation comparison: 54.09% validation accuracy with 14.38% deviation for FedCM contrasts with 55.13% and 3.25% for MoFedSAM on CIFAR-100.The validation results reveal a smaller reported deviation for MoFedSAM in this example.
  • Heterogeneity: 95.42%, 94.20%, and 92.90% are the reported FedSAM training accuracies on CIFAR-10 under IID, Dirichlet 0.6, and Dirichlet 0.3 splits.The corresponding validation accuracies are 87.36%, 82.55%, and 79.82%.
  • Heterogeneity: Increasing heterogeneity substantially degrades validation performance while not significantly affecting training accuracy.The paper reports that the proposed algorithms are less influenced by heterogeneity than the compared benchmarks.
  • Parameter impacts: Increasing the number of participating clients can improve performance, whereas increasing local epochs does not guarantee better accuracy.The effects of SAM’s ρ and momentum β depend on the algorithm and dataset.
Loading 2206.02618v1…