PaperlayerView exact revision on arXiv

Source-linked AI summary

How Edge of Stability Hinders SCAFFOLD in Federated Optimization

Anant Khandelwal, Michael Crawshaw, Mingrui Liu

arXiv:2608.25873v1cs.LG

TL;DR

SCAFFOLD’s theoretical advantage under heterogeneous data does not reliably translate to deep learning, motivating an explanation for its lackluster practical performance. The paper empirically probes progressive sharpening and Edge of Stability in federated optimization and finds that high sharpness degrades SCAFFOLD’s global-gradient estimation. It concludes that SCAFFOLD can work well in situations where Edge of Stability does not occur, while its experiments remain scoped mainly to full-batch CNN image classification.

  • Problem

    SCAFFOLD has convergence guarantees independent of data heterogeneity, yet often does not outperform FedAvg in deep learning.

  • Method

    The paper performs an extensive empirical probe of progressive sharpening, Edge of Stability, equilibrium sharpness, and update misalignment in federated optimization.

  • Results

    Both FedAvg and SCAFFOLD exhibit progressive sharpening and Edge of Stability, while SCAFFOLD’s global-gradient estimation degrades at high sharpness.

  • Takeaways & Limitations

    SCAFFOLD’s underperformance is linked to optimization dynamics associated with Edge of Stability rather than deep-network nonconvexity alone.

  • Takeaways & Limitations

    Experiments mainly use full-batch gradients for CNN image classification, and the equilibrium sharpness lacks a closed-form characterization.

Abstract

from arXiv · show

In federated learning, it is well known that heterogeneous data can (in theory) slow down optimization, and much effort has been directed at designing optimization algorithms that are unaffected by data heterogeneity, such as the SCAFFOLD algorithm. Yet, despite strong theoretical guarantees, SCAFFOLD does not usually outperform the much simpler FedAvg in practice. In this work, we propose that this gap is due to the presence of Edge of Stability (EoS) and progressive sharpening in federated optimization, supported by extensive empirical probing. First, we find that EoS-like dynamics occur with both FedAvg and SCAFFOLD under a variety of architectures and hyperparameters. We observe that the equilibrium value of the sharpness is inversely proportional to the learning rate (as in GD), and interestingly, the degree of data heterogeneity (but not the number of local steps) also affects the equilibrium value. Most importantly, we observe that SCAFFOLD's ability to estimate the gradient of the global objective is severely degraded at the EoS, as measured by the correlation between sharpness and SCAFFOLD's error in estimating the global gradient along the optimization trajectory. This suggests a mechanism for SCAFFOLD's lackluster performance in deep learning: with high sharpness at the EoS, SCAFFOLD cannot reliably estimate the global gradient.

1 Introduction

The paper asks why FedAvg can match or outperform SCAFFOLD despite SCAFFOLD’s heterogeneity-independent theoretical rate, and investigates progressive sharpening and Edge of Stability as a possible explanation.

  • Motivation: SCAFFOLD was designed to eliminate convergence slowdowns from heterogeneous data, yet often does not outperform FedAvg in deep learning.Its control variates estimate the global gradient using gradients from previous communication rounds.
  • Hypothesis: The paper conjectures that progressive sharpening and Edge of Stability degrade SCAFFOLD’s performance under high sharpness.For gradient descent, sharpness increases until equilibrating around 2/η.
  • Approach: An extensive empirical probe examines Edge-of-Stability dynamics in federated optimization.The experiments vary architectures and optimizer hyperparameters and measure both global and local sharpness.
  • Findings: Both FedAvg and SCAFFOLD exhibit progressive sharpening and Edge of Stability across image-classification architectures and optimizer hyperparameters.The phenomena occur for both local and global objectives.
  • Findings: SCAFFOLD’s update misalignment strongly correlates with sharpness, whereas FedAvg’s update misalignment is relatively unaffected by high sharpness.This supports the proposed connection between Edge of Stability and SCAFFOLD’s lackluster performance.

2 Related Work

