# Understanding Progressive Training Through the Framework of
Randomized Coordinate Descent

Rafał Szlendak  Elnur Gasanov  Peter Richtárik

###### Abstract

We propose a Randomized Progressive Training algorithm (RPT) – a stochastic
proxy for the well-known Progressive Training method (PT) (Karras et al.,
2017). Originally designed to train GANs (Goodfellow et al., 2014), PT was
proposed as a heuristic, with no convergence analysis even for the simplest
objective functions. On the contrary, to the best of our knowledge, RPT is the
first PT-type algorithm with rigorous and sound theoretical guarantees for
general smooth objective functions. We cast our method into the established
framework of Randomized Coordinate Descent (RCD) (Nesterov, 2012; Richtárik &
Takáč, 2014), for which (as a by-product of our investigations) we also
propose a novel, simple and general convergence analysis encapsulating
strongly-convex, convex and nonconvex objectives. We then use this framework
to establish a convergence theory for RPT. Finally, we validate the
effectiveness of our method through extensive computational experiments.

Progressive Training, IST, RCD

##  1 Introduction

The modern practice of supervised learning typically relies on training huge-
dimensional, overparametrized models using first-order gradient methods.
However, in most cases, using the full gradient information in the training
procedure is computationally infeasible. The gold-standard approach of
bypassing this difficulty is provided by the family of Stochastic Gradient
Descent (SGD) algorithms which use only a partial information about the
gradients for the reward of reduced compute time. Progressive Training (PT)
(Karras et al., 2017), albeit a non-stochastic method, relies on a similar
idea. In its original form, it was developed as a method for training
generative adversarial networks (GANs) (Goodfellow et al., 2014) by
progressively growing the neural network that we are training. Specifically,
for a fixed period of time we only train the first layer of our GAN while
keeping the weights corresponding to the other layers intact; then, we perform
updates with respect to the first two layers, and so on.

Karras et al. (2017) justified two main advantages of this approach: an
increased stability of the training (since the training procedure starts from
a simpler, reduced model), and reduced computation cost. The method was
proposed as a heuristic, with no convergence guarantees even for the simplest
objectives, such as quadratics or convex functions.

Observe that in the PT procedure, the domain of the parameters which we are
trying to learn is partitioned into blocks which correspond to the respective
layers of the neural network that PT aims to train. Motivated by this, in this
work we are going to be concerned with minimizing the general objective
functions $f$ whose domain will be written as
$\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}$. For example,
$\mathbb{R}^{d_{i}}$ might represent the parameters of the $i^{\text{th}}$
layer of the neural network in question.

Algorithm 1 provides formal statement of the PT procedure for the general
differentiable function
$f:\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}\rightarrow\mathbb{R}$.
The training procedure is divided into $B$ epochs; in each epoch $b\in[B]$ we
want to update the first $b$ coordinate blocks. Let
$G_{b}f(x)\in\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}$ denote the
vector

$\left(\nabla_{1}f(x),\dots,\nabla_{b}f(x),0,\dots,0\right),$

where $\nabla_{i}f$ is the gradient of $f$ with respect to the $i^{\text{th}}$
coordinate block. In each epoch $b$ we make the gradient update

$x\leftarrow x-\gamma_{b}G_{b}f(x)$

for $T_{b}$ iterations, where $\gamma_{b}>0$ is a suitably chosen stepsize
which is fixed for the entire epoch $b$.

1: Input: initial iterate
$x\in\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}$; stepsizes
$\gamma_{1},\gamma_{2},\cdots,\gamma_{n}>0$; integers
$T_{1},T_{2},\cdots,T_{n}\geq 0$

2: for $b=1,\cdots,B$ do

3: for $t=0,\dots,T_{b}-1$ do

4: $x\leftarrow x-\gamma_{b}G_{b}f(x)$

5: end for

6: end for

7: Output: $x$

Algorithm 1 Progressive Training

###  1.1 Existing literature

Wang et al. (2022) attempt to analyze the federated version of PT called
ProgFed. Unfortunately, as argued in Appendix A, their analysis is vacuous.

Recently, a novel and related research thread emerged around the idea of
_independent subnet training_ (IST) which relies on partitioning the neural
network in question into sparser subnets, training them in parallel across the
devices and periodically aggregating the model. Yuan et al. (2022) introduce
IST and provide its analysis. However, the authors make rather strong
assumptions on the objective function, such as Lipschitz continuity of the
objective function and certain assumptions on the stochastic gradient error.
Analogous technique is applied by Dun et al. (2021) for training ResNet (He et
al., 2016a). It is important to note that the IST framework differs from RPT
in several key components. Indeed, IST partitions the neural network in
question into sparser submodels which are then trained across the devices _in
parallel_. The model is then aggregated and this procedure repeats. In
contrast, RPT performs a _single update at a time_ of a random subnetwork
_chosen in a very particular way_ (which simulates the deterministic PT).
Moreover, IST is a distributed method; in contrast, in our paper we focus on
the single-node setup.

Notably, none of the aforementioned papers contains mathematically sound
analysis for general smooth functions.

###  1.2 Our contributions and organization of the paper

As we have argued, there is still a considerable gap in theoretical
understanding of the PT mechanism. _The key objective of this paper is fixing
this gap by proposing a randomized proxy for PT which mimics its behavior and
has simple and general analysis with mild assumptions._ The contributions of
our work can be summarized as follows:

  * •

In Section 2 we develop a simple and general analysis of Randomized Coordinate
Descent (RCD) for strongly-convex, convex and nonconvex objectives. Our
analysis recovers previous convergence rates of RCD and allows us to fully
characterize its iteration complexity in stronlgy-convex, convex and nonconvex
regimes via a single quantity $L_{\mathbf{P}}$.

  * •

In Sections 3 and 4 we propose a method called RPT which mimics the behavior
of PT, and, as we argue, is the first PT-type method with sound theoretical
guarantees for smooth objectives. Instead of progressively expanding the model
which we are trying to optimize, in each iteration RPT makes a choice of the
submodel according to a carefully chosen randomness, and then makes a gradient
update of this model. We then view RPT as a particular instance of RCD, which
immediately establishes its convergence. We finally identify regimes in which
RPT can be advantageous over vanilla gradient descent (GD) through a novel and
general notion of total computation complexity.

  * •

In Section 5, we present a series of empirical experiments that demonstrate
the advantages of our proposed method. These experiments are divided into
three sub-sections. The first sub-section employs synthetic data to
demonstrate that, when appropriately generated, our proposed Randomized
Progressive Training (RPT) method can converge at a faster rate than Gradient
Descent while also incurring less computation cost. In the second and third
sub-sections, we utilize real-world datasets to evaluate the performance of
the RPT method, including testing on a ridge regression problem across various
datasets and on an image classification task on the CIFAR10 dataset. The
results demonstrate that the RPT method exhibits comparable or even superior
performance in comparison to other methods.

###  1.3 Assumptions

In this section we lay out the assumptions that we are going to make
throughout the whole work. In an attempt to make the analysis as general as
possible, in this paper we focus on the problem of minimizing an arbitrary
differentiable function $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$, where
$\mathbb{R}^{d}$ is decomposed as
$\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}$ for some pre-specified
$d_{1},\dots,d_{B}$.

We first recall the notion of _matrix smoothness_ (Qu et al., 2016; Safaryan
et al., 2021) which generalizes the standard $L$-smoothness assumption and
allows one to encapsulate more curvature information of the objective
function.

######  Definition 1.

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ and $\mathbf{L}$ be a positive
semidefinite matrix. We say that $f$ is $\mathbf{L}$-smooth if

$f(x)-f(y)-\left\langle\nabla
f(y),x-y\right\rangle\leq\frac{\left\|x-y\right\|_{\mathbf{L}}^{2}}{2}$

for all $x,y\in\mathbb{R}^{d}$, where
$\left\|v\right\|_{\mathbf{L}}\stackrel{{\scriptstyle\text{def}}}{{=}}\sqrt{v^{\top}\mathbf{L}v}$.

We make the following assumption throughout the whole paper.

######  Assumption 1.

$f$ is $\mathbf{L}$-smooth.

With 1 taken for granted, we analyse RPT in each of the three standard
settings:

######  Assumption 2.

$f$ is $\mu$-strongly convex, that is, there exists a positive constant $\mu$
such that

$f(x)-f(y)-\left\langle\nabla
f(y),x-y\right\rangle\geq\frac{\mu\left\|x-y\right\|^{2}}{2}$

for all $x,y\in\mathbb{R}^{d}$.

######  Assumption 3.

$f$ is convex and has a global minimizer $x^{*}$.

######  Assumption 4.

$f$ is nonconvex and lower-bounded by $f^{*}$.

Under assumptions 2, 3 and 4 respectively we will be interesting in minimizing
the following quantities:

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|x^{t}-x^{*}\right\|^{2}\right]\text{,
}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f(x^{t})-f(x^{*})\right]\text{,
}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\nabla
f(x^{t})\right\|^{2}\right].$

##  2 New Analysis of Randomized Coordinate Descent

###  2.1 SkGD

