Sketched Adaptive Federated Deep Learning: A Sharp Convergence Analysis

## Abstract

Abstract Combining gradient compression methods (e.g., CountSketch (Charikar et al.,, 2002 ; Rothchild et al.,, 2020 ) , quantization (Tang et al.,, 2021 ) ) and adaptive optimizers (e.g., Adam (Kingma and Ba,, 2014 ) , AMSGrad (Reddi et al.,, 2019 ) ) is a desirable goal in federated learning (FL), with potential benefits on both fewer communication rounds and less per-round communication. In spite of the preliminary empirical success of sketched adaptive methods, existing convergence analyses show the communication cost to have a linear dependence on the ambient dimension (Spring et al.,, 2019 ; Tang et al.,, 2021 ) , i.e., number of parameters, which is prohibitively high for modern deep learning models. In this work, we introduce specific sketched adaptive federated learning (SAFL) algorithms and, as our main contribution, provide theoretical convergence analyses in different FL settings with guarantees on communication cost depending only logarithmically (instead of linearly) on the ambient dimension. Unlike existing analyses, we show that the entry-wise sketching noise existent in the preconditioners and the first moments of SAFL can be implicitly addressed by leveraging the recently-popularized anisotropic curvatures in deep learning losses, e.g., fast decaying loss Hessian eigen-values.
In the i.i.d. client setting of FL, we show that SAFL achieves asymptotic O ⁢ ( 1 / T ) 𝑂 1 𝑇 O(1/\sqrt{T}) italic_O ( 1 / square-root start_ARG italic_T end_ARG ) convergence, and converges faster in the initial epochs. In the non-i.i.d. client setting, where non-adaptive methods lack convergence guarantees, we show that SACFL (SAFL with clipping) algorithms can provably converge in spite of the additional heavy-tailed noise. Our theoretical claims are supported by empirical studies on vision and language tasks, and in both fine-tuning and training-from-scratch regimes. Surprisingly, as a by-product of our analysis, the proposed SAFL methods are competitive with the state-of-the-art communication-efficient federated learning algorithms based on error feedback.

## 1 Introduction

Despite the recent success of federated learning (FL), the cost of communication arguably remains the main challenge. Wang et al., () showed that a 20 Gbps network bandwidth is necessary to bring the communication overhead to a suitable scale for finetuning GPT-J-6B, which is unrealistic in distributed settings. Even with good network conditions, reduction on the communication complexity means one can train much larger models given the same communication budget.

The communication cost of FL can be represented as $O(dT)$, where $d$ is the ambient dimension of the parameter space and $T$ is the number of rounds for convergence. Various methods have been proposed to minimize $T$, e.g., local training (Stich,, ), large batch training (Xu et al.,, ).
Folklores in centralized training regimes suggest that $T$ heavily relies on the choice of optimizers, where adaptive methods usually demonstrate faster convergence and better generalization performance, especially in transformer-based machine learning models (Reddi et al.,, ). In decentralized settings, adaptive methods are also favorable due to their robustness to data heterogeneity, e.g., adaptive methods are guaranteed to converge under heavy-tailed noise while SGD does not (Zhang et al.,, ). These favorable merits, in principle, should be preserved in communication-efficient FL algorithms.

The alternative approach of reducing communication costs is to be more thrifty on the communication bits at a single round, i.e., to reduce the $O(d)$ factor, which is dominant in the communication complexity for modern neural networks where $d\gg T$. Considerable efforts have been devoted to design efficient gradient compression methods. Popular gradient compression methods include quantization (Alistarh et al.,, ; Chen et al.,, ; Reisizadeh et al.,, ; ), sparsification  (Alistarh et al.,, ; Wu et al.,, ; Rothchild et al.,, ) and sketching (Spring et al.,, ; Jiang et al.,, ; Song et al.,, ). However, most of these developments do not use adaptive methods, which involve anisotropic and nonlinear updates, and there are no easy ways to do error feedback (Tang et al.,, ) in case of discrepancies. Indeed, the design and analysis of communication-efficient adaptive FL algorithms poses non-trivial challenges.

In this work, we first introduce a family of Sketched Adaptive FL (SAFL) algorithms, with flexibility on the choice of sketching methods and adaptive optimizers, that simultaneously guarantees convergence and reduces per round bits towards improved communication efficiency. At a high level, SAFL algorithms are in principle analogous to previous attempts (Tang et al.,, ; Chen et al.,, ; Wang et al.,, ), which showed preliminary empirical success of applying gradient compression with adaptive optimizers in homogeneous data scenarios. Our SAFL algorithm adopts unbiased gradient estimators and hence eliminates the needs for error feedbacks. The linearity of gradient compressions in SAFL also avoids the extra round of compression on the server side required in sparsification (Stich et al.,, ) and quantization (Reisizadeh et al.,, ).

Despite the preliminary empirical success of combining gradient sketching and adaptive optimizers for federated deep learning, theoretical understanding of such algorithms is limited. Existing works on the theory has arguably alarming results, which do not match practice. For example, the analysis over both adaptive and non-adaptive optimizers show that the iterations $T$ needed for convergence can be inversely proportional to the compression rate (Chen et al.,, ; Song et al.,, ). For constant per-round communication bits, the bounds indicate the iteration complexity to scale as $O(d)$, i.e., linearly with the ambient dimensionality, which is prohibitively large for modern deep learning models. The mismatch between potential issues in theory vs. preliminary empirical promise and possibly also not having precise forms of such algorithms for different FL scenarios may be preventing wide adoption of such algorithms.

As a major contribution of our current work, we provide theoretical guarantees of the proposed SAFL algorithms on the convergence rate in common FL scenarios (almost i.i.d. as well as heavy tailed), which depends only on a logarithmically (instead of linearly) on the ambient dimension $d$. The central technical challenge in addressing the dimensional dependence is to handle the entry-wise sketching noise in both the preconditioners and the first moments of the adaptive optimizers, which has been acknowledged non-trivial (Tang et al.,, ; Wang et al.,, ).
Our sharper analysis leverages recent observations regarding the eigenspectrum structure of the loss Hessian in deep learning, which show the eigenvalues to be sharply decaying, with most eigenvalues being close to zero (Ghorbani et al.,, ; Zhang et al.,, ; Li et al.,, ; Liao and Mahoney,, ; ), and even conforming with a power-law decay (Xie et al.,, ; Zhang et al.,, ), while the conventional smoothness conditions assume uniform curvature in all directions which can be overly pessimistic in the context of deep learning. This specific eigenspectrum structure provides significant advantages in the sharp analysis of sketching noise in adaptive methods. Our work leverages such geometric structure, leading to the following main contributions:

(1) For the benign almost i.i.d. FL setting, we introduce the sketched adaptive FL (SAFL) framework which incorporates random sketching techniques into adaptive methods. While the preconditoner in adaptive methods morphs the shape of sketching noise, which poses challenges in leveraging the anisotropic Hessian structure, we prove that the proposed sketching method effectively balances iteration complexity and sketching dimension $b$.
We derive a high probability bound showing that a sketch size of $b=O(\log d)$ suffices to achieve an asymptotic $O(1/\sqrt{T})$ dimension-independent convergence rate in non-convex deep learning settings.

(2) For the heavy-tailed noise common in data-heterogeneous FL, where non-adaptive methods are not guaranteed to converge, we propose the Sketched Adaptive Clipped Federated Learning (SACFL) which guarantees the boundedness of the second moments. We theoretically show that SACFL can achieve optimal convergence rate under $\alpha$-moment noise with $\alpha\in(1,2]$, regardless of the extra noise introduced by random sketching.