Prior work establishes federated learning, SCAFFOLD, and Edge of Stability, while documenting a gap between SCAFFOLD’s guarantees and its deep-learning performance.

  • Federated learning and SCAFFOLD: Federated learning trains models collaboratively across machines with differing datasets, and FedAvg convergence can slow as worker data becomes more heterogeneous.FedAvg is the de facto standard federated optimization algorithm.
  • Federated learning and SCAFFOLD: SCAFFOLD uses historical gradient corrections to estimate the global gradient and has convergence rates independent of data heterogeneity under smoothness assumptions.Its premise is that gradients from previous communication rounds represent current-round gradients.
  • Federated learning and SCAFFOLD: Empirical studies report that SCAFFOLD struggles with deep neural networks, especially under high heterogeneity or without gradient clipping.Some variants apply corrections only to the final layer and improve empirically over standard SCAFFOLD.
  • Related analyses: Other work studies momentum, alternative heterogeneity assumptions, convergence guarantees, and limiting distributions for SCAFFOLD and FedAvg.FedAvg’s practical success under heterogeneity has been described as its “unreasonable effectiveness.”
  • Edge of Stability: Edge of Stability describes sharpness increasing during neural-network training until it equilibrates around 2/η, with short-term loss oscillations and longer-term decrease.The phenomenon has been observed across architectures, hyperparameters, and datasets and extended to adaptive optimizers.
  • Scope: This paper studies the occurrence and consequences of Edge of Stability in federated optimization rather than explaining the phenomenon itself.The authors position this as the scope boundary of their work.

3 Preliminaries

The preliminaries define federated objectives, FedAvg, SCAFFOLD’s gradient corrections, and the smoothness assumptions that progressive sharpening and Edge of Stability challenge.

  • Federated objectives: Federated optimization minimizes a global objective formed by averaging client-specific local objectives, whose differences constitute data heterogeneity.Large heterogeneity can slow FedAvg optimization.
  • SCAFFOLD: SCAFFOLD uses control variates and corrected local update directions to approximate the global gradient instead of relying only on local gradients.The approximation depends on historical gradients representing current gradients.
  • Sharpness and EoS: Progressive sharpening is a gradual increase in Hessian norm, while Edge of Stability is equilibration around 2/η accompanied by non-monotonic loss behavior.This conflicts with SCAFFOLD’s assumption of a fixed smoothness bound.
  • SCAFFOLD: SCAFFOLD’s gradient approximation is justified by L-smoothness and requires η ≤ O(1/KL) to control the error across K local steps.Equivalently, the sharpness must remain below O(1/ηK).
  • Sharpness and EoS: The required sharpness bound becomes smaller as K grows, while the observed equilibrium sharpness does not significantly change with the communication interval.This may hinder SCAFFOLD’s ability to estimate the global gradient.
  • Measurement: Experiments measure global sharpness at the averaged model and average local sharpness during local steps for both FedAvg and SCAFFOLD.Both methods show progressive sharpening and Edge-of-Stability-like behavior across architectures and learning rates.
  • Measurement: When sharpness is high, SCAFFOLD is conjectured to lose accuracy in estimating the global gradient.This is the proposed link between Edge of Stability and SCAFFOLD’s underperformance.

4 Experiments

Experiments find EoS-like sharpening in federated optimization and show that high sharpness is associated with SCAFFOLD update misalignment and weaker comparative performance.

  • 4.1 Progressive Sharpening Occurs In Federated Settings with High Heterogeneity: Both FedAvg and SCAFFOLD exhibit progressive sharpening and EoS-like dynamics across architectures and optimizer hyperparameters.Sharpness is measured for both global and local objectives; SCAFFOLD often shows the more pronounced increase and equilibration.
  • 4.2 Equilibrium Sharpness: Doubling the learning rate from 0.005 to 0.01 cuts equilibrium sharpness in half, while heterogeneity also influences equilibrium sharpness less strongly.For FedAvg, equilibrium sharpness decreases from about 550 to about 350 as heterogeneity increases from h = 0 to h = 0.9; SCAFFOLD is relatively insensitive except at h = 0.9.
  • 4.2 Equilibrium Sharpness: Communication interval K has almost no effect on equilibrium sharpness, despite SCAFFOLD’s theoretical threshold scaling as O(1/(ηK)).The observed insensitivity to K suggests that large communication intervals can conflict with the sharpness condition used in SCAFFOLD’s convergence theory.
  • 4.3 High Sharpness Empirically Correlates with SCAFFOLD Misalignment: SCAFFOLD’s update misalignment is strongly and positively correlated with sharpness, with detrended correlations of at least 0.54, whereas FedAvg’s correlation is near zero.The association persists across learning rates and network depth or dataset size, and is especially pronounced near the EoS threshold.
  • 4.3 High Sharpness Empirically Correlates with SCAFFOLD Misalignment: When EoS occurs, FedAvg achieves slightly lower losses than SCAFFOLD on CIFAR-10, while SCAFFOLD achieves near-zero loss and outperforms FedAvg on easier MNIST settings without EoS.The results support a dataset-complexity connection: SCAFFOLD can outperform FedAvg when sharpness remains low and EoS does not occur.

