Improving Generalization in Federated Learning by Seeking Flat Minima

Debora Caldarola; Barbara Caputo; Marco Ciccone

Source-linked AI summary

Improving Generalization in Federated Learning by Seeking Flat Minima

Debora Caldarola, Barbara Caputo, Marco Ciccone

arXiv:2203.11834v3Mar 22, 2022cs.LGcs.CV

TL;DR

Federated learning can generalize poorly under heterogeneous data, and this paper studies the problem through loss geometry and Hessian eigenspectra. It applies SAM or ASAM during local training and SWA during server aggregation, finding smoother minima and improved generalization across vision tasks. The approach is evaluated in classification, semantic segmentation, and domain generalization settings.

Problem
Heterogeneous federated training suffers from poor generalization because models converge toward sharp minima and local updates reflect limited data distributions.
Method
The paper applies SAM or ASAM on clients and SWA on the server to encourage flatter minima and average stochastic weights.
Results
The resulting models show smoother loss surfaces and improved final performance consistently across several vision tasks.
Takeaways & Limitations
Seeking flatter minima provides an effective way to improve federated generalization across heterogeneous and homogeneous vision scenarios.
Takeaways & Limitations
The analysis assumes heterogeneous clients have different data distributions and that each training round may not observe the entire distribution.

Abstract

from arXiv · show

Models trained in federated settings often suffer from degraded performances and fail at generalizing, especially when facing heterogeneous scenarios. In this work, we investigate such behavior through the lens of geometry of the loss and Hessian eigenspectrum, linking the model's lack of generalization capacity to the sharpness of the solution. Motivated by prior studies connecting the sharpness of the loss surface and the generalization gap, we show that i) training clients locally with Sharpness-Aware Minimization (SAM) or its adaptive version (ASAM) and ii) averaging stochastic weights (SWA) on the server-side can substantially improve generalization in Federated Learning and help bridging the gap with centralized models. By seeking parameters in neighborhoods having uniform low loss, the model converges towards flatter minima and its generalization significantly improves in both homogeneous and heterogeneous scenarios. Empirical results demonstrate the effectiveness of those optimizers across a variety of benchmark vision datasets (e.g. CIFAR10/100, Landmarks-User-160k, IDDA) and tasks (large scale classification, semantic segmentation, domain generalization).

1 Introduction

The paper attributes federated generalization difficulties to convergence toward sharp minima, especially under heterogeneous client data. It introduces SAM/ASAM locally and SWA server-side, reporting smoother loss surfaces and improved performance across vision tasks.

Heterogeneous client data degrades convergence and generalization because local training sees only limited portions of the underlying distribution.
The analysis links poor federated generalization to convergence toward sharp minima using loss surfaces and Hessian eigenvalues.
SAM and ASAM encourage client models to converge toward flatter minima, while SWA averages stochastic weights during server aggregation.
The combined approach produces smoother loss landscapes, lower Hessian eigenvalues, and improved generalization capacities.
Experiments cover small- and large-scale classification, domain generalization, and semantic segmentation, with comparisons against strong augmentations and state-of-the-art federated algorithms.

2 Related Works

Related work frames federated learning’s challenges around statistical heterogeneity, real-world vision settings, and generalization to unseen domains. Studies of loss-surface geometry motivate SAM and ASAM, while this work applies flatter-minimum methods and stochastic-weight averaging to these federated scenarios.

Federated Learning Challenges: Statistical heterogeneity in users’ data can cause unstable, slow convergence, suboptimal performance, and poor global-model generalization in federated learning.FedAvg remains the standard optimization method, using multiple local SGD steps per communication round.
Proposed Perspective: The paper seeks uniformly low-loss neighborhoods locally and uses cyclical server-side stochastic-weight averaging to explore broader weight-space regions and reach wider optima.It connects flatter minima and broader exploration with improved generalization without additional communication cost.
Real-World Vision Settings: Prior FL research often emphasizes algorithmic aspects, whereas real-world vision applications include large-scale classification, semantic segmentation, and domain generalization.The paper uses Landmarks-User-160k for federated large-scale classification and examines segmentation and domain shifts relevant to autonomous driving and other settings.
Real-World Vision Settings: Federated domain generalization seeks domain-agnostic models that maintain performance on unseen domains, but this setting remains poorly studied.Prior examples examine equipment-induced medical shifts and landscape or weather changes in autonomous driving.
Loss Geometry and Generalization: Research linking loss-surface geometry to generalization associates sharp minima with poor generalization and defines flatness through connected low-training-loss regions.These studies motivate examining solution geometry as a perspective on federated generalization.
Loss Geometry and Generalization: SAM seeks flatter minima by jointly minimizing loss sharpness and value, while ASAM addresses SAM’s sensitivity to parameter rescaling through adaptive sharpness.The paper asks whether these methods improve generalization in federated learning across architectures and tasks.