Randomized Coordinate Descent (RCD) is a celebrated technique of minimizing
high-dimensional objective functions $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$,
where updating all gradient coordinates is prohibitively expensive. In its
original form (Nesterov, 2012; Richtárik & Takáč, 2014) RCD used a _biased_
randomized estimator of the gradient; in each iteration it chooses a random
sub-block of coordinates and performs a gradient step with respect to that
sub-block alone.

Recently, Safaryan et al. (2021) developed an _unbiased_ version of RCD
(called SkGD) which updates the iterates through _unbiased sketches_.
Specifically, in each iteration $t$, SkGD starts from choosing a random set of
coordinates $S^{t}\subset[d]$. Denote $p_{j}=\mathbf{Prob}(j\in S^{t})$ and
assume $p_{j}>0$. Then SkGD performs an update of form

$x^{t+1}=x^{t}-\gamma\mathbf{C}^{t}\nabla f(x^{t}),$

where

$\mathbf{C}^{t}=\operatorname{Diag}(c_{1},\dots,c_{d}),\text{\ \ \
}c_{j}=\begin{cases}1/p_{j}\text{ \ \ if }j\in S^{t}\\\ 0\text{ \ \ \ \ \ \ \
\ otherwise.}\end{cases}$

It is then readily verifiable that $\mathbf{C}^{t}\nabla f(x^{t})$ is an
unbiased gradient estimator, that is

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}^{t}\nabla
f(x^{t})\right]=\nabla f(x^{t}).$

In the remainder of this paper we will always denote

$\mathbf{P}\stackrel{{\scriptstyle\text{def}}}{{=}}\operatorname{Diag}(p_{1},\dots,p_{d}).$

###  2.2 Key quantity: $L_{\mathbf{P}}$

We now define the most important component of our analysis. In the rest of the
paper $\lambda_{\max}\left(\mathbf{M}\right)$ will denote the largest
eigenvalue of matrix $\mathbf{M}$.

######  Definition 2.

Let $\mathbf{C}$ be an unbiased sketch operator and $\mathbf{P}$ its
associated probability matrix. Denote the smoothness matrix of the objective
function $f$ by $\mathbf{L}$. Then

$L_{\mathbf{P}}\stackrel{{\scriptstyle\text{def}}}{{=}}\lambda_{\max}\left(\mathbf{P}^{-1/2}\mathbf{L}\mathbf{P}^{-1/2}\right).$

It turns out that this quantity describes the convergence of RCD in _all of
the regimes in question_ (assumptions 2, 3, 4). Therefore, by estimating it
for a particular choice of sampling distribution $\mathbf{P}$ one obtains a
unified analysis of RCD in the aforementioned regimes. Moreover, in the
following lemma we provide sharp estimates for $L_{\mathbf{P}}$.

######  Lemma 1.

Let $\mathbf{L}$ be any positive semidefinite matrix, $\mathbf{C}$ – any
unbiased sketch operator, and $\mathbf{P}$ – its probability matrix. Then

$\displaystyle\lambda_{\max}({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}\mathbf{L}\mathbf{C}\right])\leq
L_{\mathbf{P}}\leq\lambda_{\max}(\mathbf{L})\lambda_{\max}(\mathbf{P}^{-1}).$

The first bound will be useful in proving the lemma on the convergence of RCD
for nonconvex function. The second one quantifies how large can
$L_{\mathbf{P}}$ be in terms of the smoothness constant of $f$ and
$\mathbf{P}$. In practice, however, $L_{\mathbf{P}}$ can be much smaller.

###  2.3 Convergence results

We begin with a preliminary lemma:

######  Lemma 2.

Let $\mathbf{W}=\operatorname{Diag}(w_{1},\dots,w_{d})$, where
$w_{1},\dots,w_{d}$ are any positive entries and assume that
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is $\mathbf{L}$-smooth. Then

$\left\|\nabla f(x)\right\|_{\mathbf{W}}^{2}\leq 2A(f(x)-f(x^{*})),$

where
$A=\lambda_{\max}\left(\mathbf{W}^{1/2}\mathbf{L}\mathbf{W}^{1/2}\right)$.

The following estimate of a second moment of the gradient estimator is then an
easy corollary of Lemma 2.

######  Lemma 3.

Let $\mathbf{C}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ be an unbiased sketch
operator. Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be differentiable and
denote $g(x)=\mathbf{C}\nabla f(x)$. Then

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|g(x)\right\|^{2}\right]=\left\|\nabla
f(x)\right\|_{\mathbf{P}^{-1}}\leq 2L_{\mathbf{P}}(f(x)-f^{*}).$

The convergence result for $f$ satisfying 2 is now an immediate consequence of
Corollary A.1 of (Gorbunov et al., 2020), which provides a unified analysis of
SGD under strong convexity.

######  Theorem 1.

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be $\mathbf{L}$-smooth and
$\mu$-strongly convex. Let $\mathbf{C}$ be an unbiased sketch with associated
probability matrix $\mathbf{P}$. Then the iterates $x^{k}$ of SkGD with
stepsize $0<\gamma\leq 1/L_{\mathbf{P}}$ and update rule
$x^{k+1}=x^{k}-\gamma\mathbf{C}\nabla f(x^{k})$ satisfy

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|x^{t}-x^{*}\right\|^{2}\right]\leq(1-\gamma\mu)^{k}\left\|x^{0}-x^{*}\right\|^{2}.$

Under assumptions 3 and 4 there also exist unified analyses of SGD
(respectively Khaled et al. (2020) and Khaled & Richtárik (2020)) where Lemma
2 is applicable. However, as we argue in Appendix D, applying these generic
results to our setting leads to _suboptimal_ convergence rates. We instead
develop a simple analysis of SkGD in convex and nonconvex regimes and present
it as Theorems 2 and 3 respectively.

######  Theorem 2.

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be $\mathbf{L}$-smoooth, convex,
and lower-bounded. Furthermore, suppose that $f$ attains its minimum at
$x^{*}$. Let $(x^{t})_{t=0}^{\infty}$ be the iterates of the SkGD with
unbiased sketch $\mathbf{C}$ and stepsize $0<\gamma\leq 1/(2L_{\mathbf{P}})$.
Denote $\bar{x}^{T}=\frac{1}{T+1}\sum_{t=0}^{T}x^{t}$. Then

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f(\bar{x}^{T})-f(x^{*})\right]\leq\frac{\left\|x^{0}-x^{*}\right\|^{2}}{\gamma(T+1)}.$

In particular, for $\gamma=1/(2L_{\mathbf{P}})$ and
$T+1\geq\frac{2L_{\mathbf{P}}\left\|x^{0}-x^{*}\right\|^{2}}{\varepsilon}$, we
have that

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f(\bar{x}^{T})-f(x^{*})\right]\leq\varepsilon.$

######  Theorem 3.

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be $\mathbf{L}$-smoooth and lower
bounded by $f^{*}$. Let $(x^{t})_{t=0}^{\infty}$ be the iterates of SkGD with
unbiased sketch $\mathbf{C}$ and stepsize $0<\gamma\leq 1/L_{\mathbf{P}}$.
Denote $\hat{x}^{T}=\operatorname*{argmin}_{t=0,\dots,T}\left\|\nabla
f(x^{t})\right\|^{2}$. Then

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\nabla
f(\hat{x}^{T})\right\|^{2}\right]\leq\frac{2\delta^{0}}{\gamma(T+1)},$

where $\delta^{0}=f(x^{0})-f^{*}$. In particular, for
$\gamma=1/L_{\mathbf{P}}$ and
$T+1\geq\frac{2L_{\mathbf{P}}\delta^{0}}{\varepsilon}$, we have that

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\nabla
f(\hat{x}^{T})\right\|^{2}\right]\leq\varepsilon.$

##  3 Randomized Progressive Training (RPT)

###  3.1 Problem formulation

For the sake of consistency with the PT framework, we shall decompose the
domain of the objective function $f$ into $B$ disjoint subblocks. Thus, we
will view the domain of $f$ as
$\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}$, where
$\sum_{i=1}^{B}d_{i}=d$. We will further assume that $f$ is
$\mathbf{L}$-smooth, where $\mathbf{L}$ is a symmetric, positive definite
matrix of size $d\times d$. We let $\mathbf{L}_{i}$ be the $i^{\text{th}}$
diagonal subblock of $\mathbf{L}$ of size $d_{i}\times d_{i}$. We put
$L_{i}\stackrel{{\scriptstyle\text{def}}}{{=}}\lambda_{\max}(\mathbf{L}_{i})$.
We will be interested in minimizing this function under assumptions 2, 3 or 4,
i.e. finding

$\operatorname*{argmin}_{x\in\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}}f(x).$

###  3.2 RPT

We now introduce the _PT-sketch operator_ , the key building block of RPT.

######  Definition 3.

Let $1=p_{1}\geq p_{2}\geq\dots\geq p_{B}>0$. We additionally put $p_{B+1}=0$.
Let $\sigma=\left(\sigma_{1},\dots,\sigma_{B}\right)$ be any permutation of
set $[B]$. For each $i\in[B]$ define the matrix
$\mathbf{C}_{i}\in\mathbb{R}^{d\times d}$ as follows:

