Year: 2019

Abstract:
Stochastic methods with coordinate-wise adaptive stepsize (such as RMSprop and Adam) have been widely used in training deep neural networks. Despite their fast convergence, they can generalize worse than stochastic gradient descent. In this paper, by revisiting the design of Adagrad, we propose to split the network parameters into blocks, and use a blockwise adaptive stepsize. Intuitively, blockwise adaptivity is less aggressive than adaptivity to individual coordinates, and can have a better balance between adaptivity and generalization. We show theoretically that the proposed blockwise adaptive gradient descent has comparable regret in online convex learning and convergence rate for optimizing nonconvex objective as its counterpart with coordinate-wise adaptive stepsize, but is better up to some constant. We also study its uniform stability and show that blockwise adaptivity can lead to lower generalization error than coordinate-wise adaptivity. Experimental results show that blockwise adaptive gradient descent converges faster and improves generalization performance over Nesterov’s accelerated gradient and Adam.

heading: 1 INTRODUCTION
text: Deep networks have achieved excellent performance in a variety of domains such as computer vision (He et al., 2016), language modeling (Zaremba et al., 2014), and speech recognition (Graves et al., 2013). The most popular optimizer is stochastic gradient decent (SGD) (Robbins & Monro, 1951), which is simple and has low per-iteration complexity. Its convergence rate is also well-established (Ghadimi & Lan, 2013; Bottou et al., 2018). However, vanilla SGD is sensitive to the choice of stepsize, and requires careful tuning.
To improve the efficiency and robustness of SGD, many variants based on coordinate-wise adaptive stepsize (Almeida et al., 1999; Duchi et al., 2011; Schaul et al., 2013; Kingma & Ba, 2015; Tieleman & Hinton, 2012; Zeiler, 2012; Zheng & Kwok, 2017; Reddi et al., 2018; Zaheer et al., 2018; Zou & Shen, 2018; Chen et al., 2019; Zou et al., 2019) have been proposed. Though this is effective in accelerating convergence, its generalization performance is often worse than SGD (Wilson et al., 2017). Recently, it is shown that coordinate-wise adaptive gradient descent is closely related to signbased gradient descent (Balles & Hennig, 2018; Bernstein et al., 2018). Theoretical arguments and empirical evidence suggest that the gradient sign impedes generalization (Balles & Hennig, 2018). To contract the generalization gap, a partial adaptive parameter for the second-order momentum is proposed (Chen & Gu, 2018). By using a smaller partial adaptive parameter, the adaptive gradient algorithm behaves less like sign descent and more like SGD. Based on the similar motivation, Luo et al. (2019) proposed to gradually transform Adam Kingma & Ba (2015) to SGD. Moreover, to avoid numerical problems in practical implementation, a small (= 10−8) parameter is typically used in these methods. This parameter controls adaptivity of the algorithm (Zaheer et al., 2018). A larger value (say, = 10−3) can reduce adaptivity and empirically helps Adam to match its generalization performance with SGD. This implies that coordinate-wise adaptivity may be too aggressive for good generalization performance.
To improve generalization performance, attempts have been made to use a layer-wise stepsize (Singh et al., 2015; You et al., 2017; Yu et al., 2017; Zhou et al., 2018), which assign different stepsizes to different layers or normalize the layer-wise gradient. However, there has been no theoretical analysis for its empirical success.
In this paper, we consider the more general case in which model parameters are partitioned into blocks, instead of simply into layers. By revisiting the derivation of Adagrad, we propose the use of a blockwise stepsize which depends on the corresponding gradient block. Such blockwise adaptivity is less aggressive than coordinate-wise adaptivity, as it adapts to parameter blocks instead of to individual parameters. This can thus have a better balance between adaptivity and generalization. Moreover, unlike coordinate-wise adaptivity, blockwise adaptivity is not sign-based gradient descent, and so does not suffer from its performance deterioration.
As in (Bernstein et al., 2018; Ghadimi & Lan, 2013; Ward et al., 2018; Zaheer et al., 2018; Zou & Shen, 2018; Zou et al., 2019), we will focus on the expected risk minimization problem:
min θ F (θ) = Ez[f(θ; z)], (1)
where f is some possibly nonconvex loss function, z is a random sample, θ is the model parameter, and the expectation is taken w.r.t. the underlying sample distribution. The expected risk measures the generalization performance on unseen data (Bottou et al., 2018), and reduces to the empirical risk when a finite training set is considered. We show theoretically that the proposed blockwise adaptive gradient descent can be faster than its coordinate-wise counterpart. Using tools on uniform stability (Bousquet & Elisseeff, 2002; Hardt et al., 2016), we also show that blockwise adaptivity has potentially lower generalization error than coordinate-wise adaptivity. Empirically, blockwise adaptive gradient descent converges faster and obtains better generalization performance than coordinatewise descent (Adam) and Nesterov’s accelerated gradient (NAG) (Sutskever et al., 2013).
Notations. For an integer n, [n] = {1, 2, . . . , n}. For a vector x, xT denotes its transpose, Diag(x) is a diagonal matrix with x on its diagonal, √ x is the element-wise square root of x, x2 is the coordinate-wise square of x, ‖x‖2 = √ xTx, ‖x‖∞ = maxi |xi|, ‖x‖2Q = xTQx, where Q is a positive semidefinite (psd) matrix, and x ≥ 0 means xi ≥ 0 for all i. For two vectors x and y, x/y, and 〈x, y〉 denote the element-wise division and dot product, respectively. For a square matrix X , X−1 is its inverse, and X 0 means that X is psd. Moreover, 1d = [1, 1, . . . , 1]T ∈ Rd.

heading: 2 RELATED WORK

heading: 2.1 ADAGRAD
text: Adagrad (shown in Algorithm 1) is an adaptive gradient method in online convex learning with coordinate-wise stepsize (Duchi et al., 2011; McMahan & Streeter, 2010). It is particularly useful for sparse learning, as parameters for the rare features can take large steps. Recently, Ward et al. (2018) established its convergence properties with a global adaptive stepsize in nonconvex optimization. It is shown that Adagrad converges to a stationary point at the optimal O(1/ √ T ) rate (up to a factor log(T )), where T is the total number of iterations.
Recall that the SGD iterate is the solution to the problem: θt+1 = arg minθ〈gt, θ〉 + 12η‖θ − θt‖ 2 2, where gt is the gradient of the loss ft at iteration t, and θt ∈ Rd is the parameter vector. To incorporate information about the curvature of sequence {ft}, the `2-norm in the SGD update can be replaced by the Mahalanobis norm, leading to (Duchi et al., 2011):
θt+1 = arg min θ 〈gt, θ〉+
1
2η ‖θ − θt‖2Diag(st)−1 , (2)
where st ≥ 0. This is an instance of mirror descent (Nemirovski & Yudin, 1983). Its regret bound has a gradient-related term ∑T t=1 ‖gt‖2Diag(st)−1 . Adagrad’s stepsize can be obtained by examining a similar objective (Duchi et al., 2011):
min s∈S T∑ t=1 ‖gt‖2Diag(s)−1 , (3)
where S = {s : s ≥ 0, 〈s, 1〉 ≤ c}, and c is some constant. At optimality, s∗,i = c‖g1:T,i‖2/ ∑d j=1 ‖g1:T,j‖2, where g1:T,i = [g1,i, . . . , gT,i]T . As st cannot depend on gj’s with j > t, this suggests st,i ∝ ‖g1:t,i‖2. Theoretically, this choice of st leads to a regret bound that is competitive with the best post-hoc optimal bound (McMahan & Streeter, 2010).
To solve the expected risk minimization problem in (1), an Adagrad variant called weighted AdaEMA is recently proposed in (Zou et al., 2019). It employs weighted averaging of g2t,i’s for stepsize. Moreover, weighted AdaEMA is a general coordinate-wise adaptive method and includes many Adagrad variants, including Adam and RMSprop, as special cases.

heading: 2.2 UNIFORM STABILITY
text: Given n samples S = {zi}ni=1 drawn i.i.d. from an underlying data distribution D, one often learns the model by minimizing the empirical risk: minθ ΦS(θ) ≡ 1n ∑n i=1 f(θ; zi). Let M(S) be the output of a possibly randomized algorithm M (e.g., SGD) running on data S. Definition 1. (Hardt et al., 2016) Let S and S′ be two data sets of size n that differ in only one sample. Algorithm M is u-uniformly stable if stab ≡ supS,S′ supz∈D EM [f(M(S); z) − f(M(S′); z)] ≤ u.
The generalization error is defined as gen ≡ ES,M [ΦS(M(S))−F (M(S))], where the expectation is taken w.r.t. the set S and randomness of M (Hardt et al., 2016). It is shown that the generalization error is bounded by the uniform stability of M , i.e., | gen| ≤ stab (Hardt et al., 2016). In other words, the more uniformly stable an algorithm is, the lower is its generalization error.

heading: 3 BLOCKWISE ADAPTIVE DESCENT
text: Consider the under-determined least squares problem:
min θ ‖Xθ − y‖22, (4)
whereX ∈ Rn×d is the input matrix (with sample size n, and dimensionality d > n), and output y ∈ Rn. We assume that XXT is invertible. As pointed out in (Zhang et al., 2017; Wilson et al., 2017), any stochastic gradient descent method on problem (4) with a global stepsize outputs a trajectory with iterates lying in the span of the rows of X . One solution of (4) is XT (XXT )−1y, which happens to be the solution with minimum `2-norm among all possible global minimizers. This minimum-norm solution has the largest margin, and maximizing margin typically leads to lower generalization error (Boser et al., 1992).
It is known that SGD converges to the minimum `2-norm solution of problem (4) (Zhang et al., 2017). On the other hand, coordinate-wise adaptive methods (such as Adagrad, RMSprop, and Adam) fail to find the minimum `2-norm solution, but converge to solutions with low `∞-norm instead (Wilson et al., 2017). Examples in (Wilson et al., 2017) show that solutions obtained by these adaptive methods can generalize arbitrarily poorly, while the SGD solution makes no error.
In Section 3.1, we first show that blockwise adaptivity, unlike coordinate-wise adaptivity, can find the minimum `2-norm solution of a nonlinear least squares problem in layer-wise training of a neural network. This motivates us to further exploit blockwise adaptivity in end-to-end neural network training (Section 3.2). To provide more analysis, Section 3.3 studies the proposed algorithm in the online convex learning setting as in (Duchi et al., 2011; Kingma & Ba, 2015; Reddi et al., 2018).