3 Behind the Curtain of Heterogeneous FL

Heterogeneous federated training produces unstable, poorly generalizing models because local updates specialize on different data distributions and can drive the global model toward sharp regions. Loss-surface and Hessian analyses connect this behavior to curvature and show that heterogeneity changes both convergence and the geometry of the solution.

Federated learning setup: Federated learning trains a global model by aggregating local client updates over communication rounds, targeting low population loss rather than training loss alone.Clients optimize local losses on privacy-protected datasets, while the server aggregates their updates.
Heterogeneous behavior: Statistical heterogeneity makes convergence noisier and slower, with a considerable performance gap relative to centralized and homogeneous settings.On Cifar100, the heterogeneous setup uses Dirichlet concentration α ∈ {0, 0.5}, while α = 1000 represents the homogeneous setting.
Heterogeneous behavior: When α = 0, the global model performs highly on one class at a time but forgets the others, whereas homogeneous training behaves similarly across clients and eventually overfits.The α = 0 local models are specialized to narrow class distributions, while α = 1k produces more uniform client behavior.
Loss-landscape analysis: Local updates overfit their client distributions and occupy high-test-error regions, causing aggregation to move the global model away from a minimum.The resulting local models specialize too much on their own data and fail to generalize to the overall distribution.
Loss-landscape analysis: Both heterogeneous and homogeneous models converge to sharp regions, while Hessian analysis uses λmax as a curvature measure and λmax/λ5 as a sharpness proxy.A larger λmax indicates greater loss change in the corresponding direction and a steeper minimum.
Hessian analysis: For α = 0, local λmax values are at most 14 but aggregation drives the global model to λmax ≃94; for α ∈ {0.5, 1k}, λmax decreases over rounds more clearly.These observations motivate explicitly searching for flatter minima to improve generalization.

4 Seeking Flat Minima in Federated Learning

The paper addresses poor federated generalization by explicitly seeking flatter minima during client optimization and averaging stochastic server-side weights. It combines SAM or ASAM locally with SWA in FedAvg and uses these methods to target smoother loss landscapes and better generalization.

Motivation: Standard optimizers minimize training loss without using curvature information, while federated heterogeneity further promotes sharp minima and poor generalization.The paper motivates flatter-neighborhood optimization as a way to address this behavior.
Client optimization: Algorithm 1 applies SAM or ASAM through sharpness-aware local updates while retaining FedAvg-style server aggregation.The procedure subsamples clients, performs local epochs and minibatch updates, then sends client updates to the server.
Overall approach: The combined approach uses client-side sharpness-aware optimization and server-side stochastic-weight averaging to seek flatter neighborhoods in federated training.The paper presents this combination as its approach for improving generalization.
Server aggregation: SWA averages stochastic weights on the server side and begins from 75% of training onward with a learning-rate schedule.The method explores high-performing regions of weight space before updating the SWA average.

5 Experiments

Experiments across federated classification, segmentation, domain generalization, and real-world settings show that SAM, ASAM, and SWA improve performance by promoting flatter solutions, especially under heterogeneity.

5.1 The Effectiveness of the Search for Flat Minima in FL: Federated models in heterogeneous settings can trail homogeneous counterparts by up to 20 percentage points, motivating explicit optimization for flat minima.The experiments attribute poor generalization partly to client overspecialization and convergence to sharp minima.
5.1 The Effectiveness of the Search for Flat Minima in FL: FedASAM + SWA reduces λmax from 93.5 with FedAvg to 24.6 in the most heterogeneous setting, indicating flatter solutions.FedSAM and FedASAM alone yield λmax values of 70.3 and 30.1.
5.1 The Effectiveness of the Search for Flat Minima in FL: Flat-minima methods outperform corresponding baselines under heterogeneity, while their positive gap is larger in federated than centralized training.Mixup and Cutout worsen federated performance alone but become beneficial when combined with the proposed methods and SWA.
5.1 The Effectiveness of the Search for Flat Minima in FL: FedASAM consistently improves accuracy by approximately 6 percentage points over the best state-of-the-art baseline on Cifar10 and Cifar100.Adding ASAM also raises the performance of other federated algorithms, including FedAvgM and SCAFFOLD.
5.2 ASAM and SWA in Real World Vision Scenarios: Across real-world tasks, client-side sharpness-aware optimization combined with server-side averaging improves final accuracy by up to 7% on Landmarks-User-160k.The broader evaluation includes large-scale classification, semantic segmentation, and domain generalization on autonomous-driving data.