$\mathbf{C}_{i}=\operatorname{Diag}\left((\xi^{i}_{\sigma_{1}})_{d_{1}},\dots,(\xi^{i}_{\sigma_{B}})_{d_{B}}\right),$

where $\operatorname{Diag}(v)$ denotes the diagonal matrix with $v$ as its
diagonal,
$(x)_{m}\stackrel{{\scriptstyle\text{def}}}{{=}}(x,\dots,x)\in\mathbb{R}^{m}$,
and

$\xi^{i}_{j}=\begin{cases}1/p_{j}&\text{ if }j\leq i\\\ 0&\text{
otherwise.}\end{cases}$

We finally define a random matrix $\mathbf{C}$ which is equal to
$\mathbf{C}_{i}$ with probability $p_{i}-p_{i+1}$ and call it a _PT-sketch_.

With this definition, we are now ready to formulate the RPT algorithm.

1: Input: initial iterates
$x^{0}\in\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}$; stepsize
$\gamma$; integer $T\geq 0$

2: for $t=1,\dots,T$ do

3: Draw the PT-sketch $\mathbf{C}^{t}$

4: $x^{t+1}=x^{t}-\gamma\mathbf{C}^{t}\nabla f(x^{t})$

5: end for

6: Output: $x^{T}$

Algorithm 2 Randomized Progressive Training

RPT is an instance of SkGD with probability matrix
$\mathbf{P}=\operatorname{Diag}\left(\left(p_{\sigma_{1}}\right)_{d_{1}},\dots,\left(p_{\sigma_{B}}\right)_{d_{B}}\right)$.
Hence, theorems 1, 2, 3 immediately apply to it. Our task is now to choose
$\mathbf{P}$ in a way that minimizes the total computation complexity.

######  Remark 1.

There are two main differences between RPT and classical approaches to
progressive training (Karras et al., 2017; Wang et al., 2022). Firstly, our
method chooses the sizes of coordinate updates in a randomized fashion. PT,
however, is a deterministic method which grows the coordinate updates
according to a pre-set schedule.

Moreover, PT assumes a particular order of importance of the coordinate blocks
– it _always_ updates the first coordinate block, then includes the second
block, and so on. However, we note that in general there is no particular
reason to prioritize the first coordinate block, and RPT allows us to decide
on the order of importance by choosing the permutation $\sigma$ in Definition
3.

######  Remark 2.

The formulation of RPT is stemming from the problem formulation described in
Section 3.1, which assumes a particular division of the domain elements into
$B$ coordinate blocks. While our analysis holds for any such division, we do
not specify how to decide on it, noting that it is often problem-specific (for
instance, a coordinate block might be a layer of a neural network (Karras et
al., 2017; Wang et al., 2022)).

##  4 Analysis of RPT

###  4.1 Cost model

In this section we will define the computation cost of optimizing the
objective function from Section 3.1. By construction, in the RPT procedure the
variable $x^{t}\in\mathbb{R}^{d_{1}}\times\dots\times\mathbb{R}^{d_{B}}$ is
divided into $B$ coordinate blocks and updated with respect to a certain
subset of these blocks. In our discussion we will assume that by carrying out
an update with respect to block $i$ the algorithm suffers a penalty $c_{i}$.
Hence, if $\mathbf{C}^{t}=\mathbf{C}_{i}$, then $\text{cost}_{t}$, the cost
incurred in the iteration $t$ is equal to $\sum_{j=1}^{i}c_{j}$. Moreover,
without loss of generality we make a simplifying normalizing assumption on the
costs $c_{i}$ for the sake of presentation.

######  Assumption 5.

$\sum_{i=1}^{B}c_{i}=1$.

Now, recall that $\mathbf{C}=\mathbf{C}_{i}$ with probability $p_{i}-p_{i+1}$.
Therefore, the expected cost of the entire optimization procedure is

$\displaystyle{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\text{cost}(T)\right]$
$\displaystyle\stackrel{{\scriptstyle\text{def}}}{{=}}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\sum_{t=1}^{T}\text{cost}_{t}\right]$
$\displaystyle=T\sum_{i=1}^{B}(p_{i}-p_{i+1})\sum_{j=1}^{i}c_{\sigma_{j}}$
$\displaystyle=T\sum_{i=1}^{B}p_{i}c_{\sigma_{i}}.$

Previous convergence results show that in order to reach the
$\varepsilon$-accurate solution in all regimes described by assumptions 2, 3,
and 4, one needs $T$ to be proportional to $L_{\mathbf{P}}$. Thus,

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\text{cost}(T)\right]\propto
c(\mathbf{P})\stackrel{{\scriptstyle\text{def}}}{{=}}L_{\mathbf{P}}\sum_{i=1}^{B}p_{i}c_{\sigma_{i}}$

regardless of the regime in question. Thus, choosing the optimal RPT sketch is
equivalent to solving the minimization problem

$\min_{\sigma}\min_{1=p_{1}\geq\dots\geq
p_{B}>0}L_{\mathbf{P}}\sum_{i=1}^{B}p_{i}c_{\sigma_{i}}.$ (1)

###  4.2 Choosing the probabilities $p_{i}$

Finding the global minimum of the quantity $c(\mathbf{P})$ in general is a
challenging task when $\mathbf{L}$ is simply any positive definite matrix;
$c(\mathbf{P})$ is a non-convex function that can intricately depend on the
structure of matrix $\mathbf{L}$. However, note that taking $p_{i}=1$ for all
$i\in[B]$ recovers gradient descent, i.e. the optimal $c(\mathbf{P})$ is at
most $L$. Thus, the solution of Problem (1) is not worse than gradient
descent. Furthermore, one can find a closed-form minimizer for the _tight
upper bound_ for $L_{\mathbf{P}}$. We first invoke Lemma 1 from Nesterov
(2012).

######  Lemma 4.

Consider an arbitrary positive semidefinite matrix
$\mathbf{M}\in\mathbb{R}^{d\times d}$ with diagonal subblocks
$\mathbf{M}_{1},\dots.\mathbf{M}_{B}$, where each
$\mathbf{M}_{i}\in\mathbb{R}^{d_{i}\times d_{i}}$. Put
$M_{i}=\lambda_{\max}(\mathbf{M}_{i})$ and $M=\lambda_{\max}(\mathbf{M})$.
Then

$\max_{i}M_{i}\leq M\leq\sum_{i=1}^{B}M_{i}.$

In Appendix C we argue that these bounds are tight.

We then have

$c(\mathbf{P})\leq\sum_{i=1}^{B}\frac{L_{\sigma_{i}}}{p_{i}}\cdot\sum_{i=1}^{B}p_{i}c_{\sigma_{i}}.$

The below lemma describes the global minimizer of this upper bound.

######  Lemma 5.

$\sum_{i=1}^{B}\frac{L_{\sigma_{i}}}{p_{i}}\cdot\sum_{i=1}^{B}p_{i}c_{\sigma_{i}}\geq\left(\sum_{i=1}^{B}\sqrt{L_{i}c_{i}}\right)^{2}.$

Moreover, this bound is met by taking $\sigma$ to be such that

$\sqrt{L_{\sigma_{1}}/c_{{\sigma_{1}}}}\geq\dots\geq\sqrt{L_{\sigma_{B}}/c_{{\sigma_{B}}}},$

and

$p_{i}=\frac{\sqrt{L_{\sigma_{i}}/c_{\sigma_{i}}}}{\max_{j}\sqrt{L_{j}/c_{j}}}.$

This choice of probabilities is similar to the one proposed by (Zhu et al.,
2016) (which, contrary to our paper, treats _accelerated version_ of RCD)
where the probabilities $p_{i}$ are proportional to $\sqrt{L_{i}}$. However,
(Zhu et al., 2016) consider the setup of non-overlapping blocks and analyse
iteration complexity only. On the contrary, our work, to the best of our
knowledge, is the first one to consider such choice of probabilities in a non-
accelerated setting and motivates this choice through the total complexity
analysis.

We summarize the discussion with the following corollary:

######  Corollary 1.

Let $\mathbf{P}$ be the optimal solution of Equation 1. Then

$c(\mathbf{P})\leq\min\left\\{L,\left(\sum_{i=1}^{B}\sqrt{L_{i}c_{i}}\right)^{2}\right\\}.$

##  5 Experiments

We conduct several experiments to evaluate the performance of Randomized
Progressive Training in comparison to other methods on various target
functions and datasets.

###  5.1 Quadratic functions on synthetic data