heading: 3.1 BLOCKWISE VS COORDINATE-WISE ADAPTIVITY
text: Consider a L-layer neural network with output φL−1(· · ·φ2(φ1(XW1)W2) · · ·WL−1)WL, where {Wl ∈ Rdl−1×dl}Ll=1 are weight matrices with d0 = d and dL = m, the output dimensionality. Assume that the nonlinear activation functions {φl}L−1l=1 are bijective with nonzero derivatives on R (e.g., tanh and leaky ReLU). For simplicity, we further assume that dl = d = m > n for all l. Training this neural network with the square loss corresponds to solving the nonlinear optimization problem: min{Wl}Ll=1 ‖φL−1(· · ·φ2(φ1(XW1)W2) · · ·WL−1)WL − Y ‖ 2 2, where Y ∈ Rn×m is the label matrix. Consider training the network layer-by-layer, starting from the bottom one. For layer l, its optimization subproblem can be rewritten as
min Wl ‖Φl(Hl−1Wl)− Y ‖22, (5)
where Φl(·) = φL−1(· · ·φl+1(φl(·)Wl+1) · · ·WL−1)WL, Hl−1 = φl−1(· · ·φ1(XW1) · · ·Wl−1) is the input activation to the lth layer, and H0 = X . Note that (4) is a special case of (5) with L = 1
and identity mapping. To minimize (5), the weights for this layer are updated as
Wt+1,l = Wt,l − ηt,lgt,l, (6)
where gt,l is a stochastic gradient evaluated at Wt,l at time t, and ηt,l is the stepsize which may be adaptive in that it depends on gt,l.
Proposition 1. Assume that Wl′ ’s (with l′ > l) are invertible. If Wl is initialized to zero, and Hl−1 has full row rank, the critical point that (6) converges to is the minimum `2-norm solution of (5) in expectation.
Another benefit of using a blockwise stepsize is that the optimizer’s extra memory cost can be reduced. Using a coordinate-wise stepsize requires an additionalO(d) memory for storing estimates of the second moment, while the blockwise stepsize only needs an extra O(B) memory, where B is the number of blocks. A deep network generally has millions of parameters but only tens of layers. If we set B to be the number of layers, memory reduction can be significant.

heading: 3.2 BLOCKWISE ADAPTIVE GRADIENT (BAG)
text: Let the gradient gt ∈ Rd be partitioned to {gt,Gb ∈ Rdb : b = 1, . . . , B}, where Gb is the set of indices in block b, and gt,Gb is the corresponding subvector of gt. Inspired by problem (3) in the derivation of Adagrad, we consider the following variant which imposes a block structure on s:
min s∈S′ T∑ t=1 ‖gt‖2Diag(s)−1 , (7)
where S ′ = {s : s = [q11Td1 , . . . , qB1 T dB ]T ≥ 0, 〈s, 1〉 ≤ c} for some qi ∈ R. We assume the indices in Gb are consecutive; otherwise, we can simply reorder the elements of the gradient. Note that reordering does not change the result, as the objective is invariant to ordering of the coordinates. It can be easily shown that at optimality of (7), qb = c‖g1:T,Gb‖2/( √ db ∑B i=1 √ di‖g1:T,Gi‖2), where g1:T,Gb = [g T 1,Gb , . . . , g T T,Gb ] T . The optimal qb is thus proportional to ‖g1:T,Gb‖2/ √ db. When st in
(2) is partitioned by the same block structure, the optimal qb suggests to incorporate ‖g1:t,Gb‖2/ √ db into st for block b at time t.
Algorithm 1 Adagrad: Adaptive gradient for online convex learning (Duchi et al., 2011).
1: Input: η > 0; > 0. 2: initialize θ1; v0 ← 0 3: for t = 1, 2, . . . , T do 4: Receive subgradient gt ∈ ∂ft(θt) 5: for i = 1, 2, . . . , d do 6: vt,i = vt−1,i + ‖gt,i‖22 7: θt+1,i = θt,i − ηgt,i/( √ vt,i + ) 8: end for 9: end for
Algorithm 2 BAG: Blockwise adaptive gradient for online convex learning.
1: Input: η > 0; > 0. 2: initialize θ1; v0 ← 0 3: for t = 1, 2, . . . , T do 4: Receive subgradient gt ∈ ∂ft(θt) 5: for b = 1, 2, . . . , B do 6: vt,b = vt−1,b + ‖gt,Gb‖22/db 7: θt+1,Gb = θt,Gb − ηgt,Gb/( √ vt,b + ) 8: end for 9: end for
The proposed procedure, which will be called blockwise adaptive gradient (BAG), is shown in Algorithm 2. Compared to Adagrad, each block, instead of each coordinate, has its own learning rate. When B = d (i.e., each block has only one coordinate), BAG reduces to Adagrad. When B = 1 (i.e., all coordinates are grouped together), Algorithm 2 produces the update: θt+1 = θt − η(gt/(‖g1:t‖2/ √ d + )) with a global adaptive learning rate, which is equivalent to AdaGrad-Norm (Ward et al., 2018).
Returning to the underdetermined least squares problem in (4), the following Proposition shows that when B > 1, BAG finds the minimum `2-norm solution in each subspace induced by the group structure. When B = 1, BAG converges to the minimum `2-norm solution of (4).
Proposition 2. Assume that for each b ∈ [B], each submatrix X:,Gb ∈ Rn×db has full row rank. BAG (with θ1 initialized to 0) converges to an optimal solution θ∗ of (4). For each b ∈ [B], the subvector θ∗,Gb of θ∗ equals X T :,Gb(X:,GbX T :,Gb) −1ub for some ub ∈ Rn and ∑ b∈[B] ub = y.

heading: 3.3 REGRET ANALYSIS
text: To further illustrate the advantages of blockwise adaptivity over coordinate-wise adaptivity, we consider the online convex learning setting. At round t, the learner picks θt, and suffers a loss ft(θt). After T rounds, the learner wants to achieve a low regret w.r.t. an optimal θ∗ = arg minθ ∑T t=1 ft(θ) in hindsight:
R(T ) ≡ T∑ t=1 ft(θt)− T∑ t=1 ft(θ∗) ≡ T∑ t=1 ft(θt)− inf θ T∑ t=1 ft(θ). (8)
We make the following assumptions. Assumption 1. Each ft in (8) is convex but possibly nonsmooth. There exists a subgradient g ∈ ∂ft(θ) such that ft(θ′) ≥ ft(θ) + 〈g, θ′ − θ〉 for all θ, θ′. Assumption 2. Each parameter block is in a ball of the corresponding optimal block throughout the iterations. In other words, for all b ∈ [B], maxt ‖θt,Gb − θ∗,Gb‖2 ≤ Db for some Db, where θ∗,Gb is the subvector of θ∗ in block b.
When B = 1, this reduces to the common assumption in online convex learning Duchi & Singer (2009). When B > 1, it naturally encodes the heterogeneity of model parameters. Theorem 1. Suppose that Assumptions 1 and 2 hold. Then,
R(T ) ≤ B∑ b=1 [ 1 2η √ db D2b + η √ db ] ‖g1:T,Gb‖2. (9)
Assume that maxt ‖θt−θ∗‖∞ ≤ D∞ for some constantD∞. WhenB = d, the above regret bound reduces to that of Adagrad (Theorem 5 of (Duchi et al., 2011)) by setting Db = D∞ for all b ∈ [B]. In the following, we show that when gradient magnitudes for elements in the same block have the same upper bound, blockwise adaptive learning can have lower regret than coordinate-wise adaptive learning. As a deep network can be naturally divided into blocks (examples will be given in Section 5.2) and parameters in the same block are likely to have gradients with similar magnitudes (which is verified empirically in Appendix B), blockwise adaptivity can be more beneficial. Corollary 1. Assume that E[g2t,i] ≤ σ2b for all i ∈ Gb. The expectation of (9) can be bounded as:
E[R(T )] ≤ B∑ b=1 σb [ 1 2η D2b + ηdb ]√ T . (10)
When B = d (which corresponds to Adagrad), Assumption 2 becomes maxt(θt,i − θ∗,i) ≤ Di for some Di. Expectation of the bound in (9) then reduces to
E[R(T )] ≤ B∑ b=1 σb
[ 1
2η ∑ i∈Gb D2i + ηdb
] √ T . (11)
Assuming that Assumption 2 is tight in the sense that D2b ≤ ∑ i∈Gb D 2 i . The bound in (10) is then smaller than that in (11). Intuitively, when gradients in a block have similar magnitudes in expectation, we have from the weak law of large numbers that vt,b = ∑ i∈Gb vt,i/db (where vt,b and vt,i are as defined in Algorithms 2 and 1, respectively) is a better estimate of E[vt,i] than vt,i for a single coordinate i. This implies using blockwise adaptivity may lead to better performance.

heading: 4 BLOCKWISE ADAPTIVE GRADIENT WITH MOMENTUM (BAGM)
text: In Algorithm 2, vt,b’s are increasing w.r.t. t. The update suffers from vanishing stepsize, making slow progress on nonconvex problems such as deep network training. To alleviate this problem, many Adagrad variants (such as RMSprop, Adam and weighted AdaEMA (Zou et al., 2019)) use weighted moving average momentum. In this paper, we extend the use of blockwise adaptive stepsize to weighted AdaEMA (Section 4.1), which includes Adam and RMSprop as special cases. The proposed algorithm will be called blockwise adaptive gradient with momentum (BAGM). Sections 4.2 and 4.3 then study its convergence and generalization properties. Note that as BAG is a special case of BAGM, the analysis there also apply to BAG.

heading: 4.1 PROPOSED ALGORITHM
text: The proposed BAGM is shown in Algorithm 4. Here, mt serves as an exponential moving averaged momentum, {βt} is a sequence of momentum parameters, and at’s assign different weights to the past gradients in the accumulation of variance. Note from Algorithm 4 that the variance estimate can be rewritten as
v̂t,b = t∑ i=1 ai At ‖gi,Gb‖22 db = 1∑t j=1 aj t∑ i=1 ai ‖gi,Gb‖22 db . (12)
In particular, we will consider the three weight sequences {at} introduced in (Zou & Shen, 2018). S.1: at = a for some a > 0; S.2: at = tτ for some τ > 0; The fraction at/At in (12) then decreases as O(1/t). S.3: at = α−t for some 0 < α < 1: It can be shown that this is equivalent to using the exponential moving average estimate: vt,b = αvt−1,b + (1− α) ‖gt,Gb‖ 2 2
db , and v̂t,b = vt,b 1−αt .
With βt = 0, weight sequence S.1, and = 0, BAGM reduces BAG. When B = d and = 0, BAGM reduces to weighted AdaEMA. As weighted AdaEMA includes many Adagrad variants, the proposed BAGM also covers the corresponding blockwise variants.
Algorithm 3 Weighted AdaEMA for stochastic nonconvex optimization.
1: Input: {ηt}; {at}; {βt}; > 0. 2: initialize θ1; v0 ← 0; m0 ← 0; A0 ← 3: for t = 1, 2, . . . , T do 4: Sample an unbiased stochastic gradient
gt 5: At = At−1 + at 6: for i = 1, 2, . . . , d do 7: vt,i = vt−1,i + atg2t,i 8: v̂t,i = vt,i/At 9: mt,i = βtmt−1,i + (1− βt)gt,i
10: θt+1,i = θt,i − ηtmt,i/ √ v̂t,i 11: end for 12: end for
Algorithm 4 BAGM: Blockwise adaptive gradient with momentum for stochastic nonconvex optimization.
1: Input: {ηt}; {at}; {βt}; > 0. 2: initialize θ1; v0 ← 0; m0 ← 0; A0 ← 0 3: for t = 1, 2, . . . , T do 4: Sample an unbiased stochastic gradient
gt 5: At = At−1 + at 6: for b = 1, 2, . . . , B do 7: vt,b = vt−1,b + at‖gt,Gb‖22/db 8: v̂t,b = vt,b/At 9: mt,Gb = βtmt−1,Gb + (1− βt)gt,Gb
10: θt+1,Gb = θt,Gb−ηtmt,Gb/( √ v̂t,b+ ) 11: end for 12: end for

heading: 4.2 CONVERGENCE ANALYSIS ON NONCONVEX PROBLEMS
text: Assumption 3. F in (1) is lower-bounded (i.e., F (θ∗) = infθ F (θ) > −∞) and L-smooth. Assumption 4. Each block of stochastic gradient has bounded second moment, i.e., Et[‖gt,Gb‖22]/db ≤ σ2b ,∀b ∈ [B],∀t, where the expectation is taken w.r.t. the random ft.
Assumption 4 implies the variance of each block of stochastic gradient is upper-bounded by dbσ2b (i.e., Et[‖gt,Gb−∇GbF (θt)‖22] = Et[‖gt,Gb‖22]−‖∇GbF (θt)‖22 ≤ dbσ2b ). This naturally encodes the notion of heterogeneous gradient in a multi-layer neural network. When B = 1, this reduces to the usual second moment bound in stochastic approximation (Shamir & Zhang, 2013; Zou et al., 2019). Assumption 5. 0 ≤ βt ≤ β for some 0 ≤ β < 1.
Assumption 5 allows us to use, for example, a constant βt = β, a decreasing sequence βt = β/tτ , or an increasing sequence βt = β(1− 1/tτ ). Assumption 6. (i) {at} is non-decreasing; (ii) at grows slowly such that {At−1/At} is nondecreasing and At/(At−1 + a1) ≤ ω for some ω ≥ 0; (iii) p ≡ limt→∞At−1/At > β2.
Assumption 6 is satisfied by the three weight sequences introduced above. Specifically, for S.1: ω = 1 and p = 1; S.2: ω = (1 + 2τ )/2 and p = 1; S.3: ω = (1 + 1/α)/2 and p = α > β2. Assumption 7. (Zou et al., 2019) The stepsize ηt is chosen such that wt = ηt/ √ at/At is “almost” non-increasing, i.e., there exists a non-increasing sequence {zt} and positive constants C1 and C2 such that C1zt ≤ wt ≤ C2zt for all t.
Assumption 7 is satisfied by the weights sequences S.1, S.2, S.3 when
ηt = η/ √ t (13)
for some η > 0. Interested readers are referred to (Zou et al., 2019) for details. Proposition 3. Suppose that Assumptions 3-7 hold. With probability at least 1 − δ2/3, min1≤t≤T ‖∇F (θt)‖22 ≤ 1δO(log(T )/ √ T ) for S.1 and S.2; and min1≤t≤T ‖∇F (θt)‖22 ≤ 1δO(1) for S.3.
Note that SGD, with the decreasing stepsize in (13), converges at a rate of O(log(T )/T ) (Ghadimi & Lan, 2013). Thus, the rates for S.1 and S.2 are as good as SGD. Though S.3 only leads to anO(1) bound, it has good performance in practice (Kingma & Ba, 2015; Zaheer et al., 2018). Moreover, recall that whenB = d and = 0, BAGM reduces to weighted AdaEMA. In this case, Proposition 3 obtains the same convergence rates as in (Zou et al., 2019).
Next, we consider B = B̃ for some B̃ 6= d (blockwise stepsize), and compare it with B = d (coordinate-wise stepsize). This requires the following assumption, which is slightly stronger than Assumption 4 (that only bounds the expectation). Assumption 8. ‖gt,Gb‖22/db ≤ G2b ,∀b ∈ [B] and ∀t.
Note that when B = d, Assumption 8 becomes g2t,i ≤ G2i for some Gi, and Assumption 4 becomes Et[g 2 t,i] ≤ σ2i for some σi. Let G̃b be the set of indices in block b when B = B̃. The following Corollary shows that BAGM has faster convergence than its coordinate-wise counterpart when {σ2i }i∈G̃b have low variability. This also agrees with our observation in Section 3.3 that blockwise adaptivity can have lower regret under this condition. Corollary 2. Suppose that Assumptions 3-8 hold. Define r1 ≡ ∑B̃ b=1 ∑ i∈G̃b log(σ2i / 2+1)∑B̃
b=1 db log(σ2b/ 2+1) , r2 ≡∑B̃
b=1 ∑ i∈G̃b
σi∑B̃ b=1 σbdb
and r3 ≡ ∑B̃ b=1 ∑ i∈G̃b σi log(σ2i / 2+1)∑B̃
b=1 σbdb log(σ2b/ 2+1) . Let C̃d(T )/δ (resp. C̃B̃(T )/δ) be the
high probability upper bound on min1≤t≤T ‖∇F (θt)‖22 when B = d (resp. B = B̃). If
min(1, r1, r2, r3)
√ maxb maxi∈G̃b G2i + 2
maxbG2b+ 2 ≥ 1, then C̃d(T ) ≥ C̃B̃(T ).
Note that
√ maxb maxi∈G̃b G2i + 2
maxbG2b+ 2 ≥ 1 when Assumption 8 is tight. By comparing the denominator
and numerator in r1, r2, r3, it can be seen that min(r1, r2, r3) is close to or greater than 1 when {σ2i }i∈G̃b have low variability. These will be verified empirically in Appendix B.

heading: 4.3 UNIFORM STABILITY AND GENERALIZATION ERROR
text: As in Definition 1, let S, S′ be two data sets of size n that differ in only one sample, and the tth iterates of BAGM on S and S′ by θt and θ′t, respectively. Let ∆t = ‖θt − θ′t‖2, and ∆̃t(z) = |f(θt; z) − f(θ′t; z)|. The following Proposition allows us to study how B affects the growth of E[∆̃t(z)], where the expectation is taken w.r.t. randomness of the algorithm. Proposition 4. Suppose that Assumptions 3-7 hold. Assume that f is γ̃-Lipschitz1, βt = 0, and the initial θ1, θ′1 values are the same. We have supS,S′ supz E[∆̃t+1(z)] ≤ 2γ̃C2 nC1 √[ w21 ∑B b=1 db log (σ 2 b/ 2 + 1) + dω ∑t k=1 η 2 k ] t + ( 1− 1n ) γ̃Wt, where Wt =
γ̃ ∑t k=1 ηkE [ maxb ∣∣∣1/(√v̂k,b + )− 1/(√v̂′k,b + )∣∣∣]+ L∑tk=1 ηkE [∆k/(√minb v̂k,b + )]. Recall that σ2b ≤ 1db ∑ i∈G̃b σ 2 i . If σ 2 b = 1 db ∑ i∈G̃b σ 2 i , the first term on the RHS of the above bound is smallest when B = d; otherwise, some B < d will make this term smallest. For the minb v̂k,b term inside Wt, this reduces to 1B̃ ∑B̃ b=1 v̂k,b when B = 1, and to minb mini∈G̃b v̂k,i when B = d.
1In other words, |f(θ; z)− f(θ′; z)| ≤ γ̃‖θ − θ′‖2 for any z.
As 1 B̃ ∑B̃ b=1 v̂k,b ≥ minb v̂k,b ≥ minb mini∈G̃b v̂k,i, this minb v̂k,b term is the smallest whenB = d, and is largest when B = 1. As for the other terms in Wt, the first term is small when B is close to d, and the second term is small when B approaches 1. Hence, for B equals some 1 < B̃ < d, supS,S′ supz E[∆̃t+1(z)], and thus the generalization error, grows slower than those of B = d and B = 1. Besides, Proposition 4 also indicates that a larger makes the bound smaller.
Note that the uniform stability bound of SGD (Theorem 3.8 in [12]) is not directly comparable with our Proposition 4. However, as SGD and B = 1 both depend on the same second moment upper bound (Assumption 4), by showing B = B̃ is better than B = 1, we expect blockwise adaptivity to be also better than SGD.

heading: 5 EXPERIMENTS
text: In Section 5.1, we first empirically validate the regret analysis results of BAG (Section 3.3) on a linear model. In Section 5.2, we run deep networks on a number of standard benchmark data sets including CIFAR-10 (Section 5.2.1), ImageNet (Section 5.3), and WikiText-2 (Section 5.4). As the focus is on deep learning, we only use BAGM (instead of BAG) in this section.

heading: 5.1 ILLUSTRATION OF THE REGRET ANALYSIS RESULTS
text: In this section, we use BAG on the linear model ft(θt) = max(0, 1−yt〈θt, xt〉), with input xt ∈ Rd generated in a blockwise manner and yt ∈ {−1, 1} is the label for xt. Assume that xt is partitioned into B̃ blocks, whose structure may be different from that of theB gradient blocks. With probability pb, each element xt,i in input block b is sampled from N (cbyt, γ2b ) for some scaling factor cb and variance γ2b , and xt,i = 0 otherwise. It can be easily shown that for elements in the same input block b, their expected gradient magnitudes have the same upper bound (E[g2t,i] ≤ pb(c2b + γ2b )). We generate a synthetic data set, with d = 100, using the above procedure. The first 50 features of xt are sampled independently fromN (10yt, 100) with probability 0.5, and zero otherwise. The last 50 features are sampled independently from N (−5yt, 25) with probability 0.4, and zero otherwise. The class label yt ∈ {−1, 1} are sampled randomly with equal probabilities. We study BAG with B = 1, 2, 3, 4 and 100. For B = 2, gradient gt is partitioned in the same way as the input. For B = 3, we form the first block using the first 35 coordinates, the second block with the next 30 coordinates, and the third block with the remaining 35 elements. For B = 4, gt is divided into four blocks each of 25 elements. We initialize θ1 to zero, fix = 10−8 and η = 0.01. Figure 1 compares the expected regret, which is estimated by averaging the regrets over 100 repetitions. BAG with B = 2 and 4 achieve lower regrets than the others (as B = 4 covers the case for B = 2). BAG with B = 3 is a little worse as
its block structure is different from the input block structure, but still performs better than B = d. For B = 1, the mismatch in block structures is severe and its performance is worst.
Figure 1: Expected regret with different B’s (note that curves for B = 2 and 4 overlap).

heading: 5.2 REAL-WORLD DATA SETS
text: We introduce four block construction strategies: B.1: Use a single adaptive stepsize for each parameter tensor/matrix/vector. A parameter tensor can be the kernel tensor in a convolution layer, a parameter matrix can be the weight matrix in a fully-connected layer, and a parameter vector can be a bias vector; B.2: Use an adaptive stepsize for each output dimension of the parameter matrix/vector in a fully connected layer, and an adaptive stepsize for each output channel in the convolution layer; B.3: Use an adaptive stepsize for each output dimension of the parameter matrix/vector in a fully
connected layer, and an adaptive stepsize for each kernel in the convolution layer; B.4: Use an adaptive stepsize for each input dimension of the parameter tensor/matrix, and an adaptive stepsize for each parameter vector. More details on the implementation can be found in Appendix A
We compare the proposed BAGM (with block construction approaches B.1, B.2, B.3, B.4) with the following baselines: (i) Nesterov’s accelerated gradient (NAG) (Sutskever et al., 2013); and (ii) Adam (Kingma & Ba, 2015). These two algorithms are widely used in deep networks. NAG provides a strong baseline with good generalization performance, while Adam serves as a fast counterpart with coordinate-wise adaptive stepsize.
As grid search for all hyper-parameters is very computationally expensive, we only tune the most important ones using a validation set and fix the rest. We use a constant βt = β (momentum parameter) and exponential increasing sequence S.3 with α = 0.999 for BAGM. For Adam, we fix its second moment parameter to 0.999 and tune its momentum parameter. Note that with such configurations, Adam is a special case of BAGM with B = d (i.e., weighted AdaEMA). For all the adaptive methods, we use = 10−3 as suggested in (Zaheer et al., 2018). All the experiments are run on a AWS p3.16 instance with 8 NVIDIA V100 GPUs.

heading: 5.2.1 CIFAR-10
text: We train a deep residual network from the MXNet Gluon CV model zoo Guo et al. (2019) on the CIFAR-10 data set. We use the 56-layer and 110-layer networks as in (He et al., 2016). 10% of the training data are carved out as validation set. We perform grid search using the validation set for the initial stepsize η and momentum parameter β on ResNet56. The obtained hyperparameters are then also used on ResNet110. We follow a similar setup as in (He et al., 2016). Details are in Appendix A.2. To reduce statistical variance, results are averaged over 5 repetitions.
top-1 error (%) top-5 error (%) NAG 20.94 5.51 Adam 21.04 5.47
BAGM-B.1 20.79 5.43 BAGM-B.2 20.90 5.39 BAGM-B.3 20.88 5.52 BAGM-B.4 20.82 5.48
Table 2: Validation set errors on ImageNet. The best results are bolded.
Table 1 shows the testing errors of the various methods. With a large = 10−3, the testing performance of Adam matches that of NAG. This agrees with (Zaheer et al., 2018) that a larger reduces adaptivity and improves generalization performance. It also agrees with Proposition 4 that the bound is smaller when is larger. As for BAGM, it outperforms Adam for all block construction schemes used. It also outperforms NAG with schemes B.1, B.2 and B.4.
Convergence of the training, testing, and generalization errors (absolute difference between training error and testing error) are shown in Figure 2.2 As can be seen, on both ResNet models, BAGM-B.1 converges to a lower training error rate than Adam. This agrees with Corollary 2 that blockwise adaptive methods can have faster convergence than their counterparts with element-wise adaptivity. Moreover, the generalization error of BAGM-B.1 is smaller than Adam, which agrees with Proposition 4 that blockwise adaptivity can have a slower growth of generalization error. On both models, BAGM-B.1 gives the smallest generalization error, while NAG has the highest generalization error on ResNet56. Hence, the proposed methods can accelerate convergence and improve generalization.

heading: 5.3 IMAGENET
text: In this experiment, we train a 50-layer ResNet model on ImageNet (Russakovsky et al., 2015). The data set has 1000 classes, 1.28M training samples, and 50,000 validation images. As the data set does not come with labels for its test set, we evaluate its generalization performance on the validation
2To reduce clutterness, we only show results of the block construction scheme BAGM-B.1, which gives the lowest testing error among the proposed block schemes. The full results are shown in Figure 3.
set. We use the ResNet50 v1d network from the MXNet Gluon CV model zoo. We train the FP16 (half precision) model on 8 GPUs, each of which processes 128 images in each iteration. More details are in Appendix A.3. As it takes long time to train on ImageNet, we only run each algorithm once.
Performance on the validation set is shown in Table 2. As can be seen, BAGM with all the block schemes (particularly BAGM-B.1) achieve lower top-1 errors than Adam and NAG. As for the top-5 error, BAGM-B.2 obtains the lowest, which is then followed by BAGM-B.1. Overall, BAGM-B.1 has the best performance on both CIFAR-10 and ImageNet.

heading: 5.4 WORD-LEVEL LANGUAGE MODELING
text: In this section, we train the AWD-LSTM word-level language model (Merity et al., 2018) on the WikiText-2 (WT2) data set (Merity et al., 2017). We use the publicly available implementation in the Gluon NLP toolkit Guo et al. (2019). We perform grid search on the initial learning rate and momentum parameter as in Section 5.2.1, and set the weight decay to 1.2×10−6 as in (Merity et al., 2018). Results are averaged over 3 repetitions. More details on the setup are in Appendix A.4. As there is no convolutional layer, B.2 and B.3 are the same.
Table 3 shows the testing perplexities, the lower the better. As can be seen, all adaptive methods achieve lower test perplexities than NAG, and BAGM-B.2 obtains the best result.

heading: 6 CONCLUSION
text: In this paper, we proposed adapting the stepsize for each parameter block, instead of for each individual parameter as in Adam and RMSprop. Regret, convergence and uniform stability analyses show that it can have lower regret, faster convergence and lower generalization error than its counterpart with coordinate-wise adaptive stepsize. Experiments on synthetic dataset, image classification and language modeling confirm these theoretical results.

heading: A EXPERIMENTAL SETUP
text: A.1 IMPLEMENTATION DETAILS
As {at} is non-decreasing, the accumulated sum At can grow significantly, which may potentially cause some numerical issue. In practice, using steps 7 and 8 in Algorithms 4, we equivalently rewrite the update in (12) as the following exponentially moving update:
v̂t,b = αtv̂t−1,b + (1− αt) ‖gt,Gb‖22 db ,
where αt = 1 − at/At. If at = α−t, then αt = α(1 − αt−1)/(1 − αt). Based on Proposition 3, this setting leads to an O(1) bound. On the other hand, if at = tτ , then we have at/At = O(1/t). This suggests that we can use polynomial-decay averaging αt = 1− (c+ 1)/(t+ c) for some c ≥ 0 (Shamir & Zhang, 2013), whereas c > 0 reduces the weight of earlier iterates compared to later ones. The larger c corresponds to the larger τ . In this case, as ∑T t=1 at = O(T γ) for some γ > 0, we have a convergence rate of O(log(T )/T ).
A.2 CIFAR-10
The CIFAR-10 data set has 50,000 training images and 10,000 testing images. As in (He et al., 2016), we employ data augmentation for training. We first pad the input picture by adding 4 pixels on each side of the image. Then, a 32 × 32 crop is randomly sampled from the padded image with random horizontal flipping. A mini-batch size of 128 is used. The stepsize is divided by 10 at the 39k-th and 59k-th iterations. We use a weight decay of 0.0001.
For NAG, the initial stepsize η is chosen from {0.01, 0.05, 0.1, 0.5, 1}. For the adaptive methods, we have η ∈ {0.0001, 0.0005, 0.001, 0.005, 0.01}. The momentum parameter is searched over {0, 0.5, 0.9}. The learning rate is multiplied by 0.1 at the 100th and 150th epochs. We perform grid search on the hyper-parameters by running each algorithm for 200 epochs on ResNet56. The hyper-parameters that give the highest accuracy on the validation set are employed. The testing performance is obtained by running each algorithm with its best hyper-parameters on full training set for 400 epochs. The same obtained hyperparameters are then used on training ResNet110. When NAG is applied to ResNet110, we use a smaller learning rate at the beginning to warm up the training. Specifically, the obtained learning rate is divided by 10 in the first 4000 iterations, and then go back to the original one and continue training. The grid search results are shown in Table 4.
Figure 3 shows that, on ResNet56, BAGM converges to a lower training error rate than Adam for all schemes used. For the deeper ResNet100 model, BAGM-B.1 and B.4 have faster convergence than Adam, while BAGM-B.2 and B.3 show the same convergence speed with Adam.
A.3 IMAGENET
In this experiment, we employ label smoothing and mixup (Zhang et al., 2018). The cosine schedule (Loshchilov & Hutter, 2017) for learning rate is used. A warmup of 5 epochs is applied. During validation, we use the center crop. Hyperparameter tuning is based on the obtained results in Section A.2. Specifically, for NAG, the initial learning rate is chosen from {0.4, 0.5}, and momentum parameter is fixed to 0.9. For Adam and BAGM, we have the initial learning rate η ∈ {0.004, 0.005}, and we use momentum parameter β = 0. A weight decay of 0.0001 is used (weight decay is not applied to bias vectors, and parameters for batch normalization layers) 3. The best learning rates for each method are shown in Table 5.
A.4 WORD LANGUAGE MODELING
In this experiment, we follow the same setting in (Merity et al., 2018). A 3-layer AWD-LSTM is considered. The model is unrolled for 70 steps, and a mini-batch of size 80 is used. We clip the norm of the gradients at 0.25. The details of the configuration used in this experiment can be found in https://github.com/dmlc/gluon-nlp/blob/master/scripts/
3The example script for running NAG with η = 0.4 can be found in https://raw. githubusercontent.com/dmlc/web-data/master/gluoncv/logs/classification/ imagenet/resnet50_v1d-mixup.sh. Details of the data augmentation can be found in https://github.com/dmlc/gluon-cv/blob/master/scripts/classification/ imagenet/train_imagenet.py.
language_model/word_language_model.py. For completeness, we show the model configuration in Table 6.
As the WikiText-2 data set comes with a validation set, we perform grid search by evaluating the performance on the validation set. For NAG, the initial stepsize is chosen from {1, 3, 10, 30}. For the adaptive methods, we select η ∈ {0.1, 0.03, 0.01, 0.003}. The momentum parameter varies in {0, 0.5, 0.9}. The learning rate is multiplied by 0.1 when the validation performance does not improve for 30 consecutive epochs. We tie the word embeddings and softmax weights. For each algorithm, we employ the iterate averaging scheme proposed in (Merity et al., 2018). The model is trained for 750 epochs. The hyper-parameters obtained by grid search are shown in Table 7. In general, B.1 and B.4 are not suitable for updating the word embedding matrix as word frequency varies a lot and thus the gradient is highly sparse. However, the gradient becomes dense when we use weight tying. In modern toolkits such as Tensorflow, MXNet, and Pytorch, the weight matrices of the LSTM gates are concatenated to speed up matrix-vector multiplication. We need to apply B.1 and B.4 to these weight matrices separately.

heading: B GRADIENTS IN PARAMETER BLOCK
text: As discussed in Corollary 2, BAGM can have faster convergence than its coordinate-wise counterpart when {σ2i }i∈G̃b have low variability. In this section, we verify this experimentally. Using the setup in Section 5.2.1, we focus on BAGM-B.1, which shows the fastest convergence. At the end of each epoch, we perform 10 full data passes with random shuffle and data augmentation (as described in Appendix A.2) to compute E[g2i ] and E[‖gG̃b‖ 2 2]/db. We then approximate σ 2 i and σ 2 b by their empirical maxima over all epochs. Figure 4 shows the coefficient of variation of {σ2i }i∈G̃b , which is defined as the ratio of the standard deviation to the mean (Everitt, 2006). As can be seen, around 86% (resp. 75%) of all blocks for ResNet56 (resp. ResNet110) have coefficient of variation smaller than 1, indicating that {σ2i }i∈G̃b have low variance and concentrate around the mean.
We also compute the empirical values of rmin ≡ min(r1, r2, r3) and √ maxb maxi∈G̃b G2i + 2
maxbG2b+ 2 in
Corollary 2. We use the two symbols (v̄T,B and C(T )) that are defined in Appendix F for the corresponding main theorem. Moreover, we estimate √ v̄T,d+ 2
v̄T,B̃+ 2 , where v̄T,d = v̄T,B=d and
v̄T,B̃ = v̄T,B=B̃ , instead of
√ maxb maxi∈G̃b G2i + 2
maxbG2b+ 2 , as C(T ) is tighter than C̃(T ). We estimate
v̄T,B̃ using max1≤t≤T maxb v̂t,b. The results are shown in Table B.
ResNet56 ResNet110√ v̄T,d+ 2
v̄T,B̃+ 2 min(r1, r2, r3)
√ v̄T,d+ 2
v̄T,B̃+ 2 min(r1, r2, r3)
B.1 3.70 1.02 3.30 1.01 B.2 1.50 1.06 1.73 1.05 B.3 1.47 1.01 1.39 1.01 B.4 2.27 1.17 2.47 1.13
Table 8: Empirical estimate of
√ v̄T,d+ 2
v̄T,B̃+ 2 and min(r1, r2, r3).
As can been seen, all min(r1, r2, r3) are larger than 1, indicating {σ2i }i∈G̃b have low variability for all the proposed blocks, which agree with Corollary 2 and explain why the proposed blockwise adaptivity leads to faster convergence. Moreover, B.1 has largest √ v̄T,d+ 2
v̄T,B̃+ 2 . Theoretically, it means
B.1 should lead to fastest convergence, which is empirically verified in Figure 3. Similarly, B.2 has second largest √ v̄T,d+ 2
v̄T,B̃+ 2 and is therefore the second fastest.

heading: C PROOF OF PROPOSITION 1
text: Proof. In this proof, we use denominator layout for matrix calculus. As all the activation functions are bijective and {Wk}Lk=l+1 are invertible, Φl is bijective and has an inverse function Φ −1 l . Specifically, Φ−1l is given by
Φ−1l (Y ) = φ −1 l (· · ·φ −1 L−2(φ −1 L−1(YW −1 L )W −1 L−1) · · ·W −1 l+1).
Then, with the assumption that Hl−1 has full row rank, the nonconvex objective (5) can be reformulated as the following convex problem:
min Wl ‖Hl−1Wl − Φ−1l (Y )‖ 2 2. (14)
It is obvious that its large margin solution is HTl−1(Hl−1H T l−1) −1Φ−1l (Y ). In the sequel, we will see that every critical point of (5) is a global optimal solution. Let hi,l−1 denotes a column vector
that is the i-th row of Hl−1 and Z:,i be the i-th column of matrix Z. The gradient of (5) is
2HTl−1 d∑ k=1 Diag(Φl(Hl−1Wl):,k − Y:,k)Gk,l = HTl−1El,
where Gk,l = [∇x=hT1,l−1WlΦl(x)k; · · · ;∇x=hTn,l−1WlΦl(x)k] ∈ R n×d and El = 2 ∑d k=1 Diag(Φl(Hl−1Wl):,k − Y:,k)Gk,l to be the error matrix. As Hl−1 has full row rank, then clearly gradient is zero if only and if El = 0. By the definition of Gk,l, we can see that El = 0 if only and if Φl(Hl−1Wl) = Y when ∇x=hTi,l−1WlΦl(x) has full row rank for all i ∈ [n]. Note that the gradient ∇x=hTi,l−1WlΦl(x) is of the following form:
∇x=hTi,l−1WlΦl(x) = (WL ◦ φ ′ L−1(h T i,L−2WL−1) T 1Td ) T · · · (Wl+1 ◦ φ′l(hTi,l−1Wl)T 1Td )T ,
where ◦ is the Hadamard product. For all k ∈ {l, . . . , L− 1}, as Wk+1 has full rank and φ′k(z) 6= 0 for any z ∈ R, we have thatWk+1◦φ′k(hTi,k−1Wk)T 1Td has full rank. Applying the fact that the multiplication of a number of invertible matrices preserves full rank, we obtain that ∇x=hTi,l−1WlΦl(x) has full rank. Therefore, every critical point satisfies Φl(Hl−1Wl) = Y and every critical point is a global optimal solution. Let it be the index chosen at iteration t and yit be the it-th row of Y . Let us define et,l = 2 ∑d k=1(Φl(h
T it,l−1Wt,l)k−yit,k)∇x=hTit,l−1Wt,lΦl(x)k. Now, we prove that if the following update
rule applied on (5) finds a critical point, then the iterate converges to the largest margin solution.
Wt+1,l = Wt,l − ηt,lhit,l−1et,l = HTl−1 − t∑ j=1 ηj,lẼj,l  , (15) where we use Wl,1 = 0, ηt,l is the stepsize for l-th layer at iteration t, and Ẽj,l is a matrix in which its ik-th row is ej,l and all the other rows are zeros. Then, the solution found by (15) lies in the span of rows of Hl−1. In other words, the solution has the following parametric form:
Wl = H T l−1αl
for some αl ∈ Rn. Thus, if (15) is converging to a critical point in expectation, then we have Wt,l → W∗,l as t → ∞, where W∗,l = HTl−1α∗,l for some optimal α∗,l. Since every critical point is an optimal solution, then W∗,l is also a solution to (14), and we have
Φ−1l (Y ) = Hl−1W∗,l = Hl−1H T l−1α∗,l.
We solve for α∗,l and obtain
α∗,l = (Hl−1H T l−1) −1Φ−1l (Y ).
Therefore, W∗,l = HTl−1(Hl−1H T l−1) −1Φ−1l (Y ).

heading: D PROOF OF PROPOSITION 2
text: Proof. Let (xit , yit) be the pair of sample selected at iteration t. The stochastic gradient of least square problem (4) at the t-th iteration is
2(xTitθt − yit)xit = X T et,
where we define et to be the error vector with value 2(xTitθt − yit) in the it-th coordinate and zeros elsewhere. For each block b, BAG with θ1 = 0 uses the following update rule:
θt+1,Gb = θt,Gb − ηt,bXT:,Gbet = X T :,Gb
( −
t∑ i=1 ηi,bei
) ,
where ηt,b = η/( √∑t
i=1 ‖XT:,Gbei‖ 2 2/db + ). Then, each subvector of the solution found by BAG
lies in the span of rows of X:,Gb . In other words, each subvector of the solution is of the following parametric form:
θGb = X T :,Gbαb
for some αb ∈ Rn. Combining with Theorem 1 and online-to-batch conversion (Cesa-Bianchi et al., 2004), BAG is converging in expectation 1t ∑t i=1 θi → θ∗ as t → ∞, where θ∗,Gb = XT:,Gbα∗,b for some optimal α∗,b. Since θ∗ is a solution to (4), we have
y = Xθ∗ = B∑ b=1 X:,GbX T :,Gbα∗,b.
Assume that each submatrix X:,Gb has full row rank, then X:,GbX T :,Gb is invertible, we can solve for α∗,b’s and obtain
α∗,b = (X:,GbX T :,Gb) −1ub for some ub ∈ Rn and ∑B b=1 ub = y.

heading: E PROOF OF THEOREM 1
text: Lemma 1. Let {θt} be the sequence generated by the Algorithm 2. Define st = [( √ vt,1 +
)1Td1 , . . . , ( √ vt,B + )1 T dB ]T . Let Ht = Diag(st). Then, for any θ, we have
ft(θt)− ft(θ) ≤ 1
2η ‖θt − θ‖2Ht −
1
2η ‖θt+1 − θ‖2Ht +
η 2 ‖gt‖2H−1t .
Proof. For any θ, the convexity of ft indicates that
ft(θt)− ft(θ) ≤ 〈gt, θt − θ〉 = 〈gt, θt+1 − θ〉+ 〈gt, θt − θt+1〉
= 1
η 〈θt+1 − θ,Ht(θt − θt+1)〉+ 〈gt, θt − θt+1〉
= 1
2η ‖θt − θ‖2Ht −
1
2η ‖θt+1 − θ‖2Ht −
1
2η ‖θt+1 − θt‖2Ht + 〈gt, θt − θt+1〉
≤ 1 2η ‖θt − θ‖2Ht − 1 2η ‖θt+1 − θ‖2Ht − 1 2η ‖θt+1 − θt‖2Ht + 1 2η ‖θt+1 − θt‖2Ht + η 2 ‖gt‖2H−1t
= 1
2η ‖θt − θ‖2Ht −
1
2η ‖θt+1 − θ‖2Ht +
η 2 ‖gt‖2H−1t ,
where the second to last inequality follows from Fenchel’s inequality applied to the conjugate functions 12η‖ · ‖ 2 Ht and η2‖ · ‖ 2 H−1t .
Lemma 2. Considering an arbitrary R-valued sequence {ai} and its vector representation a1:t = [a1, . . . , at], we have
T∑ t=1 a2t ‖a1:t‖2 ≤ 2‖a1:T ‖2.
Proof. The lemma can be proved by induction. The lemma trivially holds when T = 1. Assume the lemma holds for T − 1, we get
T∑ t=1 a2t ‖a1:t‖2 ≤ 2‖a1:T−1‖2 + a2T ‖a1:T ‖2
= 2 √ Z − x+ x√
Z ,
where we define Z = ‖a1:T ‖22 and x = a2T . As the RHS is non-increasing for x ≥ 0. We can set x = 0 to maximize the bound and obtain 2 √ Z.
Lemma 3. Let Ht be defined as in Lemma 1. Denote g1:t,Gb = [gT1,Gb , . . . , g T t,Gb ] T . We have T∑ t=1 ‖gt‖2H−1t ≤ 2 B∑ b=1 √ db‖g1:T,Gb‖2.
Proof. T∑ t=1 ‖gt‖2H−1t ≤ T∑ t=1 〈gt,Diag(st)−1gt〉
= T∑ t=1 B∑ b=1 √ db‖gt,Gb‖22 ‖g1:t,Gb‖2
= B∑ b=1 √ db T∑ t=1 ‖gt,Gb‖22 ‖g1:t,Gb‖2
≤ 2 B∑ b=1 √ db‖g1:T,Gb‖2.
where the last inequality follows from the Lemma 2 by setting ai = ‖gi,Gb‖22.
E.1 PROOF OF THEOREM 1
Proof. By summing up the equation in Lemma 1 with θ = θ∗, we obtain T∑ t=1 ft(θt)− ft(θ∗) ≤ 1 2η ‖θ1 − θ∗‖2H1 + 1 2η T−1∑ t=1 [ ‖θt+1 − θ∗‖2Ht+1 − ‖θt+1 − θ∗‖ 2 Ht ] + η 2 T∑ t=1 ‖gt‖2H−1t .
By the construction of Ht, we have that Ht+1 Ht. Then, we get ‖θt+1 − θ∗‖2Ht+1 − ‖θt+1 − θ∗‖ 2 Ht
= 〈θt+1 − θ∗,Diag(st+1 − st)(θt+1 − θ∗)〉
= B∑ b=1 ‖θt+1,Gb − θ∗,Gb‖22( √ vt+1,b − √ vt,b).
Given the above result, we have T−1∑ t=1 [ ‖θt+1 − θ∗‖2Ht+1 − ‖θt+1 − θ∗‖ 2 Ht ] =
B∑ b=1 T−1∑ t=1 ‖θt+1,Gb − θ∗,Gb‖22( √ vt+1,b − √ vt,b)
= B∑ b=1 T−1∑ t=1 ‖θt+1,Gb − θ∗,Gb‖22( √ vt+1,b − √ vt,b) + B∑ b=1 ‖θ1,Gb − θ∗,Gb‖22( √ v1,b − √ v1,b)
≤ B∑ b=1 D2b √ vT,b − B∑ b=1 ‖θ1,Gb − θ∗,Gb‖22 √ v1,b.
Recall that vT,b = ‖g1:T,Gb‖22/db. Let = 0. Combining Lemma 3 with the fact that ‖θ1−θ∗‖2H1 =∑B b=1 ‖θ1,Gb − θ∗,Gb‖22 √ v1,b, we have
T∑ t=1 ft(θt)− ft(θ∗) ≤ 1 2η B∑ b=1 D2b√ db ‖g1:T,Gb‖2 + η B∑ b=1 √ db‖g1:T,Gb‖2.
E.2 PROOF OF COROLLARY 1
Proof. Taking expectation of the gradient terms in (9), we have, for all b’s,
E[‖g1:T,Gb‖2] ≤ √√√√∑ i∈Gb T∑ t=1 E[g2t,i] ≤ σb √ dbT .
Let B = B̃. Then, (9) reduces to
E[R(T )] ≤ B̃∑ b=1 σb [ 1 2η D2b + ηdb ]√ T .

heading: F MAIN THEOREM
text: Let σt,b = √ Et[|gt,Gb‖22]. We define a sequence of virtual estimates of the second moment:
ṽt,b = 1
At (vt−1,b + at db Et[‖gt,Gb‖22])
= 1
At ( t−1∑ i=1 ai ‖gi,Gb‖22 db + at σ2t,b db ) ∀b ∈ [B]. (16)
Let v̂T,B ≡ max1≤t≤T E[maxb ṽt,b] Theorem 2. Suppose that Assumptions 3-7 hold. Let ρ = β2/p̃. We have
min 1≤t≤T
(E[‖∇F (θt)‖4/32 ])3/2 ≤ C(T ),
where4
C(T ) =
√ 2 (v̄T,B + 2)
ηTT
[ 2C2
(1− β)C1 C0 + C4
[ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )
+ B∑ b=1 C ′b
[ w1 log ( σ2b 2 + 1 ) + ω T∑ t=1 ηt √ at At ]]] , (17)
C0 = F (θ1) − F (θ∗), C4 = 2C 2 2
C21 √ Ca(1− √ ρ)(1−β) , C
′ b = LC32w1db C31Ca(1− √ ρ)2 + 2C23C2σbdb C1 , and C3 =
β/(1−β)√ CaA1/A2(1−ρ) + 1.
When B = d, the bound here is tighter than that in (Zou et al., 2019), as we exploit heterogeneous second-order upper bound (Assumption 4). Note that constants C0, C3, C4 are not related to block partition. They only depend on initial optimality gap and sequences {ηt}, {at}, {βt}. In the following, we introduce several lemmas to prove the Theorem 2.
In the sequel, we define Ht as Ht = Diag(st),
where st = [( √ v̂t,1 + )1 T d1 , . . . , ( √ v̂t,B + )1 T dB ] T .
Let δt = θt+1 − θt = −ηtmt/st. We introduce H̃t as
H̃t = Diag(s̃t),
where ṽt,b is defined in (16) and s̃t = [( √ ṽt,1 + )1 T d1 , . . . , ( √ ṽt,B + )1 T dB ] T . Assume that σt,b/ √ db ≤ σb for all t and let Σ = Diag([σ211Td1 , . . . , σ 2 B1 T dB ]T ).
4When T = 1, the second term in C(T ) (involving summation from t = 2 to T ) disappears.
Lemma 4. Let St = S0 + ∑t i=1 ai, where {at} is a non-negative sequence and S0 > 0. We have
T∑ t=1 at St ≤ log(ST )− log(S0)
Proof. The concavity of log leads to log(b) ≤ log(a) + 1a (b− a) for all a, b > 0. This suggests that
a− b a ≤ log(a)− log(b) = log (a b ) .
Hence, we have
T∑ t=1 at St = T∑ t=1 St − St−1 St ≤ T∑ t=1 log ( St St−1 ) = log(ST )− log(S0).
Lemma 5. Let {at} and {st} be two real number sequences, and let St = ∑t i=1 si. Then, we have
T∑ t=1 atst = T−1∑ t=1 (at − at+1)St + aTST .
Proof. Let S0 = 0. Expanding the summation, we obtain
T∑ t=1 atst = T∑ t=1 at(St − St−1)
= T−1∑ t=1 atSt − T−1∑ t=1 at+1St + aTST
= T−1∑ t=1 (at − at+1)St + atST
Lemma 6. Assume {at} is non-decreasing such that {At−1/At} is non-decreasing. Define wt = ηt/ √
at At . Assume wt is ”almost” non-increasing. This means there exists another non-increasing
sequence {zt} and positive constants C1 and C2 such that C1zt ≤ wt ≤ C2zt. Then,
wt ≤ C2/C1wi and ηt ≤ C2/C1ηi for all i < t.
Proof. For any i < t,
wt ≤ C2zt ≤ C2zi ≤ C2/C1wi.
Then,
ηt ≤ C2 √ at/At
C1 √ ai/Ai
ηi = C2 √
1−At−1/At C1 √ 1−Ai−1/Ai ηi ≤ C2/C1ηi.
Lemma 7. Assume that {at} is non-decreasing. For any block diagonal matrix C = Diag([c11Td1 , . . . , cB1 T dB ]T ) with cb ≥ 0 for all b, we have
T∑ t=1 E [ at At ‖H−1t gt‖2C ] ≤ B∑ b=1 cbdb [ log ( σ2b 2 + 1 ) + log ( 1 a1 T∑ i=1 ai + 1 )] .
Proof.
T∑ t=1 at At ‖H−1t gt‖2C = T∑ t=1 B∑ b=1 cb at At ‖gt,Gb‖22 ( √ v̂t,b + )2
≤ T∑ t=1 B∑ b=1 cb at At ‖gt,Gb‖22 v̂t,b + 2
= B∑ b=1 T∑ t=1 cbat‖gt,Gb‖22∑t i=1 ai‖gi,Gb‖22/db +At 2
= B∑ b=1 cbdb T∑ t=1 at‖gt,Gb‖22∑t i=1 ai‖gi,Gb‖22 + dbAt 2
≤ B∑ b=1 cbdb T∑ t=1 at‖gt,Gb‖22∑t i=1 ai‖gi,Gb‖22 + dba1 2 .
Hence,
T∑ t=1 at At ‖H−1t gt‖2C ≤ B∑ b=1 cbdb
[ log ( T∑ i=1 ai‖gi,Gb‖22 + dba1 2 ) − log(dba1 2) ] ,
where the inequality follows from Lemma 4. Using Jensen’s inequality, we get
T∑ t=1 E [ at At ‖H−1t gt‖2C ] ≤ B∑ b=1 cbdbE [ log ( T∑ i=1 ai‖gi,Gb‖22 + dba1 2 ) − log(dba1 2) ]
≤ B∑ b=1 cbdb
[ log ( T∑ i=1 aiE[‖gi,Gb‖22] + dba1 2 ) − log(dba1 2) ]
≤ B∑ b=1 cbdb
[ log ( dbσ 2 b T∑ i=1 ai + dba1 2 ) − log(dba1 2) ]
= B∑ b=1 cbdb log ( σ2b a1 2 T∑ i=1 ai + 1 ) .
Using the inequality log(1 + ab) ≤ log(1 + a+ b+ ab) = log(1 + a) + log(1 + b) for a, b ≥ 0, we have
T∑ t=1 E [ at At ‖H−1t gt‖2C ] ≤ B∑ b=1 cbdb [ log ( σ2b 2 + 1 ) + log ( 1 a1 T∑ i=1 ai + 1 )] .
Lemma 8. Assume that {at} is non-decreasing. Define wt = ηt/ √
at At . Assume wt is ”al-
most” non-increasing. This means there exists another non-increasing sequence {zt} and positive constants C1 and C2 such that C1zt ≤ wt ≤ C2zt for all t. For any block diagonal matrix C = Diag([c11Td1 , . . . , cB1 T dB ]T ) with cb ≥ 0 for all b, we have T∑ t=1 ηtE [√ at At ‖H−1t gt‖2C ] ≤ C2 C1 [ w1 B∑ b=1 cbdb log ( σ2b 2 + 1 ) + B∑ b=1 cbdb T∑ t=1 ηt √ at At At At−1 + a1 ] .
Proof. Let ξt = atAt ‖H −1 t gt‖2C , then ζt = ∑t i=1 ξi. Lemma 5 indicates that we have
T∑ t=1 ηt √ at At ‖H−1t gt‖2C = T∑ t=1 wtξt
≤ C2 T∑ t=1 ztξt
= C2 [ T−1∑ t=1 (zt − zt+1)ζt + zT ζT ] .
Define Mt = ∑B b=1 cbdb [ log ( σ2b 2 + 1 ) + log ( 1 a1 ∑t i=1 ai + 1 )] . By Lemma 7, we have E[ζt] ≤ Mt. Then,
T∑ t=1 ηtE [√ at At ‖H−1t gt‖2C ] ≤ C2 [ T−1∑ t=1 (zt − zt+1)E[ζt] + zTE[ζT ] ]
≤ C2 [ T−1∑ t=1 (zt − zt+1)Mt + zTMT ] ,
where the last inequality follows from the assumption that zt ≥ zt+1. Then, T∑ t=1 ηtE [√ at At ‖H−1t gt‖2C ] ≤ C2 [ T−1∑ t=1 (zt − zt+1)Mt + zTMT ]
= C2 [ T∑ t=1 zt(Mt −Mt−1) + z1M0 ]
= C2 [ z1
B∑ b=1 cbdb log ( σ2b 2 + 1 ) + T∑ t=1 zt B∑ b=1 cbdb log ( At + a1 At−1 + a1 )]
≤ C2 C1
[ w1
B∑ b=1 cbdb log ( σ2b 2 + 1 ) + T∑ t=1 wt B∑ b=1 cbdb log ( At + a1 At−1 + a1 )] .
As log(1 + x) ≤ x for x > −1 and the fact that At ≥ At−1, we get
log ( At + a1 At−1 + a1 ) = log ( 1 + At + a1 At−1 + a1 − 1 ) ≤ At + a1 At−1 + a1 − 1 = at At−1 + a1 .
Hence, T∑ t=1 ηtE [√ at At ‖H−1t gt‖2C ] ≤ C2 C1 [ w1 B∑ b=1 cbdb log ( σ2b 2 + 1 ) + T∑ t=1 wt B∑ b=1 cbdb log ( At + a1 At−1 + a1 )]
≤ C2 C1
[ w1
B∑ b=1 cbdb log ( σ2b 2 + 1 ) + B∑ b=1 cbdb T∑ t=1 ηt √ at At At At−1 + a1 ] .
Lemma 9. Let η̃t,b = ηt√ ṽt,b+ . For each block b and t ≥ 2, we have
( δt −
βtηt√ 1− at/Atηt−1 δt−1 ) Gb = −(1− βt)η̃t,bgt,Gb + η̃t,b at Atdb ‖gt,Gb‖22√ v̂t,b + Xt,b + η̃t,b σt,b√ db Yt,b + Zt,b,
where
Xt,b = βtmt−1,Gb√ v̂t,b + √ At−1v̂t−1,b/At + (1− βt)gt,Gb√ ṽt,b + √ v̂t,b ,
Yt,b =
at At √ db ‖gt,Gb‖2√
v̂t,b +
βtmt−1,Gb√ At−1v̂t−1,b/At +
√ at Atdb ‖gt,Gb‖2√
v̂t,b + √ At−1v̂t−1,b/At
√ at Atdb
σt,b√ ṽt,b + √ At−1v̂t−1,b/At
− at At gt,Gb√
v̂t,b + (1− βt) σt,b√db√ ṽt,b + √ v̂t,b ,
Zt,b = βtηtmt−1,Gb
( 1− √ At−1/At )
( √ At−1v̂t−1,b/At + )( √ At−1v̂t−1,b/At + √ At−1/At ) .
Proof. For any t ≥ 2,
δt − βtηt√
At−1/Atηt−1 δt−1
= −ηtmt st + βtηtmt−1√ At−1/Atst−1
= −ηt [ mt st − βtmt−1√
At−1/Atst−1
]
= − (1− βt)ηtgt√ vt + − βtηtmt−1
[ 1
√ vt + − 1√ At−1vt−1/At +
]
−βtηtmt−1
[ 1√
At−1st−1/At + − 1√ At−1st−1/At + √ At−1/At
]
= − (1− βt)ηtgt√ vt + − βtηtmt−1
[ 1
√ vt + − 1√ At−1vt−1/At +
] (18)
+βtηtmt−1
( 1− √ At−1/At )
( √ At−1vt−1/At + )( √ At−1vt−1/At + √ At−1/At )
Let expand the first term of (18) as
(1− βt)ηtgt√ vt + = (1− βt)ηtgt√ ṽt + + (1− βt)ηtgt
[ 1
√ vt + − 1√ ṽt + ] =
(1− βt)ηtgt√ ṽt + + (1− βt)ηtgt ṽt − vt ( √ vt + )( √ ṽt + )( √ ṽt + √ vt) .
For each block b, we have
(1− βt)ηtgt,Gb√ v̂t,b +
= (1− βt)ηtgt,Gb√
ṽt,b + + (1− βt)ηtgt,Gb
at Atdb (σ2t,b − ‖gt,Gb‖22) ( √ v̂t,b + )( √ ṽt,b + )( √ ṽt,b + √ v̂t,b)
= (1− βt)η̃t,bgt,Gb + η̃t,b σt,b√ db
at At gt,Gb√
v̂t,b + (1− βt) σt,b√db√ ṽt,b + √ v̂t,b − η̃t,b at Atdb ‖gt,Gb‖22√ v̂t,b + (1− βt)gt,Gb√ ṽt,b + √ v̂t,b .(19)
Then, we expand the second term of (18):
βtηtmt−1
[ 1
√ vt + − 1√ At−1vt−1/At +
]
= βtηtmt−1
√ At−1vt−1/At − √ vt
( √ vt + )( √ At−1vt−1/At + )
= βtηtmt−1 At−1vt−1/At − vt
( √ vt + )( √ At−1vt/At + )( √ vt + √ At−1vt−1/At) .
Similarly, for each block b, we have
βtηtmt−1,Gb
[ 1√
v̂t,b + − 1√ At−1v̂t−1,b/At +
]
= −βtηtmt−1,Gb at‖gt,Gb‖22/(Atdb) ( √ v̂t,b + )( √ At−1v̂t−1,b/At + )( √ v̂t,b + √ At−1v̂t−1,b/At)
= −βtηtmt−1,Gb
[ at‖gt,Gb‖22/(Atdb)
( √ v̂t,b + )( √ ṽt,b + )( √ v̂t,b + √ At−1v̂t−1,b/At)
+ at‖gt,Gb‖22/(Atdb) ( √ v̂t,b + )( √ v̂t,b + √ At−1v̂t−1,b/At)
[ 1√
At−1v̂t−1,b/At + − 1√ ṽt,b +
]]
= −η̃t,b at Atdb ‖gt,Gb‖22√ v̂t,b + βtmt−1,Gb√ v̂t,b + √ At−1v̂t−1,b/At
−η̃t,b σt,b√ db  atAt√db ‖gt,Gb‖2√ v̂t,b + βtmt−1,Gb√ At−1v̂t−1,b/At + √ at Atdb ‖gt,Gb‖2√ v̂t,b + √ At−1v̂t−1,b/At
√ at Atdb
σt,b√ ṽt,b + √ At−1v̂t−1,b/At  .(20) Combining (19) and (20) into (18), we obtain the result. Lemma 10. Suppose that {at} is a non-decreasing sequence and At = ∑t i=1 at such that
{At/At+1} is non-decreasing and limt→∞ AtAt+1 = p > 0. Let Ât,i = ∏t j=i+1 Aj−1 Aj
for 1 ≤ i < t and Ât,t = 1. For a fixed constant p̃ such that β2 < p̃ < p, we have
Ât,i ≥ Cap̃t−i ,
where Ca = (∏N
j=2 Aj−1 Aj p̃
) and N is the maximum of the indices for which Aj−1/Aj < p̃. When
there are no such indices, i.e., A1/A2 ≥ p̃, we use Ca = 1 by convention.
Proof.
Ât,i = t∏ j=i+1 Aj−1 Aj ≥  N∏ j=i+1 Aj−1 Aj  p̃t−N =  N∏ j=i+1 Aj−1 Aj p̃  p̃t−i ≥  N∏ j=2 Aj−1 Aj p̃  p̃t−i.
Lemma 11. Suppose that 0 ≤ βt ≤ β < 1 for all t. Let ρ := β 2
p̃ , where p̃ is defined in Lemma 10. Then, for all t, we have
‖mt,Gb‖22 ≤ 1
Caat/(Atdb)(1− ρ) v̂t,b,
where Ca is defined in Lemma 10.
Proof. Let β̂t,i = ∏t j=i+1 βj for i < t and β̂t,t = 1
‖mt,Gb‖22 = ∥∥∥∥∥ t∑ i=1 (1− βi)β̂t,igi,Gb ∥∥∥∥∥ 2
2
= ∥∥∥∥∥∥ t∑ i=1 (1− βi)β̂t,i√ ai Atdb √ ai Atdb gi,Gb ∥∥∥∥∥∥ 2
2
≤ ( t∑ i=1 (1− βi)2β̂2t,i ai Atdb )( t∑ i=1 ai Atdb ‖gi,Gb‖ 2 2 )
= ( t∑ i=1 (1− βi)2β̂2t,i ai Atdb ) v̂t,b. (21)
Then, with Lemma 10, we get t∑ i=1 (1− βi)2β̂2t,i ai Atdb = t∑ i=1 (1− βi)2β̂2t,i ai Aidb Ât,i ≤ 1 Caat/(Atdb) t∑ i=1 ( β2 p̃ )t−i ≤ 1 Caat/(Atdb)(1− ρ) . (22)
Then, combining (21) and (22), we obtain the result.
Lemma 12. Assume F is L-smooth, {at} is non-decreasing such that {At−1/At} is non-decreasing and limt→∞ AtAt+1 = p > 0. Let p̃ be a constant such that β 2 < p̃ < p. Assume Et[‖gt,Gb‖22] =
σ2t,b ≤ dbσ2b . Define wt = ηt/ √ at At . Assume wt is ”almost” non-increasing. This means there exists another non-increasing sequence {zt} and positive constants C1 and C2 such that C1zt ≤ wt ≤ C2zt for all t. Assume 0 ≤ βt ≤ β < 1 for all t. Define following Lyapunov function:
Mt = E[〈∇F (θt), δt〉+ L‖δt‖22]. Let C3 ≡ [
β/(1−β)√ CaA1/A2(1−ρ) + 1
] , where ρ := β 2
p̃ . Then, for any t ≥ 2, we have
Mt ≤ βtηt√
At−1/Atηt−1 Mt−1 − 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] + 2wtC 2 3E [ at At ‖H−1t gt‖2Σ1/2 ]
+LE[‖δt‖22] + βwt√
Ca(1− ρ) (√ At At−1 − 1 ) B∑ b=1 σbdb, (23)
and for t = 1, we have
M1 ≤ − 1− β1
2 η1E [ ‖∇F (θ1)‖2H̃−11 ] + 2w1C 2 3E [ a1 A1 ‖H−11 g1‖2Σ1/2 ] + LE[‖δ1‖22]. (24)
Proof. For any t ≥ 2,
E[〈∇F (θt), δt〉] = βtηt√
At−1/Atηt−1 E[〈∇F (θt), δt−1〉] + E
[〈 ∇F (θt), δt −
βtηt√ At−1/Atηt−1 δt−1
〉] . (25)
Then, for the first term of (25), we have 〈∇F (θt), δt−1〉 = 〈∇F (θt−1), δt−1〉+ 〈∇F (θt)−∇F (θt−1), δt−1〉
≤ 〈∇F (θt−1), δt−1〉+ L‖θt − θt−1‖2‖δt−1‖2 = 〈∇F (θt−1), δt−1〉+ L‖δt−1‖22,
where the first inequality follows from Schwartz inequality and the smoothness of the function F . Hence, we have
βtηt√ At−1/Atηt−1 E[〈∇F (θt), δt−1〉] ≤ βtηt√ At−1/Atηt−1 E [ 〈∇F (θt−1), δt−1〉+ L‖δt−1‖22 ] =
βtηt√ At−1/Atηt−1 Mt−1.
Now, we estimate the second term of (25). By Lemma 9, for each block b, we get
E [〈 ∇GbF (θt), δt,Gb −
βtηt√ At−1/Atηt−1 δt−1,Gb
〉]
= −(1− βt)E[〈∇GbF (θt), η̃t,bgt,Gb〉] + E [〈 ∇GbF (θt), η̃t,b
at Atdb ‖gt,Gb‖22√ v̂t,b + Xt,b
〉]
+E [〈 ∇GbF (θt), η̃t,b
σt,b√ db Yt,b
〉] + E[〈∇GbF (θt), Zt,b〉]. (26)
For the first term of (26), we have
−(1− βt)E[〈∇GbF (θt), η̃t,bgt,Gb〉] = −(1− βt)E[〈∇GbF (θt), η̃t,b∇GbF (θt)〉] = −(1− βt)η̃t,bE[‖∇GbF (θt)‖22]. (27)
For the second term of (26), we have
E [〈 ∇GbF (θt), η̃t,b
at Atdb ‖gt,Gb‖22√ v̂t,b + Xt,b
〉]
≤ E
[√ η̃t,b‖∇GbF (θt)‖2‖gt,Gb‖2/ √ db
σt,b/ √ db
√ η̃t,b
at At ‖gt,Gb‖2/
√ dbσt,b/ √ db‖Xt,b‖2√
v̂t,b +
] . (28)
Note that
√ η̃t,bσt,b/ √ db = √ ηtσ2t,b/db√ ṽt,b + ≤ √√√√ ηtσ2t,b/db√ at/Atσ2t,b/db ≤ √ ηtσb√ at/At = √ wtσb. (29)
Besides, we have
‖Xt,b‖2 = ∥∥∥∥∥ βtmt−1,Gb√v̂t,b +√At−1v̂t−1,b/At + (1− βt)gt,Gb√ṽt,b +√v̂t,b ∥∥∥∥∥
2
≤ ∥∥∥∥∥ βtmt−1,Gb√v̂t,b +√At−1v̂t−1,b/At ∥∥∥∥∥
2
+ ∥∥∥∥∥ (1− βt)gt,Gb√ṽt,b +√v̂t,b ∥∥∥∥∥
2
.
With Lemma 11, we have∥∥∥∥∥ mt−1,Gb√v̂t,b +√At−1v̂t−1,b/At ∥∥∥∥∥
2
≤ ∥∥∥∥∥ mt−1,Gb√At−1v̂t−1,b/At ∥∥∥∥∥
2
≤ 1√ CaAt−1/Atat/(Atdb)(1− ρ)
, (30)∥∥∥∥∥ gt,Gb√ṽt,b +√v̂t,b ∥∥∥∥∥
2
≤ ∥∥∥∥∥ gt,Gb√v̂t,b ∥∥∥∥∥
2
≤ ∥∥∥∥∥ gt,Gb√at/At‖gt,Gb‖22/db ∥∥∥∥∥
2
= √ db√
at/At . (31)
Then, we get
‖Xt,b‖2 ≤ βt√
CaAt−1/Atat/(Atdb)(1− ρ) +
(1− βt) √ db√
at/At
= [ βt/(1− βt)√
CaAt−1/At(1− ρ) + 1
] (1− βt)
√ db√
at/At
≤ [ β/(1− β)√
CaAt−1/At(1− ρ) + 1
] (1− βt)
√ db√
at/At
≤ [ β/(1− β)√
CaA1/A2(1− ρ) + 1
] (1− βt)
√ db√
at/At := C3
(1− βt) √ db√
at/At ,
where the last-to-second inequality follows from the assumption that βt ≤ β, and the last inequality holds as we assume {at} is chosen such that {At−1/At} is non-decreasing for all t. Hence, combining the above result with (29) and (28), we have
E [〈 ∇GbF (θt), η̃t,b
at Atdb ‖gt,Gb‖22√ v̂t,b + Xt,b
〉]
≤ E √η̃t,b‖∇GbF (θt)‖2‖gt,Gb‖2/√db σt,b/ √ db √ wtσbC3(1− βt) √ at At ‖gt,Gb‖2√ v̂t,b +  ≤ E [ 1− βt
4 η̃t,b‖∇GbF (θt)‖22‖gt,Gb‖22/db σ2t,b/db + wtσbC 2 3 (1− βt) at At ‖gt,Gb‖22 ( √ v̂t,b + )2
]
≤ E
[ 1− βt
4 η̃t,b‖∇GbF (θt)‖22Et[‖gt,Gb‖2]2/db σ2t,b/db + wtσbC 2 3 (1− βt) at At ‖gt,Gb‖22 ( √ v̂t,b + )2
]
≤ E
[ 1− βt
4 η̃t,b‖∇GbF (θt)‖22 + wtσbC23
at At ‖gt,Gb‖22 ( √ v̂t,b + )2
] , (32)
where the second inequality follows from ab ≤ a 2 2c + cb2
2 for any c > 0. Now, we estimate the third term of (26):
E [〈 ∇GbF (θt), η̃t,b
σt,b√ db Yt,b
〉] ≤ E [√ η̃t,b ‖∇GbF (θt)‖2 √ η̃t,b
σt,b√ db ‖Yt,b‖2
] .
Similarly, with (30) and (31), by expanding ‖Yt,b‖2, we have
‖Yt,b‖2 ≤ at At √ db ‖gt,Gb‖2√
v̂t,b +
βt‖mt−1,Gb‖2√ At−1v̂t−1,b/At +
√ at Atdb ‖gt,Gb‖2√
v̂t,b + √ At−1v̂t−1,b/At
√ at Atdb
σt,b√ ṽt,b + √ At−1v̂t−1,b/At
+ at At ‖gt,Gb‖2√ v̂t,b + (1− βt) σt,b√db√ ṽt,b + √ v̂t,b
≤ at At √ db ‖gt,Gb‖2√
v̂t,b +
βt√ CaAt−1/Atat/(Atdb)(1− ρ) + at At ‖gt,Gb‖2√ v̂t,b + 1− βt√ at/At
= √ at At ‖gt,Gb‖2√ v̂t,b + βt√ CaAt−1/At(1− ρ) + √ at At ‖gt,Gb‖2√ v̂t,b + (1− βt)
= √ at At ‖gt,Gb‖2√ v̂t,b + [ βt/(1− βt)√ CaAt−1/At(1− ρ) + 1 ] (1− βt)
≤ √ at At ‖gt,Gb‖2√ v̂t,b + C3(1− βt),
where C3 is the constant defined above. Hence, together with (29), we obtain
E [〈 ∇GbF (θt), η̃t,b
σt,b√ db Yt,b
〉] ≤ E √η̃t,b ‖∇GbF (θt)‖2√η̃t,b σt,b√db √ at At ‖gt,Gb‖2√ v̂t,b + C3(1− βt)  ≤ 1− βt
4 η̃t,bE
[ ‖∇GbF (θt)‖ 2 2 ] + wtσbC 2 3 (1− βt)E [ at At ‖gt,Gb‖22
( √ v̂t,b + )2
]
≤ 1− βt 4
η̃t,bE [ ‖∇GbF (θt)‖ 2 2 ] + wtσbC 2 3E [ at At ‖gt,Gb‖22
( √ v̂t,b + )2
] . (33)
The last term of (26) can be bounded as follows E[〈∇GbF (θt), Zt,b〉] ≤ E[‖∇GbF (θt)‖2‖Zt,b‖2] ≤ E[σb √ db‖Zt,b‖2],
and with (30), we get
‖Zt,b‖2 ≤ βtηt ‖mt−1,Gb‖2√ At−1v̂t−1,b/At + √ At−1/At
( 1− √ At−1/At )
( √ At−1v̂t−1,b/At + )
≤ βtηt√ CaAt−1/Atat/(Atdb)(1− ρ)
( 1− √ At−1/At ) ≤ βηt √ Atdb/at√
Ca(1− ρ) (√ At At−1 − 1 )
= βwt √ db√
Ca(1− ρ) (√ At At−1 − 1 ) .
Hence,
E[〈∇GbF (θt), Zt,b〉] ≤ βσbdbwt√ Ca(1− ρ) (√ At At−1 − 1 ) . (34)
Combining (26), (27), (32), (33), and (34), we get
E [〈 ∇GbF (θt), δt,Gb −
βtηt√ At−1/Atηt−1 δt−1,Gb
〉]
≤ −1− βt 2
η̃t,bE [ ‖∇GbF (θt)‖ 2 2 ] + 2wtσbC 2 3E [ at At ‖gt,Gb‖22
( √ v̂t,b + )2
] +
βσbdbwt√ Ca(1− ρ) (√ At At−1 − 1 ) .
Summing from b = 1 to B, we obtain
E [〈 ∇F (θt), δt −
βtηt√ At−1/Atηt−1 δt−1
〉]
≤ −1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] + 2wtC 2 3E [ at At ‖H−1t gt‖2Σ1/2 ] +
βwt√ Ca(1− ρ) (√ At At−1 − 1 ) B∑ b=1 σbdb.
Then, with (25), we have E[〈∇F (θt), δt〉]
≤ βtηt√ At−1/Atηt−1 Mt−1 − 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] + 2wtC 2 3E [ at At ‖H−1t gt‖2Σ1/2 ]
+ βwt√
Ca(1− ρ) (√ At At−1 − 1 ) B∑ b=1 σbdb.
We obtain (23) by adding the term LE[‖δt‖22] to both sides of the above equation. When t = 1, we have
M1 = E[−〈∇F (θ1), η1m1/( √ v̂1 + )〉+ L‖δ1‖22]
= E[−〈∇F (θ1), η1(1− β1)g1/( √ v̂1 + )〉+ L‖δ1‖22]. (35)
Then, following the derivation of (19), for each block b, we have
(1− β1)η1g1,Gb√ v̂1,b + = (1− β1)η̃1,bg1,Gb + η̃1,b σ1,b√ db a1 A1 g1,Gb√ v̂1,b + (1− β1) σ1,b√db√ ṽ1,b + √ v̂1,b
−η̃1,b a1 A1db ‖g1,Gb‖22√ v̂1,b + (1− β1)g1,Gb√ ṽ1,b + √ v̂1,b .
Hence, with similar argument, we get
E [ − 〈 ∇F (θ1), η1
(1− β1)g1√ v̂1 +
〉] ≤ −1− β1
2 η1E [ ‖∇F (θ1)‖2H̃−11 ] + 2w1C 2 3E [ a1 A1 ‖H−11 g1‖2Σ1/2 ] .
Combining above with (35), and adding LE[‖δ‖22], we obtain (24).
Lemma 13. With the same assumptions in Lemma 12, we have T∑ t=1 ‖δt‖22 ≤ C22/C 2 1w1 Ca(1− √ ρ)2 T∑ t=1 wt at At ‖H−1t gt‖22.
Proof. For each block b,
‖mt,Gb‖2 = ∥∥∥∥∥∥ t∑ i=1  t∏ j=i+1 βj  (1− βi)gi,Gb ∥∥∥∥∥∥
2 ≤ t∑ i=1  t∏ j=i+1 βj(1− βi)  ‖gi,Gb‖2 ≤
t∑ i=1 βt−i ‖gi,Gb‖2 .
Then,
‖mt,Gb‖2√ v̂t,b + ≤ t∑ i=1 βt−i ‖gi,Gb‖2√ v̂t,b + .
Since v̂t,b ≥ At−1v̂t−1,b/At, we have v̂t,b ≥ (∏t j=i+1Aj−1/Aj ) v̂i,b = Ât,iv̂i,b ≥ Cap̃t−iv̂i,b by Lemma 10. It follows that
‖mt,Gb‖2√ v̂t,b + ≤ t∑ i=1 βt−i ‖gi,Gb‖2√ v̂t,b + ≤ 1√ Ca t∑ i=1 ( β√ p̃ )t−i ‖gi,Gb‖2√ v̂i,b + = 1√ Ca t∑ i=1 √ ρ t−i ‖gi,Gb‖2√ v̂i,b + .
Then, as at/At = 1−At−1/At is non-decreasing, we have ‖δt‖22 = B∑ b=1 ∥∥∥∥∥ ηtmt,Gb√v̂t,b + ∥∥∥∥∥ 2
2
= B∑ b=1 ∥∥∥∥∥wt √ at/Atmt,Gb√ v̂t,b + ∥∥∥∥∥ 2
2
≤ w 2 t
Ca B∑ b=1 ( t∑ i=1 √ ρ t−i √ at/At ‖gi,Gb‖2√ v̂i,b + )2
≤ w 2 t
Ca B∑ b=1 ( t∑ i=1 √ ρ t−i √ ai/Ai ‖gi,Gb‖2√ v̂i,b + )2
≤ w 2 t
Ca B∑ b=1  t∑ j=1 √ ρ t−j 2 t∑ i=1 √ ρt−i(∑t
j=1
√ ρt−j ) √ai/Ai ‖gi,Gb‖2√ v̂i,b + 2
≤ w 2 t
Ca B∑ b=1  t∑ j=1 √ ρ t−j  t∑ i=1 √ ρ t−i ai/Ai ‖gi,Gb‖ 2 2 ( √ v̂i,b + )2 ≤ w 2 t Ca(1− √ ρ) B∑ b=1 t∑ i=1 √ ρ t−i ai/Ai ‖gi,Gb‖ 2 2 ( √ v̂i,b + )2
= w2t
Ca(1− √ ρ) t∑ i=1 √ ρ t−i ai Ai ‖H−1i gi‖ 2 2.
As wt ≤ C2/C1wi for any i ≤ t by Lemma 6, then we have
‖δt‖22 ≤ C22/C 2 1w1
Ca(1− √ ρ) t∑ i=1 √ ρ t−i wi ai Ai ‖H−1i gi‖ 2 2.
Hence,
T∑ t=1 ‖δt‖22 ≤ C22/C 2 1w1 Ca(1− √ ρ) T∑ t=1 t∑ i=1 √ ρ t−i wi ai Ai ‖H−1i gi‖ 2 2
= C22/C 2 1w1
Ca(1− √ ρ) T∑ i=1 T∑ t=i √ ρ t−i wi ai Ai ‖H−1i gi‖ 2 2
≤ C 2 2/C 2 1w1
Ca(1− √ ρ)2 T∑ i=1 wi ai Ai ‖H−1i gi‖ 2 2.
Lemma 14. With the same assumptions in Lemma 12, letMt = E[〈∇F (θt), δt〉+L‖δt‖22], we have T∑ t=1 Mt ≤ C2 C1 √ Ca(1− √ ρ) [ 2C23 T∑ t=1 wtE [ at At ‖H−1t gt‖2Σ1/2 ] + LC22/C 2 1w1 Ca(1− √ ρ)2 T∑ t=1 wtE [ at At ‖H−1t gt‖22 ]
+ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )] − 1− β 2 T∑ t=1 ηtE [ ‖∇F (θt)‖2H̃−1t ] .
Proof. Let define following quantity
Nt = 2wtC 2 3E [ at At ‖H−1t gt‖2Σ1/2 ] + LE[‖δt‖22] + βwt√ Ca(1− ρ) (√ At At−1 − 1 ) B∑ b=1 σbdb, ∀t ≥ 2,
N1 = 2w1C 2 3E [ a1 A1 ‖H−11 g1‖2Σ1/2 ] + LE[‖δ1‖22].
Then, by Lemma 12, for any t ≥ 2, we have
Mt ≤ βtηt√
At−1/Atηt−1 Mt−1 − 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] +Nt
≤ βtηt√ At−1/Atηt−1 Mt−1 +Nt
and M1 ≤ N1. Then, by recursively applying above relation, we get
Mt ≤ β̂t,1ηt√ Ât,1η1 M1 + t∑ i=2 β̂t,iηt√ Ât,iηi Ni − 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ]
≤ t∑ i=1 β̂t,iηt√ Ât,iηi Ni − 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] ,
where β̂t,i = ∏t j=i+1 βj for i < t and β̂t,t = 1 and Ât,i = ∏t j=i+1 Aj−1 Aj for i < t and Ât,t = 1. Note that β̂t,i ≤ βt−i, and ηt ≤ C2/C1ηi. By Lemma 10, we have Ât,i ≥ Cap̃t−i . Then,
Mt ≤ C2
C1 √ Ca t∑ i=1 ( β√ p̃ )t−i Ni − 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] = C2
C1 √ Ca t∑ i=1 √ ρ t−i Ni − 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] .
It can be verified that the above inequality holds for t = 1 as C2/(C1 √ Ca) ≥ 1. Then, summing from t = 1 to t = T , we obtain T∑ t=1 Mt ≤ C2 C1 √ Ca T∑ t=1 t∑ i=1 √ ρ t−i Ni − T∑ t=1 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t
] = C2
C1 √ Ca T∑ i=1 T∑ t=i √ ρ t−i Ni − T∑ t=1 1− βt 2 ηtE [ ‖∇F (θt)‖2H̃−1t ] ≤ C2
C1 √ Ca(1− √ ρ) T∑ t=1 Nt − 1− β 2 T∑ t=1 ηtE [ ‖∇F (θt)‖2H̃−1t ] . (36)
With Lemma 13, we get T∑ t=1 Nt = T∑ t=1 [ 2wtC 2 3E [ at At ‖H−1t gt‖2Σ1/2 ] + LE[‖δt‖22] ] + T∑ t=2 βwt√ Ca(1− ρ) (√ At At−1 − 1 ) B∑ b=1 σbdb
≤ 2C23 T∑ t=1 wtE [ at At ‖H−1t gt‖2Σ1/2 ] + LC22/C 2 1w1 Ca(1− √ ρ)2 T∑ t=1 wtE [ at At ‖H−1t gt‖22 ]
+ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 ) .
Combining the above with (36), we obtain the result. Lemma 15. Assume {at} is non-decreasing such that {At−1/At} is non-decreasing. Define wt = ηt/ √
at At . Assume wt is ”almost” non-increasing. This means there exists another non-increasing
sequence {zt} and positive constants C1 and C2 such that C1zt ≤ wt ≤ C2zt. We have
1
T T∑ t=1 ( E[‖∇F (θt)‖4/32 ] )3/2 ≤ √ 2 (max1≤t≤T E [maxb ṽt,b] + 2) C1/C2ηTT T∑ t=1 ηtE [ ‖∇F (θt)‖2H̃−1t ] .
Proof. By Hölder’s inequality, we have E[|XY |] ≤ (E[|X|p])1/p(E[|Y |q])1/q for any 0 < p, q < 1 with 1/p+ 1/q = 1. Taking p = 3/2, q = 3, and
X = ( ‖∇F (θt)‖22√ maxb ṽt,b + )2/3 , Y = (√ max b ṽt,b + )2/3 ,
we obtain
E[‖∇F (θt)‖4/32 ] ≤
( E [ ‖∇F (θt)‖22√ maxb ṽt,b + ])2/3( E [(√ max b ṽt,b + )2])1/3 .
Hence,( E[‖∇F (θt)‖4/32 ] )3/2 ≤ ( E [ ‖∇F (θt)‖22√ maxb ṽt,b + ])( E [(√ max b ṽt,b + )2])1/2 .
Note that
‖∇F (θt)‖22√ maxb ṽt,b +
= B∑ i=1 ‖∇GiF (θt)‖22√ maxb ṽt,b +
≤ B∑ b=1 ‖∇GbF (θt)‖22√ ṽt,b +
= ‖∇F (θt)‖2H̃−1t .
We also have
E [(√ max b ṽt,b + )2] ≤ 2E [( max b ṽt,b + 2 )] = 2 ( E [ max b ṽt,b ] + 2 ) .
Then, for any t ≤ T , we get
( E[‖∇F (θt)‖4/32 ] )3/2 ≤ √ 2 ( E [ max b ṽt,b ] + 2 ) E [ ‖∇F (θt)‖2H̃−1t ] = √ 2 (E [maxb ṽt,b] + 2)
ηt ηtE [ ‖∇F (θt)‖2H̃−1t ] ≤ √ 2 (max1≤t≤T E [maxb ṽt,b] + 2)
C1/C2ηT ηtE [ ‖∇F (θt)‖2H̃−1t ] ,
where the last inequality follows from Lemma 6. Taking average from t = 1 to T , we get
1
T T∑ t=1 ( E[‖∇F (θt)‖4/32 ] )3/2 ≤ √ 2 (max1≤t≤T E [maxb ṽt,b] + 2) C1/C2ηTT T∑ t=1 ηtE [ ‖∇F (θt)‖2H̃−1t ] .
F.1 PROOF OF THEOREM 2
Proof. As F is L-smooth, then we have
F (θt+1) ≤ F (θt) + 〈∇F (θt), θt+1 − θt〉+ L
2 ‖θt+1 − θt‖2.
Recursively applying the above relation, we get
F (θ∗) ≤ E[F (θT+1)] ≤ F (θ1) + T∑ t=1 Mt,
where Mt = E[〈∇F (θt), δt〉+ L‖δt‖22]. By Lemma 14, we have
1− β 2 T∑ t=1 ηtE [ ‖∇F (θt)‖2H̃−1t ] ≤ F (θ1)− F (θ∗) + C2
C1 √ Ca(1− √ ρ)
[ 2C23 T∑ t=1 wtE [ at At ‖H−1t gt‖2Σ1/2 ]
+ LC22/C 2 1w1
Ca(1− √ ρ)2 T∑ t=1 wtE [ at At ‖H−1t gt‖22 ] + β√ Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )]
= F (θ1)− F (θ∗) + C2
C1 √ Ca(1− √ ρ)
[ 2C23 T∑ t=1 ηtE [√ at At ‖H−1t gt‖2Σ1/2 ]
+ LC22/C 2 1w1
Ca(1− √ ρ)2 T∑ t=1 ηtE [√ at At ‖H−1t gt‖22 ] + β√ Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )] .
Applying Lemma 8, we have
1− β 2 T∑ t=1 ηtE [ ‖∇F (θt)‖2H̃−1t ] ≤ F (θ1)− F (θ∗)
+ C2
C1 √ Ca(1− √ ρ)
[ 2C23
C2 C1
[ w1
B∑ b=1 σbdb log ( σ2b 2 + 1 ) + B∑ b=1 σbdb T∑ t=1 ηt √ at At At At−1 + a1 ]
+ LC32/C 3 1w1
Ca(1− √ ρ)2
[ w1
B∑ b=1 db log ( σ2b 2 + 1 ) + B∑ b=1 db T∑ t=1 ηt √ at At At At−1 + a1 ]
+ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )] ≤ F (θ1)− F (θ∗)
+ C2
C1 √ Ca(1− √ ρ)
[ 2C23
C2 C1
[ w1
B∑ b=1 σbdb log ( σ2b 2 + 1 ) + ω B∑ b=1 σbdb T∑ t=1 ηt √ at At ]
+ LC32/C 3 1w1
Ca(1− √ ρ)2
[ w1
B∑ b=1 db log ( σ2b 2 + 1 ) + ω B∑ b=1 db T∑ t=1 ηt √ at At ]
+ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )]
= F (θ1)− F (θ∗) + C2
C1 √ Ca(1− √ ρ)
[ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )
+ B∑ b=1 [ LC33/C 3 1w1db Ca(1− √ ρ)2 + 2C23C2σbdb C1 ][ w1 log ( σ2b 2 + 1 ) + ω T∑ t=1 ηt √ at At ]] .
Combining above with Lemma 15, we have
min 1≤t≤T
( E[‖∇F (θt)‖4/32 ] )3/2 ≤ 1 T T∑ t=1 ( E[‖∇F (θt)‖4/32 ] )3/2 ≤ 2 √ 2 (max1≤t≤T E [maxb ṽt,b] + 2)
C1/C2(1− β)ηTT [F (θ1)− F (θ∗)
+ C2
C1 √ Ca(1− √ ρ)
[ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 )
+ B∑ b=1 [ LC32/C 3 1w1db Ca(1− √ ρ)2 + 2C23C2σbdb C1 ] [ w1 log ( σ2b 2 + 1 ) + ω T∑ t=1 ηt √ at At ]]]
=
√ 2 (max1≤t≤T E [maxb ṽt,b] + 2)
ηTT
[ 2C2
(1− β)C1 [F (θ1)− F (θ∗)]
+ 2C22
C21 √ Ca(1− √ ρ)(1− β)
[ β√
Ca(1− ρ) B∑ b=1 σbdb T∑ t=2 wt (√ At At−1 − 1 ) (37)
+ B∑ b=1 [ LC32w1db C31Ca(1− √ ρ)2 + 2C23C2σbdb C1 ] [ w1 log ( σ2b 2 + 1 ) + ω T∑ t=1 ηt √ at At ]]] .(38)
F.2 PROOF OF PROPOSITION 3
To prove the result, we use the following high probability bound.
Proposition 5. With probability at least 1− δ2/3, min1≤t≤T ‖∇F (θt)‖22 ≤ C(T )/δ.
Proof. By the concavity of the minimum, we have
E [ min
1≤t≤T ‖∇F (θt)‖4/32
]3/2 ≤ min
1≤t≤T
( E[‖∇F (θt)‖4/32 ] )3/2 .
Let X = min1≤t≤T ‖∇F (θt)‖22. The Theorem 2 suggests that we have E[X2/3] ≤ C(T )2/3. By Markov’s inequality, we get
P ( X2/3 > C(T )2/3
δ2/3
) ≤ E[X 2/3]
C(T )2/3 δ2/3 ≤ δ2/3.
Hence, P ( X > C(T )δ ) ≤ δ2/3, and we have P (X ≤ C(T )δ ) ≥ 1− δ 2/3.
Proof. (of Proposition 3) Recall the definition of C(T ) in (17). When at = atτ , we have At = O(t1+τ ). This suggests that
ηt √ at At ≤ η 1− β √ at tAt = O ( 1 t ) ,
wt = η
1− β̃t √ At tat ≤ η 1− β √ At tat = O (1) ,
and √ At At−1 − 1 = √ At − √ At−1√ At−1 = O ( 1 t ) .
Hence, T∑ t=1 ηt √ at At = O (log(T )) , T∑ t=1 wt (√ At At−1 − 1 ) = O (log(T )) , and C(T ) = O ( log(T )√ T ) .
On the other hand, when at = α−t, we have
ηt √ at At ≤ η 1− β √ 1− α (1− αt)t ≤ η√ t ,
wt = η
1− β̃t √ At att = η 1− β̃t √ 1− αt (1− α)t ≤ η (1− β) √ (1− α)t ,
and √ At At−1 − 1 = √ 1− αt (1− αt−1)α ≤ √ 1 + α α .
Then, we get
T∑ t=1 ηt √ at At ≤ 2η √ T , T∑ t=1 wt (√ At At−1 − 1 ) = O (√ T ) , and C(T ) = O (1) .
Combining the results with Proposition 5, we complete the proof.
F.3 PROOF OF COROLLARY 2
Proof. As ‖gt,Gb‖22/db ≤ G2b , then we have
ṽt,b = (vt−1,b + atEt[‖gt,G̃b‖ 2 2]/db])/At ≤ G2b ,
and therefore v̄T,B ≡ max1≤t≤T E [maxb ṽt,b] ≤ maxbG2b . Arranging the terms in C̃(T ), we obtain
C̃(T )
=
√ 2 (maxbG2b + 2)
ηTT
[ 2C2
(1− β)C1 [F (θ1)− F (θ∗)]
+ 2C22
C21 √ Ca(1− √ ρ)(1− β)
[[ β√
Ca(1− ρ) T∑ t=2 wt (√ At At−1 − 1 ) + 2C23C2 C1 ω T∑ t=1 ηt √ at At ] B∑ b=1 σbdb
+ LC32w 2 1
C31Ca(1− √ ρ)2 B∑ b=1 db log ( σ2b 2 + 1 ) + 2C23C2w1 C1 B∑ b=1 σbdb log ( σ2b 2 + 1 )
+ LC32w1dω
C31Ca(1− √ ρ)2 T∑ t=1 ηt √ at At ]] .
When B = d, we have
C̃d(T )
=
√ 2 ( maxb maxi∈G̃b G 2 i + 2 )
ηTT
[ 2C2
(1− β)C1 [F (θ1)− F (θ∗)]
+ 2C22
C21 √ Ca(1− √ ρ)(1− β)
[[ β√
Ca(1− ρ) T∑ t=2 wt (√ At At−1 − 1 ) + 2C23C2 C1 ω T∑ t=1 ηt √ at At ] d∑ i=1 σi
+ LC32w 2 1
C31Ca(1− √ ρ)2 d∑ i=1 log ( σ2i 2 + 1 ) + 2C23C2w1 C1 d∑ i=1 σi log ( σ2i 2 + 1 )
+ LC32w1dω
C31Ca(1− √ ρ)2 T∑ t=1 ηt √ at At ]]
=
√ 2 ( maxb maxi∈G̃b G 2 i + 2 )
ηTT
[ 2C2
(1− β)C1 [F (θ1)− F (θ∗)]
+ 2C22
C21 √ Ca(1− √ ρ)(1− β) [ β√ Ca(1− ρ) T∑ t=2 wt (√ At At−1 − 1 ) + 2C23C2 C1 ω T∑ t=1 ηt √ at At ] B̃∑ b=1 ∑ i∈G̃b σi
+ LC32w 2 1
C31Ca(1− √ ρ)2 B̃∑ b=1 ∑ i∈G̃b log ( σ2i 2 + 1 ) + 2C23C2w1 C1 B̃∑ b=1 ∑ i∈G̃b σi log ( σ2i 2 + 1 )
+ LC32w1dω
C31Ca(1− √ ρ)2 T∑ t=1 ηt √ at At ]] .
Substituting r1 := ∑B̃ b=1 ∑ i∈G̃b log(σ2i / 2+1)∑B̃
b=1 db log(σ2b/ 2+1) , r2 :=
∑B̃ b=1 ∑ i∈G̃b
σi∑B̃ b=1 σbdb
and r3 :=∑B̃ b=1 ∑ i∈G̃b
σi log(σ2i / 2+1)∑B̃
b=1 σbdb log(σ2b/ 2+1) , we get
C̃d(T )
=
√ 2 ( maxb maxi∈G̃b G 2 i + 2 )
ηTT
[ 2C2
(1− β)C1 [F (θ1)− F (θ∗)]
+ 2C22
C21 √ Ca(1− √ ρ)(1− β) [ β√ Ca(1− ρ) T∑ t=2 wt (√ At At−1 − 1 ) + 2C23C2 C1 ω T∑ t=1 ηt √ at At ] r2 B̃∑ b=1 σbdb
+ LC32w 2 1
C31Ca(1− √ ρ)2 r1 B̃∑ b=1 db log ( σ2b 2 + 1 ) + 2C23C2w1 C1 r3 B̃∑ b=1 σbdb log ( σ2b 2 + 1 )
+ LC32w1dω
C31Ca(1− √ ρ)2 T∑ t=1 ηt √ at At ]]
≥ min(1, rmin)
√ maxb maxi∈G̃b G 2 i + 2
maxbG2b + 2
CB̃(T ),
where rmin = min(r1, r2, r3). Then, C̃d(T ) C̃B̃(T ) ≥ min(1, rmin)
√ maxb maxi∈G̃b G2i + 2
maxbG2b+ 2 . When B = B̃, we assume that Assumption 4 is tight in the sense that σ2b ≤ 1db ∑ i∈G̃b σ 2 i ,
5 Thus, rmin can be larger than 1 as σ2b ≤ 1db ∑ i∈G̃b σ 2 i . Corollary 2 then indicates that blockwise adaptive stepsize will lead
to improvement if √
(maxb maxi∈G̃b G 2 i + 2)/(maxbG2b + 2) > 1rmin . Similarly, assume that the upper bound Gb is tight so that G2b ≤ 1db ∑ i∈G̃b G 2 i . Thus, maxb maxi∈G̃b G 2 i ≥ maxbG2b , and the above condition is likely to hold when rmin is close to 1 or is larger than 1. From the definitions of r1, r2 and r3, we can see that they get close to 1 or is larger than 1 when {σ2i }i∈G̃b have sufficiently low variability.
F.4 PROOF OF PROPOSITION 4
Proof. As function is γ̃-Lipschitz, we have the following result:
sup z
EM [f(M(S); z)− f(M(S′); z)] ≤ γ̃EM [‖M(S)−M(S′)‖2].
Therefore, we can consider bounding EM [‖M(S)−M(S′)‖2]. Let βt = 0 for all t.
θt+1 = θ1 − t∑
k=1
ηkH −1 k mk
= θ1 − t∑
k=1
ηkH −1 k gk
= θ1 − t∑
k=1
ηkH −1 k ∇f(θk; zik),
5Note that 1 db Et[‖gt,G̃b‖ 2 2] = 1 db ∑ i∈G̃b Et[g 2 t,i] ≤ 1db ∑ i∈G̃b σ2i . On the other hand, 1dbEt[‖gt,G̃b‖ 2 2] ≤
σ2b . Thus, this bound is tight in the sense that σ 2 b ≤ 1db ∑ i∈G̃b σ2i .
where ik ∈ [n] is the example index selected at iteration k. Then, we can bound ∆t+1 = ‖θt+1 − θ′t+1‖2 as follows
E[∆t+1] = E[‖θt+1 − θ′t+1‖2]
= E[‖θ1 − θ′1 − t∑
k=1
ηkH −1 k ∇f(θk; zik) + t∑ k=1 ηkH ′−1 k ∇f(θ ′ k; z ′ ik )‖2]
≤ E[‖θ1 − θ′1‖2] + t∑
k=1
ηkE[‖H−1k ∇f(θk; zik)−H ′−1 k ∇f(θ ′ k; z ′ ik )‖2]
= t∑ k=1 ηkE[‖H−1k ∇f(θk; zik)−H ′−1 k ∇f(θ ′ k; z ′ ik )‖2]. (39)
Note that zik = z ′ ik with probability 1−1/n. Then, we can bound each term E[‖H−1k ∇f(θk; zik)− H ′−1 k ∇f(θ′k; z′ik)‖2] as follows
E[‖H−1k ∇f(θk; zik)−H ′−1 k ∇f(θ ′ k; z ′ ik )‖2]
≤ 2 n E[‖H−1k ∇f(θk; zik)‖2] +
( 1− 1
n
) E[‖H−1k ∇f(θk; zik)−H ′−1 k ∇f(θ ′ k; zik)‖2]
≤ 2 n E[‖H−1k ∇f(θk; zik)‖2] +
( 1− 1
n
) E[‖H−1k ∇f(θk; zik)−H ′−1 k ∇f(θk; zik)‖2]
+ ( 1− 1
n
) E[‖H
′−1 k ∇f(θk; zik)−H ′−1 k ∇f(θ ′ k; zik)‖2]. (40)
The second term is bounded as
E[‖H−1k ∇f(θk; zik)−H ′−1 k ∇f(θk; zik)‖2]
≤ E[‖H−1k −H ′−1 k ‖2‖∇f(θk; zik)‖2] ≤ γ̃E[‖H−1k −H ′−1 k ‖2]
= γ̃E max b ∣∣∣∣∣∣ 1√v̂k,b + − 1√v̂′k,b + ∣∣∣∣∣∣  .
We expand the third term of (40) as
E[‖H ′−1 k ∇f(θk; zik)−H ′−1 k ∇f(θ ′ k; zik)‖2]
≤ E[‖H ′−1 k ‖2‖∇f(θk; zik)−∇f(θ ′ k; zik)‖2] ≤ LE[‖H ′−1 k ‖2‖θk − θ ′ k‖2]
≤ LE [ 1√
minb v̂k,b + ‖θk − θ′k‖2
]
= LE [ 1√
minb v̂k,b + ∆k
] .
Substituting the above results into (40) and combining with (39), we obtain
E[∆t+1] ≤ 2
n t∑ k=1 ηkE[‖H−1k ∇f(θk; zik)‖2]
+ ( 1− 1
n
) γ̃ t∑ k=1 ηkE max b ∣∣∣∣∣∣ 1√v̂k,b + − 1√v̂′k,b + ∣∣∣∣∣∣ 
+ ( 1− 1
n
) L t∑ k=1 ηkE [ 1√ minb v̂k,b + ∆k ] .
Note that if wt = ηt/ √ at/At is ”almost” non-increasing w.r.t. another non-increasing sequence {zt} and positive constant C1 and C2, then w2t is also ”almost” non-increasing w.r.t. another nonincreasing sequence {z2t } and positive constant C21 and C22 . Using Lemma 8 with C = I , we have
t∑ k=1 ηkE[‖H−1k ∇f(θk; zik)‖2]
≤ √ t √√√√ t∑ k=1 η2kE[‖H −1 k ∇f(θk; zik)‖22]
= √ t √√√√ t∑ k=1 η2k √ Ak ak E [√ ak Ak ‖H−1k ∇f(θk; zik)‖22 ]
≤ √ t √√√√C22 C21 [ w21 B∑ b=1 db log ( σ2b 2 + 1 ) + d t∑ k=1 η2k Ak Ak−1 + a1 ]
≤ C2 C1 √√√√[w21 B∑ b=1 db log ( σ2b 2 + 1 ) + dω t∑ k=1 η2k ] t.
Then, we get
E[∆t+1] ≤ 2C2 nC1 √√√√[w21 B∑ b=1 db log ( σ2b 2 + 1 ) + dω t∑ k=1 η2k ] t
+ ( 1− 1
n
) γ̃ t∑ k=1 ηkE max b ∣∣∣∣∣∣ 1√v̂k,b + − 1√v̂′k,b + ∣∣∣∣∣∣ 
+ ( 1− 1
n
) L t∑ k=1 ηkE [ 1√ minb v̂k,b + ∆k ] .
As the bound holds for any S, S′, z, it also holds for its supreme.

Review summaries:
This paper proposes blockwise adaptivity.  We divide the parameters into blocks, for example in a linear threshold unit the bias term is in a bias term block while the input weights are in an input weight block.  We then average the square norm of the gradients over each block and use the same adaptation based on this average square norm for all parameters in the block.  theoretical and experimental results are given.

In this paper, the authors propose a generalization of AdaGrad, called BAG, that operates on blocks of parameters instead of each individual parameter. The authors also propose a momentum version of BAG called BAGM. Convergence rate results are proved for the algorithm, and some uniform stability results show situations where BAG would generalize better than previous adaptive gradient methods.

The paper proposes adaptive gradient approaches where the step-size is not determined on the per-coordinate basis but rather for blocks of coordinates. Theoretical results are presented in terms of regret in online convex optimization, regarding convergence in non-convex optimization,  and with respect to uniform stability and generalization. These indicate that under certain conditions adaptivity at the block level could outperform coordinate-wise adaptivity. The approach is evaluated against alternatives on simulated and real-world problems.