(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 training-from-scratch regimes. The proposed SAFL algorithm achieves comparable performance with the full-dimensional unsketched adaptive optimizers, and are competitive with the state-of-the-art communication-efficient federated learning algorithms based on error feedback. We also validate the SACFL algorithms can achieve similar performance as the unsketched clipping algorithm when the local client gradient norms are $\alpha$-stable heavy-tailed.

## 2 Related Works

Communication-efficient federated optimization. There have been rapid advances in communication-efficient federated optimization in recent years. Local training, i.e. running SGD independently in parallel on different clients, is the off-the-shelf training mechanism which ideally reduces the frequency of communication. Stich, () shows local SGD achieves the same convergence rate as mini-batch SGD. Wang and Joshi, () study the effect of the frequency in model averaging, and propose adaptive communication strategies. Mishchenko et al., () prove the local gradient step can surprisingly accelerate the training process, which offers non-trivial advantages over SGD.

Besides the advances in efficient training mechanisms, applying gradient compression is another promising research thread in communication-efficient learning. In principle, gradient compression methods reduce the communication bits per round with negligible increase of overhead in the convergence rate. Various gradient compression methods have been proposed and exhibited preliminary improvement in practice. Quantization is one of the popular compression schemes which adopts lower bits to represent data originally represented by 32 bits on each dimension. Tang et al., () propose 1-bit Adam based on the stability of Adam’s variance term during the training time.
improve 1bit-Adam using large batch training and adaptive layerwise learning rates. Tang et al., () propose a sign-based unbiased quantization method that controls the bias of signSGD (Bernstein et al.,, ) by injecting random noise prior to the compression.

Another popular gradient compression method is sparsification, where the transmitted bits solely come from the most significant values in the model update. The communication cost is proportional to the number of non-zero elements in the sparsified gradient. Deterministic sparsification methods are simpler in practice, e.g., Random-k (Wangni et al.,, ), Top-k (Stich et al.,, ; Shi et al.,, ; ; Xu et al.,, ), deep gradient compression (Lin et al.,, ). However, the consequential biased gradient estimation reportedly hurts training performance and leads to worse generalization  (Beznosikov et al.,, ). Error compensation techniques (Zheng et al.,, ; Richtárik et al.,, ) are necessary to mitigate the effect. Rothchild et al., () first propose to apply error compensation on the server side to support sparse client participation. Wang et al., () apply targeted error compensation to specific components of the sparsified model updates.

Distinguished from all the methods above, sketching has gained increasing popularity because of several favorable properties including mergibility and unbiasedness. Sketching methods project the entire gradient vector into a tiny subspace. Our method also falls into this category and is well-compatible with adaptive optimizers and associated with a sharper convergence analysis. Ivkin et al., () utilize Count-Sketch (Charikar et al.,, ) to estimate the heavy-hitters of a gradient vector. Vargaftik et al., () estimate the coordinates of a gradient with structured random rotations in a high-dimensional sphere. Rabbani et al., () apply sketching to model weights to improve downlink communication efficiency.

Theoretical analysis on communication-efficient federated learning is also a central topic in this thread. Ivkin et al., () develop the convergence guarantees for Count-Sketch in a strongly-convex setting. Chen et al., (, ) conduct a convergence analysis for quantized Adam with error compensation. Haddadpour et al., () provide a unified convergence analysis on periodical compressed communication mechanism based on quantization and sparsification. Wang et al., () study the convergence properties of communication-efficient adaptive gradient methods under biased compressors. Song et al., () provide the first convergence result of random sketching in the non-convex setting, but the upper bound comes with a dimension dependence.

Noise in Deep Learning. In our work, we deal with noises from various sources. There have been numerous literatures discussing the noise in neural network training. However, high-probability bounds are indeed quite limited, as the mainstream of analysis of the optimization methods are conducted based on expectation. The analysis over common noise assumption, e.g. sub-Gaussian and sub-exponential is proposed by Rakhlin et al., () in the strongly-convex settings, which is subsequently improved by Harvey et al., (). Li and Orabona, () prove the high probability convergence rate
for a weighted average of the squared gradient norms of SGD assuming strong smoothness and sub-Gaussian noise. Madden et al., () prove a high probability bound under sub-Weibull noise, which generalizes sub-Gaussian and sub-exponential properties to heavier tailed distributions (Vladimirova et al.,, ).

More recently, the community finds the heavy-tailed phenomenon are prevalent in common machine learning tasks (Simsekli et al.,, ; Reddi et al.,, ). It is also observed in federated learning settings when the data are heterogeneous across clients (Yang et al.,, ). Under the heavy-tailed noise assumptions, Gorbunov et al., () prove the first high-probability convergence results
for Clip-SGD in the convex case, and is later generalized to Holder-continuous gradients in Gorbunov et al., (). In the case of non-convex problems, Cutkosky and Mehta, () provide a convergence bound for normalized clip-SGD. Subsequent works including Sadiev et al., () improve the bounds without bounded gradient assumptions.

Adaptive Learning Rates. Adaptive learning rates are the key ingredients in deep learning optimization. Adagrad is first proposed in  Duchi et al., () in aim of utilizing sparsity in stochastic gradients. Subsequent works, e.g. Adam (Kingma and Ba,, ) and AMSGrad (Reddi et al.,, ) have become the mainstream optimizers used in machine learning because of their superior empirical performance. These methods use implicit learning rates adaptive to the current iterate in the training process. In many cases, adaptive methods have been shown to converge faster than SGD, and with better generalization as well (Reddi et al.,, ). In recent literatures, adaptive methods are shown to be capable of better dealing with the noise, which partially accounts for their empirical success. Zhang et al., () show empirical connections between the noise in the gradients and Adam’s performance. On the other hand, Chezhegov et al., () demonstrate that AdaGrad and its delayed version can fail to converge in polynomial time under heavy-tailed noise, while adaptive clipping-based methods can cope with the noise with theoretical guarantees (Zhang et al.,, ). A combination of clipping-based methods and Adagrad is devised to achieve convergence under heavy-tailed noise in Chezhegov et al., ().

## 3 Sketched Adaptive FL under Mild Noise

In this section, we consider federated learning on nearly-i.i.d client data distribution. The objective is to develop communication-efficient adaptive learning algorithms. We will first propose the general framework for applying gradient compression to FedOPT (Reddi et al.,, ), and proceed with the mild-noise assumptions and convergence analysis of the algorithm.

### 3.1 Sketched Adaptive FL (SAFL)

A canonical federated learning setting involves $C$ clients, each associated with a local data distribution ${\cal D}_{c}$. The goal is to minimize the averaged empirical risks: ${\cal L}(x)=\frac{1}{C}\sum_{c=1}^{C}\mathbb{E}_{\xi\sim{\cal D}_{c}}l(x,\xi)$,
where $l$ is the loss function, $x\in\mathbb{R}^{d}$ is the parameter vector, and $\xi$ is the data sample. We denote ${\cal L}^{c}(x)=\mathbb{E}_{\xi\sim{\cal D}_{c}}l(x,\xi),c\in[C]$ as the client loss function computed over the local data distribution. We denote $g_{t,k}^{c}$ as the mini-batch gradient over ${\cal L}^{c}(x)$ at global step $t$ and local step $k$.

Figure: Algorithm 1 Sketched Adaptive Federated Learning (SAFL)

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 training steps, client $c$ sends to the server the sketched local model updates with a sketching operator sk$:\mathbb{R}^{d}\rightarrow\mathbb{R}^{b}$.
If $b\ll d$ without deteriorating the performance too much, the communication cost per round can be reduced from $O(d)$ to $O(b)$. Algorithm 2 projects the compressed updates and second moments back to the ambient dimension using a desketching operator desk$:\mathbb{R}^{b}\rightarrow\mathbb{R}^{d}$ and implements a single-step adaptive optimization. The server and clients call Algorithm 2 every epoch, i.e. communication round, to update the global model and synchronize local models.
The gradient compression steps differentiate Algorithm 1 from the subspace training methods (Gressmann et al.,, ; Wortsman et al.,, ) since we are utilizing the global gradient vector in each round rather than solely optimizing over the manifold predefined by a limited pool of parameters. The choice of server-side optimizers determines how the lossy replicates in $\mathbb{R}^{d}$ are used to update the running moments (i.e. the momentum and the second moments). The server sends the moments in $\mathbb{R}^{b}$ back to the clients so that each client can perform an identical update on its local model, which ensures synchronization as each training round starts.

###### Remark 3.1 .

(Sketching Randomness).
At each single round, the sketching operators sk’s are shared among clients, via the same random seed, which is essential for projecting the local model updates to a shared low dimensional subspace and making direct averaging reasonable. On the other hand, we use fresh sk’s at different rounds so that the model updates lie in distinct subspaces. ∎

Figure: Algorithm 2 ADA_OPT (AMSGrad)

###### Remark 3.1 .

### 3.2 Random Sketching

We will first introduce the desired characteristics of compression and then list a family of sketching algorithms which possess those properties.

###### Property 1 .

(Linearity). The compression operators are linear w.r.t the input vectors, i.e. $\texttt{sk}(\sum_{i=1}^{n}{v_{i}})=\sum_{i=1}^{n}{\texttt{sk}(v_{i})}$ and $\texttt{desk}(\sum_{i=1}^{n}{\bar{v}_{i}})=\sum_{i=1}^{n}{\texttt{desk}(\bar{v
}_{i})},~{}~{}\forall\{v_{i},\bar{v}_{i}\in\mathbb{R}^{d}\}_{i=1}^{n}$.

###### Property 2 .

(Unbiased Estimation).
For any vector $v\in\mathbb{R}^{d}$, $\mathbb{E}[\texttt{desk}(\texttt{sk}(v))]=v$.

###### Property 3 .

(Bounded Vector Products).
For any fixed vector $v,h\in\mathbb{R}^{d}$, $\mathbb{P}(|\langle\texttt{desk}(\texttt{sk}(v)),h\rangle-\langle v,h\rangle|
\geq(\frac{\log^{1.5}(d/\delta)}{\sqrt{b}})\|v\|\|h\|)\leq\Theta(\delta)$.

Property 1 and 2 guarantee the average of first moments in Algorithm 1 over clients are, in expectation, the same as those in FedOPT. Property 3 quantifies the bound on the deviation of vector products when applying compression. $\texttt{sk}(v)=Rv$ and $\texttt{desk}(\bar{v})=R^{\top}\bar{v}$, where $R\in\mathbb{R}^{b\times d}$ is a random sketching operator, satisfy all the properties above (Song et al.,, ). We denote $R_{t}$ as the sketching operator used in round $t$. Different instantiations of $R$ constitute a rich family of sketching operators, including i.i.d. isotropic Gaussian (Song et al.,, ), Subsampled Randomized Hadamard Transform (SRHT) (Lu et al.,, ), and Count-Sketch (Charikar et al.,, ), among others. The specific error bounds for these special cases can be found in Appendices A.1, A.2, and A.3 respectively.

###### Property 1 .

###### Property 2 .

###### Property 3 .

### 3.3 Convergence Analysis

We first state a set of standard assumptions commonly used in the literature of first-order stochastic methods. We focus on the mild noise assumptions, which are typically observed when the training data are nearly i.i.d. over clients. We will use $\|\cdot\|$ to denote $L_{2}$-norm throughout the work.

###### Assumption 1 .

(Bounded Global Gradients). Square norm of the gradient is uniformly bounded, i.e., $\|\nabla{\cal L}(x)\|^{2}\leq G_{g}^{2}$.

###### Assumption 2 .

(Bounded Client Gradients).
For every client, there exists a constant $G_{c}\geq 0$, such that $\|\nabla{\cal L}^{c}(x)\|^{2}\leq G_{c}^{2},~{}c\in[C]$.

For simplicity, we define $G:=\max\{\max\{G_{c}\}_{c=1}^{C},G_{g}\}$ to denote the upper bound for client and global gradient norms. We also assume the stochastic noise from minibatches is sub-Gaussian, which is widely adopted in first-order optimization (Harvey et al.,, ; Mou et al.,, ).

###### Assumption 3 .

(Sub-Gaussian Noise).
The stochastic noise $\|\nabla{\cal L}^{c}(x)-g^{c}(x)\|$ at each client is a $\sigma$-sub-Gaussian random variable, i.e. $\mathbb{P}(\|\nabla{\cal L}^{c}(x)-g^{c}(x)\|\geq t)\leq 2\exp(-t^{2}/\sigma^{
2})$, for all $t\geq 0$.

Besides, we have assumptions on the Hessian eigenspectrum $\{\lambda_{i},v_{i}\}_{i=1}^{d}$ of the loss function ${\cal L}$.

###### Assumption 4 .

(Hessian Matrix Eigenspectrum)
The smoothness of the loss function ${\cal L}$, i.e. the largest eigenvalue of the loss Hessian $H_{\cal L}$ is bounded by $L$, $\max_{i}\lambda_{i}\leq L$. The sum of absolute values of $H_{\cal L}$ eigenspectrum is bounded by $D$, i.e. $\sum_{i=1}^{d}|\lambda_{i}|\leq D$.

The assumption of bounded sum of eigenspectrum has been validated by several recent literatures, in the context of deep learning, where the eigenspectrum is observed to sharply decay (Ghorbani et al.,, ; Li et al.,, ; ), have bulk parts concentrate at zero (Sagun et al.,, ; Liao and Mahoney,, ) or conform with a power-law distribution (Xie et al.,, ; Zhang et al.,, ). We quote their plots in Appendix D for completeness. Our empirical verification under the setting of FL can also be found in Fig. 6 in Appendix D.

###### Remark 3.2 .

(Three types of noises in Algorithm 1).
One of the key technical contributions of this work is to theoretically balance the noises of different sources and derive a reasonable convergence rate which is independent of the ambient dimension. The noise in the training process stems from the mini-batch training, the client data distribution, and the compression error due to sketching. The stochastic error of mini-batch training is sub-Gaussian by Assumption 3. We will adopt $\delta_{g}$ to control the scale of the sub-Gaussian noise. The i.i.d. data distribution leads to the bounded gradient assumption (Assumption 2).
The sketching error depends on the specific choice of sketching methods, but is always controlled by the bounded property on vector products (Property 3) with a universal notation $\delta$. All three kinds of noises are unbiased and additive to the gradient, though may have sequential dependencies. Therefore, for the analysis (Appendix B), we will introduce a martingale defined over the aggregated noise, using which we can derive a high-probability concentration bound for the variance. We denote $\nu$ as the scale of the $\psi_{2}$-norm (Vershynin,, ) in the martingale. ∎

Now we characterize the convergence of Algorithm 1 in Theorem 3.1. All technical proofs for this section are in Appendix B and we provide an outline of the proof techniques in Section 3.4.

###### Theorem 3.1 .

Suppose the sequence of iterates $\{x_{t}\}_{t=1}^{T}$ is generated by Algorithm 1 (SAFL) with a constant learning rate $\eta_{t}\equiv\eta$. Under Assumptions 1-4, for any $T$ and $\epsilon>0$, with probability $1-\Theta(\delta)-O(\exp(-\Omega(\nu^{2})))-\delta_{g}$,

$$ $\displaystyle\left(\sqrt{1+\frac{\log^{1.5}(CKd^{2}T^{2}/\delta)}{\sqrt{b}}} \eta KG+\epsilon\right)^{-1}\kappa\eta K\sum_{t=1}^{T}\|\nabla{\cal L}(x_{t}) \|^{2}\leq{\cal L}(z_{1})+\frac{1}{\epsilon}\kappa\eta^{2}LK^{2}G^{2}T$ $\displaystyle+\nu\kappa\eta K\sqrt{T}(\frac{\log^{1.5}(CKTd/\delta)}{\sqrt{b}} \frac{G^{2}}{\epsilon}+\frac{\sigma}{\epsilon}\log^{\frac{1}{2}}(\frac{2T}{ \delta_{g}}))+\eta^{2}\kappa^{2}T(1+\frac{\log^{1.5}(CKdT^{2}/\delta)}{\sqrt{b }})^{2}\frac{8}{(1-\beta_{1})^{2}}\frac{DK^{2}G^{2}}{\epsilon^{2}},$ $$

where $\delta,\delta_{g}$, and $\nu$ are the randomness of sketching, sub-Gaussian noise, and martingales respectively.

A non-asymptotic convergence bound of training with practical decaying learning rates can be found in Theorem B.2 in appendix. Given that we only introduce logarithmic factors on $d$ in the iteration complexity and the per-round communication $b$ is a constant, the total communication bits of training a deep model till convergence is also logarithmic w.r.t $d$. To better understand Theorem 3.1, we can investigate different regimes based on the training stages. For the asymptotic regime, where $T$ is sufficiently large, we can achieve an $O(1/\sqrt{T})$ convergence rate in Corollary 1.

###### Corollary 1 .

With the same condition as in Thereom 3.1,
for sufficiently large $T\geq\frac{G^{2}}{\epsilon^{2}}$, with probability $1-\Theta(\delta)-O(\exp(-\Omega(\nu^{2})))-\delta_{g}$,

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\|\nabla{\cal L}(x_{t})\|^{2}$ $\displaystyle\leq\frac{2{\cal L}(z_{1})\epsilon}{\kappa\sqrt{T}}+\frac{2}{ \epsilon}\frac{LG^{2}}{\sqrt{T}}+\nu\frac{2}{\sqrt{T}}(\frac{\log^{1.5}(CKTd/ \delta)}{\sqrt{b}}G^{2}+\sigma\log^{\frac{1}{2}}(\frac{2T}{\delta_{g}}))$ $\displaystyle+\kappa\frac{1}{\sqrt{T}}(1+\frac{\log^{1.5}(CKdT^{2}/\delta)}{ \sqrt{b}})^{2}\frac{16}{(1-\beta_{1})^{2}}\frac{DG^{2}}{\epsilon},$ $$

where $\delta,\delta_{g}$ and $\nu$, are the randomness of sketching, sub-Gaussian noise and martingales respectively.

More interestingly, for the near-initialization regime, where $T$ is relatively small, we can observe that the coefficient of $\|\nabla{\cal L}(x_{t})\|^{2}$ on the left hand side in Theorem 3.1 and B.2 is approximately a constant, given that $\epsilon$ is tiny. Therefore, SAFL can achieve an $O(1/T)$ convergence near initialization, which accounts for the empirical advantages over non-adaptive methods.

###### Corollary 2 .

With the same condition as in Thereom 3.1, set $b\geq\log^{3}(CKd^{2}T^{2}/\delta)$ and constant $J_{1}>\sqrt{2}G$, then for any $T\leq\frac{J_{1}-\sqrt{2}G}{\epsilon^{2}}$, with probability $1-\Theta(\delta)-O(\exp(-\Omega(\nu^{2})))-\delta_{g}$,

$$ $\displaystyle\frac{1}{J_{1}T}\sum_{t=1}^{T}\|\nabla{\cal L}(x_{t})\|^{2}$ $\displaystyle\leq\frac{{\cal L}(z_{1})\epsilon}{\kappa T}+\frac{1}{\epsilon} \frac{LG^{2}}{T}+\frac{\nu}{T}(G^{2}+\sigma\log^{\frac{1}{2}}(\frac{2T}{\delta _{g}}))+\frac{\kappa}{T}\frac{32}{(1-\beta_{1})^{2}}\frac{DG^{2}}{\epsilon},$ $$

where $\delta,\delta_{g}$ and $\nu$ are the randomness of sketching, sub-Gaussian noise and martingales respectively.

###### Assumption 1 .

###### Assumption 2 .

###### Assumption 3 .

###### Assumption 4 .

###### Remark 3.2 .

###### Theorem 3.1 .

###### Corollary 1 .

###### Corollary 2 .

### 3.4 Technical Results and Proof Sketch

In this section, we provide a sketch of the proof techniques behind the main results. We focus on the proof of Theorem 3.1, and the proof of Theorem B.2 shares the main structure. The proof of Theorem 3.1 contains several critical components, which are unique to adaptive methods. We follow the common proof framework of adaptive optimization, and carefully deal with the noise introduced by random sketching in the momentum. We adopt AMSGrad (Alg. 2) as the server optimizer and it would be straightforward to extend the analysis to other adaptive methods.

We first introduce the descent lemma for AMSGrad. For conciseness, we denote the precondtioner matrix ${\rm diag}((\sqrt{\hat{v}_{t}}+\epsilon)^{2})$ as $\hat{V}_{t}$. Define an auxiliary variable $z_{t}=x_{t}+\frac{\beta_{1}}{1-\beta_{1}}(x_{t}-x_{t-1})$. The trajectory of ${\cal L}$ over $\{z_{t}\}_{t=1}^{T}$ can be tracked by the following lemma.

###### Lemma 3.2 .

For any round $t\in[T]$, there exists function $\Phi_{t}\geq 0$, and $\Phi_{0}\leq G$ such that

$$ $\displaystyle{\cal L}(z_{t+1})$ $\displaystyle\leq{\cal L}(z_{t})+\Phi_{t}-\Phi_{t+1}-\frac{\kappa\eta}{C}\sum_ {c=1}^{C}\sum_{k=1}^{K}\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}R_{t} ^{\top}R_{t}g_{t,k}^{c}+(z_{t}-x_{t})^{\top}H_{{\cal L}}(\hat{z}_{t})(z_{t+1}- z_{t}),$ $$

where $H_{{\cal L}}(\hat{z}_{t})$ is the loss Hessian at some $\hat{z}_{t}$ within the element-wise interval of $[x_{t},z_{t}]$.

Our objective henceforth is to bound the first-order descent term and the second-order quadratic term on the right hand side respectively.

Second-Order Quadratic Term. Denote $\{\lambda_{j},v_{j}\}_{j=1}^{d}$ as the eigen-pairs of $H_{{\cal L}}(\hat{z}_{t})$. The quadratic term can be written as $(z_{t}-x_{t})^{\top}H_{{\cal L}}(\hat{z}_{t})(z_{t+1}-z_{t})=\sum_{j=1}^{d}
\lambda_{j}\langle z_{t+1}-z_{t},v_{j}\rangle\langle z_{t}-x_{t},v_{j}\rangle.$
The inner product terms can be viewed as a projection of the updates onto anisotropic bases. Since $z_{t+1}-z_{t}$ and $z_{t}-x_{t}$ can both be expressed by $x_{t+1}-x_{t}$ and $x_{t}-x_{t-1}$, we can bound the quadratic term using the following lemma.

###### Lemma 3.3 .

For any $t\in[T]$, $|\langle x_{t}-x_{t-1},v_{j}\rangle|\leq\kappa\eta(1+\frac{\log^{1.5}(CKtd/
\delta)}{\sqrt{b}})\frac{KG}{\epsilon}$, with probability $1-\delta$.

Bounding the inner-product term is non-trivial since $z_{t}$ contains momentum information which depends on the randomness of previous iterations. A proof of a generalized version of this statement is deferred to the appendix, where induction methods are used to address the dependence. Combining Lemma 3.3 with Assumption 4 yields a dimension-free bound on the second-order quadratic term.

###### Remark 3.3 .