In the first experiment, we show how fast Randomized Progressive Training can
be compared to vanilla Gradient Descent. When running RPT with probabilities
proposed in Lemma 5, theory suggests that total computation cost of
$O((\sum\limits_{i=1}^{B}\sqrt{L_{i}c_{i}})^{2}/\varepsilon)$ is required to
achieve $\varepsilon$-accuracy. At the same time, standard theory for GD under
5 (meaning that a cost of one GD iteration is one) suggests $O(L/\varepsilon)$
total cumulative cost. We vary block smoothness constants $L_{i}$ and
computation costs $c_{i}$ so that the ratio
$(\sum\limits_{i=1}^{B}\sqrt{L_{i}c_{i}})^{2}/L$ is relatively large. In this
vein, we generate nine quadratic functions
$f(x)=\frac{1}{2}x^{\top}\mathbf{A}x$ with varying dimensions and properties
of the diagonal matrix $\mathbf{A}$. The number of blocks is three in all
experiments. We consider the computation cost $c_{i}$ of block $i$ to be
proportional to its size. As shown on Figure 1, when the size of the block
with the smallest smoothness constant is much larger than one of the largest
$L_{i}$, RPT shows tremendous speedup over GD. Table 1 compares speedups
predicted by theory with actual ones.

Figure 1: Comparison of Randomized Progressive Training (RPT) Gradient Descent (GD) on quadratic functions. Target function is $f(x)~{}=~{}\frac{1}{2}x^{\top}\mathbf{A}x$, where the diagonal matrix $\mathbf{A}\in\mathbb{R}^{l\times l}$ and its dimension $l$ differ in each plot. For RPT, the iterate is separated into three blocks. ’blocks’ in the plot title indicates the block sizes, ’Ls’ indicates smoothness constants within the corresponding block. Numbers in ’blocks’ sum up to $l$. The block computation cost $c_{i}$ is the size of block $i$ divided by $l$, and the computation cost of a single GD iteration is one. Probabilities for RPT are set as stated in Lemma 5.  Setup | Theory | Real
---|---|---
blocks | 10 | 10 | 10 | 1.1 | 1.9
Ls | 272 | 53.3 | 11
blocks | 10 | 50 | 250 | 3.4 | 19.8
Ls | 270.5 | 55 | 11
blocks | 10 | 100 | 1000 | 5.7 | 26.9
Ls | 256.7 | 54.8 | 11
blocks | 10 | 500 | 25000 | 12.7 | 36.6
Ls | 274.8 | 55 | 11
blocks | 10 | 10 | 10 | 1.5 | 2.2
Ls | 1066.5 | 108.4 | 11
blocks | 10 | 50 | 250 | 6.4 | 17.6
Ls | 1080.4 | 109.6 | 11
blocks | 10 | 100 | 1000 | 12.3 | 64.9
Ls | 1093.9 | 109.7 | 11
blocks | 10 | 500 | 25000 | 36.6 | 128.1
Ls | 1066.7 | 110 | 11
blocks | 10 | 10 | 10 | 2.2 | 2.8
Ls | 27414 | 547 | 11
blocks | 10 | 50 | 250 | 15.4 | 20
Ls | 27166 | 550 | 11
blocks | 10 | 100 | 1000 | 40.1 | 62.4
Ls | 26020 | 549 | 11
blocks | 10 | 500 | 25000 | 282.5 | 407.8
Ls | 27363 | 550 | 11
Table 1: Comparison of theoretical and actual speedup of Randomized
Progressive Training versus Gradient Descent in a synthetic experiment using
quadratic functions.

###  5.2 Ridge regression on real datasets

In our second set of experiments, we assess the convergence of Randomized
Progressive Training, Gradient Descent, and Cyclic Block Coordinate Descent
(Sun & Hong, 2015) on the ridge regression problem

$\displaystyle\min f(x)=\frac{1}{n}\sum_{i=1}^{n}(\langle a_{i},x\rangle-
b_{i})^{2}+\lambda\|x\|^{2}_{2},$

where $a_{i}\in\mathbb{R}^{d}$, $b_{i}\in\mathbb{R}$ and $\lambda$ is a
regularization parameter.

We consider four regression datasets: California housing (Kelley Pace & Barry,
1997), prostate cancer, Los Angeles ozone (Hastie et al., 2009), and white
wine quality (Cortez et al., 2009). Each dataset feature is normalized, and a
bias vector is added. Then, heuristically, features are permuted so that the
diagonal of the Hessian matrix is sorted in descending order. As in the
previous experiment, we consider $c_{i}$ proportional to the corresponding
block size.

As can be seen from Figure 2, although the theoretical speedup of RPT over GD,
according to Table 2, is close to 1, the actual speedup is still noteworthy.
However, both algorithms are slower than Cyclic Block Coordinate Descent.

Figure 2: Comparison of Randomized Progressive Training (RPT) vs. Gradient Descent (GD) and Cyclic Block Coordinate Descent (CBCD) on the ridge regression problem. The block computation cost $c_{i}$ is the size of block $i$ divided by the dimension of the problem, and the computation cost of one GD iteration is 1. Probabilities of block sampling for RPT are set as stated in Lemma 5. Dataset | Theoretical | Actual
---|---|---
speedup | speedup
california housing | 1 | 1.3
prostate cancer | 1 | 2.1
los angeles ozone | 1 | 2.4
white wine quality | 1.1 | 1.3
Table 2: Comparison of theoretical and actual speedup of Randomized
Progressive Training versus Gradient Descent for ridge regression problems.

###  5.3 Image classification on CIFAR10

In our final computational experiment, we trained a Deep Residual Neural
Network (He et al., 2016b) with 18 layers (ResNet18) for the task of image
classification on the CIFAR10 dataset (Krizhevsky, 2009).

Implementation. The CIFAR10 dataset comprises 50’000 training samples and
10’000 test samples, which are distributed evenly among 10 classes.. The
samples are RGB images of size $32\times 32$. To facilitate hyperparameter
tuning, we further divided the training dataset into training and validation
datasets with a ratio $9:1$. During the training process, we applied a padding
of 4, randomly cropped to size $32\times 32$, and randomly flipped images
horizontally. We trained the neural network with three algorithms: Stochastic
Gradient Descent, Progressive Training (Wang et al., 2022), and our own
method. We implemented all settings using the Python packages JAX (Bradbury et
al., 2018), Flax (Heek et al., 2020), and Optax.

Fine-tuning. Fine-tuning is a crucial step in the process of training neural
networks, as the lack of smoothness and convexity properties inherent in these
models often renders the application of theoretical results ineffective
(Szegedy et al., 2014). To address this, a common approach for hyperparameter
tuning is to utilize a grid search method. In our case, for the RPT method, we
fine-tuned the learning rate and sampling probabilities. A grid search was
conducted for the learning rate over the values $[1,0.5,0.1,0.05,0.01,0.001]$,
and for sampling probabilities, we evaluated two different strategies. The
first strategy used exponentially decaying probabilities with rates
$[0.999,0.995,0.99,0.95,0.9,0.85]$, e.g., for a rate of $0.9$, the
probabilities were calculated as
$(p_{1},\dots,p_{22})=(1,0.9,0.9^{2},\dots,0.99^{21})$. The second strategy
used evenly spaced probabilities between $p_{1}=1$ and $p_{22}$, where
$p_{22}$ ran over the list $[0.7,0.6,0.5,0.4,0.3,0.2]$. The best results were
obtained with a learning rate of $1$ and exponentially decaying sampling
probabilities with a rate of $0.9$. For Progressive Training (PT), we only
fine-tuned the learning rate and utilized the round sizes $T_{1},\dots,T_{22}$
suggested in (Wang et al., 2022) as the best strategy. Specifically, we set
$T_{1}=\dots=T_{21}=\frac{T}{44}$, and $T_{22}=\frac{23T}{44}$, where $T$ is
the total number of steps. The grid search for the learning rate was the same
as for RPT. Additionally, Stochastic Gradient Descent parameters such as the
learning rate, weight decay, and momentum were set according to the guidelines
in (Moreau et al., 2022). The batch size in all three settings was 128.

(a) Training loss versus computation cost

(b) Test accuracy versus computation cost

(c) Training loss versus epoch

(d) Test accuracy versus epoch

Figure 3: Comparison of Randomized Progressive Training (RPT), Stochastic
Gradient Descent (SGD), and Progressive Training (PT) on an image
classification task using a ResNet18 model. The block computation cost $c_{i}$
is proportional to the size of block $i$, with a computation cost of 1
assigned to SGD.

Results. The results of our image classification task are summarized in Figure
3. While it may appear that RPT converges more slowly and has a lower test
accuracy than SGD, based on Figures 3(c) and 3(d), it is important to note
that this comparison is not entirely fair as one RPT step can take
significantly less time than one SGD step. To more accurately compare the
performance of RPT to its competitors, we have evaluated the performance
versus computation cost, as shown in Figures 3(a) and 3(b). From these
figures, it is clear that RPT outperforms both competitors.

