ScaLA: Accelerating Adaptation of Pre-Trained Transformer-Based Language Models via Efficient Large-Batch Adversarial Noise

In the part, we present results that are not included in the main text due to the space limit.

For all configurations, we fine-tune against the GLUE datasets and set the maximum number of epochs to 6. We use a linear learning rate decay schedule with a warm-up ratio of 0.1. For ScaLA, we set $\lambda=1$, perturbation clipping radius $\omega=10^{-5}$, step size $\rho=10^{-4}$, and $t_{s}$={3,5}. These values worked well enough that we did not feel the need to explore more. For fairness, we perform a grid search of learning rates in the range of $\{$1e-5, 3e-5, 5e-5, 7e-5, 9e-5, 1e-4, 3e-4$\}$ for small batch sizes and $\{$5.6e-5, 8e-5, 1e-4, 1.7e-4, 2.4e-4, 2.8e-4, 4e-4, 5.6e-4, 1e-3$\}$ for large batch sizes. We keep the remaining hyperparameteres unchanged.

In this part, we investigate how large-batch adaptation affects the generalizability of transformer networks on downstream tasks. As there are various heuristics for setting the learning rates (Smith et al., 2018; Goyal et al., 2017; Smith, 2018; You et al., 2019a), and because few work studies the learning rate scaling effects on adapting pre-trained Transformer networks, we perform a grid search on learning rates $\{$1e-4, 3e-4,5e-4, 7e-4, 9e-4, 1e-3, 3e-3$\}$ and batch sizes $\{$1K, 2K, 4K, 8K$\}$ while keeping the other hyperparameters the same to investigate how ScaLA affects the hyperparameter tuning effort.

Table 6 shows the results of using the square root scaling rule to decide the learning rates for large batch sizes vs. accuracy results with tuned learning rate results, without and with ScaLA. The first row represents the best accuracy found through fine-tuning with a small batch size 32. The next two rows correspond to fine-tuning with batch size 1024 using tuned learning rates vs. using the scaling rule. The last two rows represent fine-tuning using ScaLA with batch size 1024, also using tuned learning rates vs. the scaling rule. Even with square-root scaling, the large-batch baseline still cannot reach the small-batch accuracy (88.7 vs. 89.4). Moreover, although tuning the learning rates lead to better results on some datasets such as MNLI-m (84.9 vs. 85.1) and SST-2 (92.9 vs. 93.5), the square-root scaling rule leads to better results on other tasks such as QNLI (90.8 vs. 90) and QQP (91.4/88.4 vs. 90.9/87.7). So the best learning rates on fine-tuning tasks are not exactly sqrt. However, given that ScaLA with square-root learning rate scaling achieves on average better results than the grid search of learning rates (89.4 vs. 89.7), we suggest to use sqrt scaling for learning rates to simplify the hyperparameter tuning effort for ScaLA.

In this section, we provide the formal statements and detailed proofs for the convergence rate. The convergence analysis builds on techniques and results in (Davis & Drusvyatskiy, 2018; You et al., 2019b). We consider the general problem of a two-player sequential game represented as nonconvex-nonconcave minimax optimization that is stochastic with respect to the outer (first) player playing $x\in\mathbb{X}$ while sampling $\xi$ from $Q$ and deterministic with respect to the inner (second) player playing $y\in\mathbb{Y}$, i.e.,

Since finding the Stackelberg equilibrium, i.e., the global solution to the saddle point problem, is NP-hard, we consider the optimality notion of a local minimax point (Jin et al., 2020). Since maximizing over $y$ may result in a non-smooth function even when $f$ is smooth, the norm of the gradient is not particularly a suitable metric to track the convergence progress of an iterative minimax optimization procedure. Hence, we use the gradient of the Moreau envelope (Davis & Drusvyatskiy, 2019) as the appropriate potential function. Let $\mu\in\mathbb{R}_{+}^{h}$. The $\mu$-Moreau envelope for a function $g:\mathbb{X}\to\mathbb{R}$ is defined as $g_{\mu}(x):=\min_{z}g(z)+\sum_{i=1}^{h}\frac{1}{2\mu^{i}}\|x^{i}-z^{i}\|^{2}$. Another reason for the choice of this potential function is due to the special property (Rockafellar, 2015) of the Moreau envelope that if its gradient $\nabla_{x}[g_{\mu}(x)]$ almost vanishes at $x$, such $x$ is close to a stationary point of the original function $g$.