6 Conclusions

The paper connects poor federated generalization to sharp minima and introduces SAM, ASAM, and server-side SWA to encourage flatter solutions. These methods improve performance across vision tasks, with the strongest results when combined.

Conclusions: Heterogeneous federated training can converge toward sharp minima associated with degraded generalization and slower training.The analysis links poor generalization to the geometry of the loss surface and convergence toward sharp solutions.
Conclusions: SAM, ASAM, and SWA produce smoother loss surfaces and improve final performance across several vision tasks.Each modification is effective individually, while combining them yields the best performance.
Conclusions: SAM and ASAM modify local client optimization to seek parameters surrounded by neighborhoods of uniformly low training loss.SAM uses a neighborhood-based sharpness objective, while ASAM adapts sharpness to parameter scaling.
Conclusions: SWA averages high-performing stochastic weights on the server side, using learning-rate schedules to explore weight-space regions.In federated training, the server maintains a standard model and an SWA model, updating the latter periodically during late training.
Conclusions: Client-side data augmentation uses either Mixup or Cutout to improve learned-model generalization.Mixup interpolates images and labels, whereas Cutout masks randomly located square image regions.

B Training in Heterogeneous Scenarios - Additional Material

Additional analysis shows that statistical heterogeneity causes unstable, client-dependent model behavior, unlike the more uniform behavior in homogeneous training. Output-feature norms reveal this changing attention pattern over rounds.

Training in Heterogeneous Scenarios: Under statistical heterogeneity, model predictions fluctuate and generalization and performance deteriorate.The analysis describes oscillations and loss in performance when client distributions differ.
Training in Heterogeneous Scenarios: Higher output-feature norms indicate greater network attention to the corresponding class or client distribution.The feature norms are computed from the current global model on local client data.
Training in Heterogeneous Scenarios: For α = 0, the model’s attention shifts among client distributions and can focus on a single class at each round.This behavior varies over time without a constant trend, contrasting with the homogeneous setting.
Training in Heterogeneous Scenarios: For α = 1000, the model gives similar attention to each distribution, producing more uniform features.The homogeneous scenario shows substantially more consistent behavior across clients.

C.1 Datasets and Models

The experiments cover federated image classification, robustness, landmark recognition, and semantic segmentation under uniform and heterogeneous client partitions. They use benchmark-specific datasets, models, preprocessing, and tuned training settings.

Datasets and Models: CIFAR10 and CIFAR100 are split across 100 clients using a label-based Dirichlet process with concentration parameter α.Each client receives 500 images, and larger α produces more locally diverse and homogeneous class distributions.
Models and preprocessing: The CIFAR experiments use a LeNet5-like CNN, while CIFAR100-PAM uses a modified ResNet18 with GroupNorm.CIFAR inputs use standard cropping, flipping, and normalization; CIFAR100-PAM follows the CIFAR LDA preprocessing.
Datasets and Models: CIFAR100-PAM extends the setup to 500 clients with 100 images each using Pachinko Allocation followed by LDA.The partition reflects both coarse and fine label structure.
Datasets and Models: Landmarks-User-160k contains 164,172 training images across 2,028 landmarks and uses an ImageNet-pretrained MobileNetV2 with GroupNorm.The reported ImageNet-pretrained model reaches approximately 68% top-1 validation accuracy, with limited hyperparameter tuning noted.
Datasets and Models: Cityscapes and IDDA evaluate semantic segmentation with BiSeNetv2 under uniform or domain-heterogeneous federated splits.IDDA contains 105 domains spanning cities, viewpoints, and atmospheric conditions, including unseen countryside and rainy settings.
Hyper-parameters: The training configuration is dataset-specific, with CIFAR experiments selecting one local epoch and a client learning rate of 0.01.The server-side optimizer comparison identifies SGD with learning rate 1 as the best choice in the described setup.