## References

  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Cortez et al. (2009) Cortez, P., Cerdeira, A., Almeida, F., Matos, T., and Reis, J.  Modeling wine preferences by data mining from physicochemical properties.  _Decision Support Systems_ , 47(4):547–553, 2009\.  ISSN 0167-9236.  doi: https://doi.org/10.1016/j.dss.2009.05.016.  URL https://www.sciencedirect.com/science/article/pii/S0167923609001377.  Smart Business Networks: Concepts and Empirical Evidence.
  * Dun et al. (2021) Dun, C., Wolfe, C. R., Jermaine, C. M., and Kyrillidis, A.  Resist: Layer-wise decomposition of resnets for distributed training.  _CoRR_ , abs/2107.00961, 2021.  URL https://arxiv.org/abs/2107.00961.
  * Goodfellow et al. (2014) Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y.  Generative adversarial nets.  In _Advances in neural information processing systems_ , pp. 2672–2680, 2014.
  * Gorbunov et al. (2020) Gorbunov, E., Hanzely, F., and Richtarik, P.  A unified theory of sgd: Variance reduction, sampling, quantization and coordinate descent.  In Chiappa, S. and Calandra, R. (eds.), _Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics_ , volume 108 of _Proceedings of Machine Learning Research_ , pp. 680–690. PMLR, 26–28 Aug 2020.  URL https://proceedings.mlr.press/v108/gorbunov20a.html.
  * Hastie et al. (2009) Hastie, T., Tibshirani, R., and Friedman, J.  _The Elements of Statistical Learning: Data Mining, Inference, and Prediction_.  Springer series in statistics. Springer, 2009.  ISBN 9780387848846.  URL https://books.google.com.sa/books?id=eBSgoAEACAAJ.
  * He et al. (2016a) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 770–778, 2016a.  doi: 10.1109/CVPR.2016.90.
  * He et al. (2016b) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2016b.
  * Heek et al. (2020) Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepierre, B., Steiner, A., and van Zee, M.  Flax: A neural network library and ecosystem for JAX, 2020.  URL http://github.com/google/flax.
  * Karras et al. (2017) Karras, T., Aila, T., Laine, S., and Lehtinen, J.  Progressive growing of gans for improved quality, stability, and variation, 2017.  URL https://arxiv.org/abs/1710.10196.
  * Kelley Pace & Barry (1997) Kelley Pace, R. and Barry, R.  Sparse spatial autoregressions.  _Statistics & Probability Letters_, 33(3):291–297, 1997.  ISSN 0167-7152.  doi: https://doi.org/10.1016/S0167-7152(96)00140-X.  URL https://www.sciencedirect.com/science/article/pii/S016771529600140X.
  * Khaled & Richtárik (2020) Khaled, A. and Richtárik, P.  Better theory for sgd in the nonconvex world, 2020.  URL https://arxiv.org/abs/2002.03329.
  * Khaled et al. (2020) Khaled, A., Sebbouh, O., Loizou, N., Gower, R. M., and Richtárik, P.  Unified analysis of stochastic gradient methods for composite convex and smooth optimization, 2020.  URL https://arxiv.org/abs/2006.11573.
  * Krizhevsky (2009) Krizhevsky, A.  Learning multiple layers of features from tiny images.  Technical report, 2009.
  * Moreau et al. (2022) Moreau, T., Massias, M., Gramfort, A., Ablin, P., Bannier, P.-A., Charlier, B., Dagréou, M., Dupré la Tour, T., Durif, G., F. Dantas, C., Klopfenstein, Q., Larsson, J., Lai, E., Lefort, T., Malézieux, B., Moufad, B., T. Nguyen, B., Rakotomamonjy, A., Ramzi, Z., Salmon, J., and Vaiter, S.  Benchopt: Reproducible, efficient and collaborative optimization benchmarks.  In _NeurIPS_ , 2022.  URL https://arxiv.org/abs/2206.13424.
  * Nesterov (2012) Nesterov, Y.  Efficiency of coordinate descent methods on huge-scale optimization problems.  _SIAM Journal on Optimization_ , 22(2):341–362, 2012.  doi: 10.1137/100802001.  URL https://doi.org/10.1137/100802001.
  * Qu et al. (2016) Qu, Z., Richtarik, P., Takac, M., and Fercoq, O.  Sdna: Stochastic dual newton ascent for empirical risk minimization.  In Balcan, M. F. and Weinberger, K. Q. (eds.), _Proceedings of The 33rd International Conference on Machine Learning_ , volume 48 of _Proceedings of Machine Learning Research_ , pp. 1823–1832, New York, New York, USA, 20–22 Jun 2016. PMLR.  URL https://proceedings.mlr.press/v48/qub16.html.
  * Richtárik & Takáč (2014) Richtárik, P. and Takáč, M.  Iteration complexity of randomized block-coordinate descent methods for minimizing a composite function.  _Mathematical Programming_ , 144(1):1–38, 2014\.  doi: 10.1007/s10107-012-0614-z.  URL https://doi.org/10.1007/s10107-012-0614-z.
  * Safaryan et al. (2021) Safaryan, M., Hanzely, F., and Richtárik, P.  Smoothness matrices beat smoothness constants: Better communication compression techniques for distributed optimization, 2021.  URL https://arxiv.org/abs/2102.07245.
  * Sun & Hong (2015) Sun, R. and Hong, M.  Improved iteration complexity bounds of cyclic block coordinate descent for convex problems.  In Cortes, C., Lawrence, N., Lee, D., Sugiyama, M., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 28. Curran Associates, Inc., 2015.  URL https://proceedings.neurips.cc/paper/2015/file/96b9bff013acedfb1d140579e2fbeb63-Paper.pdf.
  * Szegedy et al. (2014) Szegedy, C., Zaremba, W., Sutskever, I., Bruna, J., Erhan, D., Goodfellow, I., and Fergus, R.  Intriguing properties of neural networks.  In _International Conference on Learning Representations_ , 2014.  URL http://arxiv.org/abs/1312.6199.
  * Wang et al. (2022) Wang, H.-P., Stich, S., He, Y., and Fritz, M.  ProgFed: Effective, communication, and computation efficient federated learning by progressive training.  In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 23034–23054. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/wang22y.html.
  * Yuan et al. (2022) Yuan, B., Wolfe, C. R., Dun, C., Tang, Y., Kyrillidis, A., and Jermaine, C.  Distributed learning of fully connected neural networks using independent subnet training.  _Proc. VLDB Endow._ , 15(8):1581–1590, jun 2022\.  ISSN 2150-8097.  doi: 10.14778/3529337.3529343.  URL https://doi.org/10.14778/3529337.3529343.
  * Zhu et al. (2016) Zhu, Z. A., Qu, Z., Richtárik, P., and Yuan, Y.  Even faster accelerated coordinate descent using non-uniform sampling.  In _ICML_ , 2016.

##  Appendix A Issues with the theory for ProgFed by Wang et al. (2022)

We will first briefly outline the notation that Wang et al. (2022) use in
their paper.

We denote by $\mathcal{M}$ the neural network that we aim to train. Put
$\mathcal{M}^{s}$ to be network $\mathcal{M}$ truncated to its first $s$
layers. Let $\mathcal{L}$ denote loss function of interest and put $f$ and
$f^{s}$ to be $\mathcal{L}\circ\mathcal{M}$ and
$\mathcal{L}\circ\mathcal{M}^{s}$ respectively. Denote $\mathbf{x}^{s}$ to be
the parameter vector of $\mathcal{M}^{s}$. Finally, for $s\leq i\leq S$ we put
$\mathbf{x}^{i}_{|E_{s}}$ and $\nabla f^{i}(\mathbf{x}^{i})_{|E_{s}}$ to be
the projections of vectors $\mathbf{x}^{i}$ and $\nabla f^{i}(\mathbf{x}^{i})$
onto the parameter space of $\mathcal{M}^{s}$.

The first result by Wang et al. (2022) proves that

$\frac{1}{T}\sum_{t=0}^{T-1}\alpha_{t}^{2}\left\|\nabla
f^{s}(\mathbf{x}_{t}^{s})_{|E_{s}}\right\|<\varepsilon$

after
$\mathcal{O}\left(\frac{\sigma^{2}}{\epsilon^{2}}+\frac{1}{\epsilon}\right)\cdot
LF_{0}$ iterations. Here $\mathbf{x}_{t}^{s}$ is the $t^{\text{th}}$ iterate
of ProgFed,

$\alpha_{t}=\min\left\\{1,\frac{\left\langle\nabla
f(\mathbf{x}_{t})_{|E_{s}},\nabla
f^{s}(\mathbf{x}_{t}^{s})_{|E_{s}}\right\rangle}{\left\|\nabla
f^{s}(\mathbf{x}_{t}^{s})\right\|^{2}}\right\\},$

$L$ is the smoothness constant of $f$, and $F_{0}$ is the function
suboptimality. Thus, the numbers $\alpha_{t}$ depend on the iterates $x_{t}$.
No bounds on $\alpha_{t}$ are provided. In particular, the theory does not
exclude the possibility that $\alpha_{t}^{2}$ are extremely close to 0. This
renders the result meaningles.

The second result that Wang et al. (2022) attempts to bound
$\frac{1}{T}\sum_{t=0}^{T-1}\left\|\nabla
f\left(\mathbf{x}_{t}\right)\right\|^{2}$. Unfortunately, it suffers from a
similar issue. The authors claim that
$\frac{1}{T}\sum_{t=0}^{T-1}\left\|\nabla
f\left(\mathbf{x}_{t}\right)\right\|^{2}$ after

$\mathcal{O}\left(\frac{q^{4}\sigma^{2}}{\epsilon^{2}}+\frac{q^{2}}{\epsilon}\right)\cdot
LF_{0},$

