SA-PEF: Step-Ahead Partial Error Feedback for Efficient Federated Learning

TL;DR

SA-PEF introduces a novel step-ahead partial error feedback method for federated learning, improving convergence speed and stability under non-IID data and partial client participation.

Contribution

It combines step-ahead correction with partial error feedback, providing theoretical convergence guarantees and practical acceleration over traditional error feedback methods.

Findings

SA-PEF achieves faster convergence to target accuracy.

Theoretical analysis confirms convergence under non-convex objectives.

Experimental results show consistent performance improvements.

Abstract

Biased gradient compression with error feedback (EF) reduces communication in federated learning (FL), but under non-IID data, the residual error can decay slowly, causing gradient mismatch and stalled progress in the early rounds. We propose step-ahead partial error feedback (SA-PEF), which integrates step-ahead (SA) correction with partial error feedback (PEF). SA-PEF recovers EF when the step-ahead coefficient $\alpha=0$ and step-ahead EF (SAEF) when $\alpha=1$. For non-convex objectives and $\delta$-contractive compressors, we establish a second-moment bound and a residual recursion that guarantee convergence to stationarity under heterogeneous data and partial client participation. The resulting rates match standard non-convex Fed-SGD guarantees up to constant factors, achieving $O((\eta,\eta_0TR)^{-1})$ convergence to a variance/heterogeneity floor with a fixed inner step size. Our analysis reveals a step-ahead-controlled residual contraction $\rho_r$ that explains the observed acceleration in the early training phase. To balance SAEF's rapid warm-up with EF's long-term stability, we select $\alpha$ near its theory-predicted optimum. Experiments across diverse architectures and datasets show that SA-PEF consistently reaches target accuracy faster than EF.