D Results on Corrupted CIFAR10 and CIFAR100

On corrupted CIFAR10-C and CIFAR100-C, ASAM provides the strongest reported generalization, either alone or combined with SWA, as α varies.

Results on Corrupted CIFAR10 and CIFAR100: ASAM, alone or combined with SWA, achieves the best generalization across the corrupted CIFAR10-C and CIFAR100-C experiments.The comparison includes FedAvg, FedSAM, FedASAM, and their SWA variants across varying α.
Ablation: The section reports an ablation comparing client learning-rate cycling and server-side weight averaging.The ablation examines which component contributes more to stability and model performance.

E.1 Ablation Study on Server-Side Optimizers

Among tested server-side optimizers, FedAvg is the strongest baseline for acceptable performance, especially under heterogeneity. The SAM/ASAM hyperparameter study shows that smaller SAM neighborhoods are safer, while ASAM tolerates larger neighborhoods.

FedAvg is the best server-side choice for acceptable Cifar100 performance in both homogeneous and heterogeneous settings.The comparison used 5 clients per round and varying learning rates at 20k rounds.
High SAM ρ values rapidly reduce performance, indicating that SAM handles smaller neighborhoods better.This sensitivity pattern holds regardless of the data distribution in the Cifar100 study.
ASAM remains robust when expanding the neighborhood up to ρ = 0.5.The authors describe this as providing more freedom than SAM in the tested setting.
ASAM performance improves linearly as η increases, with η balancing stability and adaptivity.

E.3 Ablation Study on SWA

The SWA ablations separate the effects of cyclical learning rates and stochastic-weight averaging, while also showing that the timing of SWA activation affects its effectiveness. Additional comparisons cover optimizer settings and centralized versus federated gains.

SWA combines stochastic-weight averaging with a cyclical learning rate transmitted to clients during federated training.The learning rate decreases from γ1 to γ2 according to the cycle length c.
Server-side optimizer results are reported with varying learning rates on Cifar100 at 20k rounds and 5% client participation.
The ablation compares cyclical learning rates, constant learning rates, and whether server-side stochastic-weight averaging is applied.The comparison uses c > 1 for cyclical learning and c = 1 for a constant learning rate.
Longer pre-training before activating SWA leads to greater effectiveness on FedAvg with 5 clients per round.

F.2 Data Augmentations with CIFAR10

On CIFAR10, strong data augmentations do not improve generalization in the federated setting and can inhibit learning. The comparison includes FedAvg, SAM, ASAM, and SWA across augmentation conditions.

Mixup and Cutout do not improve federated generalization on CIFAR10 and can inhibit learning.The authors report that these augmentations sometimes produce worse results than FedAvg.
The broader comparison reports improvements in centralized and federated scenarios for α ∈{0.5, 1k} with 5 clients on Cifar100.
The CIFAR10 comparison evaluates FedAvg, SAM, ASAM, and SWA with strong Mixup and Cutout augmentations.

G Figures Omitted in the Main Text

The omitted figures examine convergence, loss surfaces, Hessian spectra, and client-wise curvature across heterogeneous federated settings. Together, they document optimizer stability, low-loss neighborhoods, and curvature measurements.

The omitted-figure section notes that plots are best viewed in color.
Convergence plots: With server-side momentum, SAM and ASAM reach convergence in the most heterogeneous Cifar10 and Cifar100 scenarios.FedAvgM uses server-side momentum β = 0.9.
Convergence plots: SWA improves stability and performance when compared with FedAvg, FedSAM, and FedASAM in the hardest setting of α = 0 with 5 clients per round.
Loss Surfaces: ASAM reaches the best local minimum among FedAvg, FedSAM, and FedASAM in heterogeneous Cifar100 scenarios.The comparison covers α ∈{0, 0.5}.
Hessian Eigenvalues: The Hessian analysis reports the top 50 global-model eigenvalues and complete client-wise eigenvalue comparisons over training rounds.The client-wise analysis introduces the values of λk.
Loss Surfaces: The loss-surface figures compare test error and training loss for local models and federated convergence points across heterogeneous settings.They include α = 0, α = 0.5, and α = 1000.