A straightforward application of smoothness to the second-order term yields a quadratic term $\|R^{\top}Rg\|^{2}$, which is linearly proportional to $d$ in scale (Rothchild et al.,, ; Song et al.,, ). We avoid this dimension dependence by combining Property 3 of sketching and Assumption 4.∎

First-Order Descent Term. The first-order term in the descent lemma can be decomposed into three components, which we will handle separately:

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}R_{t}^{\top}R_{ t}g_{t,k}^{c}=$ $\displaystyle\underbrace{\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2} \nabla{\cal L}^{c}(x_{t})}_{{\cal D}_{1}}+\underbrace{\nabla{\cal L}(x_{t})^{ \top}{\hat{V}_{t-1}}^{-1/2}(R_{t}^{\top}R_{t}g_{t,k}^{c}-\nabla{\cal L}^{c}(x_ {t,k}^{c}))}_{{\cal D}_{2}}$ $\displaystyle+\underbrace{\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}( \nabla{\cal L}^{c}(x_{t,k}^{c})-\nabla{\cal L}^{c}(x_{t}))}_{{\cal D}_{3}}.$ $$

First, ${\cal D}_{3}$ can be reduced to a second-order term by smoothness over ${\cal L}$, $\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}(\nabla{\cal L}^{c}(x_{t,k}^
{c})-\nabla{\cal L}^{c}(x_{t}))=-\eta\sum_{\tau=1}^{k}\nabla{\cal L}(x_{t})^{
\top}{\hat{V}_{t-1}}^{-1/2}\hat{H}_{\cal L}^{c}g_{t,\tau}^{c}.$
Note that this term does not involve any stochasticity from random sketching, hence we can directly derive the upper bound by Cauchy-Schwartz. Next, since $\frac{1}{{C}}\sum_{c=1}^{C}\nabla{\cal L}^{c}(x_{t})=\nabla{\cal L}(x_{t})$, ${\cal D}_{1}$ composes a scaled squared gradient norm. Applying element-wise high probability bound on random sketching yields the lower bound for the scale.

###### Lemma 3.4 .

For ${\hat{V}_{t-1}}^{-1/2}$ generated by Algorithm 1
(SAFL), with probability $1-\delta$,

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}\nabla{\cal L}( x_{t})\geq\left(\sqrt{1+\frac{\log^{1.5}(CKtd^{2}/\delta)}{\sqrt{b}}}\eta KG+ \epsilon\right)^{-1}\|\nabla{\cal L}(x_{t})\|^{2}.\vspace*{-2mm}$ $$

Martingale for zero-centered noise. ${\cal D}_{2}$ contains a zero-centered noise term $R_{t}^{\top}R_{t}g_{t,k}^{c}-\nabla{\cal L}^{c}(x_{t,k}^{c})$, where the randomness is over $R_{t}$ and the mini-batch noise at round $t$. Although $x_{t,k}^{c}$ has temporal dependence, the fresh noise due to mini-batching and sketching-desketching at round $t$ is independent of the randomness in the previous iterations. Therefore, the random process defined by the aggregation of the zero-centered noise terms over time forms a martingale. The martingale difference can be bounded with high probability under our proposed sketching method. Then by adapting Azuma’s inequality on a sub-Gaussian martingale, we have

###### Lemma 3.5 .

With probability $1-O(\exp(-\Omega(\nu^{2})))-\delta-\delta_{g}$,

$$ $\displaystyle\sum_{t=1}^{T}\left|\frac{1}{C}\sum_{c=1}^{C}\sum_{k=1}^{K}\nabla {\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}(R_{t}^{\top}R_{t}g_{t,k}^{c}- \nabla{\cal L}^{c}(x_{t,k}^{c}))\right|\leq\nu\sqrt{T}(\frac{\log^{1.5}(CKTd/ \delta)}{\sqrt{b}}\frac{KG^{2}}{\epsilon}+\frac{\sigma}{\epsilon}\log^{\frac{1 }{2}}(\frac{2T}{\delta_{g}})).$ $$

Finally, applying union bounds to these parts and telescoping the descent lemma leads to Theorem 3.1.

###### Lemma 3.2 .

###### Lemma 3.3 .

###### Remark 3.3 .

###### Lemma 3.4 .

###### Lemma 3.5 .

## 4 Sketched Clipping FL for Heavy-tailed Noise

In this section, we study the performance of Sketched Adaptive Clipped FL (SACFL) defined in Algorithm 3 in the context of heavy-tailed noise over gradient norms. This is arguably the more challenging setting and requires carefully addressing the noises with clipping.

### 4.1 Heavy-tailed Noise and Sketched Clipping-based Adaptive Methods

We start with the key bounded $\alpha$-moment assumption for the heavy-tailed stochastic first-order oracle.

###### Assumption 5 .

(Bounded $\alpha$-Moment). There exists a real number $\alpha\in(1,2]$ and a constant $G\geq 0$, such that $\mathbb{E}[\|\nabla{\cal L}^{c}(x,\xi)\|^{\alpha}]\leq G^{\alpha}$, $\forall c\in[C],x\in\mathbb{R}^{d}$, where $\xi$ is the noise from the minibatch.

Assumption 5 implies that the noise can have unbounded second moments when $\alpha<2$, which is much weaker compared to Assumption 2. This assumption can be satisfied by a family of noises including the Pareto
distribution (Arnold,, ) and $\alpha$-stable Levy distribution (Nolan,, ), both of which have unbounded variances. Heavy-tailed noises have detrimental effects on most of existing optimization theories, while, at the same time, being prevalent in FL due to data heterogeneity, i.e., non-i.i.d. client data distributions. This phenomenon has been shown in (Charles et al.,, ), and Assumption 5 has been adopted in existing theoretical analysis (Zhang et al.,, ; Yang et al.,, ). Clipping-based methods (Koloskova et al.,, ), a mainstream approach to handle exceedingly large gradient norms, use adaptive learning rates to normalize the gradient. These methods have empirically demonstrated the capability under heavy-tailed scenarios and are also proven to have optimal convergence rates (Zhang et al.,, ; Liu et al.,, ).

Figure: Algorithm 3 Sketched Adaptive Clipped Federated Learning (SACFL)

Our goal is to apply the sketching techniques to the clipping-based adaptive methods. This is indeed a challenging task. As we have already shown in Section 3, random sketching introduces a significant amount of noise to the client updates. It is unknown whether these noises additionally introduced by sketching will affect the behavior of clipping methods, given that the intrinsic noises are already heavy-tailed due to data heterogeneity.

To address this open question, we propose the Sketched Adaptive Clipped Federated Learning (SACFL) in Algorithm 3. In each round, besides sketching the local updates, the client directly sends the $L_{2}$-norm of the update to the server. The $L_{2}$-norm can be viewed as a global second moment specific to clipping methods. Notably, the $L_{2}$-norm is a scalar value and does not require any compression. Upon receiving the running moments $\bar{m}_{t}^{c},\bar{v}_{t}^{c}$ from clients, the server averages the sketched local updates and the $L_{2}$-norms respectively, and then updates the global model by $x_{t}=x_{t-1}-\kappa\min\{\frac{\tau}{\bar{v}_{t}},1\}\texttt{desk}(\bar{m}_{t})$, where $\kappa$ is the learning rate and $\tau$ is a horizon-dependent clipping threshold. When the averaged gradient norm exceeds $\tau$, i.e., when the gradient is in the heavy-tailed regime, clipping is enabled to downscale the gradient. Otherwise, the recovered gradients are directly used to update the global model.

###### Assumption 5 .

### 4.2 Convergence Analysis

Next, we state the convergence guarantees of SACFL under Assumption 5. All technical proofs for this section are in Appendix C. We start with the descent lemma for clipping methods.

###### Lemma 4.1 .

If the sketching dimension $b$ satisfies $b\geq 4\log^{3}(d/\delta)$, taking expectation over the stochasticity of gradients yields, with probability $1-\Theta(\delta)$,

$$ $\displaystyle\mathbb{E}[{\cal L}(x_{t+1})]-{\cal L}(x_{t})+\frac{\kappa\eta K} {4}\|\nabla{\cal L}(x_{t})\|^{2}$ $\displaystyle\leq\frac{\kappa\eta K}{2}\|\nabla{\cal L}(x_{t})-\frac{1}{K} \frac{1}{{C}}\sum_{c=1}^{C}\mathbb{E}[\tilde{\Delta}_{t}^{c}]\|^{2}+\frac{ \kappa^{2}\eta^{2}}{2}\mathbb{E}[(\frac{1}{C}R^{\top}R\sum_{c=1}^{C}\tilde{ \Delta}_{t}^{c})^{\top}H_{\cal L}(\hat{x}_{t})(\frac{1}{C}R^{\top}R\sum_{c=1}^ {C}\tilde{\Delta}_{t}^{c})],$ $$

where $\tilde{\Delta}_{t}^{c}=\min\{1,\frac{\tau}{\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta
_{t}^{c}\|}\}\Delta_{t}^{c}$, and $\Delta_{t}^{c}=\sum_{k=1}^{K}g_{t,k}^{c}$, $H_{{\cal L}}(\hat{x}_{t})$ is the loss Hessian at some $\hat{x}_{t}$ within the element-wise interval of $[x_{t},x_{t+1}]$.

Intuitively, the first-order terms are barely affected by the heavy-tailed noise since we assume $\alpha>1$ and these terms do not involve the potentially-unbounded second moments, although special attention for the first-order terms is necessary to achieve the optimal convergence rate, which are deferred to the appendix. Next, we show how to deal with the second-order term. With probability $1-\Theta(\delta)$,

$$ $\displaystyle\mathbb{E}[(\frac{1}{C}R^{\top}R\sum_{i=1}^{C}\tilde{\Delta}_{t}^ {c})^{\top}\hat{H}_{\cal L}(\frac{1}{C}R^{\top}R\sum_{i=1}^{C}\tilde{\Delta}_{ t}^{c})]=\mathbb{E}[\sum_{j=1}^{d}\lambda_{j}\langle\frac{1}{C}R^{\top}R\sum_{ i=1}^{C}\tilde{\Delta}_{t}^{c},v_{j}\rangle^{2}]$ $\displaystyle\leq\mathbb{E}\left[\sum_{j=1}^{d}\lambda_{j}1_{\lambda_{j}\geq 0 }\left(\frac{\tau M}{\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta_{t}^{c}\|}\frac{1}{{C }}\sum_{c=1}^{C}\|\Delta_{t}^{c}\|\right)^{2-\alpha}\left(\frac{M}{C}\sum_{i=1 }^{C}\|\Delta_{t}^{c}\|\right)^{\alpha}\right]\leq M\sum_{j=1}^{d}\lambda_{j}1 _{\lambda_{j}\geq 0}K^{2}\tau^{2-\alpha}G^{\alpha},$ $$

where $M:=(1+\frac{\log^{1.5}(d/\delta)}{\sqrt{b}})$. The first equality follows by using the same eigen-decomposition as in the previous section where $\{\lambda_{j},v_{j}\}$ are the eigenpairs of $H_{\cal L}(\hat{x}_{t})$ and the second order term can be reduced to a squared inner product term. The primary trick thereafter (in the first inequality) is to split the inner product terms into two parts, which can be handled by the two-sided adaptive learning rates respectively. By resorting to the bounded second moment of random sketching, we find the first part with order $2-\alpha$ is contained in a $(1+\frac{\log^{1.5}(d/\delta)}{\sqrt{b}})\tau-$ball with high probability, and the second part with order $\alpha$ is bounded by applying Assumption 5.

###### Remark 4.1 .

The bound is high-probability w.r.t the randomness of sketching functions, while the expectation is over other randomness, including local stochastic noises and the heavy-tailed noises.

Finally, we can derive the convergence rate for SACFL by combining the analysis.

###### Theorem 4.2 .

If the sketch size $b$ satisfies $b\geq 4\log^{3}(dT/\delta)$, then under Assumption 4 and 5, the sequence $\{x_{t}\}$ generated by Algorithm 3 (SACFL) satisfies:

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla{\cal L}(x_{t})]\|^{2 }\leq\frac{4({\cal L}(x_{1})-{\cal L}(x_{T}))}{\kappa\eta KT}$ $\displaystyle+3K^{2}(L^{2}\eta^{2}G^{2}+G^{2\alpha}\tau^{-2(\alpha-1)}+L\eta G ^{1+\alpha}\tau^{1-\alpha})$ $\displaystyle+3D\kappa\eta(KG^{\alpha}\tau^{2-\alpha}),~{}~{}w.p.~{}1-\Theta(\delta)$ $$

With a proper choice of hyper-parameters with $\kappa=K^{\frac{3\alpha-6}{3\alpha-2}}T^{-\frac{1}{3\alpha-2}}$, $\eta=T^{\frac{1-\alpha}{3\alpha-2}}K^{\frac{4-4\alpha}{3\alpha-2}}$ and $\tau=(K^{4}T)^{\frac{1}{3\alpha-2}}$, we achieve $\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla{\cal L}(x_{t})\|^{2}]\leq O(T^{
\frac{2-2\alpha}{3\alpha-2}}K^{\frac{4-2\alpha}{3\alpha-2}}),~{}~{}w.p.~{}1-
\Theta(\delta).$

###### Remark 4.2 .

The convergence rate depends on the noise level $\alpha$. When $\alpha=2$, i.e. the bounded variance case, the convergence rate is $O(1/\sqrt{T})$, which matches the rate in Theorem 3.1. We also claim that the iteration complexity matches the optimal bound in the heavy-tailed case (Yang et al.,, ; Zhang et al.,, ). ∎

###### Lemma 4.1 .

###### Remark 4.1 .

###### Theorem 4.2 .

###### Remark 4.2 .

## 5 Empirical Studies

In this section, we instantiate the algorithmic framework of SAFL in Section 3 and SACFL in Section 4 to demonstrate the effect of sketching in different settings.

Figure: Figure 1: Test Error on CIFAR-10 with ResNet of 42M parameters. The plot starts from the 10th epoch for better demonstration. Optimizers: ADA_OPT $\in$ {SGD, SGDm, Adam}, FetchSGD and 1bit-Adam with sketch size $b\in\{4e5,1e6,4e6\}$; Rightmost: ADA_OPT is Adam. The legend 4e7 represents training in the ambient dimension without sketching. Adam optimizer consistently outperforms other optimizers. Larger sketch sizes improves the convergence rate and test errors.
Refer to caption: x1.png

Figure: Figure 2: Test Error on CIFAR-10. We finetune a ViT-base model (with 86M parameters) from the pretrained backbone checkpoint (Dosovitskiy et al.,, ). SGDm, Adam, FetchSGD are compared under sketch size $b\in\{8e4,8e5,8e6\}$. 1Bit-Adam has comparable compression rates with $b=8e5$. Sketched Adam optimizer consistently outperforms other communication-efficient algorithms.
Refer to caption: x2.png

Figure: Figure 3: Test Error on SST2 (GLUE) with BERT of 100M parameters. Left: sketch size $b=2e5$; Middle: $b=2e6$; Right: ADA_OPT is Adam, with sketch size $b\in\{2e4,2e5,2e6\}$. The legend $1e8$ represents training in the ambient dimension without sketching. Adam achieves faster convergence and lower test errors across different sketch sizes. Larger sketch sizes mainly improves the convergence rate and achieve comparable test errors at the end of training.
Refer to caption: x3.png