5 Conclusion

The paper attributes SCAFFOLD’s underperformance in deep federated learning to progressive sharpening and Edge of Stability, which degrade its global-gradient estimates. The evidence is primarily from full-batch CNN image-classification settings, while broader extensions remain open.

  • 5 Conclusion: Progressive sharpening and Edge of Stability occur for both FedAvg and SCAFFOLD in federated deep-learning experiments.These dynamics appear across network architectures and optimizer hyperparameters.
  • 5 Conclusion: Equilibrium sharpness is inversely proportional to learning rate, depends less strongly on data heterogeneity, and has negligible dependence on local-step count.With many local steps, this prevents SCAFFOLD from meeting its theoretical sharpness requirement.
  • 5 Conclusion: SCAFFOLD’s update misalignment is strongly correlated with sharpness, and it fails to estimate the global gradient at the Edge of Stability.Update misalignment quantifies SCAFFOLD’s error in estimating the global gradient.
  • 5 Conclusion: SCAFFOLD can work well in deep learning in situations where Edge of Stability does not occur.This suggests the relevant distinction is optimization dynamics rather than deep-network nonconvexity alone.
  • 5 Conclusion: The study focuses mainly on full-batch FedAvg and SCAFFOLD for CNN image-classification tasks and does not derive closed-form equilibrium sharpness values.Stochastic gradients, Transformers, larger image datasets, and language tasks are identified as possible extensions.

A.1 Neural Network Architectures

The appendix specifies fixed CNN and MLP architectures for MNIST and CIFAR-10, using ReLU activations and ten-output linear classifiers. CNN depth and channel layouts vary by dataset and architecture family.

  • A.1 Neural Network Architectures: All networks output ten unnormalized logits without a final softmax or other activation.Convolutional and fully connected layers include biases unless stated otherwise.
  • A.1 Neural Network Architectures: ReLU is used throughout the CNNs and CIFAR-10 MLP, mapping negative inputs to zero and leaving positive inputs unchanged.It is piecewise linear and non-saturating for positive inputs.
  • A.1 Neural Network Architectures: CIFAR-10 CNNs use 3 × 3 stride-one padded convolutions, 32 channels after the first layer, ReLU after each convolution, and pooling after convolutional blocks.For L ≥ 3, the architecture repeats two-convolution blocks, with an extra pair when L is odd.
  • A.1 Neural Network Architectures: The listed L = 3 CIFAR-10 CNN maps 3 × 32 × 32 inputs through five convolutions and a 2048-input linear classifier to ten logits.This architecture matches the cited Cohen et al. (2021) model.
  • A.1 Neural Network Architectures: MNIST CNNs increase channels as 1 → 32 → 64 → 128, then preserve 128 channels, with 3 × 3 convolutions, ReLUs, and pooling after every second convolution.The listed MNIST CNN maps to a 6272-input classifier after two pooling operations.
  • A.1 Neural Network Architectures: The CIFAR-10 MLP uses 3072 inputs, one 200-unit ReLU hidden layer, and a ten-output linear classifier, while the MNIST MLP uses two 1000-unit hidden layers.The MNIST MLP maps 784 inputs through two ReLU layers to ten logits.

A.2 Global and Local Sharpness

