Beyond Scaffold: A Unified Spatio-Temporal Gradient Tracking Method

TL;DR

This paper introduces ST-GT, a unified spatio-temporal gradient tracking method for distributed optimization that reduces communication costs and handles data heterogeneity effectively, achieving linear convergence in certain settings.

Contribution

ST-GT unifies gradient tracking with a novel spatio-temporal approach, improving convergence and communication efficiency in federated learning over time-varying graphs.

Findings

Achieves linear convergence for strongly convex problems.

Attains the first linear speed-up in communication complexity with respect to local updates.

Reduces the topology-dependent noise term from σ² to σ²/τ.

Abstract

In distributed and federated learning algorithms, communication overhead is often reduced by performing multiple local updates between communication rounds. However, due to data heterogeneity across nodes and the local gradient noise within each node, this strategy can lead to the drift of local models away from the global optimum. To address this issue, we revisit the well-known federated learning method Scaffold (Karimireddy et al., 2020) under a gradient tracking perspective, and propose a unified spatio-temporal gradient tracking algorithm, termed ST-GT, for distributed stochastic optimization over time-varying graphs. ST-GT tracks the global gradient across neighboring nodes to mitigate data heterogeneity, while maintaining a running average of local gradients to substantially suppress noise, with slightly more storage overhead. Without assuming bounded data heterogeneity, we prove that ST-GT attains a linear convergence rate for strongly convex problems and a sublinear rate for nonconvex cases. Notably, ST-GT achieves the first linear speed-up in communication complexity with respect to the number of local updates per round $\tau$ for the strongly-convex setting. Compared to traditional gradient tracking methods, ST-GT reduces the topology-dependent noise term from $\sigma^2$ to $\sigma^2/\tau$, where $\sigma^2$ denotes the noise level, thereby improving communication efficiency.