Experimental Configurations.
We adopt three distinct experimental settings, from vision to language tasks, and in finetuning and training-from-scratch regimes. For the vision task, we train a ResNet101 (Wu and He,, ) with a total of 42M parameters from scratch and finetune a ViT-Base (Dosovitskiy et al.,, ) with 86M parameters on the CIFAR-10 dataset (Krizhevsky et al.,, ). For the language task, we adopt SST2, a text classification task, from the GLUE benchmark (Wang et al.,, ). We train a BERT model (Devlin,, ) which has 100M parameters.
The client optimizer is mini-batch SGD. At each round, the client trains one single epoch (iterate over the client dataset). For other hyperparameters used in the training process, please refer to the appendix.

Sharp-Decaying Hessian Eigenspectrum. As a key technical cornerstone of the theory, Assumption 4 states that the Hessian matrix has a sharp-decaying eigenspectrum. While this assumption has been repeatedly validated in the previous works, it is unknown if the property holds in the context of federated deep learning. We show an affirmative verification in Fig. 6 in the Appendix.

Sketched Adaptive FL. We adopt Adam as the base adaptive optimizer at the server side, and make comparison with the sketched non-adaptive optimizers SGD, SGDm (SGD with momentum). We also involve the state-of-the-art communication-efficient algorithms FetchSGD (Rothchild et al.,, ) and 1bit-Adam (Tang et al.,, ), which are based on biased sparsification and quantization respectively. In the i.i.d client setting, the data are uniformly distributed over 5 clients.
Fig. 1 depicts the test errors on CIFAR-10 when training ResNet(40M) with sketch sizes $b\in\{4e5,1e6,4e6\}$.
We can see for all sketch sizes, our sketched Adam consistently outperforms other optimizers in the convergence rate and the test error. The compression rate of 1bit-Adam is fixed at 97%, which is comparable with the compression rate 99% achieved at $b=4e5$. 1bit-Adam is plotted separately because it takes remarkably longer to converge. More interestingly, even with distinct sketch sizes, the iterations needed for convergence in Adam are almost the same. The test performance degrades slightly with smaller sketch sizes. This is anticipated and totally acceptable considering that the communication cost has been drastically reduced. With the same budget of communication bits, using a lower compression rate facilitates larger model training, which has the potential in better generalization performance. For results on extremely large compression rates, please refer to Appendix D.

We also present results on finetuning a ViT-Base model (80M) in Fig. 2. The sketch size $b\in\{8e4,8e5,8e6\}$. We see, in the finetuning regime, the sketched Adam optimizer also achieves competitive performance with the baseline methods. Similar phenomenon is observed in the language task. Fig. 3 shows the test errors of training SST2 with BERT (100M). The sketch sizes are selected from $\{2e4,2e5,2e6,1e8\}$. We observe sketched Adam converges faster and achieves a slightly better test performance than other optimizers. Note that the sketch size of $2e4$ is tiny, given that the ambient dimension is 100M. It is quite thrilling that using an extremely high compression rate (99.98%), the model can still achieve comparable performance as trained in the ambient dimension.

Figure: (a) Heavy tailed gradient norms.
Refer to caption: x4.png

Sketched Clipping Method. Next, we study the performance of the sketched clipping methods. In Section 4, we claim that SACFL excels in the heavy-tailed regime even when interfered with the noise from random sketching. To show empirical evidence, we first build an (extremely pessimistic) environment of heavy-tailed noise on the CIFAR-10 dataset. Specifically, the data categories are extremely imbalanced among 80 clients. Each client has 4 distinct majority classes which occupy 80% of its entire client dataset, while the remaining data samples are the minority categories. In this data heterogeneous setting, the features are hardly learned. We run sketched clipping methods in this environment and fit the local gradient norms with an $\alpha$-stable Levy distribution in Fig. 4(a). We select the clipping threshold $\tau$ in $\{0.2,0.4\}$ such that the clipping operator is in effect in most rounds. Fig. 4(b) depicts the test errors in the first 200 epochs under distinct sketch sizes, where we observe the trajectories significantly overlap. Hence, we can verify that the sketching operator has minor effects on the clipping method, while providing the benefits of lower communication costs. For completeness, we present results on the BERT model trained with SST2 dataset in Fig. 8 in Appendix, where we can also observe the sketched clipping method preserves the original convergence guarantees.

## 6 Conclusion

In this paper, we investigated sketched adaptive methods for FL. While the motivation behind combining sketching and adaptive methods for FL is clear, there is limited understanding on its empirical success due to the inherent technical challenges. We consider both mild-noise and heavy-tailed noise settings, propose corresponding adaptive algorithms for each, and show highly promising theoretical and empirical results. Inspired by the recently observations on heterogeneity in weights across neural network layers (Zhang et al.,, ), an important future direction is to independently sketch layer-wise gradients, rather than sketching the concatenated gradient vectors. We believe our novel work can form the basis for future advances on the theme.

Acknowledgements. The work was supported by the National Science Foundation (NSF) through awards IIS 21-31335, OAC 21-30835, DBI 20-21898, as well as a C3.ai research award.

## Appendix A Lemma for Random Sketching

For completeness, we provide the following lemmas that give high probability bounds on the inner products.

###### Lemma A.1 .

(SRHT)[Same as Lemma D.23 Song et al., ()]
Let $R\in\mathbb{R}^{b\times d}$ denote a subsample randomized Hadamard transform or AMS sketching matrix. Then for any fixed vector $h\in\mathbb{R}$ and any fixed vector $g\in\mathbb{R}$ the following properties hold:

$$ $\displaystyle\mathbb{P}\left[|\langle g^{\top}R^{\top}Rh-g^{\top}h|\geq\frac{ \log^{1.5}(d/\delta)}{\sqrt{b}}\|g\|_{2}\|h\|_{2}\right]\leq\Theta(\delta).$ $$

###### Lemma A.2 .

(Gaussian)[Same as Lemma D.24 Song et al., ()]
Let $R\in\mathbb{R}^{b\times d}$ denote a random Gaussian matrix. Then for any fixed vector $h\in\mathbb{R}$ and any fixed vector $g\in\mathbb{R}$ the following properties hold:

$$ $\displaystyle\mathbb{P}\left[|\langle g^{\top}R^{\top}Rh-g^{\top}h|\geq\frac{ \log^{1.5}(d/\delta)}{\sqrt{b}}\|g\|_{2}\|h\|_{2}\right]\leq\Theta(\delta).$ $$

###### Lemma A.3 .

(Count-Sketch)[Same as Lemma D.25 Song et al., ()]
Let $R\in\mathbb{R}^{b\times d}$ denote a count-sketch matrix. Then for any fixed vector $h\in\mathbb{R}$ and any fixed vector $g\in\mathbb{R}$ the following properties hold:

$$ $\displaystyle\mathbb{P}\left[|\langle g^{\top}R^{\top}Rh-g^{\top}h|\geq\log(1/ \delta)\|g\|_{2}\|h\|_{2}\right]\leq\Theta(\delta).$ $$

###### Lemma A.1 .

###### Lemma A.2 .

###### Lemma A.3 .

## Appendix B Proof of Theorem 3.1

### B.1 Proof of Lemma 3.2

Let

$$ $\displaystyle z_{t}=x_{t}+\frac{\beta_{1}}{1-\beta_{1}}(x_{t}-x_{t-1})=\frac{1 }{1-\beta_{1}}x_{t}-\frac{\beta_{1}}{1-\beta_{1}}x_{t-1}.$ $$

Then, the update on $z_{t}$ can be expressed as

$$ $\displaystyle z_{t+1}-z_{t}$ $\displaystyle=\frac{1}{1-\beta_{1}}(x_{t+1}-x_{t})-\frac{\beta_{1}}{1-\beta_{1 }}(x_{t}-x_{t-1})$ $\displaystyle=-\frac{1}{1-\beta_{1}}\kappa\hat{V_{t}}^{-1/2}\cdot m_{t}+\frac{ \beta_{1}}{1-\beta_{1}}\kappa\hat{V_{t-1}}^{-1/2}\cdot m_{t-1}$ $\displaystyle=-\frac{1}{1-\beta_{1}}\kappa\hat{V_{t}}^{-1/2}\cdot(\beta_{1}m_{ t-1}+(1-\beta_{1})\cdot R_{t}^{\top}\bar{m}_{t})+\frac{\beta_{1}}{1-\beta_{1}} \kappa{\hat{V}_{t-1}}^{-1/2}\cdot m_{t-1}$ $\displaystyle=\frac{\beta_{1}}{1-\beta_{1}}\left(\kappa{\hat{V}_{t-1}}^{-1/2}- \kappa\hat{V_{t}}^{-1/2}\right)m_{t-1}-\frac{\kappa}{C}\hat{V_{t}}^{-1/2}R_{t} ^{\top}\sum_{c=1}^{C}\bar{m}_{t}^{c}$ $\displaystyle=\frac{\beta_{1}}{1-\beta_{1}}\left(\kappa{\hat{V}_{t-1}}^{-1/2}- \kappa\hat{V_{t}}^{-1/2}\right)m_{t-1}-\frac{\kappa}{C}\hat{V_{t}}^{-1/2}R_{t} ^{\top}\sum_{c=1}^{C}R_{t}(x_{t,0}^{c}-x_{t,K}^{c})$ $\displaystyle=\frac{\beta_{1}}{1-\beta_{1}}\left(\kappa{\hat{V}_{t-1}}^{-1/2}- \kappa\hat{V_{t}}^{-1/2}\right)m_{t-1}-\frac{\kappa\eta}{C}\hat{V_{t}}^{-1/2} \sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c}$ $$

By Taylor expansion, we have

$$ $\displaystyle{\cal L}(z_{t+1})$ $\displaystyle={\cal L}(z_{t})+\nabla{\cal L}(z_{t})^{\top}(z_{t+1}-z_{t})+ \frac{1}{2}(z_{t+1}-z_{t})^{\top}\hat{H}_{{\cal L}}(z_{t+1}-z_{t})$ $\displaystyle={\cal L}(z_{t})+\nabla{\cal L}(x_{t})^{\top}(z_{t+1}-z_{t})+( \nabla{\cal L}(z_{t})-\nabla{\cal L}(x_{t}))^{\top}(z_{t+1}-z_{t})+\frac{1}{2} (z_{t+1}-z_{t})^{\top}\hat{H}_{{\cal L}}(z_{t+1}-z_{t}).$ $$

Bounding the first-order term

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}(z_{t+1}-z_{t})$ $\displaystyle=$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}\left(\frac{\beta_{1}}{1-\beta_{1}} \left(\kappa{\hat{V}_{t-1}}^{-1/2}-\kappa\hat{V_{t}}^{-1/2}\right)m_{t-1}- \frac{\kappa\eta}{C}\hat{V_{t}}^{-1/2}\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top} R_{t}g_{t,k}^{c}\right)$ $\displaystyle\leq$ $\displaystyle\frac{\beta_{1}}{1-\beta_{1}}\|\nabla{\cal L}(x_{t})\|_{\infty}( \|\kappa{\hat{V}_{t-1}}^{-1/2}\|_{1,1}-\|\kappa\hat{V_{t}}^{-1/2}\|_{1,1})\|m_ {t-1}\|_{\infty}$ $\displaystyle-\frac{\eta}{C}\nabla{\cal L}(x_{t})^{\top}(\kappa\hat{V_{t}}^{-1 /2}-\kappa{\hat{V}_{t-1}}^{-1/2})\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t} g_{t,k}^{c}-\frac{\kappa\eta}{C}\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{- 1/2}\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c}$ $\displaystyle\leq$ $\displaystyle\left(\frac{\beta_{1}}{1-\beta_{1}}\|m_{t-1}\|_{\infty}+\frac{ \eta}{C}\|\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c}\|_{\infty} \right)\|\nabla{\cal L}(x_{t})\|_{\infty}(\|\kappa{\hat{V}_{t-1}}^{-1/2}\|_{1, 1}-\|\kappa\hat{V_{t}}^{-1/2}\|_{1,1})$ $\displaystyle-\frac{\kappa\eta}{C}\sum_{c=1}^{C}\sum_{k=1}^{K}\nabla{\cal L}(x _{t})^{\top}{\hat{V}_{t-1}}^{-1/2}R_{t}^{\top}R_{t}g_{t,k}^{c}.$ $$

The quadratic terms can be written as

$$ $\displaystyle(\nabla{\cal L}(z_{t})-\nabla{\cal L}(x_{t}))^{\top}(z_{t+1}-z_{t })=(z_{t}-x_{t})^{\top}\hat{H}_{{\cal L}}(\frac{1}{1-\beta_{1}}(x_{t+1}-x_{t}) -\frac{\beta_{1}}{1-\beta_{1}}(x_{t}-x_{t-1})),$ $$

where $\hat{H}_{\cal L}$ is a second-order Taylor remainder. So the quadratic term can be further seen as a quadratic form over $z_{t+1}-z_{t}$ and $z_{t}-x_{t}$, denote as ${\cal Q}(z_{t+1}-z_{t},z_{t}-x_{t})$. For the same reason, the term $\frac{1}{2}(z_{t+1}-z_{t})^{\top}\hat{H}_{{\cal L}}(z_{t+1}-z_{t})$ can also be written into a quadratic form ${\cal Q}(z_{t+1}-z_{t},z_{t+1}-z_{t})$. Putting the two terms together yields a quadratic form of ${\cal Q}(z_{t+1}-z_{t},z_{t}-x_{t})$.

### B.2 Proof of Lemma B.1 (Generalized version of Lemma 3.3 )

###### Proof.

We can prove by induction. For $t=0$, since $m_{0}=0$, the inequality holds. Suppose we have for $h\in\mathbb{R}^{d}$, s.t. $\|h\|\leq H$, with probability $1-\Theta((t-1)\delta)$,

$$ $\displaystyle|m_{t-1}^{\top}h|\leq(1+\frac{\log^{1.5}(CKd/\delta)}{\sqrt{b}})G$ $$

Then by the update rule,

$$ $\displaystyle|m_{t}^{\top}h|$ $\displaystyle=|(\beta_{1}\cdot m_{t-1}+(1-\beta_{1})\cdot\frac{\eta}{C}\sum_{c =1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c})^{\top}h|$ $\displaystyle\leq\beta_{1}|m_{t-1}^{\top}h|+\frac{(1-\beta_{1})\eta}{C}\sum_{c =1}^{C}\sum_{k=1}^{K}|\langle R_{t}^{\top}R_{t}g_{t,k}^{c},h\rangle|$ $\displaystyle\leq\beta_{1}|m_{t-1}^{\top}h|+(1-\beta_{1})(1+\frac{\log^{1.5}( CKd/\delta)}{\sqrt{b}})\eta\sum_{k=1}^{K}\|g_{t,k}^{c}\|_{2}\|h\|_{2}$ $\displaystyle\leq(1+\frac{\log^{1.5}(CKd/\delta)}{\sqrt{b}})\eta KGH,~{}~{}w.p .~{}1-\Theta(t\delta).$ $$

Let $h=\hat{V}_{t}^{-1/2}v_{i}$. Then $\|h\|_{2}\leq 1/\epsilon$. We have

$$ $\displaystyle|(\hat{V}_{t}^{-1/2}m_{t})^{\top}v_{i}|$ $\displaystyle\leq(1+\frac{\log^{1.5}(CKd/\delta)}{\sqrt{b}})\eta KG/\epsilon$ $$

∎

###### Proof.

### B.3 Proof of Lemma 3.4