Assumptions: We assume that $\mathbb{X}=\bigsqcup_{i=1}^{h}\mathbb{X}^{i}$ is partitioned into $h$ disjoint groups , i.e., in terms of training a neural network, we can think of the network having the parameters partitioned into $h$ (hidden) layers. The measure $Q$ characterizes the training data. Let $\widehat{\nabla}_{x}f(x,y)$ denote the noisy estimate of the true gradient $\nabla_{x}f(x,y)$. We assume that the noisy gradients are unbiased, i.e., $\mathbb{E}[\widehat{\nabla}_{x}f(x,y)]=\nabla_{x}f(x,y)$. For each group $i\in[h]$, we make the standard (groupwise) boundedness assumption (Ghadimi & Lan, 2013) on the variance of the stochastic gradients, i.e., $\mathbb{E}\|\widehat{\nabla}^{i}_{x}f(x,y)-\nabla^{i}_{x}f(x,y)\|^{2}\leq\sigma^{2}_{i}$, $\forall i\in[h]$. We assume that $f(x,y)$ has Lipschitz continuous gradients. Specifically, let $f(x,y)$ be $\alpha$-smooth in $x$ where $\alpha:=(\alpha_{1},\ldots,\alpha_{h})$ denotes the $h$-dimensional vector of (groupwise) Lipschitz parameters, i.e., $\|\nabla^{i}_{x}f(x_{a},y)-\nabla^{i}_{x}f(x_{b},y)\|\leq\alpha_{i}\|x_{a}^{i}-x_{b}^{i}\|$, $\forall i\in[h]$ and $x_{a},x_{b}\in\mathbb{X},y\in\mathbb{Y}$. Let $\kappa_{\alpha}:=\frac{\max_{i}\alpha_{i}}{\min_{i}\alpha_{i}}$.

Super-scripts are used to index into a vector ($i$ denotes the group index and $j$ denotes an element in group $i$). For any $c\in\mathbb{R}$, the function $\nu:\mathbb{R}\rightarrow[\mathcal{L},\mathcal{U}]$ clips its values, i.e., $\nu(c):=\max(\mathcal{L},\min(c,\mathcal{U}))$ where $\mathcal{L}<\mathcal{U}$. Let $\|.\|$, $\|.\|_{1}$ and $\|.\|_{\infty}$ denote the $\ell_{2}$, $\ell_{1}$, and $\ell_{\infty}$ norms. We assume that the true gradients are bounded, i.e., $\|\nabla_{x}f(x,y)\|_{\infty}\leq\mathcal{G}$.

First, we begin with relevant supporting lemmas. The following lemma characterizes the convexity of an additive modification of $g$.

The following property of the Moreau envelope relates it to the original function.

We now present the formal version of Theorem 4.1 in Theorem D.3. Note that Lemma D.2 facilitates giving the convergence guarantees in terms of the gradient of the Moreau envelope. Recall that $t\in[T]$ denotes the epochs corresponding to the outer maximization. Without loss of generality, we set the delay parameter for injection of the adversarial perturbation in Algorithm 1 as $t_{s}=0$. Here, we assume that the PGA provides an $\epsilon$-approximate maximizer.