where $q:=\max_{t\in[T]}\left(q_{t}:=\frac{\left\|\nabla
f\left(\mathbf{x}_{t}\right)\right\|}{\alpha_{t}\left\|\nabla
f^{s}\left(\mathbf{x}_{t}^{s}\right)_{\mid E_{s}}\right\|}\right)$. Wang et
al. (2022) do not provide any bound on $q$, which does not depend on the
iterates. Thus, it is not possible to decide a priori how many iterations do
we need in order to reach the $\varepsilon$-accurate solution, because the
theory is dependent on the iterates _produced during the run of the method_.
In other words, their second convergence result (in its current state) does
not exclude an infinite runtime, so in fact it does not show the convergence
at all.

##  Appendix B Proofs

###  B.1 Proof of Lemma 1

See 1

###### Proof.

We start from proving the first inequality. Let $x\in\mathbb{R}^{d}$. Then

$\displaystyle
x^{\top}\mathbf{P}^{1/2}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}\mathbf{L}\mathbf{C}\right]\mathbf{P}^{1/2}x$
$\displaystyle={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\mathbf{L}^{1/2}\mathbf{C}\mathbf{P}^{1/2}x\right\|^{2}\right]$
$\displaystyle\leq\lambda_{\max}(\mathbf{L}){\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\mathbf{C}\mathbf{P}^{1/2}x\right\|^{2}\right]$
$\displaystyle=\lambda_{\max}(\mathbf{L}){\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\sum_{i=1}^{n}\frac{1}{p_{i}}x_{i}^{2}\chi_{\\{\mathbf{C}=\mathbf{C}_{i}\\}}\right]$
$\displaystyle=\lambda_{\max}(\mathbf{L})\sum_{i=1}^{n}x_{i}^{2}\frac{1}{p_{i}}p_{i}$
$\displaystyle=\lambda_{\max}(\mathbf{L})\left\|x\right\|^{2}.$

Thus,

$\lambda_{\max}(\mathbf{P}^{1/2}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}\mathbf{L}\mathbf{C}\right]\mathbf{P}^{1/2})\leq\lambda_{\max}(\mathbf{L}),$

i.e.

$\lambda_{\max}({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}\mathbf{L}\mathbf{C}\right])\leq\lambda_{\max}(\mathbf{P}^{-1/2}\mathbf{L}\mathbf{P}^{-1/2}).$

For the second inequality, similarly,

$\displaystyle x^{\top}\mathbf{P}^{-1/2}\mathbf{L}\mathbf{P}^{-1/2}x$
$\displaystyle=\left\|\mathbf{L}^{1/2}\mathbf{P}^{-1/2}x\right\|^{2}$
$\displaystyle\leq\lambda_{\max}(\mathbf{L})\left\|\mathbf{P}^{-1/2}x\right\|^{2}$
$\displaystyle\leq\lambda_{\max}(\mathbf{L})\lambda_{\max}(\mathbf{P}^{-1})\left\|x\right\|^{2}.$

Thus,

$\lambda_{\max}(\mathbf{P}^{-1/2}\mathbf{L}\mathbf{P}^{-1/2})\leq\lambda_{\max}(\mathbf{L})\lambda_{\max}(\mathbf{P}^{-1}).$

∎

###  B.2 Proof of Lemma 2

See 2

###### Proof.

Fix $\alpha\in\mathbb{R}$ and $x\in\mathbb{R}^{d}$. Since $x^{\star}$ is a
minimizer of $f$ and since $f$ is $\mathbf{L}$-smooth, we have

$\displaystyle f\left(x^{\star}\right)$ $\displaystyle\leq
f(x-\alpha\mathbf{W}\nabla f(x))$ $\displaystyle\leq\quad
f(x)-\alpha\langle\nabla f(x),\mathbf{W}\nabla
f(x)\rangle+\frac{\alpha^{2}}{2}\|\mathbf{W}\nabla f(x)\|_{\mathbf{L}}^{2}$
$\displaystyle=f(x)-\alpha\langle\nabla f(x),\mathbf{W}\nabla
f(x)\rangle+\frac{\alpha^{2}}{2}\langle\nabla
f(x),\mathbf{W}\mathbf{L}\mathbf{W}\nabla f(x)\rangle$

By definition of $A$, we have
$\mathbf{W}^{1/2}\mathbf{L}\mathbf{W}^{1/2}\preceq A\mathbf{I}_{d}$, which
implies that $\mathbf{W}\mathbf{L}\mathbf{W}\preceq A\mathbf{W}$. Plugging
this estimate into the above inequality, we get

$\displaystyle f\left(x^{\star}\right)$ $\displaystyle\leq f(x)-\alpha\|\nabla
f(x)\|_{\mathbf{W}}^{2}+\frac{\alpha^{2}A}{2}\|\nabla f(x)\|_{\mathbf{W}}^{2}$
$\displaystyle=f(x)-\left(\alpha-\frac{\alpha^{2}A}{2}\right)\|\nabla
f(x)\|_{\mathbf{W}}^{2}$

Since the above inequality holds for any $\alpha\in\mathbb{R}$, it also holds
for $\alpha^{\star}=\frac{1}{A}$; this $\alpha$ maximizes the expression
$\alpha-\frac{\alpha^{2}A}{2}$. It only remains to plug $\alpha^{\star}$ into
the above expression and rearrange the inequality. ∎

###  B.3 Proof of Lemma 3

See 3

###### Proof.

$\displaystyle{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|g(x)\right\|^{2}\right]$
$\displaystyle={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\mathbf{C}\nabla
f(x)\right\|^{2}\right]$
$\displaystyle={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\nabla
f(x)^{\top}\mathbf{C}^{2}\nabla f(x)\right]$ $\displaystyle=\nabla
f(x)^{\top}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}^{2}\right]\nabla
f(x)$ $\displaystyle=\nabla f(x)^{\top}\mathbf{P}^{-1}\nabla f(x)$
$\displaystyle=\left\|\nabla f(x)\right\|_{\mathbf{P}^{-1}}.$

The last inequality follows directly from Lemma 2. ∎

###  B.4 Proof of Theorem 2

See 2

###### Proof.

Taking conditional expectations and expanding the brackets yields

$\displaystyle{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|x^{t+1}-x^{*}\right\|^{2}|x^{t}\right]$
$\displaystyle=\left\|x^{t}-x^{*}\right\|^{2}-2\gamma{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\langle\mathbf{C}\nabla
f(x^{t}),x^{t}-x^{*}\right\rangle|x^{t}\right]+\gamma^{2}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\mathbf{C}\nabla
f(x^{t})\right\|^{2}|x^{t}\right]$
$\displaystyle\leq\left\|x^{t}-x^{*}\right\|^{2}-2\gamma\left\langle\nabla
f(x^{t}),x^{t}-x^{*}\right\rangle+2\gamma^{2}L_{\mathbf{P}}\left(f(x^{t})-f(x^{*})\right)$
$\displaystyle\leq\left\|x^{t}-x^{*}\right\|^{2}-2\gamma\left(f(x^{t})-f(x^{*})\right)+\gamma\left(f(x^{t})-f(x^{*})\right)$
$\displaystyle=\left\|x^{t}-x^{*}\right\|^{2}-\gamma\left(f(x^{t})-f(x^{*})\right),$

where the first inequality follows from Lemma 3, and the second one from
convexity of $f$ and $0<\gamma\leq 1/(2L_{\mathbf{P}})$. Rearranging the above
inequality, taking the expectation of both sides and using the tower property
gives

$\displaystyle\gamma{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f(x^{t})-f(x^{*})\right]\leq{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|x^{t}-x^{*}\right\|^{2}-\left\|x^{t+1}-x^{*}\right\|^{2}\right].$

Thus, by convexity of $f$, we finally arrive at

$\displaystyle{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f(\bar{x}^{T})-f(x^{*})\right]$
$\displaystyle\leq\frac{1}{\gamma(T+1)}\sum_{t=0}^{T}\left(f(x^{t})-f(x^{*})\right)$
$\displaystyle\leq\frac{1}{\gamma(T+1)}\left(\left\|x^{0}-x^{*}\right\|^{2}-{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|x^{T+1}-x^{*}\right\|^{2}\right]\right)$
$\displaystyle\leq\frac{\left\|x^{0}-x^{*}\right\|^{2}}{\gamma(T+1)}.$

∎

###  B.5 Proof of Theorem 3

See 3 Before carrying out the proof of Theorem 3 we note that in this theorem,
one can replace $L_{\mathbf{P}}$ by

$L_{\mathbf{M}}\stackrel{{\scriptstyle\text{def}}}{{=}}\lambda_{\max}({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}\mathbf{L}\mathbf{C}\right])$

which is _at most_ $L_{\mathbf{P}}$ by Lemma 1. However, the constant
$L_{\mathbf{M}}$ is difficult to interpret qualitatively. For the sake of
unified presentation of the RPT analysis in all regimes in question, we
therefore chose to replace it by $L_{\mathbf{P}}$.

###### Proof.

By $\mathbf{L}$-smoothness we have