We first prove the element-wise lower bound of the diagonal matrix $\hat{V}_{t-1}^{-1/2}$. Denote $(\hat{V}_{t-1}^{-1/2})_{i}$ as the $i$-th element on the diagonal of $\hat{V}_{t-1}^{-1/2}$. By the update rule,

$$ $\displaystyle(\hat{V}_{t-1}^{-1/2})_{i}\geq(\max_{t-1}(\sqrt{v_{t,i}})+ \epsilon)^{-1}\geq(\sqrt{1+\frac{\log^{1.5}(CKtd/\delta)}{\sqrt{b}}}\eta KG+ \epsilon)^{-1},~{}~{}w.p.~{}1-\Theta(\delta)$ $$

where the last inequality follows by letting $h$ as a one-hot vector $h_{i}$ in Lemma A.1, observing that the elements can be transformed to an inner product form $v_{t,i}=v_{t}^{\top}h_{i}$. Then the scaled gradient norm can be lower bounded as

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}\nabla{\cal L}( x_{t})$ $\displaystyle\geq\min_{i}(\hat{V}_{t-1}^{-1/2})_{i}\sum_{i=1}^{d}[\nabla{\cal L }(x_{t})]_{i}^{2}$ $\displaystyle\geq(\sqrt{1+\frac{\log^{1.5}(CKtd/\delta)}{\sqrt{b}}}\eta KG+ \epsilon)^{-1}\|\nabla{\cal L}(x_{t})\|^{2},~{}~{}w.p.~{}1-\Theta(d\delta)$ $$

which completes the proof by applying union bounded on the dimension $d$.

### B.4 Proof of Lemma 3.5

Since the noise is zero-centered, we view the random process of

$$ $\displaystyle\{Y_{t}=\sum_{\tau=1}^{t}\frac{1}{{C}}\sum_{c=1}^{C}\sum_{k=1}^{K }\nabla{\cal L}(x_{\tau})^{\top}{\hat{V}_{\tau-1}}^{-1/2}(R_{\tau}^{\top}R_{ \tau}g_{\tau,k}^{c}-g_{\tau,k}^{c})\}_{t=1}^{T}$ $$

as a martingale. The difference of $|Y_{t+1}-Y_{t}|$ is bounded with high probability

$$ $\displaystyle|Y_{t+1}-Y_{t}|=|\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/ 2}(R_{t}^{\top}R_{t}g_{t,k}^{c}-g_{t,k}^{c})|\leq\frac{\log^{1.5}(d/\delta)}{ \sqrt{b}}G\|\hat{V}_{t}^{-1/2}\nabla{\cal L}(x_{t})\|_{2},~{}~{}w.p.~{}1- \Theta(\delta)$ $$

Then by Azuma’s inequality,

$$ $\displaystyle\mathbb{P}(|Y_{T}|\geq\nu\sqrt{\sum_{t=1}^{T}\left(\frac{\log^{1. 5}(d/\delta)}{\sqrt{b}}G\|\hat{V}_{t}^{-1/2}\nabla{\cal L}(x_{t})\|_{2}\right) ^{2}})=O({\rm exp}(-\Omega(\nu^{2})))+T\delta$ (1) $$

Note that the original Azuma’s is conditioned on a uniform bound of the difference term, while our bound here is of high probability. Hence, we need another union bound. A similar bound can be achieved for the sub-Gaussian noise in stochastic gradient. Let

$$ $\displaystyle Z_{t}=\sum_{\tau=1}^{t}\frac{1}{{C}}\sum_{c=1}^{C}\sum_{k=1}^{K} \nabla{\cal L}(x_{\tau})^{\top}{\hat{V}_{\tau-1}}^{-1/2}(g_{\tau,k}^{c}-\nabla {\cal L}^{c}(x_{t,k}^{c})).$ $$

Then

$$ $\displaystyle\mathbb{P}(|Z_{T}|\geq\nu\sqrt{\sum_{t=1}^{T}\frac{\sigma^{2}}{ \epsilon^{2}}\log(\frac{2T}{\delta_{g}})})=O({\rm exp}(-\Omega(\nu^{2})))+ \delta_{g}$ $$

Combining the two bounds by union bound completes the proof.

### B.5 Proof of Theorem 3.1

After applying Lemma 3.2. The second order quadratic forms in the descent lemma can be written as

$$ $\displaystyle(\nabla{\cal L}(z_{t})-\nabla{\cal L}(x_{t}))^{\top}(z_{t+1}-z_{t})$ $\displaystyle=$ $\displaystyle(z_{t}-x_{t})^{\top}\hat{H}_{{\cal L}}(\frac{1}{1-\beta_{1}}(x_{t +1}-x_{t})-\frac{\beta_{1}}{1-\beta_{1}}(x_{t}-x_{t-1}))$ $\displaystyle=$ $\displaystyle-\kappa\frac{\beta_{1}}{1-\beta_{1}}(\hat{V}_{t-1}^{-1/2}m_{t-1}) ^{\top}\hat{H}_{{\cal L}}(\frac{1}{1-\beta_{1}}(-\kappa\hat{V}_{t}^{-1/2}m_{t} )-\frac{\beta_{1}}{1-\beta_{1}}(-\kappa\hat{V}_{t-1}^{-1/2}m_{t-1}))$ $\displaystyle=$ $\displaystyle\kappa^{2}\frac{\beta_{1}}{(1-\beta_{1})^{2}}(\hat{V}_{t-1}^{-1/2 }m_{t-1})^{\top}\hat{H}_{{\cal L}}(\hat{V}_{t}^{-1/2}m_{t})-\kappa^{2}\frac{ \beta_{1}^{2}}{(1-\beta_{1})^{2}}(\hat{V}_{t-1}^{-1/2}m_{t-1})^{\top}\hat{H}_{ {\cal L}}(\hat{V}_{t-1}^{-1/2}m_{t-1}),$ $$

and

$$ $\displaystyle(z_{t+1}-z_{t})^{\top}\hat{H}_{\cal L}(z_{t+1}-z_{t})$ $\displaystyle=$ $\displaystyle(\frac{1}{1-\beta_{1}}(x_{t+1}-x_{t})-\frac{\beta_{1}}{1-\beta_{1 }}(x_{t}-x_{t-1}))^{\top}\hat{H}_{{\cal L}}(\frac{1}{1-\beta_{1}}(x_{t+1}-x_{t })-\frac{\beta_{1}}{1-\beta_{1}}(x_{t}-x_{t-1}))$ $\displaystyle=$ $\displaystyle\frac{1}{(1-\beta_{1})^{2}}(x_{t+1}-x_{t})^{\top}\hat{H}_{{\cal L }}(x_{t+1}-x_{t})-\frac{2\beta_{1}}{(1-\beta_{1})^{2}}(x_{t+1}-x_{t})^{\top} \hat{H}_{{\cal L}}(x_{t}-x_{t-1})$ $\displaystyle+\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}(x_{t}-x_{t-1})^{\top} \hat{H}_{{\cal L}}(x_{t}-x_{t-1}),$ $$

which is essentially a quadratic form defined on $\hat{V}_{t}^{-1/2}m_{t}$ and $\hat{V}_{t-1}^{-1/2}m_{t-1}$. Hence, we provide a generalized version of Lemma 3.3, as follows.

###### Lemma B.1 .

With probability $1-\Theta(t\delta)$, for eigenvector $v_{i}$ of the Hessian matrix, $|(\hat{V}_{t}^{-1/2}m_{t})^{\top}v_{i}|\leq(1+\frac{\log^{1.5}(CKd/\delta)}{
\sqrt{b}})\eta KG/\epsilon$.

Note that $v_{i}$ can be any basis and is constant throughout the training process. Then the sum of quadratic forms is written as

$$ $\displaystyle(\nabla{\cal L}(z_{t})-\nabla{\cal L}(x_{t}))^{\top}(z_{t+1}-z_{t})$ $\displaystyle\leq$ $\displaystyle\kappa^{2}\frac{\beta_{1}}{(1-\beta_{1})^{2}}(\hat{V}_{t-1}^{-1/2 }m_{t-1})^{\top}\hat{H}_{{\cal L}}(\hat{V}_{t}^{-1/2}m_{t})-\kappa^{2}\frac{ \beta_{1}^{2}}{(1-\beta_{1})^{2}}(\hat{V}_{t-1}^{-1/2}m_{t-1})^{\top}\hat{H}_{ {\cal L}}(\hat{V}_{t-1}^{-1/2}m_{t-1}),$ $\displaystyle=$ $\displaystyle\kappa^{2}\frac{\beta_{1}}{(1-\beta_{1})^{2}}\sum_{i=1}^{d} \lambda_{i}(\hat{V}_{t-1}^{-1/2}m_{t-1})^{\top}(v_{i}v_{i}^{\top})\hat{V}_{t}^ {-1/2}m_{t}-\kappa^{2}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}\sum_{i=1}^{d} \lambda_{i}(\hat{V}_{t-1}^{-1/2}m_{t-1})^{\top}(v_{i}v_{i}^{\top})\hat{V}_{t-1 }^{-1/2}m_{t-1}$ $\displaystyle\leq$ $\displaystyle\kappa^{2}\frac{\beta_{1}}{(1-\beta_{1})^{2}}\sum_{i=1}^{d}| \lambda_{i}||(\hat{V}_{t-1}^{-1/2}m_{t-1})^{\top}v_{i}||(\hat{V}_{t}^{-1/2}m_{ t})^{\top}v_{i}|+\kappa^{2}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}\sum_{i=1}^{ d}|\lambda_{i}||(\hat{V}_{t-1}^{-1/2}m_{t-1})^{\top}v_{i}|^{2}$ $\displaystyle\leq$ $\displaystyle\kappa^{2}\frac{2}{(1-\beta_{1})^{2}}D(1+\frac{\log^{1.5}(CKd/ \delta)}{\sqrt{b}})^{2}\eta^{2}K^{2}G^{2}/\epsilon^{2},$ $$

where the last inequality is by $\beta_{1}\leq 1$ and Lemma. B.1.

First-Order Descent Term. The first-order term in the descent lemma can be decomposed into three components, which we will handle separately.

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}R_{t}^{\top}R_{ t}g_{t,k}^{c}=$ $\displaystyle\underbrace{\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2} \nabla{\cal L}^{c}(x_{t})}_{{\cal D}_{1}}+\underbrace{\nabla{\cal L}(x_{t})^{ \top}{\hat{V}_{t-1}}^{-1/2}(R_{t}^{\top}R_{t}g_{t,k}^{c}-\nabla{\cal L}^{c}(x_ {t,k}^{c}))}_{{\cal D}_{2}}$ $\displaystyle+\underbrace{\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}( \nabla{\cal L}^{c}(x_{t,k}^{c})-\nabla{\cal L}^{c}(x_{t}))}_{{\cal D}_{3}}.$ $$

First, ${\cal D}_{3}$ can be reduced to a second-order term by smoothness over ${\cal L}$,

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}(\nabla{\cal L} ^{c}(x_{t,k}^{c})-\nabla{\cal L}^{c}(x_{t}))=\nabla{\cal L}(x_{t})^{\top}{\hat {V}_{t-1}}^{-1/2}\hat{H}_{L}^{c}(x_{t,k}^{c}-x_{t})$ $\displaystyle=$ $\displaystyle-\eta\sum_{\tau=1}^{k}\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}} ^{-1/2}\hat{H}_{\cal L}^{c}g_{t,\tau}^{c}$ $\displaystyle\leq$ $\displaystyle\frac{1}{\epsilon}L\|\nabla L\|\sum_{\tau=1}^{k}\|g_{t,\tau}^{c} \|\leq\frac{1}{\epsilon}\eta LKG^{2}.$ $$

Note that this term does not involve any stochasticity with regard to random sketching, which means we can directly derive the upper bound by Cauchy-Schwartz in the last inequality.

Next observing that $\frac{1}{{C}}\sum_{c=1}^{C}\nabla{\cal L}^{c}(x_{t})=\nabla{\cal L}(x_{t})$, ${\cal D}_{1}$ composes a scaled squared gradient norm. Applying element-wise high probability bound on random sketching yields the lower bound for the scale.
By Lemma 3.4, we can derive the lower bound for ${\cal D}_{1}$. Note that applying union bound to ${\cal D}_{1}$ does not introduce another $T$ dependence, since $\hat{v}_{t,i}$ is monotonically non-decreasing.

Martingale for zero-centered noise. ${\cal D}_{2}$ contains a zero-centered noise term $R_{t}^{\top}R_{t}g_{t,k}^{c}-\nabla{\cal L}^{c}(x_{t,k}^{c})$, where the randomness is over $R_{t}$ and the mini-batch noise at round $t$. Despite $x_{t,k}^{c}$ has temporal dependence, the fresh noise at round $t$ is independent of the randomness in the previous iterations. Hence, the random process defined by the aggregation of these norm terms over time forms a martingale. By Lemma 3.5, we can bound this term ${\cal D}_{2}$.

Finally, putting these parts together by union bound over $[T]$ and telescoping the descent lemma leads to Theorem 3.1.

###### Lemma B.1 .

### B.6 Proof of Corollary 1

In the aysmptotic regime, with sufficiently large $T$, the term $\sqrt{1+\frac{\log^{1.5}(CKd^{2}T^{2}/\delta)}{\sqrt{b}}}\eta KG$ approaches $\epsilon$, so the denominator on the LHS can be replaced with $2\epsilon$. Then the derivation is straightforward by just substituting $\eta=\frac{1}{\sqrt{T}K}$ into Theorem 3.1.

### B.7 Proof of Corollary 2

We first develop the convergence bound in Theorem 3.1 under the condition $b\geq\log^{3}(CKd^{2}T^{2}/\delta)$,

$$ $\displaystyle\left(\sqrt{2}\eta KG+\epsilon\right)^{-1}\kappa\eta K$ $\displaystyle\sum_{t=1}^{T}\|\nabla{\cal L}(x_{t})\|^{2}\leq{\cal L}(z_{1})+ \frac{1}{\epsilon}\kappa\eta^{2}LK^{2}G^{2}T$ $\displaystyle+\nu\kappa\eta K\sqrt{T}(\frac{G^{2}}{\epsilon}+\frac{\sigma}{ \epsilon}\log^{\frac{1}{2}}(\frac{2T}{\delta_{g}}))+\eta^{2}\kappa^{2}T\frac{3 2}{(1-\beta_{1})^{2}}\frac{DK^{2}G^{2}}{\epsilon^{2}},$ $$

The condition on $T\leq\frac{J_{1}-\sqrt{2}G}{\epsilon^{2}}$ is equivalent to

$$ $\displaystyle\frac{\sqrt{2}\eta KG+\epsilon}{\eta K}\leq J_{1},$ $$

since $\eta=\frac{1}{\sqrt{T}K}$. Then scaling the coefficient on the left hand side and substituting $\frac{1}{\sqrt{T}K}$ for $\eta$, we derive

$$ $\displaystyle\frac{1}{J_{1}T}\sum_{t=1}^{T}\|\nabla{\cal L}(x_{t})\|^{2}$ $\displaystyle\leq\frac{{\cal L}(z_{1})\epsilon}{\kappa T}+\frac{1}{\epsilon} \frac{LG^{2}}{T}+\frac{\nu}{T}(G^{2}+\sigma\log^{\frac{1}{2}}(\frac{2T}{\delta _{g}}))+\frac{\kappa}{T}\frac{32}{(1-\beta_{1})^{2}}\frac{DG^{2}}{\epsilon},$ $$