• Proof.

  In this proof, for brevity, we define the vector $\nabla_{t}:=\nabla_{x}f(x,y)$, i.e., the gradient of the objective with respect to $x$, evaluated at the outer step $t$. Since evaluating gradients using mini-batches produces noisy gradients, we use $\widehat{\nabla}$ to denote the noisy version of a true gradient $\nabla$, i.e., $\widehat{\nabla}=\nabla+\Delta$ for a noise vector $\Delta$. For any outer step $t$, we have $f(x_{t},\widehat{y})\geq g(x_{t})-\epsilon$ where $\widehat{y}$ is an $\epsilon$-approximate maximizer. For any $\widetilde{x}\in\mathbb{X}$, using the smoothness property (Lipschitz gradient) of $f$, we have

  	$\displaystyle g(\widetilde{x})$	$\displaystyle\geq f(\widetilde{x},y_{t})$
  		$\displaystyle\geq f(x_{t},y_{t})+\sum_{i=1}^{h}\langle\nabla^{i}_{t},\widetilde{x}^{i}-x^{i}_{t}\rangle-\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widetilde{x}^{i}-x^{i}_{t}\|^{2}$
  		$\displaystyle\geq g(x_{t})-\epsilon+\sum_{i=1}^{h}\langle\nabla^{i}_{t},\widetilde{x}^{i}-x^{i}_{t}\rangle-\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widetilde{x}^{i}-x^{i}_{t}\|^{2}$		(4)

  Let $\phi_{\mu}(x,z):=g(z)+\sum_{i=1}^{h}\frac{1}{2\mu^{i}}\|x^{i}-z^{i}\|^{2}$. Recall that the $\mu$-Moreau envelope for $g$ is defined as $g_{\mu}(x):=\min_{z}\phi_{\mu}(x,z)$ and its gradient is the groupwise proximal operator given by $\nabla_{x}[g_{\mu}(x)]=\left[\frac{1}{\mu^{1}}\left(x^{1}-\arg\min_{z^{1}}\phi_{\mu}(x,z)\right),\ldots,\frac{1}{\mu^{h}}\left(x^{h}-\arg\min_{z^{h}}\phi_{\mu}(x,z)\right)\right]$.

  Now, let $\widehat{x}_{t}=\arg\min_{x}\phi_{1/2\alpha}(x_{t},x)=\arg\min_{x}\left(g(x)+\sum_{i=1}^{h}\alpha_{i}\|x^{i}_{t}-x^{i}\|^{2}\right)$. Then, plugging in the update rule for $x$ at step $t+1$ in terms of quantities at step $t$, using the shorthand $\nu^{i}_{t}:=\nu(\|x^{i}_{t}\|)$ and conditioning on the filtration up to time $t$, we have

  	$\displaystyle g_{1/2\alpha}(x_{t+1})$	$\displaystyle\leq g(\widehat{x}_{t})+\sum_{i=1}^{h}\alpha_{i}\|x^{i}_{t+1}-\widehat{x}^{i}_{t}\|^{2}$
  		$\displaystyle\leq g(\widehat{x}_{t})+\sum_{i=1}^{h}\alpha_{i}\left\|x^{i}_{t}-\eta_{t}\nu^{i}_{t}\frac{\widehat{\nabla}^{i}_{t}}{\|\widehat{\nabla}^{i}_{t}\|}-\widehat{x}^{i}_{t}\right\|^{2}$
  		$\displaystyle\leq g(\widehat{x}_{t})+\sum_{i=1}^{h}\alpha_{i}\left\|x^{i}_{t}-\widehat{x}^{i}_{t}\right\|^{2}+\sum_{i=1}^{h}2\alpha_{i}\eta_{t}\left\langle\nu^{i}_{t}\frac{\widehat{\nabla}^{i}_{t}}{\|\widehat{\nabla}^{i}_{t}\|},\widehat{x}^{i}_{t}-x^{i}_{t}\right\rangle+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\leq g_{1/2\alpha}(x_{t})+\sum_{i=1}^{h}2\alpha_{i}\eta_{t}\left\langle\nu^{i}_{t}\frac{\widehat{\nabla}^{i}_{t}}{\|\widehat{\nabla}^{i}_{t}\|},\widehat{x}^{i}_{t}-x^{i}_{t}\right\rangle+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\leq g_{1/2\alpha}(x_{t})+2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\sum_{j=1}^{d_{i}}\left(\frac{\widehat{\nabla}^{i,j}_{t}}{\|\widehat{\nabla}^{i}_{t}\|}-\frac{{\nabla^{i,j}_{t}}}{\|{\nabla^{i}_{t}}\|}+\frac{{\nabla^{i,j}_{t}}}{\|{\nabla^{i}_{t}}\|}\right)\times(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\leq g_{1/2\alpha}(x_{t})+2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\sum_{j=1}^{d_{i}}\left(\frac{{\nabla^{i,j}_{t}}}{\|{\nabla^{i}_{t}}\|}\right)\times(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})$
  		$\displaystyle+2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\sum_{j=1}^{d_{i}}\left(\frac{\widehat{\nabla}^{i,j}_{t}}{\|\widehat{\nabla}^{i}_{t}\|}-\frac{{\nabla^{i,j}_{t}}}{\|{\nabla^{i}_{t}}\|}\right)\times(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\leq g_{1/2\alpha}(x_{t})+2\eta_{t}\sum_{i=1}^{h}\frac{\alpha_{i}\nu^{i}_{t}}{\|{\nabla^{i}_{t}}\|}\left\langle{\nabla^{i}_{t}},\widehat{x}^{i}_{t}-x^{i}_{t}\right\rangle$
  		$\displaystyle+2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\sum_{j=1}^{d_{i}}\left(\frac{\nabla^{i,j}_{t}+\Delta^{i,j}_{t}}{\|\nabla^{i}_{t}+\Delta^{i}_{t}\|}-\frac{{\nabla^{i,j}_{t}}}{\|{\nabla^{i}_{t}}\|}\right)\times(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$

  		$\displaystyle\leq g_{1/2\alpha}(x_{t})+2\eta_{t}\mathcal{U}\sum_{i=1}^{h}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left\langle{\nabla^{i}_{t}},\widehat{x}^{i}_{t}-x^{i}_{t}\right\rangle$
  		$\displaystyle+2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\sum_{j=1}^{d_{i}}\left(\frac{\|{\nabla^{i}_{t}}\|(\nabla^{i,j}_{t})(\nabla^{i,j}_{t}+\Delta^{i,j}_{t})-\|\nabla^{i}_{t}+\Delta^{i}_{t}\|(\nabla^{i,j}_{t})^{2}}{\|\nabla^{i}_{t}+\Delta^{i}_{t}\|\|{\nabla^{i}_{t}}\|}\right)\times\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}$
  		$\displaystyle+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\stackrel{{\scriptstyle E_{1}}}{{\leq}}g_{1/2\alpha}(x_{t})+2\eta_{t}\mathcal{U}\max_{i}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle+2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\left(\frac{\langle\nabla^{i}_{t},\nabla^{i}_{t}+\Delta^{i}_{t}\rangle-\|\nabla^{i}_{t}+\Delta^{i}_{t}\|\|\nabla^{i}_{t}\|}{\|\nabla^{i}_{t}+\Delta^{i}_{t}\|}\right)+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\leq g_{1/2\alpha}(x_{t})+2\eta_{t}\mathcal{U}\max_{i}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle-2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\left(\frac{\|\nabla^{i}_{t}+\Delta^{i}_{t}\|\|\nabla^{i}_{t}\|-\|\nabla^{i}_{t}+\Delta^{i}_{t}\|^{2}+\langle\Delta^{i}_{t},\nabla^{i}_{t}+\Delta^{i}_{t}\rangle}{\|\nabla^{i}_{t}+\Delta^{i}_{t}\|}\right)$
  		$\displaystyle+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$		(5)
  		$\displaystyle\leq g_{1/2\alpha}(x_{t})+2\eta_{t}\mathcal{U}\max_{i}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle-2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\left(\|\nabla^{i}_{t}\|-\|\nabla^{i}_{t}+\Delta^{i}_{t}\|-\frac{|\langle\Delta^{i}_{t},\nabla^{i}_{t}+\Delta^{i}_{t}\rangle|}{\|\nabla^{i}_{t}+\Delta^{i}_{t}\|}\right)+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\stackrel{{\scriptstyle E_{2}}}{{\leq}}g_{1/2\alpha}(x_{t})+2\eta_{t}\mathcal{U}\max_{i}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle-2\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\left(\|\nabla^{i}_{t}\|-\|\nabla^{i}_{t}+\Delta^{i}_{t}\|-\|\Delta^{i}_{t}\|\right)+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  		$\displaystyle\stackrel{{\scriptstyle E_{3}}}{{\leq}}g_{1/2\alpha}(x_{t})+2\eta_{t}\mathcal{U}\max_{i}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle-4\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\|\Delta^{i}_{t}\|+\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$
  	$\displaystyle g_{1/2\alpha}(x_{T})$	$\displaystyle\stackrel{{\scriptstyle E_{4}}}{{\leq}}g_{1/2\alpha}(x_{0})+2\mathcal{U}\sum_{t=0}^{T-1}\eta_{t}\max_{i}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle-4\sum_{t=0}^{T-1}\eta_{t}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\|\Delta^{i}_{t}\|+\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}\eta_{t}^{2}(\nu^{i}_{t})^{2}$

  where we have used Hölder’s inequality along with bound (4) in $E_{1}$, Cauchy-Schwarz inequality in $E_{2}$, triangle inequality in $E_{3}$, telescoping sum in $E_{4}$. Rearranging and using $\eta_{t}=\eta$ in $E_{5}$ along with Hölder’s inequality,

  	$\displaystyle\frac{1}{2\eta\mathcal{U}}\left(g_{1/2\alpha}(x_{T})-g_{1/2\alpha}(x_{0})\right)$	$\displaystyle\leq\sum_{t=0}^{T-1}\max_{i}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle-\frac{2}{\mathcal{U}}\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\|\Delta^{i}_{t}\|+\frac{\eta}{2\mathcal{U}}\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}(\nu^{i}_{t})^{2}$
  	$\displaystyle\frac{1}{2\eta\mathcal{U}}\left(g_{1/2\alpha}(x_{T})-g_{1/2\alpha}(x_{0})\right)$	$\displaystyle\stackrel{{\scriptstyle E_{5}}}{{\leq}}\max_{i,t}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}\sum_{t=0}^{T-1}\left(g(\widehat{x}_{t})-g(x_{t})+\epsilon+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$
  		$\displaystyle-\frac{2}{\mathcal{U}}\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\|\Delta^{i}_{t}\|+\frac{\eta}{2\mathcal{U}}\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}(\nu^{i}_{t})^{2}$

  Dividing by $T$ and rearranging,

  	$\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\left(g(x_{t})-g(\widehat{x}_{t})-\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$	$\displaystyle\leq\epsilon-\frac{1}{2\eta\mathcal{U}\zeta T}\left(g_{1/2\alpha}(x_{T})-g_{1/2\alpha}(x_{0})\right)$
  		$\displaystyle-\frac{2}{\mathcal{U}\zeta T}\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}\nu^{i}_{t}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\|\Delta^{i}_{t}\|$
  		$\displaystyle+\frac{\eta}{2\mathcal{U}\zeta T}\sum_{i=1}^{h}\alpha_{i}\sum_{t=0}^{T-1}(\nu^{i}_{t})^{2}$

  where we define $\zeta:=\max_{i,t}\frac{\alpha_{i}}{\|{\nabla^{i}_{t}}\|}$. Defining $\mathcal{D}:=\left(g_{1/2\alpha}(x_{0})-\mathbb{E}(\min_{x}g(x))\right)$ and taking expectation with respect to $\xi$ on both sides, we have

  	$\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left(g(x_{t})-g(\widehat{x}_{t})-\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}\right)$	$\displaystyle\leq\epsilon+\frac{\mathcal{D}}{2\eta\mathcal{U}\zeta T}$
  		$\displaystyle-\frac{2\mathcal{L}}{\mathcal{U}\zeta T}\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\mathbb{E}\|\Delta^{i}_{t}\|+\frac{\eta\mathcal{U}\|\alpha\|_{1}}{2\zeta}$
  		$\displaystyle\stackrel{{\scriptstyle E_{6}}}{{\leq}}\epsilon+\frac{\mathcal{D}}{2\eta\mathcal{U}\zeta T}$
  		$\displaystyle-\frac{2\mathcal{L}}{\mathcal{U}\zeta T}\sum_{t=0}^{T-1}\sum_{i=1}^{h}\alpha_{i}\max_{j}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\frac{\sigma_{i}}{\sqrt{b}}+\frac{\eta\mathcal{U}\|\alpha\|_{1}}{2\zeta}$
  		$\displaystyle\stackrel{{\scriptstyle E_{7}}}{{\leq}}\epsilon+\frac{\mathcal{D}}{2\eta\mathcal{U}\zeta T}$
  		$\displaystyle-\frac{2\mathcal{L}\|\alpha\|_{1}}{\mathcal{U}\zeta\sqrt{b}}\max_{i,j,t}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\sigma_{i}+\frac{\eta\mathcal{U}\|\alpha\|_{1}}{2\zeta}$
  		$\displaystyle\stackrel{{\scriptstyle E_{8}}}{{=}}\epsilon+\frac{\mathcal{D}}{2\eta\mathcal{U}\zeta T}-\frac{2\mathcal{L}\|\alpha\|_{1}\mathcal{Z}}{\mathcal{U}\zeta\sqrt{b}}+\frac{\eta\mathcal{U}\|\alpha\|_{1}}{2\zeta}$		(6)

  where we have used the assumption on the variance of stochastic gradients in $E_{6}$, Hölder’s inequality in $E_{7}$ and we define $\mathcal{Z}:=\max_{i,j,t}\frac{(\widehat{x}^{i,j}_{t}-x^{i,j}_{t})}{(\nabla^{i,j}_{t})}\sigma_{i}$ in $E_{8}$; $b$ denotes batch size. Now, we lower bound the left hand side using the convexity of the additive modification of $g$.

  	$\displaystyle g(x_{t})-g(\widehat{x}_{t})-\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}$
  	$\displaystyle\geq g(x_{t})+0-g(\widehat{x}_{t})-\sum_{i=1}^{h}\alpha_{i}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}$
  	$\displaystyle\geq\left(\left(g(x_{t})+\sum_{i=1}^{h}\alpha_{i}\|x^{i}_{t}-x^{i}_{t}\|^{2}\right)-\min_{x}\left(g(x_{t})+\sum_{i=1}^{h}\alpha_{i}\|x^{i}-x^{i}_{t}\|^{2}\right)\right)+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}$
  	$\displaystyle\geq\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}+\sum_{i=1}^{h}\frac{\alpha_{i}}{2}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}{=}\sum_{i=1}^{h}\frac{4\alpha_{i}^{2}}{4\alpha_{i}}\|\widehat{x}^{i}-x^{i}_{t}\|^{2}$
  	$\displaystyle\stackrel{{\scriptstyle E_{9}}}{{\geq}}\frac{1}{4\max_{i}\alpha_{i}}\|\nabla g_{1/2\alpha}(x_{t})\|^{2}$		(7)

  where we have used the expression for the gradient of the Moreau envelope in $E_{9}$. Combining the inequalities from Equation (7) and Equation (6), we have

  $\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left(\frac{1}{4\max_{i}\alpha_{i}}\|\nabla g_{1/2\alpha}(x_{t})\|^{2}\right)$

  $\displaystyle\leq\epsilon+\frac{\mathcal{D}}{2\eta\mathcal{U}\zeta T}+\left(\frac{\eta\mathcal{U}}{2\zeta}-\frac{2\mathcal{L}\mathcal{Z}}{\mathcal{U}\zeta\sqrt{b}}\right)\|\alpha\|_{1}$

  Setting the learning rate as $\eta=\frac{1}{\mathcal{U}\sqrt{T}}$ and batch size as $b=\frac{16T\mathcal{L}^{2}\mathcal{Z}^{2}}{\mathcal{U}^{2}}$,

  $\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\|\nabla g_{1/2\alpha}(x_{t})\|^{2}\right]$

  $\displaystyle\leq 4\epsilon\max_{i}\alpha_{i}+\frac{2\mathcal{D}\max_{i}\alpha_{i}}{\zeta\sqrt{T}}$

  Now, to simplify $\zeta$, using the inequality that $\max_{k}(a_{k}\cdot b_{k})\geq\min_{k_{a}}a_{k_{a}}\cdot\min_{k_{b}}b_{k_{b}}$ for two finite sequences $\{a,b\}$ with positive values, along with the bounded gradients assumption, we have

  $\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\|\nabla g_{1/2\alpha}(x_{t})\|^{2}\right]$

  $\displaystyle\leq 4\epsilon\max_{i}\alpha_{i}+\frac{2\mathcal{DG}\max_{i}\alpha_{i}}{\sqrt{T}\min_{i}\alpha_{i}}=4\epsilon\|\alpha\|_{\infty}+\frac{2\kappa_{\alpha}\mathcal{DG}}{\sqrt{T}}$

  where $\kappa_{\alpha}:=\frac{\max_{i}\alpha_{i}}{\min_{i}\alpha_{i}}$. ∎