$\displaystyle f(x^{t+1})$ $\displaystyle\leq f(x^{t})+\left\langle\nabla
f(x^{t}),x^{t+1}-x^{t}\right\rangle+\frac{\left\|x^{t+1}-x^{t}\right\|_{\mathbf{L}}^{2}}{2}$
$\displaystyle=f(x^{t})-\gamma\left\langle\nabla f(x^{t}),\mathbf{C}\nabla
f(x^{t})\right\rangle+\frac{\gamma^{2}}{2}\left\|\mathbf{C}\nabla
f(x^{t})\right\|_{\mathbf{L}}^{2}$
$\displaystyle=f(x^{t})-\gamma\left\langle\nabla f(x^{t}),\mathbf{C}\nabla
f(x^{t})\right\rangle+\frac{\gamma^{2}}{2}\left\langle\nabla
f(x^{t}),\mathbf{C}\mathbf{L}\mathbf{C}\nabla f(x^{t})\right\rangle.$

Thus, taking expectation conditional on $x^{t}$ yields

$\displaystyle{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f(x^{t+1})|x^{t}\right]$
$\displaystyle\leq f(x^{t})-\gamma\left\|\nabla
f(x^{t})\right\|^{2}+\frac{\gamma^{2}}{2}\nabla
f(x^{t})^{\top}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\mathbf{C}\mathbf{L}\mathbf{C}\right]\nabla
f(x^{t})$ $\displaystyle\leq f(x^{t})-\gamma\left\|\nabla
f(x^{t})\right\|^{2}+\frac{\gamma^{2}}{2}L_{\mathbf{M}}\left\|\nabla
f(x^{t})\right\|^{2}$ $\displaystyle\leq f(x^{t})-\gamma\left\|\nabla
f(x^{t})\right\|^{2}+\frac{\gamma^{2}}{2}L_{\mathbf{P}}\left\|\nabla
f(x^{t})\right\|^{2},$

i.e.

$\displaystyle\left(\gamma-\frac{\gamma^{2}L_{\mathbf{P}}}{2}\right)\left\|\nabla
f(x^{t})\right\|^{2}\leq
f(x^{t})-{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f(x^{t+1})|x^{t}\right].$

By taking expectations of both sides, using tower property and taking the
average over $t$, we get

$\displaystyle\frac{\gamma}{2}\min_{t=0,\dots,T}\left\|\nabla
f(x^{t})\right\|^{2}$
$\displaystyle\leq\frac{\gamma}{2(T+1)}\sum_{t=0}^{T}\left\|\nabla
f(x^{t})\right\|^{2}$
$\displaystyle\leq\left(\frac{\gamma}{2}+\frac{\gamma}{2}\left(1-\gamma
L_{\mathbf{P}}\right)\right)\frac{1}{T+1}\sum_{t=0}^{T}\left\|\nabla
f(x^{t})\right\|^{2}$
$\displaystyle=\left(\gamma-\frac{\gamma^{2}L_{\mathbf{P}}}{2}\right)\frac{1}{T+1}\sum_{i=0}^{T}\left\|\nabla
f(x^{t})\right\|^{2}$
$\displaystyle\leq\frac{1}{T+1}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[f\left(x^{0}\right)-f\left(x^{T+1}\right)\right]$
$\displaystyle\leq\frac{1}{T+1}\delta^{0}.$

Rearranging gives

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\nabla
f(\hat{x}^{T})\right\|^{2}\right]\leq\frac{2\delta^{0}}{\gamma(T+1)}.$

∎

###  B.6 Proof of Lemma 4

See 4

###### Proof.

The second inequality is just Lemma 1 from (Nesterov, 2012). Let us now prove
the first inequality. Without loss of generality we may assume that
$\max_{i}L_{i}=L_{1}$. Denote $e_{i}$ to be the $i^{\text{th}}$ standard unit
vector of $\mathbb{R}^{d}$. Let $S=\text{Span}(e_{1},\dots,e_{d_{1}})$. Now,

$\displaystyle L_{1}$
$\displaystyle=\sup_{\begin{subarray}{c}x\in\mathbb{R}^{d_{1}},\\\
\left\|x\right\|=1\end{subarray}}x^{\top}\mathbf{L}_{1}x$
$\displaystyle=\sup_{\begin{subarray}{c}x\in S,\\\
\left\|x\right\|=1\end{subarray}}x^{\top}\mathbf{L}x$
$\displaystyle\leq\sup_{\begin{subarray}{c}x\in\mathbb{R}^{d},\\\
\left\|x\right\|=1\end{subarray}}x^{\top}\mathbf{L}x.$

This concludes the proof. ∎

###  B.7 Proof of Lemma 5

See 5

###### Proof.

By Cauchy-Schwarz inequality,

$\displaystyle\sum_{i=1}^{B}\frac{L_{\sigma_{i}}}{p_{i}}\cdot\sum_{i=1}^{B}p_{i}c_{\sigma_{i}}$
$\displaystyle\geq\left(\sum_{i=1}^{B}\sqrt{L_{\sigma_{i}}/p_{i}}\cdot\sqrt{p_{i}c_{\sigma_{i}}}\right)^{2}$
$\displaystyle=\left(\sum_{i=1}^{B}\sqrt{L_{i}c_{i}}\right)^{2}.$

Since we are free to choose permutation $\sigma$, we may clearly choose it in
such a way that

$\sqrt{L_{\sigma_{1}}/c_{{\sigma_{1}}}}\geq\dots\geq\sqrt{L_{\sigma_{B}}/c_{{\sigma_{B}}}}.$

Plugging in this choice of permutation $\sigma$ and probablities
$p_{1},\dots,p_{B}$ verifies that the lower bound from Lemma 5 is attainable.
∎

##  Appendix C Discussion about tightness of Lemma 4

Here we consider the 2-block case argue that the maximum eigenvalue
$\lambda_{\max}(\mathbf{M})$ can be arbitrarily close to
$\lambda_{\max}(\mathbf{M}_{1})+\lambda_{\max}(\mathbf{M}_{2})$. Extension of
the below result to an arbitrary number of blocks is very straightforward.

######  Proposition 1.

Let
$\mathbf{M}_{1}\in\mathbb{S}_{++}^{d_{1}},\mathbf{M}_{2}\in\mathbb{S}_{++}^{d_{2}}$,
and
$\mathbf{C}=\sqrt{\lambda_{\max}(\mathbf{M}_{1})(\lambda_{\max}(\mathbf{M}_{2})-\varepsilon)}\cdot
vu^{\top}$, where $\varepsilon\in(0,\lambda_{\max}(\mathbf{M}_{2}))$,
$u\in\mathbb{R}^{d_{1}}$ and $v\in\mathbb{R}^{d_{2}}$ are unit eigenvectors
corresponding to maximum eigenvalues of matrices $\mathbf{M}_{1}$ and
$\mathbf{M}_{2}$, respectively. Then the matrix
$\mathbf{M}=\begin{bmatrix}\mathbf{M}_{1}&\mathbf{C}^{\top}\\\
\mathbf{C}&\mathbf{M}_{2}\end{bmatrix}$ is positive definite, and it holds
that

$\lambda_{\max}(\mathbf{M})\geq\lambda_{\max}(\mathbf{M}_{1})+\lambda_{\max}(\mathbf{M}_{2})-\varepsilon.$

###### Proof.

First, let us check the matrix
$\mathbf{M}=\begin{bmatrix}\mathbf{M}_{1}&\mathbf{C}^{\top}\\\
\mathbf{C}&\mathbf{M}_{2}\end{bmatrix}$ is positive definite. Since
$\mathbf{M}_{1}$ is positive definite, it suffices to check that
$\mathbf{M}_{2}-\mathbf{C}\mathbf{M}_{1}^{-1}\mathbf{C}^{\top}$ is positive
definite. The second term can be rewritten as

$\displaystyle\mathbf{C}\mathbf{M}_{1}^{-1}\mathbf{C}^{\top}=\lambda_{\max}(\mathbf{M}_{1})(\lambda_{\max}(\mathbf{M}_{2})-\varepsilon)v\underbrace{u^{\top}\mathbf{M}_{1}^{-1}u}_{=\frac{1}{\lambda_{\max}(\mathbf{M}_{1})}}v^{\top}=(\lambda_{\max}(\mathbf{M}_{2})-\varepsilon)vv^{\top}.$

Now, through eigenvalue decomposition of $\mathbf{M}_{2}$, we see that
eigenvalues of $\mathbf{M}_{2}-\mathbf{C}\mathbf{M}_{1}^{-1}\mathbf{C}^{\top}$
are the same as eigenvalues of $\mathbf{M}_{2}$ except for the maximum one,
which is replaced by $\varepsilon>0$. Thus,
$\mathbf{M}_{2}-\mathbf{C}\mathbf{M}_{1}^{-1}\mathbf{C}^{\top}$ is positive
definite.