### B.8 A non-asymptotic bound on practical learning rates

We first state a convergence bound on using practical learning rates, which decays as the optimization procedure.

###### Theorem B.2 .

Suppose the sequence of iterates $\{x_{t}\}_{t=1}^{T}$ is generated by Algorithm 1 with a decaying learning rate $\eta_{t}=\frac{1}{\sqrt{t+T_{0}}K}$, where $T_{0}=\lceil\frac{1}{1-\beta_{1}^{2}}\rceil$. Under Assumptions 1-4, for any $T$ and $\epsilon>0$, with probability $1-\Theta(\delta)-O(\exp(-\Omega(\nu^{2})))-\delta_{g}$,

$$ $\displaystyle\sum_{t=1}^{T}\left(\sqrt{1+\frac{\log^{1.5}(CKd^{2}T^{2}/\delta) }{\sqrt{b}}}\eta_{t}JKG+\epsilon\right)^{-1}\kappa\eta_{t}\|\nabla{\cal L}(x_{ t})\|^{2}\leq{\cal L}(z_{1})+\frac{1}{\epsilon}\kappa LG^{2}\log T$ $\displaystyle+\nu\kappa\log T(\frac{\log^{1.5}(CKTd/\delta)}{\sqrt{b}}\frac{G^ {2}}{\epsilon}+\frac{\sigma}{\epsilon}\log^{\frac{1}{2}}(\frac{2T}{\delta_{g}} ))+\kappa^{2}\log T(1+\frac{\log^{1.5}(CKdT^{2}/\delta)}{\sqrt{b}})^{2}\frac{8 }{(1-\beta_{1})^{2}}\frac{DG^{2}}{\epsilon^{2}},$ $$

where $\delta,\delta_{g}$, and $\nu$ are the randomness from sketching, sub-Gaussian stochastic noise and martingales respectively, and $J$ is a constant defined in Lemma. B.3.

Alike the analysis in the constant learning rate case, we first define auxiliary variables $z_{t}$

$$ $\displaystyle z_{t}=x_{t}+\frac{\beta_{1}}{1-\beta_{1}}(x_{t}-x_{t-1})=\frac{1 }{1-\beta_{1}}x_{t}-\frac{\beta_{1}}{1-\beta_{1}}x_{t-1}.$ $$

Then, the update on $z_{t}$ can be expressed as

$$ $\displaystyle z_{t+1}-z_{t}$ $\displaystyle=\frac{1}{1-\beta_{1}}(x_{t+1}-x_{t})-\frac{\beta_{1}}{1-\beta_{1 }}(x_{t}-x_{t-1})$ $\displaystyle=\frac{\beta_{1}}{1-\beta_{1}}\left(\kappa{\hat{V}_{t-1}}^{-1/2}- \kappa\hat{V_{t}}^{-1/2}\right)m_{t-1}-\frac{\kappa\eta_{t}}{C}\hat{V_{t}}^{-1 /2}\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c}$ $$

By Taylor expansion, we have

$$ $\displaystyle{\cal L}(z_{t+1})$ $\displaystyle={\cal L}(z_{t})+\nabla{\cal L}(z_{t})^{\top}(z_{t+1}-z_{t})+ \frac{1}{2}(z_{t+1}-z_{t})^{\top}\hat{H}_{{\cal L}}(z_{t+1}-z_{t})$ $\displaystyle={\cal L}(z_{t})+\nabla{\cal L}(x_{t})^{\top}(z_{t+1}-z_{t})+( \nabla{\cal L}(z_{t})-\nabla{\cal L}(x_{t}))^{\top}(z_{t+1}-z_{t})+\frac{1}{2} (z_{t+1}-z_{t})^{\top}\hat{H}_{{\cal L}}(z_{t+1}-z_{t}).$ $$

Bounding the first-order term

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}(z_{t+1}-z_{t})$ $\displaystyle=$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}\left(\frac{\beta_{1}}{1-\beta_{1}} \left(\kappa{\hat{V}_{t-1}}^{-1/2}-\kappa\hat{V_{t}}^{-1/2}\right)m_{t-1}- \frac{\kappa\eta_{t}}{C}\hat{V_{t}}^{-1/2}\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{ \top}R_{t}g_{t,k}^{c}\right)$ $\displaystyle\leq$ $\displaystyle\frac{\beta_{1}}{1-\beta_{1}}\|\nabla{\cal L}(x_{t})\|_{\infty}( \|\kappa{\hat{V}_{t-1}}^{-1/2}\|_{1,1}-\|\kappa\hat{V_{t}}^{-1/2}\|_{1,1})\|m_ {t-1}\|_{\infty}$ $\displaystyle-\frac{\eta_{t}}{C}\nabla{\cal L}(x_{t})^{\top}(\kappa\hat{V_{t}} ^{-1/2}-\kappa{\hat{V}_{t-1}}^{-1/2})\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R _{t}g_{t,k}^{c}-\frac{\kappa\eta_{t}}{C}\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{ t-1}}^{-1/2}\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c}$ $\displaystyle\leq$ $\displaystyle\left(\frac{\beta_{1}}{1-\beta_{1}}\|m_{t-1}\|_{\infty}+\frac{ \eta_{t}}{C}\|\sum_{c=1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c}\|_{ \infty}\right)\|\nabla{\cal L}(x_{t})\|_{\infty}(\|\kappa{\hat{V}_{t-1}}^{-1/2 }\|_{1,1}-\|\kappa\hat{V_{t}}^{-1/2}\|_{1,1})$ $\displaystyle-\frac{\kappa\eta_{t}}{C}\sum_{c=1}^{C}\sum_{k=1}^{K}\nabla{\cal L }(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}R_{t}^{\top}R_{t}g_{t,k}^{c}.$ $$

The quadratic terms can be written as

$$ $\displaystyle(\nabla{\cal L}(z_{t})-\nabla{\cal L}(x_{t}))^{\top}(z_{t+1}-z_{t })=(z_{t}-x_{t})^{\top}\hat{H}_{{\cal L}}(\frac{1}{1-\beta_{1}}(x_{t+1}-x_{t}) -\frac{\beta_{1}}{1-\beta_{1}}(x_{t}-x_{t-1})),$ $$

where $\hat{H}_{\cal L}$ is a second-order Taylor remainder.

To bound the quadratic term, the counterpart of Lemma B.1 can be stated as

###### Lemma B.3 .

With learning rate $\eta_{t}=O(\frac{1}{\sqrt{t+T_{0}}})$, where $T_{0}=\lceil\frac{1}{1-\beta_{1}^{2}}\rceil$. Denote $J=\frac{1-\beta_{1}}{\sqrt{T_{0}+1}}/(\frac{1}{\sqrt{T_{0}+1}}-\frac{\beta_{1}
}{\sqrt{T_{0}}}).$ Then with probability $1-\Theta(t\delta)$,

$$ $\displaystyle|m_{t-1}^{\top}h|\leq(1+\frac{\log^{1.5}(CKd/\delta)}{\sqrt{b}})JKGH$ $$

###### Proof.

For $t=0$, since $m_{0}=0$, the inequality holds. Suppose we have for $h\in\mathbb{R}^{d}$, s.t. $\|h\|\leq H$, with probability $1-\Theta((t-1)\delta)$,

$$ $\displaystyle|m_{t-1}^{\top}h|\leq(1+\frac{\log^{1.5}(CKd/\delta)}{\sqrt{b}})JKGH$ $$

By the update rule,

$$ $\displaystyle|m_{t}^{\top}h|$ $\displaystyle=|(\beta_{1}\cdot m_{t-1}+(1-\beta_{1})\cdot\frac{\eta}{C}\sum_{c =1}^{C}\sum_{k=1}^{K}R_{t}^{\top}R_{t}g_{t,k}^{c})^{\top}h|$ $\displaystyle\leq\beta_{1}|m_{t-1}^{\top}h|+\frac{(1-\beta_{1})\eta}{C}\sum_{c =1}^{C}\sum_{k=1}^{K}|\langle R_{t}^{\top}R_{t}g_{t,k}^{c},h\rangle|$ $\displaystyle\leq\beta_{1}|m_{t-1}^{\top}h|+(1-\beta_{1})(1+\frac{\log^{1.5}( CKd/\delta)}{\sqrt{b}})\eta_{t}\sum_{k=1}^{K}\|g_{t,k}^{c}\|_{2}\|h\|_{2}$ $\displaystyle\leq(1+\frac{\log^{1.5}(CKd/\delta)}{\sqrt{b}})\eta_{t}JKGH,~{}~{ }w.p.~{}1-\Theta(t\delta).$ $$

∎

By exactly the same as in Sec. B.3, we can lower bound the scaled gradient term by

$$ $\displaystyle\nabla{\cal L}(x_{t})^{\top}{\hat{V}_{t-1}}^{-1/2}\nabla{\cal L}( x_{t})$ $\displaystyle\geq\min_{i}(\hat{V}_{t-1}^{-1/2})_{i}\sum_{i=1}^{d}[\nabla{\cal L }(x_{t})]_{i}^{2}$ $\displaystyle\geq(\sqrt{1+\frac{\log^{1.5}(CKtd/\delta)}{\sqrt{b}}}\eta_{t}KG+ \epsilon)^{-1}\|\nabla{\cal L}(x_{t})\|^{2},~{}~{}w.p.~{}1-\Theta(d\delta).$ $$

On the martingale of zero-centered noises, we can simply incorporate the learning rate $\eta_{t}$ into the martingale. Define the random process of sketching noise as

$$ $\displaystyle\{Y_{t}=\sum_{\tau=1}^{t}\frac{\eta_{\tau}}{C}\sum_{k=1}^{K} \nabla{\cal L}(x_{\tau})^{\top}{\hat{V}_{\tau-1}}^{-1/2}(R_{\tau}^{\top}R_{ \tau}g_{\tau,k}^{c}-g_{\tau,k}^{c})\}_{t=1}^{T}$ $$

as a martingale. The difference of $|Y_{t}-Y_{t-1}|$ is bounded with high probability

$$ $\displaystyle|Y_{t}-Y_{t-1}|$ $\displaystyle=|\frac{\eta_{t}}{C}\sum_{c=1}^{C}\sum_{k=1}^{K}\nabla{\cal L}(x_ {t})^{\top}{\hat{V}_{t-1}}^{-1/2}(R_{t}^{\top}R_{t}g_{t,k}^{c}-g_{t,k}^{c})|$ $\displaystyle\leq\frac{\log^{1.5}(d/\delta)}{\sqrt{b}}\eta_{t}KG\|\hat{V}_{t}^ {-1/2}\nabla{\cal L}(x_{t})\|_{2},~{}~{}w.p.~{}1-\Theta(CK\delta).$ $$

Then by Azuma’s inequality,

$$ $\displaystyle\mathbb{P}(|Y_{T}|\geq\nu\sqrt{\sum_{t=1}^{T}\left(\frac{\log^{1. 5}(d/\delta)}{\sqrt{b}}\eta_{t}KG\|\hat{V}_{t}^{-1/2}\nabla{\cal L}(x_{t})\|_{ 2}\right)^{2}})=O({\rm exp}(-\Omega(\nu^{2})))+T\delta$ (2) $$

A similar bound can be achieved for the sub-Gaussian noise in stochastic gradient. Let

$$ $\displaystyle Z_{t}=\sum_{\tau=1}^{t}\frac{\eta_{\tau}}{C}\sum_{k=1}^{K}\nabla {\cal L}(x_{\tau})^{\top}{\hat{V}_{\tau-1}}^{-1/2}(g_{\tau,k}^{c}-\nabla{\cal L }^{c}(x_{t,k}^{c})).$ $$

Then

$$ $\displaystyle\mathbb{P}(|Z_{T}|\geq\nu\sqrt{\sum_{t=1}^{T}(\frac{\eta_{t} \sigma}{\epsilon})^{2}\log(\frac{2T}{\delta_{g}})})=O({\rm exp}(-\Omega(\nu^{2 })))+\delta_{g}$ $$

Combining the two bounds by union bound completes the proof.

###### Theorem B.2 .

###### Lemma B.3 .

###### Proof.

## Appendix C Proof of Theorem 4.2

### C.1 Proof of Lemma 4.1

Denote $\Delta_{t}^{c}=\sum_{k=1}^{K}g_{t,k}^{c}$, $\tilde{\Delta}_{t}^{c}=\min\{1,\frac{\tau}{\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta
_{t}^{c}\|}\}\Delta_{t}^{c}$. Then $x_{t+1}-x_{t}=-\kappa\eta R^{\top}R\frac{1}{{C}}\sum_{c=1}^{C}\tilde{\Delta}_{
t}^{c}$.

###### Proof.

Taking the expectation of randomness in stochastic gradient yields

$$ $\displaystyle\mathbb{E}[{\cal L}(x_{t+1})]-{\cal L}(x_{t})=-\kappa\eta\langle \nabla{\cal L}(x_{t}),\frac{1}{C}R^{\top}R\sum_{c=1}^{C}\mathbb{E}[\tilde{ \Delta}_{t}^{c}]\rangle+\frac{\kappa^{2}\eta^{2}}{2}\mathbb{E}[(\frac{1}{C}R^{ \top}R\sum_{c=1}^{C}\tilde{\Delta}_{t}^{c})^{\top}\hat{H}_{\cal L}(\frac{1}{C} R^{\top}R\sum_{c=1}^{C}\tilde{\Delta}_{t}^{c})]$ $\displaystyle=-\kappa\eta\langle\nabla{\cal L}(x_{t}),\frac{1}{C}\sum_{c=1}^{C }R^{\top}R\mathbb{E}[\tilde{\Delta}_{t}^{c}]-\mathbb{E}[\tilde{\Delta}_{t}^{c} ]\rangle-\kappa\eta\langle\nabla{\cal L}(x_{t}),\frac{1}{C}\sum_{c=1}^{C} \mathbb{E}[\tilde{\Delta}_{t}^{c}]\rangle$ $\displaystyle+\frac{\kappa^{2}\eta^{2}}{2}\mathbb{E}[(\frac{1}{C}R^{\top}R\sum _{c=1}^{C}\tilde{\Delta}_{t}^{c})^{\top}\hat{H}_{\cal L}(\frac{1}{C}R^{\top}R \sum_{c=1}^{C}\tilde{\Delta}_{t}^{c})]$ $\displaystyle\leq\frac{\kappa\eta K}{2}\frac{\log^{1.5}(d/\delta)}{\sqrt{b}}\| \nabla{\cal L}\|^{2}+\frac{\kappa\eta}{2K}\frac{\log^{1.5}(d/\delta)}{\sqrt{b} }\|\frac{1}{{C}}\sum_{c=1}^{C}\mathbb{E}[\tilde{\Delta}_{t}^{c}]\|^{2}-\frac{ \kappa\eta K}{2}\|\nabla{\cal L}\|^{2}-\frac{\kappa\eta}{2K}\|\frac{1}{{C}} \sum_{c=1}^{C}\mathbb{E}[\tilde{\Delta}_{t}^{c}]\|^{2}$ $\displaystyle+\frac{\kappa\eta K}{2}\|\nabla{\cal L}-\frac{1}{K}\frac{1}{{C}} \sum_{c=1}^{C}\mathbb{E}[\tilde{\Delta}_{t}^{c}]\|^{2}+\frac{\kappa^{2}\eta^{2 }}{2}\mathbb{E}[(\frac{1}{C}R^{\top}R\sum_{c=1}^{C}\tilde{\Delta}_{t}^{c})^{ \top}\hat{H}_{\cal L}(\frac{1}{C}R^{\top}R\sum_{c=1}^{C}\tilde{\Delta}_{t}^{c})]$ $\displaystyle\leq-(1-\frac{\log^{1.5}(d/\delta)}{\sqrt{b}})\frac{\kappa\eta K} {2}\|\nabla{\cal L}\|^{2}+\frac{\kappa\eta K}{2}\|\nabla{\cal L}-\frac{1}{K} \frac{1}{{C}}\sum_{c=1}^{C}\mathbb{E}[\tilde{\Delta}_{t}^{c}]\|^{2}$ $\displaystyle+\frac{\kappa^{2}\eta^{2}}{2}\mathbb{E}[(\frac{1}{C}R^{\top}R\sum _{c=1}^{C}\tilde{\Delta}_{t}^{c})^{\top}\hat{H}_{\cal L}(\frac{1}{C}R^{\top}R \sum_{c=1}^{C}\tilde{\Delta}_{t}^{c})]$ $\displaystyle\leq-\frac{\kappa\eta K}{4}\|\nabla{\cal L}\|^{2}+\frac{\kappa \eta K}{2}\|\nabla{\cal L}-\frac{1}{K}\frac{1}{{C}}\sum_{c=1}^{C}\mathbb{E}[ \tilde{\Delta}_{t}^{c}]\|^{2}$ $\displaystyle+\frac{\kappa^{2}\eta^{2}}{2}\mathbb{E}[(\frac{1}{C}R^{\top}R\sum _{i=1}^{C}\tilde{\Delta}_{t,i}^{c})^{\top}\hat{H}_{\cal L}(\frac{1}{C}R^{\top} R\sum_{i=1}^{C}\tilde{\Delta}_{t,i}^{c})],~{}~{}w.p.~{}1-\Theta(\delta)$ $$

