Distributed Learning over Arbitrary Topology: Linear Speed-Up with Polynomial Transient Time

Abstract

We study a distributed learning problem in which $n$ agents, each with potentially heterogeneous local data, collaboratively minimize the sum of their local cost functions via peer-to-peer communication. We propose a novel algorithm, \emph{Spanning Tree Push-Pull} (STPP), which employs two spanning trees extracted from a general communication graph to distribute both model parameters and stochastic gradients. Unlike prior approaches that rely heavily on spectral gap properties, STPP leverages a more flexible topological characterization, enabling robust information flow and efficient updates. Theoretically, we prove that STPP achieves linear speedup and polynomial transient iteration complexity -- up to $\mathcal{O}(n^7)$ for smooth nonconvex objectives and $\tilde{\mathcal{O}}(n^3)$ for smooth strongly convex objectives -- under arbitrary network topologies. Moreover, compared with existing methods, STPP achieves faster convergence rates on sparse and non-regular topologies (e.g., directed rings) and reduces communication overhead on dense networks (e.g., static exponential graphs). Numerical experiments further demonstrate the strong performance of STPP across various graph architectures.