Sketched Adaptive Distributed Deep Learning: A Sharp Convergence Analysis

Combining gradient compression with adaptive optimizers is a highly desirable goal in distributed learning, with potential benefits in both fewer communication rounds and less per-round communication. In spite of preliminary empirical promise, certain major challenges in the convergence analysis of such methods have stayed open: handling compression based approximation of both first and second moments (pre-conditioner) which appear as a ratio; avoiding dependence on the number of parameters, which is extremely large in modern deep models; and providing high-probability guarantees instead of in-expectation, which can hide high variance behavior. In this work, we introduce a family of Sketched Adaptive Distributed Learning (SADL) algorithms which can use suitable unbiased gradient sketching for compression with suitable adaptive optimization algorithms. As our main contribution, we provide theoretical convergence guarantees of SADL algorithms which addresses all of the existing challenges. In particular, our guarantees hold with high probability, picks up only a logarithmic dependence on the number of parameters, and the first and second moment approximation is handled precisely yielding a dependence on the intrinsic dimension of the loss Hessian, which is significantly smaller than the full dimensionality of deep learning models. Empirically, the SADL algorithms are shown to be competitive with and often outperform baselines on both vision and language tasks, in both supervised fine-tuning and training-from-scratch regimes. Further, the SADL algorithms are also competitive with the state-of-the-art communication-efficient distributed learning algorithms based on error feedback. C C c=1 mc t ;

Update parameters and moments: x t , m t , v t , vt = ADA_OPT(x t-1 , m t-1 , v t-1 , vt-1 , mt ). end for do not involve computing the Hessian eigenspectrum, which is only used for the convergence analysis. Further, our results do not follow from the standard toolbox for sketching methods, such as Johnson-Lindenstrauss (Kpotufe and Sriperumbudur, 2020), which ensures nearby vectors stay close after sketching, but has no direct implications in an optimization setting, especially with adaptive optimizers which will sketch both the first and second moment, and work with their ratio. Based on such analysis, our work has the following main contributions:

(1) We introduce SADL algorithms which combine random sketching and adaptive methods. While the preconditoner in adaptive methods morphs the shape of sketching noise, posing challenges in leveraging the fast-decaying Hessian eigenstructure, we prove that the proposed sketching effectively balances iteration complexity and sketching dimension b. We derive a high probability bound showing that a sketch dimension of b = O(log d) suffices to achieve an O(1/ √ T ) convergence rate depending only on the intrinsic dimension in non-convex deep learning settings.

(2) Unlike existing works (Reddi et al., 2020;Xie et al., 2020), we provide a general convergence analysis without assuming uniformly bounded gradient norms on either the server or client side. We prove that SADL automatically generates bounded gradients along the entire optimization trajectory, where the involvement of mini-batch stochasticity and multiple local training steps calls for careful analysis on connecting the noisy local training steps with the global loss.

(3) We validate our theoretical claims with empirical evidence on deep learning models from vision (ResNet, Vision Transformer) and language (BERT) tasks. We cover both fine-tuning and trainingfrom-scratch regimes. SADL in general achieves comparable performance with adaptive methods without any gradient compression, using only 0.1% of the full dimensioanlity. Further, SADL is competitive with the state-of-the-art algorithms based on error feedback and adaptive methods.

A canonical distributed learning setting involves C clients, each associated with a local data distribution D c . The goal is to minimize the averaged empirical risk: L(x) = 1 C C c=1 E ξ∼Dc l(x, ξ), where l is the loss function, x ∈ R d is the parameter vector, and ξ is the data sample. We denote L c (x) = E ξ∼Dc l(x, ξ), c ∈ [C] as the client loss computed over the local distribution. We denote g c t,k as the mini-batch gradient over L c (x) at global step t and local step k. Algorithm 1 presents a generic framework of communication-efficient adaptive methods, which calls adaptive optimizers as subroutines. We denote T as the total training rounds. At each round, after K local SGD steps, client c sends to the server the sketched local model updates with a sketching operator sk: R d → R b . If b ≪ d without deteriorating the performance too much, the communication