where the first inequality is directly from Lemma A.1. The second and last inequalities are from the condition of $b\geq 4\log^{3}(d/\delta)$.
∎

###### Proof.

### C.2 Proof of Theorem 4.2

The first order term in Lemma 4.1 can be handled by

$$ $\displaystyle\|\nabla{\cal L}-\frac{1}{K}\frac{1}{{C}}\sum_{c=1}^{C}\mathbb{E} [\tilde{\Delta}_{t}^{c}]\|\leq\|\nabla{\cal L}-\frac{1}{K}\frac{1}{{C}}\sum_{c =1}^{C}\mathbb{E}[\Delta_{t}^{c}]\|+\frac{1}{K}\|\frac{1}{{C}}\sum_{c=1}^{C} \mathbb{E}[\Delta_{t}^{c}]-\frac{1}{{C}}\sum_{c=1}^{C}\mathbb{E}[\tilde{\Delta }_{t}^{c}]\|$ $\displaystyle\leq\frac{\eta L}{KC}\sum_{c=1}^{C}\sum_{i=1}^{K}\mathbb{E}[\| \nabla{\cal L}_{t,k}^{c}\|]+\frac{1}{KC}\sum_{c=1}^{C}\mathbb{E}[\|\Delta_{t}^ {c}\|1_{\{\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta_{t}^{c}\|\geq\tau\}}]$ $\displaystyle\leq\eta KLG+K^{\alpha-1}G^{\alpha}\tau^{1-\alpha},$ $$

where the last inequality follows by

$$ $\displaystyle\frac{1}{C}\sum_{c=1}^{C}\mathbb{E}[\|\Delta_{t}^{c}\|1_{\{\frac{ 1}{{C}}\sum_{c=1}^{C}\|\Delta_{t}^{c}\|\geq\tau\}}]=\mathbb{E}[\frac{1}{{C}} \sum_{c=1}^{C}\|\Delta_{t}^{c}\|1_{\{\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta_{t}^{ c}\|\geq\tau\}}]$ $\displaystyle=\mathbb{E}[(\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta_{t}^{c}\|)^{ \alpha}(\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta_{t}^{c}\|)^{1-\alpha}1_{\{\frac{1} {{C}}\sum_{c=1}^{C}\|\Delta_{t}^{c}\|\geq\tau\}}]\leq(KG)^{\alpha}\tau^{1- \alpha}.$ $$

The second order term can be handled as follows.
With probability $1-\Theta(\delta)$,

$$ $\displaystyle\mathbb{E}[(\frac{1}{C}R^{\top}R\sum_{i=1}^{C}\tilde{\Delta}_{t,i }^{c})^{\top}\hat{H}_{\cal L}(\frac{1}{C}R^{\top}R\sum_{i=1}^{C}\tilde{\Delta} _{t,i}^{c})]=\mathbb{E}[\sum_{j=1}^{d}\lambda_{j}\langle\frac{1}{C}R^{\top}R \sum_{i=1}^{C}\tilde{\Delta}_{t,i}^{c},v_{j}\rangle^{2}]$ $\displaystyle\leq\mathbb{E}[\sum_{j=1}^{d}\lambda_{j}1_{\lambda_{j}\geq 0} \langle\frac{1}{C}R^{\top}R\sum_{i=1}^{C}\tilde{\Delta}_{t,i}^{c},v_{j}\rangle ^{2-\alpha}\langle\frac{1}{C}R^{\top}R\sum_{i=1}^{C}\tilde{\Delta}_{t,i}^{c},v _{j}\rangle^{\alpha}]$ $\displaystyle\leq\mathbb{E}\left[\sum_{j=1}^{d}\lambda_{j}1_{\lambda_{j}\geq 0 }\left((1+\frac{\log^{1.5}(d/\delta)}{\sqrt{b}})\frac{\tau}{\frac{1}{{C}}\sum_ {c=1}^{C}\|\Delta_{t}^{c}\|}\frac{1}{{C}}\sum_{c=1}^{C}\|\Delta_{t}^{c}\| \right)^{2-\alpha}\left(\frac{1}{C}(1+\frac{\log^{1.5}(d/\delta)}{\sqrt{b}}) \sum_{i=1}^{C}\|\Delta_{t}^{c}\|\right)^{\alpha}\right]$ $\displaystyle\leq(1+\frac{\log^{1.5}(d/\delta)}{\sqrt{b}})DK^{2}\tau^{2-\alpha }G^{\alpha},$ $$

where the first equation follows by using the eigen-decomposition of $\hat{H}_{\cal L}$ and the second order term can be reduced to a squared inner product term. The primary trick thereafter (in the first inequality) is to split the inner product terms into two parts, which can be handled by the two-sided adaptive learning rates respectively. By applying the bounded second moment of random sketching, we find the first part with order $2-\alpha$ is contained in a $(1+\frac{\log^{1.5}(d/\delta)}{\sqrt{b}})\tau-$ball with high probability, and the second part with order $\alpha$ is bounded by applying Assumption 5. Then Theorem 4.2 follows by combining the first order term and the second-order term by union bounds, as well as applying the condition of $b\geq 4\log^{3}(d/\delta)$.

## Appendix D Experimental Details and Additional Results

Figure: Figure 5: The power-law structure of the Hessian spectrum on LeNet. Quoted from Fig.1 Xie et al., ().
Refer to caption: extracted/5993506/arxiv/imgs/hessian_ping_li.png

Figure: Figure 6: Eigenspectrum density every 5 epochs. The model is ViT-Small and trained on CIFAR10. The majority of eigenvalues concentrates near 0 and the density enjoys a super fast decay with the absolute values of eigenvalues, indicating a summable eigenspectra.
Refer to caption: x6.png

Figure: Figure 7: Comparing the performance of tiny sketch sizes on ResNet and BERT. The experiment settings are the same as in Fig. 1 and Fig. 3
Refer to caption: x7.png

Figure: (a) Heavy tailed gradient norms.
Refer to caption: x9.png

Aside from the experimental configurations described in the main paper, we provide additional details. We use Cross Entropy with label smoothing as the loss function. The parameter for label smoothing is 0.1. We use a cosine learning rate scheduler on the server optimizer, with the minimal learning rate is $1e-5$. Client batch size is 128, and weight decay is $1e-4$. For SGD and SGDm methods, the learning rate is 1.0. For SGDm, the momentum is 0.9. For Adam optimizer, the learning rate is 0.01, and the momentum is 0.9. The learning rates are tuned to achieve the best performance.

Our experiments were conducted on a computing cluster with AMD EPYC 7713 64-Core Processor and NVIDIA A100 Tensor Core GPU.

To verify Assumption 4, we plot the full Hessian eigenspectrum throughout the training process in Fig. 6. We used stochastic lanczos algorithm implemented by the pyHessian library Yao et al., () to approximate the distribution of the full eigenspectrum. Our main claim in Assumption 4 is that the Hessian eigenspectrum at an iterate is summable and the sum is independent of the ambient dimension, which can be satisfied by common distributions, like power-laws. We run testing experiments on ViT-small and train on CIFAR-10 dataset, with sketched Adam optimizer. In Fig. 6, we see the majority of eigenvalues concentrates near 0. The density enjoys a super fast decay with the absolute values of eigenvalues. The decay also holds throughout the training process. This empirical evidence shows the validity of our assumption.

In the main body of the paper, we have achieved 99.9% compression rate and 99.98% compression rate for ResNet and BERT respectively. We further include the results on smaller $b$ in Fig. 7. In principle, an extremely tiny sketch size (with 400 in vision tasks and 2000 in language tasks) still converges but generates an unfavorable local minima that hardly generalizes.

In the following, we present another empirical result on the clipping method. The goal here is to show the superiority of (sketched) clipping methods over the plain SGD optimizer. We run BERT model on SST2 dataset. The dataset is split among 5 distinct clients in an i.i.d way. The normalization factor in the clipping method is set as 0.03. In Figure. 8, we show that (sketched) clip SGD method has better performance in convergence and generalization, while the plain SGD method fails to converge. It is also observed that sketching does not cause degrade in testing error.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: See Weakness

+ Error Feedback 
Can you additionally compare w sketching + EF 

+ Overall I feel like the paper is simply analysis of 
Adaptive FL + Sketching + clipping ~ 
which was later packaged as a new algorithm. 
Combining different techniques => Low Novelty.
soundness: 2
strengths: The paper analyzes Adaptive Opt + Grad Sketching. 

While adaptive optimizers (e.g., Adam) and sketching techniques have been individually studied in FL, their combined use has been relatively unexplored. 

Convergence analyses showing that communication costs scale logarithmically with model dimensionality, which contrasts with prior work where communication costs are linearly dependent on model size. 

This logarithmic dependency is critical for scalability and suggests that SAFL could be feasible for modern deep learning models with millions or billions of parameters. 

Notably, the analysis leverages recent findings on the anisotropic Hessian spectrum in deep networks, showcasing an innovative application of this structural property in theoretical work.

The dual versions of the framework, SAFL (for i.i.d. or mild noise) and SACFL (for heavy-tailed noise), allow it to address both homogeneous and heterogeneous data environments, which is essential for federated scenarios with variable client data. This flexibility increases its potential applicability to real-world FL problems, where client data distributions often deviate from the i.i.d. assumption.
Experimental Validation Across Domains:

The paper presents a range of experiments on tasks in both vision (CIFAR-10) and NLP (SST2), validating that the framework performs well on distinct model architectures (ResNet, Vision Transformer, BERT) and data regimes. The empirical results show that SAFL maintains comparable performance to full-dimensional optimizers, which supports its practical relevance.
summary: This paper introduces the Sketched Adaptive Federated Learning (SAFL) framework, aiming to alleviate communication burdens in federated learning (FL) by leveraging gradient sketching alongside adaptive optimizers. The authors propose that by combining techniques like CountSketch and quantization with adaptive optimizers (e.g., Adam, AMSGrad), FL can achieve improved communication efficiency. The framework is extended to handle non-i.i.d. settings via a clipped variant, SACFL, which uses gradient clipping to ensure stability even with heavy-tailed noise. The paper’s convergence results are claimed to scale logarithmically with respect to model size, which is theoretically significant for applications with large model parameters. Additionally, experiments in vision and NLP domains provide empirical validation, showing that SAFL and SACFL can maintain comparable accuracy while achieving higher compression rates.
weaknesses: Dependency on Assumptions Regarding Hessian Eigen Spectrum Decay:

The convergence guarantees rely heavily on the assumption that the Hessian’s eigen spectrum decays sharply (anisotropy in curvature). Although this phenomenon is observed in many neural network architectures, it is not universally applicable, particularly in certain tasks, models, or data distributions. For instance, models used on structured data or non-standard architectures may not exhibit this property, limiting SAFL’s applicability in such cases. A more in-depth discussion on these limitations, or an empirical analysis of SAFL’s performance under weakly decaying Hessian spectrums, would clarify the practical applicability of the theoretical results.
Insufficient Baseline Comparisons and Ablation Studies:

The experimental comparisons are mostly limited to other communication-efficient algorithms like FetchSGD and 1bit-Adam, without sufficient benchmarking against non-adaptive FL methods (e.g., FedAvg) or ablation studies isolating the contributions of sketching and adaptivity. This makes it challenging to discern the individual impacts of gradient sketching and adaptive optimizers on communication efficiency and convergence stability. Including a wider range of baselines and detailed ablation studies would better contextualize the advantages of SAFL.

Scalability and Stability in Real-World FL Environments:

The experiments focus on setups with small client numbers and controlled conditions. Real-world FL often involves thousands of clients with highly heterogeneous data, variable network conditions, and diverse computational resources. The scalability of SAFL to large client pools, as well as its robustness under more realistic conditions, remains uncertain. A thorough examination of SAFL’s performance under these conditions, particularly with varied network latencies and client-side hardware capacities, would enhance confidence in its practical relevance.

Trade-offs Between Compression and Model Quality:

The paper briefly notes that SAFL performs well under high compression rates, but a deeper analysis of the trade-offs between compression rates and model performance is missing. High compression could potentially degrade convergence rates or accuracy, especially for tasks that are sensitive to gradient fidelity. A more comprehensive exploration of how compression affects different aspects of training—such as convergence speed, model accuracy, and robustness—would be informative for practitioners aiming to apply SAFL in resource-constrained environments.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The power law decay of the spectrum is only valid after a few epochs of SGD. Given this, is assuming that assumption 4 is true throughout the training process valid? Figure 6 shows an eigenvalue decay but does not necessarily confirm that the decay follows a power law.

2. Typically, SGD outperforms Adam in vision tasks with convnets. However, this does not seem to be the case in SAFL. Why is this occurring?

3. How feasible would it be to extend the proof to relax assumptions 1 and 2 to affine conditions, i.e. the gradient is bounded by an affine function?

4. On real data, we can see that Adam outperforms SGD when used with this framework. However, does this still occur on synthetic data when comparing AMSGrad to SGD to show that the improvement in corollary 2 is due to adaptivity? Or is there a more straightforward explanation?

5. Could the authors include the sketching algorithm used in the experiment?

6. If FetchSGD can be used with Adam, does SAFL still outperform the new “FetchAdam” algorithm?

(I am willing to raise the score if the weakness referred to above and questions are well addressed)
soundness: 4
strengths: The paper provides theoretical results on the convergence rates of the algorithm in both the i.i.d. and non-i.i.d. client settings by utilizing a novel assumption observed in other recent work. It is interesting to see strong assumptions validated in practice, and the resulting theory is of independent interest to other stochastic optimization results for deep learning. The proposed framework and analysis are highly flexible and natural and can be used with various optimizers and sketching algorithms. The framework is tested on different types of models and datasets.