Sharpness is defined through the largest Hessian eigenvalue, with separate global and local measurements. Global sharpness evaluates the averaged model on the combined objective, whereas local sharpness tracks client models during local training.

  • A.2 Global and Local Sharpness: Sharpness is the largest eigenvalue of an objective’s Hessian evaluated at model parameters w.This quantity measures the objective’s local curvature.
  • A.2 Global and Local Sharpness: Global sharpness evaluates the averaged model at the end of round t on the global objective.It measures curvature of the combined training loss at the model produced by federated averaging.
  • A.2 Global and Local Sharpness: Local sharpness averages participating clients’ local-objective sharpness at their respective local models and local step k of round t.It follows curvature as client models move away from the shared model during local training.
  • A.2 Global and Local Sharpness: The leading Hessian eigenvalue is estimated with Lanczos iterations using Hessian–vector products from automatic differentiation.This implementation applies to the sharpness measurements.

A.3 Heterogeneity Protocol

The experiments simulate client heterogeneity with a label-skew protocol controlled by h, while comparing FedAvg and SCAFFOLD under federated local-training procedures. Each client receives both shared and label-skewed data shards.

  • A.3 Heterogeneity Protocol: The heterogeneity parameter h ranges from 0 for fully i.i.d. data to 1 for maximal data heterogeneity.Larger h induces more non-i.i.d. label structure.
  • A.3 Heterogeneity Protocol: The protocol splits each class between a shuffled global i.i.d. pool and an ordered non-i.i.d. pool whose contiguous chunks are label-concentrated.The fraction assigned to the non-i.i.d. pool is controlled by h.
  • A.3 Heterogeneity Protocol: FedAvg iterates client local training for K steps, sends model updates to the server, and averages those updates into the next global model.The supplied algorithm description identifies this as federated averaging.
  • A.3 Heterogeneity Protocol: SCAFFOLD is included as the controlled-averaging comparison with a server control variate and client state.The algorithm specification names its initial global model, server control variate, and client state.
  • A.3 Heterogeneity Protocol: Both pools are partitioned evenly across clients, giving every client the same shared and label-skewed proportions while h controls skew severity.At h = 0, all client datasets have i.i.d. labels; increasing h specializes clients to different class subsets.

A.4 Algorithm Pseudocode

FedAvg performs distributed gradient descent with periodic model averaging, while SCAFFOLD adds control variates to reduce data heterogeneity and steer updates toward global gradients.

  • FedAvg runs gradient descent across multiple datasets and periodically averages models at a specified communication interval.
  • SCAFFOLD reduces heterogeneity by subtracting previous local gradients and adding previous global gradients to adjust model updates.

B Additional Experimental Results

Additional experiments show that sharpness dynamics and their relationship with SCAFFOLD updates depend on dataset simplicity and algorithm. SCAFFOLD exhibits strong sharpness–misalignment alignment, while FedAvg generally converges faster and more stably.

  • FashionMNIST: On FashionMNIST, both algorithms show a modest sharpness increase followed by rapid decrease as models converge and overfit, although some runs remain at a threshold or never sharpen.Sharpness remains strongly inversely correlated with learning rate.
  • MNIST: On MNIST, sharpness rises modestly and then declines as architectures nearly perfectly fit the data.The authors link these lower-sharpness trajectories to stronger relative SCAFFOLD performance and argue that sharpness, rather than convexity, explains failures.
  • Update misalignment: Across SCAFFOLD trajectories, normalized local sharpness and update misalignment consistently correlate above 0.5 across model depths and dataset sizes.Their oscillations also qualitatively match in the reported trajectories.
  • Update misalignment: FedAvg shows little or negative correlation between update misalignment and sharpness, with opposing growth directions and no qualitative similarity in oscillations.The reported sharpness decreases during training because initialization sharpness exceeds 500.
  • Loss trajectories: FedAvg consistently reaches lower loss faster than SCAFFOLD across CNN depths and dataset sizes, with less instability during training.The authors characterize FedAvg as easier and more consistent to train and suggest high sharpness makes SCAFFOLD updates less accurate and stable.
Loading 2608.25873v1…