Let
$\mathbf{M}_{1}=\mathbf{\Lambda}_{1}\mathbf{\Sigma}_{1}\mathbf{\Lambda}_{1}^{\top}$
and
$\mathbf{M}_{2}=\mathbf{\Lambda}_{2}\mathbf{\Sigma}_{2}\mathbf{\Lambda}_{2}^{\top}$
be eigenvalue decompositions of the matrices, and without the loss of
generality assume that the eigenvalues are sorted in a descending order. We
also set
$a~{}\stackrel{{\scriptstyle\text{def}}}{{=}}~{}\sqrt{\lambda_{\max}(\mathbf{M}_{1})(\lambda_{\max}(\mathbf{M}_{2})-\varepsilon)}$.
Then

$\displaystyle\mathbf{M}^{\prime}\stackrel{{\scriptstyle\text{def}}}{{=}}\operatorname{Diag}(\mathbf{\Lambda}_{1},\mathbf{\Lambda}_{2})^{\top}\mathbf{M}\operatorname{Diag}(\mathbf{\Lambda}_{1},\mathbf{\Lambda}_{2})=\begin{bmatrix}\mathbf{\Sigma}_{A}&\operatorname{Diag}(a,0,\dots,0)^{\top}\\\
\operatorname{Diag}(a,0,\dots,0)&\mathbf{\Sigma}_{2}\end{bmatrix}$

has the same set of eigenvalues as $\mathbf{M}$. We now have

$\displaystyle\lambda_{\max}(\mathbf{M}^{\prime})$
$\displaystyle=\max_{\|u^{\prime}\|^{2}+\|v^{\prime}\|^{2}=1}(u^{\prime},v^{\prime})^{\top}\mathbf{M}^{\prime}(u^{\prime},v^{\prime})$
$\displaystyle=\max_{\left\|u^{\prime}\right\|^{2}+\left\|v^{\prime}\right\|^{2}=1}u^{\prime\top}\mathbf{\Sigma}_{1}u^{\prime}+v^{\prime\top}\mathbf{\Sigma}_{2}v^{\prime}+2u^{\prime\top}\operatorname{Diag}(a,0,\dots,0)^{\top}v^{\prime}$
$\displaystyle\geq\max_{x^{2}+y^{2}=1}x^{2}\lambda_{\max}(\mathbf{M}_{1})+y^{2}\lambda_{\max}(\mathbf{M}_{2})+2xy\sqrt{\lambda_{\max}(\mathbf{M}_{1})(\lambda_{\max}(\mathbf{M}_{2})-\varepsilon)}$
$\displaystyle\geq\max_{x^{2}+y^{2}=1}x^{2}\lambda_{\max}(\mathbf{M}_{1})+y^{2}(\lambda_{\max}(\mathbf{M}_{2})-\varepsilon)+2xy\sqrt{\lambda_{\max}(\mathbf{M}_{1})(\lambda_{\max}(\mathbf{M}_{2})-\varepsilon)}$
$\displaystyle=\max_{x^{2}+y^{2}=1}\left(x\sqrt{\lambda_{\max}(\mathbf{M}_{1})}+y\sqrt{\lambda_{\max}(\mathbf{M}_{2})-\varepsilon}\right)^{2}$
$\displaystyle\geq\left(x\sqrt{\lambda_{\max}(\mathbf{M}_{1})}+y\sqrt{\lambda_{\max}(\mathbf{M}_{2})-\varepsilon}\right)^{2}\Big{|}_{x=\sqrt{\frac{\lambda_{\max}(\mathbf{M}_{1})}{\lambda_{\max}(\mathbf{M}_{1})+\lambda_{\max}(\mathbf{M}_{2})-\varepsilon}},y=\sqrt{\frac{\lambda_{\max}(\mathbf{M}_{2})-\varepsilon}{\lambda_{\max}(\mathbf{M}_{1})+\lambda_{\max}(\mathbf{M}_{2})-\varepsilon}}}$
$\displaystyle=\lambda_{\max}(\mathbf{M}_{1})+\lambda_{\max}(\mathbf{M}_{2})-\varepsilon.$

It remains to note that
$\lambda_{\max}(\mathbf{M}^{\prime})=\lambda_{\max}(\mathbf{M})$, since the
sets of eigenvalues coincide. ∎

##  Appendix D Discussion about the unified analysis in nonconvex and convex
settings

###  D.1 Nonconvex setting

Recent theoretical breakthrough of Khaled & Richtárik (2020) offers generic
result on the convergence of SGD under very mild assumptions. However, in this
section we argue that the rates promised by this result is suboptimal for the
setting of our paper, which is why we derived a separate analysis of RCD in
this regime.

######  Assumption 6 (Assumption 2, Khaled & Richtárik (2020)).

Let $g(x)$ be an unbiased estimator of $\nabla f(x)$ for all
$x\in\mathbb{R}^{d}$. We say that $g$ satisfies an _ABC assumption_ if

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|g(x)\right\|^{2}\right]\leq
2A\left(f(x)-f^{\inf}\right)+B\left\|\nabla f(x)\right\|^{2}+C.$

for all $x\in\mathbb{R}^{d}$.

######  Theorem 4.

Suppose that $f$ is lower-bounded and $L$-smooth. Let $g$ be an unbiased
gradient estimator satisfying 6 with $B=C=0$. Let
$0<\gamma\leq\frac{1}{\sqrt{LAK}}$. Denote $f(x^{0})-f^{\inf}$ by
$\delta_{0}$. Then, the iterates $x^{t}$ of SGD using the estimator $g$ and
stepsize $\gamma$ satisfy

$\min_{0\leq t\leq
T-1}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\nabla
f(x^{t})\right\|^{2}\right]\leq\frac{2(1+L\gamma^{2}A)^{T}\delta_{0}}{\gamma
T}.$

In particular, for

$T\geq\frac{144\delta_{0}^{2}LA}{\varepsilon^{2}}$

we have $\min_{0\leq t\leq
T-1}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|\nabla
f(x^{t})\right\|^{2}\right]\leq\varepsilon.$

Notice that for nonconvex smooth problems, by Lemma 3, Assumption 6 is
satisfied for $A=L_{\mathbf{P}}$ and $B=C=0$. Moreover,

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|g(x)\right\|^{2}\right]=\left\|\nabla
f(x)\right\|_{\mathbf{P}^{-1}}^{2}\leq\lambda_{\max}\left(\mathbf{P}^{-1}\right)\left\|\nabla
f(x)\right\|^{2},$

so 6 is also satisfied for $B=\lambda_{\max}\left(\mathbf{P}^{-1}\right)$ and
$A=C=0$. There seem to be no other ways in which the RPT estimator can satisfy
6.

Plugging in both of these cases into Theorem 4 yields suboptimal convergence
rates.

Indeed, for $A=L_{\mathbf{P}}$ and $B=C=0$ he rate of RPT guaranteed by
Theorem 4 is $O\left(1/\sqrt{T}\right)$, which is suboptimal (we have proved
that it is in fact $O(1/T)$).

For $B=\lambda_{\max}\left(\mathbf{P}^{-1}\right)$ and $A=C=0$ the rate is
$O(1/T)$ but it is proportional to

$L\lambda_{\max}\left(\mathbf{P}^{-1}\right)=\lambda_{\max}(\mathbf{L})\lambda_{\max}\left(\mathbf{P}^{-1}\right)\geq
L_{\mathbf{P}},$

by Lemma 4.

We manage to obtain faster rate by exploiting $\mathbf{L}$-smoothness and the
fact that
${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|g(x)\right\|_{\mathbf{L}}^{2}\right]$
can be expressed as a suitable matrix norm.

###  D.2 Convex setting

Applying generic result of Khaled et al. (2020) yields somewhat good
convergence guarantee ($O(1/T)$ rate is recovered); however, it still suffers
from an unnecessary dependency on function suboptimality.

######  Theorem 5 (Adapted from Khaled et al. (2020)).

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be convex, differentiable and
lower bounded with global minimizer $x^{*}$. Suppose that an unbiased gradient
estimator $g$ satisfies

${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathbb{E}}\left[\left\|g(x)\right\|^{2}\right]\leq
2A(f(x)-f(x^{*})).$

Choose stepsize

$\gamma\leq\min\left\\{\frac{1}{4A},\frac{1}{2L}\right\\}.$

Then the SGD procedure with gradient estimator $g$ and stepsize $\gamma$
satisfies

$\mathbb{E}\left[f\left(\bar{x}^{T}\right)-f\left(x^{*}\right)\right]\leq\frac{2\gamma\delta^{0}+\left\|x^{0}-x^{*}\right\|^{2}}{\gamma
T},$

where
$\bar{x}^{T}\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{1}{t}\sum_{i=0}^{T-1}x^{i}$
and $\delta^{0}\stackrel{{\scriptstyle\text{def}}}{{=}}f(x^{0})-f(x^{*})$.

As discussed previously, the assumption from Theorem 5 is satisfied for
$A=L_{\mathbf{P}}\geq L$. Thus, the optimal theoretical stepsize is
$\gamma=1/(4L_{\mathbf{P}}$). Plugging in this stepsize yields the rate

$O\left(\frac{\delta^{0}+L_{\mathbf{P}}\left\|x^{0}-x^{*}\right\|^{2}}{T}\right),$

whereas we obtained the rate of

$O\left(\frac{L_{\mathbf{P}}\left\|x^{0}-x^{*}\right\|^{2}}{T}\right).$