The paper reads clearly and is well-motivated.
summary: This paper proposes a framework for sketched adaptive federated learning (SAFL), a method for communication-efficient federated learning. In the (nearly) i.i.d. client setting, they prove a 1/sqrt(T) convergence rate for SAFL with a 1/T rate near initialization, with low (logarithmic) communication costs by using a novel assumption on the decay of the eigenspectrum of Hessians in deep neural networks. In the non i.i.d. setting, they prove convergence, even in the presence of heavy-tailed noise.

Experiments show improvements using SAFL over existing communication-efficient federated learning methods when tested on the CIFAR-10 dataset and the SST2 task.
weaknesses: Assumptions 1 and 2, while standard, are strong and hard to ensure in practice. The claim that SAFL enjoys better performance near initialization due to adaptivity feels unmotivated. It is also unclear if performance improvements are primarily due to using Adam over SGD, which is used in FetchSGD. Finally, the sketching algorithm used for the experiments is unclear.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 4
questions: 1. What are the technical challenges of reducing the dependence of $d$ from polynomial to logarithmic, other than applying Assumption 4? Does the proof require any new techniques compared to previous work?
2. What is the motivation for simultaneously assuming bounded gradients (Assumption 1) and decaying Hessian eigenspectrum (Assumption 4)? Do you have any natural examples of ML training objectives that satisfy these two assumptions simultaneously?
3. It seems possible that $\hat{L}$ depends on $d$. Are there any experiments in the literature investigating how $\hat{L}$ changes with $d$?
4. What is the communication cost of FetchSGD compared to the other evaluated algorithms? Did you control for communication cost when comparing your algorithm to theirs?
5. Line 079: What does almost iid mean? This phrase is used a couple times without definition, then is not referenced again.
6. Line 234: "The i.i.d. data assumption leads to the bounded gradient assumption". I don't understand this claim. Many applications have i.i.d. data and unbounded gradients, for example if the global objective is quadratic.
7. What was the protocol for hyperparameter tuning? Is it possible that some algorithms were tuned more than others? How many random seeds were used to evaluate each algorithm?
8. What does Figure 4(a) show exactly? Local stochastic gradients sampled at the same point? Local stochastic gradients sampled at different points? Or local gradients sampled at different points?
soundness: 3
strengths: 1. The paper is well-written. The theoretical motivation and contributions of this work are all clearly stated.
2. The high-level problem of reducing communication cost in federated learning is important, especially when training deep networks.
summary: This paper introduces Sketched Adaptive Federated Learning (SAFL), which combines gradient sketching methods in FL with adaptive optimizers (e.g. Adam). It is proven that the proposed framework can converge with (1) sub-gaussian noise and (2) heavy-tailed noise (with a clipped variant), and the dependence on the dimension $d$ is only logarithmic. Experiments show that Sketched Adam is on par or better than other sketched baselines for training ResNet, ViT, and BERT.
weaknesses: 1. The improved theoretical results compared to prior work may be a direct result of additional assumptions.

    1a. The biggest assumption here compared to prior work is the decay of the Hessian eigenspectrum (Assumption 4), and it seems that the reduced dependence on $d$ (from polynomial to logarithmic) may come as an immediate consequence of this assumption (Lines 299-313). The reduced $d$ dependence is the main theoretical improvement claimed by the authors, but I'm not sure whether there are any real technical challenges to achieve this reduced dependence, other than applying Assumption 4.

    1b. It is unclear whether the constant $\hat{L}$ changes with $d$ in practice, which seems very possible to me. If so, then there is really some hidden dependence on $d$ in the constant $\hat{L}$.

    1c. The combination of assumptions is not well motivated, neither theoretically nor empirically. Specifically, the paper assumes bounded gradients (Assumption 1) together with Hessian eigenspectrum decay (Assumption 4). Deep neural networks may satisfy Assumption 4, but Assumption 1 definitely does not hold for such networks. On the other hand, some classic optimization problems in ML satisfy Assumption 1 (e.g. logistic regression), but such examples fall outside of deep learning, which is the motivating example for Assumption 4. I don't see any reason to consider these assumptions together other than ease of analysis.

2. The experiments don't contain significant contributions.

    2a. Sketched Adam does outperform other sketched optimizers (e.g. sketched sgdm), but I don't interpret this as a contribution of this paper, since this is essentially a comparison of Adam against SGDM and SGD. The improvement against prior work (1-bit Adam, FetchSGD), is a minor contribution to me, since the novelty of Sketched Adam is somewhat low.

    2b. The experiments fail to compare against any prior work from the last three years. 1-bit Adam and FetchSGD are both old algorithms at this point. The other parts of the paper cite Efficient-Adam (Chen et al, 2022), Compressed AMSGrad (Wang et al, 2022), CocktailSGD (Wang et al, 2023), and many others, but recent baselines are missing from the experimental comparison.

    2c. The methodology around evaluation is somewhat vague. Line 1351 says "The learning rates are tuned to achieve the best performance", and this is all of the information about hyperparameter tuning. It seems possible that there was no strict hyperparameter protocol, which invites the possibility that some algorithms were tuned more than others. Also, there is no mention of multiple random seeds, so it seems that all experiments include results over a single random seed.

3. Weak motivation for heavy-tailed noise in federated learning. A constant theme throughout the paper is the importance of heavy-tailed noise in federated learning, as observed in (Yang et al, 2022). However, I am skeptical about this connection. (Yang et al, 2022) observes a heavy-tailed distribution of \textit{local client updates aggregated across an entire round}, but they use this as justification to assume a heavy tailed distribution of stochastic gradient noise \textit{at a single point}. I don't agree with this justification at all, which appears to confuse the distribution of local client updates aggregated across an entire round with the distributino of stochastic gradient noise at each individual point. The current submission uses this observation as inspiration for the connection between FL with heterogeneous data and heavy-tailed noise, e.g. on Line 048, where it is claimed that adaptive methods are better for FL with heterogeneous data due to heavy-tailed noise. These claims are not justified, in my opinion, and this severely undermines the motivation for a large part of this submission.

4. Unclear/missing comparisons with previous work.

    4a. The paper contains very little comparison of the current theoretical work against prior work. The only such comparison is the claim that previous works have a linear dependence on $d$, compared to the logarithmic dependence on $d$ achieved by this paper. However, the convergence rates achieved by the previous papers are not explicitly written here, and the explanation provided for the $d$ dependence of prior work (Lines 074-076) relies on the condition of constant per-round communication cost as $d$ scales. It is unclear whether this is an apples-to-apples comparison against the convergence rate of the current work. I recommend that the authors include a more clear, direct comparison of their convergence rate against that of prior works.

    4b. The experiments are missing many baselines (see Weakness #2).

    4c. The discussion of related work on heavy tailed noise (Lines 355-365) should be amended. First, I did not see any mention of heavy tailed noise in (Charles et al, 2021), so I recommend a more specific discussion of how this work related to yours. Also, I recommend clarifying the sentence at the end of this paragraph, specifically the claim of optimality. (Zhang, 2020) indeed show optimality of certain methods in the centralized case, but they don't consider the distributed case. (Liu, 2022) consider gradient clipping in the distributed case, but they do not consider heavy tailed noise or make claims of optimality.

Formatting/typos:
- The abstract should be limited to one paragraph, as written in the submission instructions.
- Line 081: "depends only on a logarithmically on the ambient dimension $d$". Should say "depends only logarithmically"?
- Line 133: "Update paramters" should say "Update parameters"
- Line 193: In property 2, I recommend modifying the definition of $\texttt{sk}$ and $\texttt{desk}$ to include $R$ and $R_t$ as parameters over which we can take expectation. Right now the expectation in property 2 is not very rigorously defined, since there are no random variables appearing in property 2 over which we could take expectation.
- Figure 1 is a little confusing to read, since the first two plots show one compression rate, the third plot shows a different compression rate, and the last plot shows results for a single algorithm with many compression rates. It would be helpful if this information were more organized. Also, I recommend using colors consistently to represent a single algorithm across different plots, as much as possible. For example, your Figure 1 uses green/blue for Adam in the first/second plot, respectively, then all colors for Adam in the last plot. Also, orange is used for sgdm in the first plot, then for 1 bit Adam in the second plot.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: See weaknesses. Also, how does bounded gradient norm or bounded $\alpha$-moment play a role in the analysis? Is it possible to remove them while retaining nuclear norm bound assumption, and achieve a result similar to [Song et al., 2023] by replacing the $dL$ factor with $\hat L$?
soundness: 3
strengths: This paper greatly extends the theoretical analysis of [Rothchild et al., 2020; Song et al., 2023] that uses sketching to improve the communication cost. Specifically, it gives a very general analysis in which adaptive optimizers are allowed and it could incorporate elementwise second moment of the sketched gradient. The analysis is involved as there are 3 types of errors: sketching error, local steps taken by client training and client data distribution. In particular, allowing client to do adaptive optimization coupled with SGD faces many complications caused by sketching. This paper gives a good error decomposition and uses a Martingale to handle the challenge. It further considers the setting where the bounded gradient norm (square) is relaxed to bounded $\alpha$-moment for $\alpha\in (1, 2]$. It is important to improve the dimension dependence for the convergence in the previous work, as what is observed in practice is that the sketching dimension has a far less impact on the overall convergence than causing a linear in dimension blowup. Finally, experiments are performed to show the effectiveness of the theory.
summary: This paper considers sketching-based federated learning algorithms, the main contributions are theoretical improvements over [Rothchild et al., 2020; Song et al., 2023] and empirically validations of their theory. It proposes an FL algorithm that utilizes sketching to compress the communication while enabling advanced optimizers beyond SGD. From a convergence rate perspective, it shows that under stronger assumptions, it is possible to achieve ${\rm poly}\log d$ dependence in the iteration complexity instead of $d$ in previous works. It also performs a thorough study on the performance of mixing sketching with stochastic gradient plus noisy second order information. Empirically, it shows superior performance over [Rothchild et al., 2020] on CIFAR10 and GLUE.
weaknesses: It is worth noting that removal of linear dependence on $d$ is a consequence of stronger assumption rather than an advancement of analysis, as in previous works, they only assume basic smoothness and bounded gradient (note that [Song et al., 2023] didn't assume bounded gradient) which only says that the spectral norm of the Hessian is at most $L$, while in this work the Hessian matrix eigenspectrum assumption in addition requires the nuclear norm of the Hessian is at most some $\hat L$. Note that in the worst case, $\hat L=d L$, which I believe is the result of [Song et al., 2023]. Arguably, in many scenarios it is true that the Hessian only has a few large eigenvalues, but I think it's important to point out the impact of the assumption, and compare \& contrast with previous works about them. I would encourage the authors to further clarify and compare the results about assumptions and convergence in a table with previous works (e.g., for comparing with [Rothchild et al., 2020], one could use gradient instead of stochastic gradient, take the optimizer to be the identity mapping, then does the convergence improves upon it due to nuclear norm assumption? For comparing with [Song et al., 2023], one might consider further remove the bounded gradient assumption and see whether the analysis still goes through).

Meta-review: This paper introduces the Sketched Adaptive Federated Learning (SAFL) framework, aiming to alleviate communication burdens in federated learning (FL) by leveraging gradient sketching alongside adaptive optimizers. The authors propose that by combining techniques like CountSketch and quantization with adaptive optimizers (e.g., Adam, AMSGrad), FL can achieve improved communication efficiency. The framework is extended to handle non-i.i.d. settings via a clipped variant, SACFL, which uses gradient clipping to ensure stability even with heavy-tailed noise. The paper’s convergence results are claimed to scale logarithmically with respect to model size, which is theoretically significant for applications with large model parameters. Additionally, experiments in vision and NLP domains provide empirical validation, showing that SAFL and SACFL can maintain comparable accuracy while achieving higher compression rates.


Main criticism:

1. Sketching belongs to the class of unbiased compressor with conically bounded variance, and hence all papers (hundreds of works) on communication compression using such compressors includes sketching as a special case. The authors do not compare to the key works in this very closely related literature. For example, for smooth nonconvex problems, the Marina method with the so-called permutation (PermK) compressor of Szlendak et al, Permutation Compressors for Provably Faster Distributed Nonconvex Optimization, arXiv:2110.03300, 202, achieves O(1/T) rate for the convergence of the square norm of the gradient, where the rate is independent of the dimension in the case when the number of workers ($C$) is larger than the number of dimensions ($d$)! In FL, the number of workers is very large. On the other hand, if $d>C$, the constant in the rate is $d/C$, and this will be equal to $\log d$ if $C = d/\log(d)$. This is true in the nonstochastic case - extensions to the stochastic case were considered in Gorbunov et al, MARINA: Faster Non-Convex Distributed Learning with Compression, ICML 2021 (see their VR-MARINA method which obtains $O(1/T)$ rate even in the stochastic case, unlike the $O(1/\sqrt{T})$ rate in the reviewed paper).

2. The authors claim, in several parts of the paper, including the abstract, that their results are competitive with the state of the art error feedback results. However, this claim is not substantiated in the paper, and, moreover, the authors do not even cite the SOTA EF methods, which are based on the EF21 method of Richtarik, Sokolov and Fatkhullin, EF21: A new, simpler, theoretically better, and practically faster error feedback, NeurIPS 2021. See, for example (there are many more follow-up works), 
- Fatkhullin et al, EF21 with bells & whistles: practical algorithmic extensions of modern error feedback
- Richtarik et al, 3PC: Three point compressors for communication-efficient distributed training and a better theory for lazy aggregation, ICML 2022
- Fatkhullin et al, Momentum provably improves error feedback!, NeurIPS 2023
- Richtarik et al, Error feedback reloaded: from quadratic to arithmetic mean of smoothness constants, ICLR 2024

Granted, the above method are not adaptive. However, the claim that the communication complexity of all know methods has linear dependence on $d$ is not true; the question depends on the number of clients $C$ as well. In subsequent works, modern EF methods have been analyzed in the nonsmooth nonconvex case as well, with the assumption of bounded gradients, as in the reviewed paper.

3. The authors claim that error feedback is not needed when unbiased compressors are used. This is true, but in this case, one case use DIANA-type variance reduction [1, 2 & many follow-up works], which is just a different type of error-feedback applicable in this setting. The theoretical results of methods based on such VR mechanisms are vastly superior to the EF-based theory, which is necessarily worse since it handles a larger class of compressors. 

- [1] Mishchenko et al, Distributed learning with compressed gradient differences, arXiv:1901.09269, 2019
- [2] Horvath et al, Stochastic distributed learning with gradient quantization and double variance reduction, arXiv:1904.05115, 2019

4. It is interesting that the results obtained by the authors depend in the Hessian spectrum. Since Newton's method can be seen as a prototypical "adaptive" method, and since communication-efficient variant of Newton method was proposed in the context of the FL literature before [3], with interesting theoretical (rates independent of condition number) and empirical results, comparison to this line of work seems to be called for:
- [3] Safaryan et al, FedNL: Making Newton-type methods applicable to federated learning, ICML 2022.

5. Other pieces of criticism were raised by the reviewers, but I am now running out of the character limit for this metareview.

The scores were mixed (3,3,6,6). In summary, I believe the paper is interesting, but needs a major revision.

Decision: Reject