MAST: model-agnostic sparsified training
Authors: Yury Demidovich
Grigory Malinovsky
Egor Shulgin
Peter Richtárik
Venue: ICLR 2025 Poster (2025)
Keywords: dropout theory
sparse training
Primary area: optimization

Abstract: We introduce a novel optimization problem formulation that departs from the conventional way of minimizing machine learning model loss as a black-box function. Unlike traditional formulations, the proposed approach explicitly incorporates an initially pre-trained model and random sketch operators, allowing for sparsification of both the model and gradient during training. We establish insightful properties of the proposed objective function and highlight its connections to the standard formulation. Furthermore, we present several variants of the Stochastic Gradient Descent (SGD) method adapted to the new problem formulation, including SGD with general sampling, a distributed version, and SGD with variance reduction techniques. We achieve tighter convergence rates and relax assumptions, bridging the gap between theoretical principles and practical applications, covering several important techniques such as Dropout and Sparse training. This work presents promising opportunities to enhance the theoretical understanding of model training through a sparsification-aware optimization approach.

Published as a conference paper at ICLR 2025
MAST: MODEL-AGNOSTIE SPARSIFIED TRAINING
Yury Demidovich Grigory Malinovsky Egor Shulgin Peter Richtárik King Abdullah University of Science and Technology (KAUST), Saudi Arabia"
ABSTRACT
We introduce novel optimization problem formulation that departs from the con- ventional way of minimizing machine learning model loss as a black-box function. Unlike traditional formulations, the proposed approach explicitly incorporates an initially pre-trained model and random sketch operators, allowing for sparsifica- tion of both the model and gradient during training. We establish the insightful properties of the proposed objective function and highlight its connections to the standard formulation. Furthermore, we present several variants of the Stochastic Gradient Descent (SGD) method adapted to the new problem formulation, includ- ing SGD with general sampling, a distributed version, and SGD with variance reduction techniques. We achieve tighter convergence rates and relax assumptions, bridging the gap between theoretical principles and practical applications. covering several important techniques such as Dropout and Sparse training. This work presents promising opportunities to enhance the theoretical understanding of model training through sparsification optimization approach.
INTRODUCTION
Efficient optimization methods have played an essential role in the advancement of modern Machine Learning (ML), given that many supervised problems ultimately reduce to the task of minimizing an abstract loss function -a process that can be formally expressed as:
min f(xx). =ER
where f: Rd - is loss function of a model with parameters/weights æ E Rd, The problem (1) has been comprehensively analyzed within the domain of optimization literature. Within the ML community, substantial attention has been directed towards studying this problem, particularly in the context of the Stochastic Gradient Descent (SGD) method (Robbins & Monro, 1951). SGD stands as a foundational and highly effective optimization algorithm within the realm of ML (Bottou et al., 2018). The pervasive use of SGD in the field attests its versatility and success in training a diverse array of models (Goodfellow et al., 2016). Notably, contemporary deep learning practices owe a substantial debt to SGD, is the cornerstone for many state-of-the-art training techniques (Sun, 2020).
While problem (1) has been of primary interest in mainstream optimization research, this approach is not always best suited for representing recent ML techniques, such as sparse/quantized training (Wu al., 2016; Hoefler al., 2021), resource-constraaned distributed learning (Caldas al., 2018; Wang et al., 2018), 8), model personalization (Smith et al., 2017; Hanzely & Richtárik, 2020; Mansour et al., 2020), and meta-learning (Schmidhuber, 1987; Finn al., 2017). Although various attempts have been made to analyze some of these settings (Khaled & Richtárik, 2019; Lin et al., 2019; Mohtashami et al., 2022), there is still no satisfactory optimization theory that can explain the success these techniques in deep learning. Previous works analyze variants of SGD trying to solve problem (1), which often results in vacuous convergence bounds and/or overly restrictive assumptions on the class of functions and algorithms involved (Shulgin & Richtárik, 2024). We assert that these issues arise due to mismatches between the method used and the problem formulation being solved.
In this work, address the issues mentioned above, propose a new optimization problem formulation called Model-Agnostic Sparsified Training (MAST):
zERd [f(()) def E [[s()))]
min
(2)
Contacte: Iyory.denidee grigori ia inouskis, egor.ahulgin, peter peter.richterikigee

---

Published as a conference paper at ICLR 2025
where (s(x) def [(v + S(x - v)) for E Rd (e.g., a pre-trained model, possibly compressed), S e R2xd is a random matrix (i.c., a sketch) sampled from distribution D.
When O, the function takes the following form: fz( = f(Sx). This scenario may be considered process in which the architecture requires training from scratch, with quantization consideration. Such an approach involves a substantial increase in training time and the necessity of hyperparameter tuning to achieve effectiveness. In terms of theory, formulation (2) can be viewed as a nested (stochastic) composition optimization (Bertsekas, 1977; Polyak, 1979; Ermoliev, 1988). Moreover, our formulation is partially related to the setting of splitting methods (Condat et al., 2023). However, due to their generality, other sctups do not consider the problem instance's 's peculiarities and focus on other applications. Conversely, when v is non-zero and pre-trained weights are utilized, the newly formulated approach can be interpreted as acquiring a "meta-model" x. Solving problem (2) then ensures that the sketched model u S(x u) exhibits strong performance on average. This interpretation shares many similarities with Model-Agnostic Meta-Learning (Finn et al., 2017).
central
We argue that this framework may be better suited for modeling various practical ML techniques, as discussed below. Furthermore, our proposed framework facilitates thorough theoretical analysis and holds potential independent interest for a broader audience. While our analysis and algorithms work for quite general class of sketches, we focus applications that are relevant for sparse S.
1.1 MOTIVATING EXAMPLES
Dropout (Hanson, 1990; Hinton et al., 2012; Frazier-Logue & Hanson, 2018) is regularization technique initially introduced to prevent overfitting in neural networks by dropping some of the model's units (activations) during training. Later, this approach was generalized to incorporate Gaussian noise (Wang & Manning, 2013; Srivastava et al., 2014) (instead of Bernoulli masks) and zeroing the weights via DropConnect (Wan et al., 2013) or entire layers (Fan et al., 2019). was also observed (Srivastava et al., 2014) that training with Dropout induces sparsity in the activations. In addition, using the Dropout-like like technique (Gomez et al., 2019; LeJeune et al., 2021) can make the resulting network more amenable to subsequent sparsification (via pruning) before deployment. Modern DL has seen a huge increase in the size of models (Villalobos et al., 2022). This has resulted in growing energy and computational costs, necessitating optimizations of neural networks' training pipeline (Yang et al., 2017). Among others, pruning and sparsification were proposed as effective techniques due to the overparametrization properties of large models (Chang et al., 2021).
Sparse training algorithms (Mocanu et al., 2016; Guo et al., 2016; Mocanu et al., 2018), in particular, suggest working with smaller subnetwork during every optimization step. This increases efficiency by reducing the model size (via compression), which naturally brings memory and computation acceleration benefits to the inference stage. Moreover, sparse training has recently been shown to speed up optimization by leveraging sparse backpropagation (Nikdan et al., 2023).
On-device learning also creates a need for sparse or submodel computations due to the memory. energy, and computational constraints of edge devices. In settings like cross-device federated learning (Konečný et al., 2016; McMahan et al., 2017: 7; Kairouz et al., 2021), models are trained in distributed way across population of heterogeneous nodes, such mobile phones. The heterogeneity of the clients' hardware makes it necessary to adapt the (potentially large) server model to the needs of low-tier devices (Caldas al., 2018; Bouacida al., 2021; Horváth et al., 2021).
Contributions. The main results of this work include:
A rigorous formalization of a new optimization formulation, as shown in Equation (2), which can encompass various important practical settings as special cases, such as Dropout and Sparse training.
In-depth theoretical characterization of the proposed problem's properties, highlighting its connec- tions to the standard formulation in Equation (1). Notably, our problem is efficiently solvable with practical methods.
The development of optimization algorithms that naturally emerge from the formulation in Equation (2), along with insightful convergence analyses in non-convex and (strongly) convex settings.
The generalization of the problem and methods to the distributed scenario, expanding the range of applications even further, including scenarios like IST and Federated Learning.
2

---

Published as a conference paper at ICLR 2025
An experimental study of the proposed algorithms and MAST properties in the context of machine learning, highlighting the advantages of the proposed approach.
Paper organization. We introduce the basic formalism and discuss sketches in Section 2. The properties of the suggested formulation (2) are analyzed in Section Section 4 contains convergence results for full-batch, stochastic, and variance-reduced methods for solving problem (2). Extensions to the distributed case are presented Section 5. Section 6 describes some of our experimental results. The last Section 7 concludes the paper and outlines potential directions of future work. All proofs are provided in the Appendix.
SKETCHES
(Stochastic) gradient-base based methods are mostly used to solve problem (1). Applying this paradigm to our formulation (2) requires computing the gradient of fo = E [fs] In the case of the general distribution S ~ DD such computation may not be possible due to expectation E. Therefore, practical algorithms typically rely on gradient estimates. An elegant property of the MAST problem (2) is that the gradient estimator takes the form
Vss(2) = S S S(x
(3)
due to the chain rule, as matrix S is independent z. Sketches/matrices S are random and sampled from distribution D. Note that estimator (3) sketches both the model x and the gradient of f. Next, we explore some of the sketches' important properties and give practical examples. Assumption 1. The sketching matrix S satisfies:
E I,
and
E [SS is finite,
(4)
where I is the identity matrix. Note that fT)( an unbiased estimator of the gradient f
Denote Ls del Amax (S (STS), us Amin (ST S), and Lo Amax (E [S S) Amin (E [S ST]
where Amax and Amin represent the largest and smallest eigenvalues. Clearly, Ls 2 us 0 and
LD 2 HD 0. If Assumption I is satisfied, then E |SSS > I, which means that up 2 1 and
| S pD ||]||
EE |Sxx1]]
S Lp ||=||
2.1 DIAGONAL SKETCHES
Let CI, .2,...,, Ca be a collection of random variables and define matrix with c-s on the diagonal
S = Diag(e1, Ca),..,
(
which satisfies Assumption 1 when E [c:] = 1 and E [c]] is finite for every i.
The following example illustrates how our framework can be used model Dropout.
Example 1. The independent Bernoulli sparsification operator is defined as a diagonal sketch (5), where every 6 is (independent) scaled Bernoulli random variable:
{ 1 1/ps,
with probability pi
(6)
O,
with probability - pi
for ps E (0, 1] and E [d] def (1.....d) d).
It can be shown that for independent Bernoulli sparsifiers,
LD maxp;" def Pminn
def
μD = min ps
Pmax-
(7)
Notice that when p: p. [d]. and 0, gradient estimator (3) results in sparse update as S (Sxz)) drops out d(1 - p) (on average) components of model weights and the gradient. The difference from the Dropout described by Hinton et al. (2012) is that they do not scaling 1/p; in equation (6) during training to ensure unbiasedness. Experimental comparison is presented in Appendix H.2.4.
Next, we show another practical example of a random sketch often used for reducing communication costs in distributed learning (Konečný al., 2016; Wangni et al., 2018; Stich et al., 2018).
3

---

Published as a conference paper at ICLR 2025
Example 2. Random K sparsification (in short, Rand-K for K [d] operator is defined by def d
S Saand-K KK
Σ Eee,
(8)
K
LES
where ei,...,ee. ed E RRd are standard unit basis vectors, and a random subset of [d] sampled from the uniform distribution over all subsets of [d] with cardinality K.
Rand-K K belongs to the class of diagonal sketches (5). St preserves K non-zero coordinates out of total d coordinates, Since E [S S] = I d/K/K, this sketch satisfies Lp = μrD d/ K. This example is suitable for modeling fixed budget sparse training when the proportion (K/d /)) of network parameters being updated remains constant during optimization (Mocanu et al., 2018; Evci al., 2020).
PROBLEM PROPERTIES
We show that the proposed formulation (2) inherits the smoothness and convexity properties of the original problem (1). Let us introduce the most standard assumptions the optimization field. Assumption 2. Function f is differentiable and Ly-smooth, i.e., there is Ly > O such that
f(x + S f(x) + (VI(z), h) + Lf | d
2
We also require f to be lower bounded by finl E R. Assumption 3. Function f is differentiable and p stronnyyyly convex, i.e., there is py > such that f(z + 2) (Vf(z) h) + Pf |/bl1 Vz,h e Rd.
2
Next, we show how the choice of the sketch S affects the smoothness parameters of fs and fp- Lemma (Consequences of sysoooothessss. If f is Ly-smooth, then
(i) fs is Lis -smooth with is S LsLs.
(ii) fo is L with Lsp S LpLs.
(ii)) fo(z)= s f(x) + (1D=1)42 - r|]
Va Rd.
In particular, property (Hii) in Lemma demonstrates that the gap between the sketched loss fo and the original function depends on the model weights and the smoothness parameter of function f.
Lemma (Consequence of Convexity). If f is convex, then fo is convex and fD(x) 2 f(x).
It shows that the convexity of f is preserved, and the "sketched" loss is always greater than the original loss. Moreover, Lemma (along with other results in this section) offers a huge advantage of the proposed problem formulation over the sparsification sparification-promotigg alternatives based on fo-norm regularization (Louizos et al., 2018; Peste et al., 2021), that make the problem hard to solve.
Lemma 3 (Consequences of μ ( If f is p-covee,x then
(i) fs Pfs -onvex with jis 2 PsP
(ii) fo is pro convex with P 2 UDHf.
(ii) fo(x) 2 f(x) (um-1)mr || m||
Va Rdd.
As consequence, we get the following result for the condition number of the proposed problem:
del Lsp < LoLs
Lo
Rio
Kzg-
UfD
PDHl
PDD
Therefore, Ki SKg LTD as HD 2 1. Thus, the resulting condition number may increase, which indicates that fo may be harder to optimize, which agrees with the intuition that compressed training harder (Evci al., 2019). In addition, for independent Bernoulli sparsifiers (6),
KD
del
TDD/HD Pmnax/
Punin; Omaxx/Pmii;
(10)
4

---

Published as a conference paper at ICLR 2025
which shows that the upper bound on the ratio / can be made as large possible by choosing a small enough Pmin- At the same time, Kip = 1 for classical Dropout: p. = p. E [d], indicating that training with Dropout may be no harder than optimizing the original model.
Relation between f and fo minima. Let X be the solutions to problem 1, and X3 the solutions of the new MAST problem (2). We now show that solution xip E X of (2) is approximate solution of the original problem (1).
Theorem 1. Let Assumptions and hold, and let xi E Xi and X Then
f(x') s f(xb)) 1 f(x)) (LD -1)LB ||0" vll? (UD 1)ps
| u11?
2
2
f(x*)+ + (HD (uD-1)jus |0 - u11² S fo((z) s f(x")) + (LD
1)L
1| - 211
2
2
Consider Rand-K as a sketch. If = (1 E)d for some E E [0,1),, 1). which corresponds to dropping
roughly an share of coordinates, then LD 1 uD 1 =/K states that
5/(1-z), Theorem I then
S f(x*)
T.p 1|= - -112- - Ms 11 - -v(l() -ol²²):
2(1
[(x)) + 2(1 - | - 01 S fo(ab) [(x')( + 2(1
EPI
EL,
||x0 u1|²
If e is small "light" sparsification). small loss, comparable to the loss of the optimal model xt; æ will be close to the pre-trained model v; MAST loss will be small comparable to the loss of the optimal uncompressed model x*
if the pre-trained model is close x², then aip will have
INDIVIDUAL NODE SETTING
In this section, we discuss the properties of the SGD-type algorithm applied to problem (2)
r+ = x' vV fs (2') ='
(z') - (S') T ff('')
(1I)
for y' = --- S'x'''''''- v). where St is sampled from D. One advantage of the proposed formulation
(2) is that it naturally gives rise to the described method, which generalizes standard (Stochastic) Gradient Descent. As noted before, due to the properties of the gradient estimator (3), recursion (11)
defines an Algorithm (I) that sketches both the model and the gradient of f.
Let us introduce notation frequently used in our convergence results. This quantity is determined by the spectral properties of the sketches being used.
L
max def sup Ls sup Amax (STS),
(12)
S
where sups L.s represents the tightest constant such that L.s S sups I.s almost surely. For indepen- dent random sparsification sketches (6): T.g. 1/r2... If Rand-K (8) then L. T.max d² K2
Our convergence analysis relies on the following Lemma:
Lemma 4. Assume that f is -smooth (2) and S satisfies Assumption 1. Then we have that Vr E RR
E
[
< 2L max (fo(x)- - pran),
where the expectation taken with respect to S.
It generalizes standard property of smooth functions often used in the non-convex analysis of SGD (Khaled& & Richtárik, 2023). Now, we are ready to present our first convergence result.
Theorem 2. Assume that ff is Ly-smooth (2), h strsonyygly convex (3), and S satisfies Assumption I. Then, for stepsize < 1/ Lg ) the iterates of Algorithm 1 (I) satisfy
s (1 7PDD1)) 1100 - =tll + 2 L max
E |||T - xbol"
(Jbal - plal).
(13)
PDpD
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Double Sketched (S)GD
1: Parameters: learning rate > O; distribution D; initial model and shift x v E R.
2: for t = 0, 2.....o do
3: Sample a sketch: S ~ D
4: Form gradient estimator:
g Ss (x') > exact (I)
< (stochastic) inexact (II) 5: Perform gradient-type step: x/+1 = x² "
9Sa (x')
6: end for
This theorem establishes linear convergence rate with constant stepsize up neighborhood of the MAST problem (2) solution. Our result is similar to the convergence of SGD (Gower et al., 2019) for standard formulation (1) with two differences, which are discussed below.
1. Both terms of the upper bound (13) depend not only on the smoothness and convexity parameters of the original function f. but also on the spectral properties of sketches S. Thus, for independent Bernuuli sparsification sketches (6) with p; (or Rand-K), the linear convergence term deterio- rates and the neighborhood size is increased by 1/p² (d²) /K² respectively). Therefore, we conclude that higher sparsity makes optimization harder.
Interestingly, the neighborhood size of (13) depends on the difference between the minima of fo and in contrast to the variance of stochastic gradients at the optimum typical for SGD (Gower et al., 2019). Thus, the method may even linearly converge to the exact solution when finf = fint which we refer to the interpolation condition. This condition may naturally hold when the original and sketched models are sufficiently overparametrized (allowing minimization of the loss to zero). Notable examples when similar phenomena have been observed in practice are training with Dropout (Srivastava et al., 2014) and the "lottery ticket hypothesis" (Frankle & Carbin, 2018).
Next, we provide results in the non-convex setting.
Theorem 3. Assume that I is Ly-smooth (2) and S satisfies Assumption 1. Then, for the stepsize S LLLLLyLLLLgaee the iterates of Algorithm 1 (I) satisfy
3 (fo(x²) -
min E [[VSD(()]]] 0StT
+DDLLDLg
max (fbe - pint).
IT
This theorem shows an OLL/ VT) convergence rate for reaching a stationary point. Our result shares similarities with the theory of SGD (Khaled & Richtárik, 2023) for problem (l), with the main difference that the rate depends on the distribution on sketches as y /IpL.a.. Moreover, the second term depends on the difference between the minima of fo and f. as in the strongly convex case. However, the first term depends on the gap between the initialization and the lower bound of the loss function, which more common in non-convex settings (Khaled & Richtárik, 2023).
- -Siot Corollary 1 (Informal). Fix & > 0 and denote 8 def E [Fo(x*)
def E [V
2 Then, for the stepsize 7 min VDT 2D( where D def L LJV LDDLL max Algorithm 1 (1)
needs T Z 128" max {350, finé - finty iterations reach stationary point minogrer rf s
ga which is order (se-1) optimal (Ghadimi & Lan, 2013; Drori & Shamir, 2020).
In the Appendix, we also provide a general convex analysis.
4.1 (STOCHASSICC) INEXACT GRADIENT
Algorithm 1 (I) is probably the simplest approach for solving the MAST problem (2). Analyzing this algorithm isolates and highlights the unique properties of the proposed problem formulation. Algorithm (I) requires exact (sketched) gradient computations every iteration, which may not be feasible/effcieett for modern MI. applications. Hence, we consider the following generalization:
x¹ 7gs (x²),
(14)

---

Published as a conference paper at ICLR 2025
where 9s (5: (z") is gradient estimator satisfying
E [as (1:) Vzs(2))
(15)
E (11)(())) S 2A (a) 1) B+VV/s(>)]] +C,
(16)
for V/ E IRd and some constants A, 2
The first condition (15) is an unbiasedness assumption standard for analyzing SGD-type methods.
The second (so-called "ABC") inequality (16) is one of the most general assumptions covering bounded stochastic gradient variance, uubsampling/minibaccing of data, and gradient compression (Khaled & Richtárik, 2023; Demidovich et al., 2023). Note that the expectation in (15) and (16) is taken with respect only to the randomness of the stochastic gradient estimator and not the sketch S.
Algorithm (II) describes the resulting method in detail. We state the convergence result for it.
Theorem Assume that f is L--smooth (2), satisfies Assumption I, and the gradient es- timator g(x) satisfies conditions 15, 16. Denote DA,B A + B L Lmax then, for stepsize
1/11v LLLLLDDD the iterales of Algorithm 1 (II) satisfy
min E [IV 3 (fo(xz) ost<T
- fint)
TLyLD
{c+ 2DA,B (frr }int)}
+
7T
2
Similarly Theorem this result establishes an O(I/VT) convergence rate. However, the upper bound in (4) is affected by constants A, B, C due to the inexactness of the gradient estimator, The case of A C 0, B sharply recovers our previous Theorem 3. When 1, C. a², we obtain convergence of SGD with bounded (by o²) variance of stochastic gradients. Moreover, when
the loss is represented as a finite sum
f0(1) Σ E [ -))],
(17)
where each fa is f smooth and lower-bounded by fint, then A = max; Lf., C = 2A 4 Ei= fint - ffitt) if losses fi are sampled uniformly at every iteration. Finally, our result (4) guarantees optimal OEe complexity in a similar manner to Corollary 1.
We direct the reader to the Appendix for the convex and strongly convex results.
4.2 DISCUSSION OF RELATED WORKS
Compressed (sparse) model training. To our knowledge, the first work that analyzed convergence of (full batch) Gradient Descent with compressed iterates (model updates) is the work of Khaled & Richtárik (2019). They considered general unbiased compressors and, in the strongly convex setting, showed linear convergence to the irreducible neighborhood, depending on the norm of the model at the optimum |a==113 In addition, their analysis requires the variance of the compressor (d/K for Rand-K) to be lower than the inverse condition number of the problem μ1 Ly. which basically means that the compressor has to be close the identity mapping practical settings. These results were extended using modified method to distributed training with compressed model updates (Chraibi al., 2019; Shulgin & Richtárik, 2022). Lin al. (2019) consider dynamic pruning with feedback inspired by the Error Feedback mechanism (Seide et al., 2014; Alistarh et al., 2018; Stich Karimireddy, 2020). Their result is similar to (Khaled & Richtárik, 2019), as the method also converges only to the irreducible neighborhood, the size of which is proportional to the norm of model weights. However, (Lin et al., 2019), the norm of stochastic gradients is required to be uniformly upper-bounded, narrowing the class of losses. The partial SGD method proposed in (Mohtashami et al., 2022) allows general perturbations of the model weights where the gradient (additionally sparsified) is computed. Unfortunately. their analysis (Wang al., 2022) was recently shown to be vacuous (Szlendak et al., 2024).
Dropout convergence analysis. Despite the wide empirical success of Dropout, there is limited theoretical understanding of its behavior and success. A few recent works (Mianjy & Arora, 2020) suggest convergence analysis of this technique. However, these attempts typically focus on certain
7

---

Published as a conference paper at ICLR 2025
Algorithm 2 Distributed Double Sketched GD
1: Parameters: learning rate > 0; sketch distributions Dı.... Dai; initial model and shift
x², 0 E R2
2: for 1= 0, 1, 1......d do
3: Sample sketches: S
4: Compute 3/
(x v) for [MM] and broadcast to corresponding nodes
5:
for 1, .............................. .M in parallel do
6: Compute local gradient: f(((v)
7: Send gradient (S!) (v(( to the server
8: end for
9: Aggregate messages and make gradient-type step: x²+1 x² 2 E (S)) Vf((!))
10: end for
class of models, such as shallow linear Neural Networks (NNs) (Senen-Cerda & Sanders, 2022) or deep NNs with ReLU activations (Senen-Cerda & Sanders, 2020). Moreover, Liao & Kyrillidis (2022) analyzed overparameterized single-hidden layer perceptron with regression loss in the context of Dropout. In contrast, our approach is model-agnostic and requires only mild assumptions like the smoothness of the loss (2) (and convexity (3) for some of the results).
5 DISTRIBUTED SETTING
Consider f being finite sum over number of participants, i.e., in the distributed setup:
1 M M
min ZER
Σ (x),
(18)
111
where fr.D, def E (B (39) as every node has its own distribution of sketches S, ~ Dr Every machine performs local computations with a model of different size, which is crucial for scenarios with heterogeneous computing hardware. The shift model is shared across all fiP, We solve (18) with the method
2)}} This setting is more general than problem
M
Σ St) V((())). M
7
(19)
where = Algorithm describes the proposed approach in detail. Local gradients can be computed for sketched (sparse) model weights, which decreases the computational load on the computing nodes. Moreover, the local gradients are sketched as well, which brings communication efficiency in the case of sparsifiers S,.
Recursion (19) is closely related to the distributed Independent Subnetwork Training (IST) framework (Yuan et al., 2022). At every iteration of IST, large model zf is decomposed into submodels Sf for independent computations (e.g., local training), which are then aggregated on the server to update the whole model. IST efficiently combines model and data parallelism, allowing the training of huge models that cannot fit onto a single device. IST was shown to be very successful for range of DL applications (Dun al., 2022; Wolfe al., 2023). Shulgin Richtárik (2024) analyzed the convergence of IST for a quadratic model decomposed with permutation sketches (Szlendak al., 2022), which satisfy Assumption They also showed that naively applying IST standard distributed optimization problems ((18) for S; = results in biased method and may not converge.
Resource-constrained Federated Learning (FI.) (Kairouz al., 2021: Konečný et al., 2016; McMahan et al., 2017) is another important practical scenario covered by Algorithm 2. In cross-device FL., local computations are typically performed by edge devices (e.g.. mobile phones), which have limited memory, computational power, and energy (Caldas et al., 2018). Thus, this forces practitioners rely on smaller (potentially less capable) models or use techniques such as Dropout in distributed setting (Alam et al., 2022; Bouacida et al., 2021: Charles et al., 2022; Chen et al., 2022: Diao et al., 2021; Horváth et al., 2021; Jiang et al., 2022; Qiu et al., 2022: Wen et al., 2022; Yang et al., 2022; Dun al., 2023). Despite extensive experimental studies of this problem setting, the principled theoretical
8

---

Published as a conference paper at ICLR 2025
understanding remains minimal. Our work can be considered the first rigorous analysis in the most general setting without restrictive assumptions.
Theorem Assume that each S is L -smooth (2) and sketches S, satisfy Assumption 1. Let
D Dmax
dof
max; Lg.x) Then, for $ Dmax. the iterates of Algorithm 2 satisfy
3 (fo(2²) - )bb") + 7/max
1 M
M
min OSET
fint
pinf
This result resembles our previous Theorem 3 in the single-node setting. Namely, Algorithm reaches stationary point at rate ((1/vT), However, due the distributed setup, convergence depends Mmax expressed as the maximum product of local smoothness and constants related to the sketches' properties. Thus, clients with more aggressive sparsification may slow down the method, given the same local smoothness constant f. Yet, "easier" local problems (with smaller Ly,) can allow the use of "harsher" sparsifiers (with larger LD, Lg. Lmax) without negatively affecting the convergence.
5.1 DISCUSSION OF RELATED WORKS
notable distinction between our result and the theory of methods like Distributed Compressed
Gradient Descent (Khirirat et al., 2018) 8) lies in the second convergence term of Theorem 5. Instead of relying on the variance of local gradients at the optimum, given by 82 Σ f I|? our result depends on the average difference between the lower bounds of the global and local losses:
f - fiaf This term measures heterogeneity within a distributed setting (Khaled et al., 2020).
Furthermore, our findings may provide better explanation of the empirical efficacy of distributed methods. Namely, 8² is less likely to be equal to zero, unlike our term, which can be very small when models over-vvvaneee parameterizedd allowing local losses to be minimized to zero.
Yuan et al. (2022) analyzed convergence Independent Subnetwork Training in their original work using the framework of Khaled & Richtárik (2019). Their analysis was performed in the single-node setting and required additional assumptions on the gradient estimator, which were recently shown to be problematic (Shulgin & Richtárik, 2024). In federated setting, Zhou al. (2022) suggested method that combines model pruning with local compressed Gradient Descent steps. They provided non-convex convergence analysis relying bounded stochastic gradient assumption, which results in "pathological" bounds (Khaled et al., 2020) for a heterogeneous distributed case.
6 EXPERIMENTS
To empirically validate our theoretical frame- work and its implications, we focus on care- fully controlled settings that satisfy the assump- tions of our work. Specifically, we consider an fzrregularized logistic regression optimization problem with the a5a dataset from the LibSVM repository (Chang & Lin, 2011). See Appendix H for further details and Appendix H.2 for more results on other methods and sketches.
0.85
0.80
Test Accuracies
0.75
0.70
0.65
0.60
In Figure 1, we compare the test accuracy of sparsified solutions for the standard (ERM) prob- lem (1) and introduced MAST formulation (2). Visualization is performed using the boxplot method from Seaborn (version 0.11.0) library (Waskom, 2021) with default parameters. For ERM, we find the exact (up machine pre- cision) optimum, which subsequently used for the accuracy evaluation. For the MAST op- timization problem, we run DSGD with exact sketched gradient V f for every sparsity level.
0.55
ERM
MAST
0.50
70
90
Sparsity
Figure 1: Test accuracies distributions of sparsified solutions for the ERM formulation (1) and MAST problem (2). "Sparsity" corresponds to the percent- age of zeroed weights.

---

Published as a conference paper at ICLR 2025
After the ERM and MAST models (x") are obtained, apply partition sketches (47) to model weights and evaluate the test accuracy the sparsified solutions (Sx").
Figure reveals that models obtained using the MAST approach exhibit greater robustness to random pruning compared to their ERM counterparts given the same sparsity. Moreover, the ERM model suffers from greater aceuracy variability, while the median test accuracies of the MAST models are markedly higher. Increasing the sparsity leads to the degradation of the performance of both approaches.
Neural network results. Next we present a subset of our distributed deep learning results (full details provided in Appendix H.2.3). Our experimental setup closely follows that of Liao & Kyrillidis (2022), which is based on the ResNet-50 model (He et al., 2016). We study the Algorithm 2 with Bernoulli sketches (6) and pt p for the standard (ERM) loss (18) with S
High sparsity, p 0.5
Mid sparsity, p
Low sparsity, p 0.9
10
y=1
0.0.5
y= 0.5
0.0.1
10"
Loss
10
500
1000
1500
2000
0
500
1000
1500
2000
0
500
1000
1500
2000
Iteration number
Iteration number
Iteration number
Figure 2: Performance of Algorithm 2 with Bernoulli sketches (6) on standard loss (18) (for S, I)
Figure illustrates the impact of sparsity level (p) and step size (7) on the method's performance. Across all sparsity levels, we observe an optimal "sweet spot" (r = 0.5) for the step size, beyond which increasing Y results slower convergence. Crucially, a nuanced interplay of y with sparsity level exists. Namely, at y L, convergence slows down for p 0.9, while for 0.7, performance degrades due to high variance, eventually being outperformed by smaller step size.
Notably, high sparsity (p- = 0.5) leads to quick loss stagnation even with a small step size y 0.01 in contrast to E {0.7, 0.9]. Remarkably, the left plot in Figure 2 illustrates that an excessively large step size may even lead to divergence of the method. This indicate that high sparsity significantly alters the minimized loss, confirming that Sparse/Dropout training indeed optimizes a formulation distinct from standard ERM. In general, larger step sizes and more aggressive sparsification (lower p) result in increased loss variance, aligning with our theoretical predictions from Sections 2 and 4.
One of the key practical insights derived from theoretical analysis that the step size y (learning rate) must be decreased for sparse optimization and training with Dropout. Our results demonstrate that this insight applies not only convex models (Figure 5) but also broader range of neural networks.
7 CONCLUSIONS AND FUTURE WORK
This work introduced novel theoretical framework for sketched model learning. We rigorously formalized a new optimization paradigm that captures practical scenarios like Dropout and Sparse training. Efficient optimization algorithms tailored the proposed formulation were developed and analyzed in multiple settings. We expanded this methodology to distributed environments, encompassing areas such as IST and Federated Learning, underscoring its broad applicability.
In future research, it would be interesting to expand the class of linear matrix sketches to encompass other compression techniques, particularly those exhibiting conic variance property (contractive compressors). Such extension might offer insights into seggntuud-bbssede pruning methods and quantized training. Nevertheless, a potential challenge to be considered is the non-differentiabiiity of such compression techniques.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We would like to thank anonymous reviewers for their helpful comments and suggestions to improve the manuscript.
The work was supported by funding from King Abdullah University of Science and Technology (KAUST): i) KAUST Baseline Research Scheme, ii) Center of Excellence for Generative AI, under award number 5940, iii) SDAIA-KAUST Center of Excellence in Artificial Intelligence and Data Science.
REFERENCES
Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. FedRolex: Model-hetterogeneuus federated learning with rolling sub-model extraction. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: openreview net forum? 1dxotdyytUUxUBBEE..... (Cited on page 8)
Dan Alistarh, Torsten Hoefler, Mikael Johansson, Sarit Khirirat, Nikola Konstantinov, and Cédric Renggli. The convergence of sparsified gradient methods. In Advances Neural Information Processing Systems (NeurIPS), 2018. (Cited on page 7)
Dimitri P Bertsekas. Approximation procedures based on the method of multipliers. Journal of Optimization Theory and Applications, 23(4):487-510, 1977. (Cited on page 2)
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223-311, 2018. (Cited on page
Nader Bouacida, Jiahui Hou, Hui Zang, and Xin Liu. Adaptive federated dropout: Improving communication efficiency and generalization for federated learning. In IEEE INFOCOM 2021 IEEE Conference Computer Communications Workshops (INFOCOM WKSHPS), pp. 1-6, 2021. (Cited on pages and 8)
Sebastian Caldas, Jakub Konečný, H Brendan McMahan, and Ameet Talwalkar, Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXiv:1812.077 10, 2018. (Cited on pages 1, 2, and 8)
Chih-Chung Chang and Chih-Jen Lin. LIBSVM: a library for support vector machines. ACM Transactions Intelligent Systems and Technology (TIST), 2(3):1-27, 2011. (Cited page 9)
Xiangyu Chang, Yingcong Li, Samet Oymak, and Christos Thrampoulidis. Provable benefits of overparameterization in model compression: From double descent to pruning neural networks. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 6974-6983, 2021. (Cited on page 2)
Zachary Charles, Kallista Bonawitz, Stanislav Chiknavaryan, Brendan McMahan, et al. Federated select: A primitive for communication-ant memory- efficcent federated learning. arXiv preprint arXiv:2208.094: 09432, 2022. (Cited on page 8)
Yuanyuan Chen, Zichen Chen, Pengcheng Wu, and Han Yu. FedOBD: Opportunistic block dropout for efficiently training large-scale neural networks through federated learning. arXiv preprint 2wXii:2208.05174, 2022. (Cited page 8)
Sélim Chraibi, Ahmed Khaled, Dmitry Kovalev, Peter Richtárik, Adil Salim, and Martin Takáč. Distributed fixed point methods with compressed iterates. arXiv preprint arXiv:2102.07245. 2019. (Cited on page
Laurent Condat, Daichi Kitahara, Andrés Contreras, and Akira Hirabayashi. Proximal splitting algorithms for convex optimization: tour of recent advances, with new twists. SIAM Review, 65 (2):375-435 2023. (Cited on page
Yury Demidovich, Grigory Malinovsky, Igor Sokolov, and Peter Richtárik. A guide through the zoo of biased sgd. Advances Neural Information Processing Systems, 1122315--23177 2023. (Cited
on page
11

---

Published as a conference paper at ICLR 2025
Enmao Diao, Jie Ding, and Vahid Tarokh. HeteroFL: Computation and communication efficient feder- ated learning for heterogeneous clients. In International Conference on Learning Representations, 2021. URL https: http///ppwweeeeeeeeeepeeeeeeeeep net wwwpee/iwwweeeeettotnnoooeeeeeet wwwrrummnddINNPPBBEPP ootp://eee.....eooooooo.o.......
kXg. (Cited page 8)
Yoel Drori and Ohad Shamir. The complexity of finding stationary points with stochastic gradient descent. In International Conference on Machine Learning, pp.
PMLR, 2020. (Cited
on page
Chen Dun, Cameron R Wolfe, Christopher M Jermaine, and Anastasios Kyrillidis. ResIST: Layer- wise decomposition of resnets for distributed training. In Uncertainty in Artificial Intelligence, pp. 610-620. PMLR, 2022. (Cited page 8)
Chen Dun, Mirian Hipolito, Chris Jermaine, Dimitrios Dimitriadis, and Anastasios Kyrillidis. Eff- cient and wightt federated learning via asynchronous distributed dropout. In International Conference Artificial Intelligence and Statistics, pp- 6630-6660 PMLR, 2023. (Cited on page
Yuri Ermoliev. Stochastic quasigradient methods. numerical techniques for stochastic optimization. Springer Series in Computational Mathematics, (10):141-185, 1988. (Cited on page 2)
Utku Evci, Fabian Pedregosa, Aidan Gomez, and Erich Elsen. The difficulty of training sparse neural networks. arXiv preprini arXiv: /906. 10732, 2019. (Cited page 4)
Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning, pp. PMLR, 2020. (Cited on page
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth demand with structured dropout. In International Conference on Learning Representations, 2019. (Cited on
page 2)
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, pp. 1135. PMLR, 2017. (Cited on pages and 2)
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2018. (Cited on page 6)
Noah Frazier-Logut and Stephen José Hanson. Dropout is a special case of the stochastic delta rule: Faster and more accurate deep learning. arXiv preprint arXiv:1808.033: 03578, 2018. (Cited on page 2)
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):3341 2013. (Cited on page 6)
Aidan N Gomez, Ivan Zhang, Siddhartha Rao Kamalakara, Divyam Madaan, Kevin Swersky, Yarin Gal, and Geoffrey Hinton. Learning sparse networks using targeted dropout, arXiv preprint arXiv: 13678, 2019. (Cited on page 2)
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016. (Cited on page 1)
Robert M Gower, Mark Schmidt, Francis Bach, and Peter Richtárik. Variance-reduced methods for machine learning. Proceedings of the JEEE, 108(11):1988-198 2020. (Cited on page 40)
Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General analysis and improved rates. Proceedings of the 36th International
Conference on Machine Learning, Long Beach, California, 2019. (Cited on page 6)
Yiwen Guo, Anbang Yao, and Yurong Chen. Dynamic network surgery for efficient dnns. Advances in Neural Information Processing Systems, 29, 2016. (Cited on page 2)
Stephen José Hanson. A stochastic version of the delta rule. Physica D: Nonlinear Phenomena, 42 (1-3):265-22) 1990. (Cited on page 2)
Filip Hanzely and Peter Richtárik. Federated learning of a mixture of global and local models. arXiv preprint 20XXx:2002.005/4 2020. (Cited on page
12

---

Published as a conference paper at ICLR 2025
Charles R Harris, K Jarrod Millman, Stéfan J Van Der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel Smith, et al. Array program- ming with numpy. Nature, 585(7825):3577336) 2020. (Cited on page-77
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Putlern Recognition (CVPR), pp. 770-778, 2016. (Cited on pages 10 and 50)
Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Πya Sutskever, and Ruslan R Salakhutdinov. Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprint urXiv:/207.05555 2012. (Cited on pages 2, 3, and 52)
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. The Journal of Machine Learning Research, 22(1): (2()):00822-1005 2021. (Cited on page 1)
Thomas Hofmann, Aurelien Lucchi, Simon Lacoste-Julien, and Brian McWilliams. Variance reduced stochastic gradient descent with neighbors. Advances in Neural Information Processing Systems, 28, 2015. (Cited on page 40)
Samuel Horváth, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Lane. FjORD: Fair and accurate federated learning under heterogeneous targets with ordered dropout. Advances in Neural Information Processing Systems. 01112277-1288 2021.
(Cited on pages 2 and 8)
Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin Leung, and Leandros Tassiulas. Model pruning enables efficient federated learning on edge devices. IEEE Transactions on Neural Networks and Learning Systems, 2022. (Cited on page 8)
Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael L. Hubert Hichner, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaïd Harchaoui, Chaoyang He, Lie He. Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konečný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi. Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song. Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Praneeth Vepakomma, Jianyu Wang. Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Found. Trends Mach. Learn., 1411-)::1-211 2021. doi: 10. 10.162200000088. URL https: http:///o....gg//////// 1ww.1116222000000000 (Cited pages and 8)
Ahmed Khaled and Peter Richtárik. Gradient descent with compressed iterates. NeurIPS 2019 Workshop on Federated Learning for Data Privacy and Confidentiality, 2019. (Cited on pages 1, 7, and 9)
Ahmed Khaled and Peter Richtárik. Better theory for SGD in the nonconvex world. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https: /oppnrevieww... net/ Survey Certification. (Cited pages 5, 6, and 7)
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local SGD on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pp. 4519-4529. PMLR, 2020. (Cited on page 9)
Sarit Khirirat, Hamid Reza Feyzmahdavian, and Mikael Johansson. Distributed learning with compressed gradients. arXiv preprint 1TXiv:1806.0577 2018. (Cited on page 9)
Jakub Konečný, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. NIPS Private Multi-Party Machine Learning Workshop, 2016. (Cited on pages and
13

---

Published as a conference paper at ICLR 2025
Dmitry Kovalev, Samuel Horváth, and Peter Richtárik. Don't jump through hoops and remove those loops: SVRG and Katyusha are better without the outer loop. In Algorithmic Learning Theory, pp. 451-467. PMLR, 2020. (Cited page 40)
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar- 10 and cifar- 100 datasets. URI: https:///ww. cs. toronto. edu/kriz/cifar html, 6(1):1, 2009, (Cited on page 50)
Daniel LeJeune, Hamid Javadi, and Richard Baraniuk. The flip side of the reweighted coin: duality of adaptive dropout and regularization. Advances in Neural Information Processing Systems, 34: 2021. (Cited on page 2)
Zhize Li. Hongyan Bao, Xiangliang Zhang, and Peter Richtárik. PAGE: simple and optimal probabilistic gradient estimator for nonconvex optimization. In International Conference on Machine Learning (ICML), 2021. arXiv:2008. 10898. (Cited on page 44)
Fangshuo Liao and Anastasios Kyrillidis. On the convergence shallow neural network training with randomly masked neurons. Transactions on Machine Learning Research, 2022. URL https: /ttp:///ppereee /pppneeeeeeeeeeeeeeeeeeeeeeeeeee net forum? (ideem?id=e? matimnddd=enmmdyyyeemmmmm 7mYYMSy 2H. (Cited on pages 8, 10, and 50)
Tao Lin, Sebastian U Stich, Luis Barba, Daniil Dmitriev, and Martin Jaggi. Dynamic model pruning with feedback. In International Conference on Learning Representations, 2019, (Cited on pages and 7)
Christos Louizos, Max Welling, and Diederik P Kingma. Learning sparse neural networks through I_ regularization. In International Conference on Learning Representations, 2018. (Cited on
page 4)
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning, arXiv preprint 2.00:::0022//6//99 2020. (Cited on page 1)
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas Communication-efficieee learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282. PMLR, 2017. (Cited on pages 2 and 8)
Poorya Mianjy and Raman Arora. On convergence and generalization of dropout training. Advances in Neural Information Processing Systems, 33:21 151-21161, 2020. (Cited on page 7)
Konstantin Mishchenko, Ahmed Khaled, and Peter Richtárik. Random reshuffling: Simple analysis with vast improvements. Advances in Neural Information Processing Systems, 33:17309-17320, 2020. (Cited on page 19)
Decebal Constantin Mocanu, Elena Mocanu, Phuong H Nguyen, Madeleine Gibescu, and Antonio Liotta. A topological insight into restricted boltzmann machines. Machine Learning, 104:243-270, 2016. (Cited page 2)
Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H Nguyen, Madeleine Gibescu, and Antonio Liotta. Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science. Nature Communications, 9(1::2383, 2018. (Cited on pages 2 and
Amirkeivan Mohtashami, Martin Jaggi, and Sebastian Stich. Masked training of neural networks with partial gradients. In International Conference Artificial Intelligence and Statistics, pp. 5876-5890. PMLR, 2022. (Cited on pages and 7)
Yurii Nesterov. A method of solving convex programming problem with convergence rate o(1/2²). Doklady Akademii Nauk USSR, 26443)543,,547 1983. (Cited on page 47)
Mahdi Nikdan, Tommaso Pegolotti, Eugenia Iofinova, Eldar Kurtic, and Dan Alistarh. SparseProp: Efficient sparse backpropagation for faster training of neural networks at the edge. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan
Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 26215-26221 26227. PMLR, 23-29 Jul 2023. wtt..eeeeeemmmmm..............ee / URL https: /ttpc///occeeee///eoooooooocooot ing.idiiii....... ml w22ss//v22//n/11 kdaiiddaa22333 html. (Cited on page
14

---

Published as a conference paper at ICLR 2025
Alexandra Peste, Eugenia lofinova, Adrian Vladu, and Dan Alistarh. AC/DC: Alternating com- pressed/decompprsseed training of deep neural networks. Advances in neural information processing systems, 34:8557-8570, 2021. (Cited on page
Polyak. On the bertsekas' method for minimization of composite functions. In International Symposium on Systems Optimization and Analysis, pp. 178-186 Springer Berlin/Heidelbrrg. 1979. (Cited on page 2)
Xinchi Qiu, Javier Fernandez-Marques, Pedro PB Gusmao, Yan Gao, Titouan Parcollet, and Nicholas Donald Lane. ZeroFL: Efficient on-device training for federated learning with lo- cal sparsity. In International Conference Learning Representations, 2022. URL https: // //openreview net wwtp/eeoumiid-22eeeeeee www.rrrmiddd2soo__mmmmo hmnM. (Cited on page 8)
Peter Richtárik, Igor Sokolov, and Iyas Fatkhullin. EF21: A new, simpler, theoretically better, and practically faster error feedback. Advances in Neural Information Processing Systems, 34: 4384-4396, 2021. (Cited on page 19)
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, pp. 400-407, 1951. (Cited on page 1)
Jürgen Schmidhuber, Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-... hook. PhD thesis, Technische Universität München, 1987. (Cited on page
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs. In Fifteenth Annual Conference of the International Speech Communication Association, 2014. (Cited on page 7)
Albert Senen-Cerda and Jaron Sanders. Almost sure convergence of dropout algorithms for neural networks. arXiv preprint arXiv:2002.02247, 2020. (Cited page 8)
Albert -Ceeda and Jaron Sanders. Asymptotic convergence rate of dropout on shallow linear neural networks. Proceedings of the ACM on Measurement and Analysis of Computing Systems, 6 (2):1-53 53, 2022. (Cited on page 8)
Egor Shulgin and Peter Richtárik. Shifted compression framework: Generalizations and improve- ments. The 38th Conference on Uncertainty in Artificial Intelligence, 2022. (Cited on page 7)
Egor Shulgin and Peter Richtárik. Towards better theoretical understanding of independent subnetwork training. In International Conference on Machine Learning, pp. 45285. PMLR, 2024. (Cited pages 1, 8, and 9)
Virginia Smith, Chao-Kat Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning, Advances in neural information processing systems, 30, 2017. (Cited page 1)
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15())119291998, 2014. (Cited on pages 2 and 6)
Sebastian U. Stich and Sai Praneeth Karimireddy. The error-feedback framework: SGD with
delayed gradients. Journal of Machine Learning Research, 21(237): 2020. URL http: // httpg//pprss/vveee 1 html. (Cited on page
Sebastian U Stich, Jean-Baptiste Cordonnier, and Martin Jaggi. Sparsified SGD with memory. In
Advances in Neural Informaionn Processing Systems. pp. 4447-4458, 2018. (Cited page 3)
Ruo-Ye Yu Sun. Optimization for deep learning: An overview. Journal of the Operations Research Society of China, 8(2):249-29 2020. (Cited page 1)
Rafat Szlendak, Alexander Tyurin, and Peter Richtárik. Permutation compressors for provably
faster distributed nonconvex optimization. In International Conference Learning Representa- tions, 2022. URL https:/ htttts:///ppereeeeee //tp///vepe/e..//////// htt://orummmiii wtt.:///oooommiieooon2e2eeze (Cited on pages
and 47)
15

---

Published as a conference paper at ICLR 2025
Rafat Szlendak, Elnur Gasanov, and Peter Richtarik. Understanding progressive training through the framework randomized coordinate descent. In International Conference on Artificial Intelligence and Statistics, pp. 2161-2169 PMLR. 2024. (Cited page
Pablo Villalobos, Jaime Sevilla, Tamay Besiroglu, Lennart Heim, Anson Ho, and Marius Hobbhahn. Machine learning model sizes and the parameter gap. arXiv preprint arXiv:2207.02833 022222 2022. (Cited on page
Li Wan, Matthew Zeiler, Sixin Zhang, Yann Le Cun, and Rob Fergus. Regularization of neural networks using dropconnect. In International Conference Machine Learning, pp. 1058-1066. PMLR, 2013, (Cited on page 2)
Hui-Po Wang, Sebastian Stich, Yang He, and Mario Fritz. Progfed: effective, communication, and computation efficient federated learning by progressive training. In International Conference Machine Learning, pp. 23034-23054. PMLR, 2022. (Cited on page 7)
Shigiang Wang, Tiffany Tuor, Theodoros Salonidis, Kin K Leung, Christian Makaya, Ting He, and Kevin Chan. When edge meets learning: Adaptive control for resource-constrained distributed machine learning. In IEEE INFOCOM 2018-IEEE Conference on Computer Communications, pp- 63-71. IEEE, 2018. (Cited on page 1)
Sida Wang and Christopher Manning. Fast dropout training. In International Conference on Machine Learning, pp. 126. PMLR, 2013. (Cited on page
Jianqiao Wangni, Jialei Wang, Ji Liu, and Tong Zhang. Gradient sparsification for communication-
efficient distributed optimization. In Advances in Neural Information Processing Systems, pp. 1299-1309, 2018. (Cited on page 3)
Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software, 6(60): 3021, 2021. doi: 10.21 0,211005/jos.003221 URL https:t///s/oiii.. or0/10..00000000 21105/ 03000 03021.
(Cited on page 9)
Dingzhu Wen, Ki-Jun Jeon, and Kaibin Huang. Federated dropout- a simple approach for enabling federated learning on resource constrained devices. IEEE Wireless Communications Letters, 11(5): 923-927, 2022. (Cited on page 8)
Cameron R Wolfe, Jingkang Yang, Fangshuo Liao, Arindam Chowdhury, Chen Dun, Artun Bayer, Santiago Segarra, and Anastasios Kyrillidis. GIST: Distributed training for large-scale graph convolutional networks. Journal of Applied and Computational Topology, pp. 1-53, 2023. (Cited on page 8)
Jiaxiang Wu, Cong Leng, Yuhang Wang, Qinghao Hu, and Jian Cheng. Quantized convolutional neural networks for mobile devices. In Proceedings of the IEEE Conference Computer Viston and Puttern Recognition, pp. 4820-4828, 2016. (Cited on page 1)
Tien-Ju Yang, Yu-Hsin Chen, Joel Emer, and Vivienne Sze. A method to estimate the energy consumption of deep neural networks. In 2017 51st Asilomar Conference on Signals, Systems, and Computers, pp. 1916-1920. IEEE, 2017. (Cited page 2)
Tien-Ju Yang, Dhruv Guliani, Françoise Beaufays, and Giovanni Motta. Partial variable training for efficient on-device federated learning. ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4348-4352. IEEE, 2022. (Cited on page 8)
Binhang Yuan, Cameron R Wolfe, Chen Dun, Yuxin Tang, Anastasios Kyrillidis, and Chris Jer- maine. Distributed learning of fully connected neural networks using independent subnet training. Proceedings of the VLDB Endowment, 1588)1155111590 2022. (Cited on pages 8 and
Hanhan Zhou, Tian Lan, Guru Venkataramani, and Wenbo Ding. On the convergence of heterogeneous federated learning with arbitrary adaptive online model pruning. arXiv preprint arXiv:2001.11803 2022. URL https::////c /ttp::///ppeevveeeeeeeeee net formm? d= (Cited on page 9)
16

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
1.1
Motivating examples
2 Sketches
2.1 Diagonal sketches
3
Problem properties
4 Individual node setting
4.1 (Stochastic) inexact gradient
4.2 Discussion of related works
Distributed setting
5.1 Discussion of related works
Experiments
7 Conclusions and future work
10
A Basic facts
19
Auxiliary facts about functions fD(() and fs(x)
20
B.1
Consequences of Tmy-smothsesss
20
B.2
Consequences of convexity
B.3
Consequences of convexxy
21
C Relation between minima of f and fp.
23
C.I
Consequences of Lipschitz continuity of the gradient
23
D Double sketched GD
24
Nonconvex analysis: proof theorem
D.2 Strongly convex analysis: proof of theorem 2
26
D.3 Convex analysis
27
E (Stochastic) inexact gradient
29
E.1 Nonconvex analysis: proof of theorem
29
E.2 Strongly convex analysis
31
E,3 Convex analysis
32
Distributed setting
34
Fl Nonconvex analysis: proof of theorem 5
34
F2 Strongly convex analysis
36
17

---

Published as a conference paper at ICLR 2025
F33 Convex analysis
37
Variance reduction
40
G.1 L-SVRDSG: strongly convex analysis
40
G.1.1 Convex analysis
43
G.2 S-PAGE: nonconvex analysis
44
H Additional experiments and details
47
H.1 Experimental details
47
H.2
Additional experiments
48
H.2.1
MAST loss trajectory
48
H.2.2
Rand-K sketches
50
H.2.3
Neural networks
50
H.2.4
Standard vs Unbiased Dropout
52
18

---

Published as a conference paper at ICLR 2025
A BASIC FACTS
For all a, E Rd and a > 0,p E (0,1] the following relations hold:
2(a, b) = ||a| + ||6] - ||a b112
(20)
||a bl22 v (1 + |a||2|1 (1 a
|b|1
(21)
- -Ila b112 v 1
||a| + +=11l ||bl1
(22)
1 ") S 2' p20. 2
(1 p)
(23)
Lemma 5 (Lemmaa from (Mishchenko et al., 2020)). Let , E Rd be fixed vectors, X def - X, be their average. Fix any (l, X X.,,, be sampled uniformly without replacement from { { and X be their average. Then, the sample average and variance are given by
E X =X
E X x -
x
n - 1
x
k(n 1)
1=1
Lemma (Lemma from (Richtárik et al., 2021)) Let a, > 0. 1f0 S S
voto
1 then
a²² + < 1. The bound is tight up to the factor of 2 since
Vatb S min
Smn
1 11 vo bJ
Vata"
Proposition Nonzero eigenvalues of SST and S coincide.
Proof. Indeed, suppose À an eigenvalue of T with eigenvector v E Rd, then is an eigenvalue of SS with an eigenvector Se.
Lemma 7, Suppose that f(x) is Lg smooth, differentiable, and bounded from below by fint Then
|V)())) 2T, (f(x) - Vo TRd
(24)
Proof. Let Let =x- by (xx), then using the I, T.y-mooohhness of S. we obtain
L1 || x=|2 2
f(x f(x + (VI(x), xt x)
Since piul < f(x+) and the definition of - we have,
pint < f(a S f(x)
1 Ly [V(( + 2L,
[VS(()) f(x)) 2 [V(())]
Rearrangement of the terms provides the claimed result.
19

---

Published as a conference paper at ICLR 2025
B AUXILIARY FACTS ABOUT FUNCTIONS fD(x) AND fs(x)
For differentiable function f: Rd + R and z.y R Bregman divergence associated with f is Dy(m,y) def = f(x) f(y) (Vf .z v).
Lemma 8 (Bregman divergence), If f is continuously differentiable, then D (F.y) = E [1
Proof. Since f is continuously differentiable, we interchange integration and differentiation. The result follows from the linearity expectation.
B. CONSEQUENCES OF SSMOOTIIESS
Recall the Ly-smoothness assumption. Assumption 2. Function f is differentiable and -y-ooh,, i.e., there is Ly such that Vax, h R
Lif ||A112
f(x + + h)
We also require f to be lower bounded by fint Lemma 9 (Consequences of T.y-smoothness). If f is L.g-mmothh, then
(i) Is is Lifs -smooth with I.fs S LsI.g. That is,
L.g.,
fs(o + h) - fs( + (Vis(a),)) h) +
||]12, V
2
(ii) fo is Li smooth with Lip LpLs. That is,
folx + h)) fo(a) (VSb(a), 13 LoLI 11122 Vx, e Rd.
(iii)
fo(x) S f(x) + (LD 1))
|| u11?, Vz eR.
(25)
Proof.
For any x, Rd, we have
fs(x + h) =
+ h u))
f(u+S(rr Sh)
S f(u+S(r - v)) + (V/(v ((z Sh) + Lg |Shh]
2
= J(u+ - w)) + // + S(x r)), + Ls
(s Sh
2
[s(x) - (V/s(t), + LsLs
S
||A112
2
ii) For any x, h Rd, we have
fo( + h)
= =
Elf(o+((r-0)+S5h)
E + + Lg 2
|Sh]
Eff(u ( -(z)))))))))))))
LE [|MH|]]]
+
2
fo(Φ) + (V/o(x), + L, (B(STS) h,
2
s fo(ø) + (z) h) LoL
||112.
2
20

---

Published as a conference paper at ICLR 2025
(iii For any R., we have
[D(x)
E[S(v + S(x [[[u+S(rr-)))
Asn. 2
u)) + Lg
S| v) - (0 2)12
S E (V/())
2
= (V)),
(x-- w)])
LE 1)
+
2
(
f(x) + (Vf(z),0)
L/f (x vz- (E S - 1) (x v)
2
f(x) + (ID 1)
| u||
B.2 CONSEQUENCES OF CONVEXITY
We do not assume differentiabilityy of f here. Recall that function f is convex if, for all a. E Rd and [0, we have that f(x + (1 aa)y)= <af()) + (1 a)f(y)
Lemma 10. f is convex, then fo is conver and fD(x) ( all Rd.
Proof.
ii Let E Rd and a E [0,1]. Then
fo(ax + a)g)
( + - (v v))) + (1-)))((1++++ S(y r)]] all [f e))] + (1- a) + -(y-
afo(z) + -1- (1-))
Alternative proof: Each fs is obviously convex, and expectation of convex functions is convex function.
Fix x E Rd and let g E 0z(z) be a subgradient of f at x. Then
fo(x)
@
e[f(++S(a - r))]
[f((c) + (g. S(x v) -x)(x))))
f(z) + (g. [S(z v) (xx-(x-i)) -)])
4
f(x) + (g. 0) f(z).
Alternative proof: Using Jensen's inequality, f(u S(x v)) [u S(x v)] 2 f(B (v + S(x v)]) fz)).
B.3 CONSEQUENCES OF CONVEXITY
Recall the μ strong convexity (or, for simplicity, /y-convexity) assumption. Assumption 3. Function f is differentiable and stronnyy convex, i.c., there μu > 0 such that
Vx, Rd
f(x ( h) f(x) + (Vf(z), 1f |A|]?.
2
Lemma 11 (Consequences co-coxiei)ty). If f is coocvee, then
(i) fs is Pfs -Cover with M' 2 pupn- That is,
fs(x + h) fs(z) + (V + USKI ||b|12 Vo,h
2
21

---

Published as a conference paper at ICLR 2025
(ii) fo is coneex with Pfr 2 HDPj- That is, fol + h) fD(2) (VfD(z), PDHS ||bl|
Vz, h E Rd.
(iil)
|1 - e11? Vo E RR.. (ND 1)ll
fr(x) 2 f(x)+ +
(26)
2
Proof.
(i) For any x. h E RRd we have
Js(a - h)
f(v ([v+S(za) S(r+
Sh)
Asn. Z
f(++ S(x u)) + VV + -(z- + Pf |Sh|
2
f(r+ S(x - u)) + (S V + -(x-
+ 14
(S (S Sh. h)
2
[s(x) + (V/s(<)) h) + /SHl |1112
2
(ii For any x. h E Rdd we have
fD( h))
Ejs(c S(x+h(()+-))))
= [[(o+S(z-0)) Shi]
Asn. a
E [[( [[(u+S(r-2)) u)) + (ff
"3 [|h]|²]
2
E [is(o + S(x- (1 o)]] ./)
=
E [|SH|]]]
+
2
SD(x) + (VSD(I), + PS (B S] h,
2
N SD(x) + (V/D(1, HDH ||A11
2
(iii) For any x Rdd we have
fo (x)
E[ + SSx u])l
Asn.
(Vf()) v) (z - u)) +
P$
||[(x v) (x -)]l²
2
= v) -(x v)))
1/ 1) 2)]
+
2
(
f(z) + (Vf(z))0) + 14 (x-v) -) (S - (x u)
2/2
f(x) + (ND - 1)pf
x - 01
2
22

---

Published as a conference paper at ICLR 2025
C RELATION BETWEEN MINIMA OF f AND fD.
Theorem 6. Let Assumptions and hold, and let xi E X and x² X Then
f(x*) s f(x/)) - f(x))= (LD -1 |20 vll² (UD 1)pss
| u112
2
2
f(x*) + (UD - 11 | =11² & fo(s)) S f(x*)+++++)+)))+))))))) + (LD - 1)L1
1| u/12
2
Proof. To obtain the result, combine inequalities (25) and (26):
(26) s fv(x)) [(x)) + (ND 1) 1r w1l2 S s
(25)
1|x - 1/12. S f(x)) (LD 1)Ls
2
Theorem 7. Let Assumption 2 hold, and let wh E Xx and x² E X Then
f(x) S fD(x)) f(x)) (LD 1)Ls || ull?
Proof. To obtain the result, inequality (25). Also, note that, since for every S D, have [s(x) fc v)) 2 f(x"), for all x E RRd we can conclude that fo(x)) = IF [s(((())) >
=[f(=')
C.1 CONSEQUENCES OF LIPSCHITZ CONTINUITY OF THE GRADIENT
The gradient of f(x) is Ly-Lisschizz if, for all x,y E Rd we have that [V/(()- VS(i)| s Ty || yll-
Lemma 12. V Lipschht,, then fD L Lipschitz with
Lip S L S |||]]]
Proof. We have that
-
| S -/ S V
E 8 (/(==+((r-)))
v =))
E S S(c S(y++((y
E
[[| 45 ||x Syl]
[||T [Is]
|| =ll
23

---

Published as a conference paper at ICLR 2025
D DOUBLE SKETCHED GD
Recall that L mmax sups { {Amax (ST = sups {max (SST)} (we used Proposition 1).
D.I NONCONVEX ANALYSIS: PROOF OF THEOREM 3
The following lemma is restated Lemma 4 from the main part of the paper.
Lemma 13. For all a Rd, we have that
[[V[s(())]] & 2 max (fo(x) - fjint.
where the expectation taken with respect to S.
Proof. Due to -mmoothnsss of f, we have that
[(()) || Vf
(S"Vf())
S(x-
E
V/(y), V/(z)) ((STV/
E Amax
(SS
V
L
V f 9 S(r-n) 11"]
(24)
< 2L E [s (u S(x u)) piul]
2 max (Jp(x) -in)).
All convergence results in the nonconvex scenarios rely on the following key lemma: Lemma 14. The iterates {z' he2 of SGD satisfy
yrr $ 21+ 8r 5 + 7 M2,
(27)
where M1 and M2 are non non-negative constants, 8 def [ff(xt) - f and def [V
w Fix w-1 and, for all 1, 2 0, define wi I47 satisfy
Then, for any T 2 1, the iterates {a²) Je2
T-
T-1
Σ
wor <
W1 8 WT-1 8t M2 Σ Wr.
Proof. Multiplying both sides of (27) by we obtain
7
Wt-1 8t
We gt+1 + M2.
wor
For every Osts 1, sum these incqualities. We arrive at
T
T-1 Σ
w w-1 8 WT 8T + M22
w
<
wi-
1=0
1=0
Recall that D def
L
LoL
Theorem 8. Let Assumptions I and 2 hold. For every 2 0, put 8t def "E [fo((a') finf and
p def
[VD((>))1
Then, for any T 2 1, the iterates {T of Algorithum satisfy
(1+ D2z2))
min p² 0St<T
50 + D22 (frrt finl).
(28)
7T
24

---

Published as a conference paper at ICLR 2025
Proof. Due to sssoothhsssss of f. we have that
f(s (( - 5)) - (s + S (x'- s))
+
V (1) (=s++ s (1ltl -
(o'
+ 41 || (xf+1 -) -) gr+1
S (x' 8]
2
-f(s ((s (x'- - -)) -7( V (f(() (v) b-o0st+10 1 1 Ss: (x-
||st_||/// (s:))=1)
+ 2
S 5s(s SSt 8)) - -7 V (u'), , Vis (a+))
1,/72/ + 2
( (5++1) T s (Vs''s) Viss((z)))
Taking the expectation with respect to S yields
7 V (vssst (2)(
fo fo (x s fo (x²)
+ 2
B ( s Vss (x*), (s')
2
$ fo (x') - Vfo (z'). (is) I [xt))
(v
Conditioned on x", take expectation with respect to S
E [/D (z )) |x*]] S Jo (x*) - V/o (x²))|
E [[SSs
2
From Lemma 4. we obtain that
E [fo (a++)) |a4] S fo (z') - > |/ (1²)
L (2L, muax (fr(z) - ) + 2 max (fial - finl))
+ 2
Subtract fine from both sides, take expectations on both sides, and use the tower property:
E [Jo (z
| - (1+ + D2-²) E [fo (x*) - sbul -E (+V56
+ D222 (Jirl plul).
We obtain that
22) 7r S + 8 8++1 + D22 (rint- -int).
Notice that the iterates {x*) J120 of Algorithm 1 satisfy condition (27) of Lemma 14 with Mj 102
and M2 = 2 (fial fial). Therefore, we can conclude that, for any T 2 1, the iterates {s²) of Algorithm I satisfy
T-1
T-1 Σ
T
wor
W-11 8 WT-1 8T + +D2 (sba plun)
Σ Wt-
7
7
t=0
Divide both sides by Et- Wy. From IT-1 Wi > -WTT = Twi
"14Day2 we can conclude that
(1+ D2n2) T
min /¹ <
8D + D2- (frrl piml).
7T
25

---

Published as a conference paper at ICLR 2025
Corollary 2. Fix > 0. Choose the stepsize y 0 aas
8²
min
DVT 2D (f fini)
Then, provided that
128 D
T
max { (30", fime - pint)
e²
have
[1V//[(>')| S
min E. 0StST
Proof. Since Y S 2Da( -int) we obtain
g²
Dºy (Jbs - s ?²
2
1 Since 7 DVT
(1+ S exp (TD???) S exp(1) 3.
y DVT
then, since
36 (6D) D
T 2
g
we have 35
2 Further, i
22
S 7T
2D 20" finry then, since
2
128 D2
(Jinf rlal)
T
g
we have 35" 7T 5. Combining it with (28), we arrive at minosteT E | fol''\)') SE
D.2 STRONGLY CONVEX ANALYSIS: PROOF THEOREM
Theorem 9. Let Assumptions 1, 2, and 3 hold. Let n² def x' - D 2 O. Choose a stepsize
0<7S Lmg Then the iterates {x" J120 of Algorithm I satisfy
<1 -1DD))) E [|]-1]]] 2>² L, Lmax (Sbl plul). E [1142111]]
(29)
Proof. Let nd del xt x.. We get
| = || (a' 7V fsr (x')) - 2/1 - =||=- = V fs: (x'))2
= |r'| - (r fss: + 7² V (-+) ||?
Now we compute expectation of both sides of the inequality, conditional on x²
E |a'] 1|r""| (r E/V fs: (I (x'1) + 2 [| (2*) (x²)]| 112 Taking the expectation with respect St, using the fact that f is continuously differentiable and using Lemma 4, we obtain that
< ||-1 27(7'. ViD(z') + 27 Ly ((fo(z') - sint) (frnt finf))
Since fo is convex, we conclude that (1a,, Vfo(z')) 2 fo(x*)~ - cinf + -innt fore,
yaaa || There-
E
111t1111 s (1 7HDH))
||, 2y
-27 (fo(x*) fbf (1 y L max
+ 2. L (fiuf - pint).
26

---

Published as a conference paper at ICLR 2025
Since
L, Lig. taking expectation and using the tower property get
E
[[ (1 (HDNy) [|>01|]]]
+ 27 Ls LS - (sor fiuf).
Unrolling the recurrence, get
(1-TwDur)) |101| 2y 2 Lgax (frr - fint)
E [|r""]]] s (1-7444
MDHs
Corollary 3. Fix Choose the stepsize y
pDDUE |r0
7 min
L
2L Lga (Sbr fim))
Then, provided that
L max
1, ( finf)
12
log 8'
1Dbuy #
have E [1041]]]] s 25 ||0
MD 11112
Proof. Since
S
have that
2L, min 1 fintt
Lg mmx (fief sint)
S 8 1|=0
PDHF
1 2L, Lg (f j
then, since
2L 2 Lmmax (fbal riul)
1 8'
log
15130 110
we obtain that
(1- (D S exp (-PPDHH/H) s
Further, y LyLgax then, since
L LgL mmax
t t2
log 8
PDHS
we obtain that
(DDD))) S exp expp->yDDupp
Thus, combining it with (29), we arrive at E [1=-1|]] S 26 1|
D.3 CONVEX ANALYSIS
Assumption 4. A set & (abb fo (x+) S fD (x) Vr E Rd) is nonempty.
Theorem 10. Let x² def x' v Let Assumptions 1, and 4 hold. Let f be convex. Choose
stepsize 0 < 7 < 2L,L Fix T 2 and let āf chosen uniformly from the iterates x0 ....... T T-1
of Algorithm I. Then
||:0|12 7T
E [f piul <
L, Lg
max (fitt -
(30)
where e def se 2t. E (0, 10.....7 1}.
27

---

Published as a conference paper at ICLR 2025
Proof. Let start by analyzing the behavior of ||a -3b1² 1: By developing the squares, obtain
|| - - || = ||= - - - 27 (>(Vss('') - = + + +
Hence, after taking the expectation with respect to conditioned on x², we can use the convexity of and Lemma 4 to obtain:
E || =|l |ot|
= -
-1 z51l
+ 2 (T(V)) x²) + (z))(("))13
- |=" - 2 + 27 ( L 1) (fo(a*) -
+ 27 LS (sbrf f
Rearranging, taking expectation and taking into account the condition on the stepsize, we have
T| [fo(z') - fb < L |||*||]
E
[|,*
+ 2, (fbr
finl).
Summing over 1. ,,...,T - and using telescopic cancellation gives:
E [f((x')
1 = ||r0||2 - E [11-11]] L, max (sbe - plut).
1=0
Since E
|| > 0, dividing both sides by 7I' gives:
1 2
7 (E
(E [f((' 1b S 7T
+ Ly (fB fint).
We treat the ( Jensen's 's inequality gives
:)
as it is probability vector. Indeed, using that fn is convex together with
|0011
E [f((T fbl S T LyL max (fiof pint).
Corollary 4. Fix 8 > 0. Choose the stepsize >> 0 as
0 2 (fint fint)) I.
7 = min
Then, provided that
2 m
1, rinl pinf
T 2
max
8
81 11
we have E [fx(T') - [bll] < 26 1
8 11
Proof. Since 7 S
we have that
24, A fi fimtt)
2yL (fint - flul)
I 8 10
2
If 21, A fo forj. then, since
T 2T. S (firf plul)
82 11
11 S 8 2 Further, 2L then, since T > 2L,2 we have
we obtain that
x5 T < 66|1 11 Thus, combining it with (30), we arrive at E [fx(x*) - -ia < 28 1|1 O
28

---

Published as a conference paper at ICLR 2025
E (STOCHASTIC) INEXACT GRADIENT
E. NONCONVEX ANALYSIS: PROOF OF THEOREM
We solve the problem (2) with the method
x²+1 =''-' 79'
(31)
where g' = g(x") is the gradient estimator that satisfies
B[g(x)] (x))
Vo E R
(32)
[|{(=)]1]] I 2A (fs(()) - f
V C, V E E
(33)
Recall that DA,B A + BL, 1Sax.
Theorem 11. Ler Assumptions 1, 2, 15 and 16 hold. For every t 2 0, put 8 d E [fx(x*) 1
and nd def
f/x'')) Then, for any T 2 1, the iterates {x" Ir-o of Algorithm 14 satisfy
(1+ + DA. 80 LALDY (2 (2 - sint) DA, +0)
min y² 0S1CT
(34)
TT
Proof. Due to smmmothhsssss of f, we have that
[(s + (1 - 8)) S [(s + S (x' -))
V (v) (1=)) S ( (s -)
-5 (x' -s)
+
by ||st (1'+1-s) -5 (x² - s)(" ?
+
2
= f(s + S (x' - -s) - 7 (f(() b= S't.gg
Lgn²
+ 2 || 9912
S fss S [s) -7 Visctt (o'),
+ - / (st+l) T St+1 "stig'ggt
2
Taking the expectation with respect to S1+1, we obtain that
L E ((S42)
f S ff(x' 7 (foo(=")
s g.
L
$ - > (fD(a*))
||g/||
Taking the expectation with respect to St conditional on x", using Lemma 4 and (16) we have that
L, E ||10112 [ot]
E [f S fo(z') 7
v V.
2
= fo((') 7 [>x
L E 2A (Js, (x²)- - finf) B
V fs,
(2+)]| (s.)''>)'''
c|
2
< fo(x') 7 [>D(>')|l )1 ALL (fo(x') - sint)
Ly
(2L Lg (So(x') -fatt + 2L Lg max (fint - )in())
2
L
+ AL (fint -E pint
+
2
29

---

Published as a conference paper at ICLR 2025
Substitute fiot from both sides, take expectation on both sides and use the tower property:
E [fo(x't1) - 1b I E [SD(x') - fbb (1 + DA.BI,IIo')) B [VD(at)]
L (2 (2( (lo DA,B C..
We obtain that
yr* s (L+DBBL 8 - LyLDr (2 (fr - fith DA C).
Notice that the iterates h2o of Algorithm 14 satisfy condition (27) of Lemma 14 with M1 DA,DL.I.
LSLD (2 (fbfff fint) DA
M2
Therefore, for any T 2 the iterates {2²) of Algorithm 14 satisfy
T-1
Σ wer'
w-1 80
WT-1 8T
7
I.gb17 L
T-1
(2A fin - ( + L
max (firt - pint) + +C)
Σ
we.
2
Divide both sides by T-1
Wi. From ΣT=
wt 2 TwT-1
Tw 1+Da,BL,LD9
conclude that
(1+ (ABLLBL))
80 L
min
rl
(2 (f - flut) DA.
Ost<T
IT
Corollary 5. Fix > O. Choose the stepsize 0as
22
min
LLLLDDA LDLD (2 (f fint) DA, DA.B+O)
Then, provided that
68 L T
DA.B, 2 2 pint) DA.B
max
we have
| S
min E OSEST
Proof. Since 7 S
2 LgLo (2( - nut Da,)+C)'
we obtain
LILDT (2 rin rint) DA, +0) 5²
2
Since 7 "l
deduce that
(1+ (A+DABL
s exp (TL < exp(11 < 3.
If VlyboDA.
then, since
36 (60)²
Lg LD.
T>
we have
38 77
< S. Further, if
2
fint DD,)+C) then, since
L, Lo (2 ntt
68 LD (2 (fint T2
finft
DAB+CC)
gl
we have 35t
& < 3. Combining it with (34), we arrive at ming E [IVb((>))]] SE?
77
30

---

Published as a conference paper at ICLR 2025
E.2 STRONGLY CONVEX ANALYSIS
15 and 16 hold. Let r² 5 x² IT, t 2
Theorem 12. Let Assumptions
Choose a stepsize
0 S 1_ Then the iterates {ar h120 of Algorithm 14 satisfy
DA,"'
2 (2 5 (1- ( ||/0|/
fint) DA.B +C)
(35)
HDHS
Proof. We get
|| = (1(2' (g6) z111 =1|x² - a/ 7g"l; = 1|1211 2y
, +² ||g²11?.
Now compute expectation of both sides of the inequality with respect to St, conditioned on x², use the fact that is continuously differentiable and use (16):
E
||
1|a
||/111 2y E
E
|| lx**
I |r1 - -2 E
2
[22 (fs, (x))
( B V fs. (z ||( + c)]
|| -2 ffz(=")) + 2A7² (fo(z') 1bb + 72 BE [| V (< ]|²]
+ + (C + 24 fb - {
Since fo is HDH we conclude that (r", VfD(z')) 2 fo(x*) - ginf + "DHL ||,*| 1ا There- fore, taking the expectation and using the tower property, we obtain
E
S (1 7UDP)) ||r111 2y (fo(z)- (bx)) (1 7DA,B)
+72 (2 (firt - pint) D
Since 7 S DA,
we get
s (1- E [|]- +2² ( (2 fiul) DA. CO).
E
Unrolling the recurrence, get
E [|>\1|]] (1->HDuy)) 11+112++ (2 fint) DA.B +0)
PDus
Corollary 6. Fix > Choose the stepsize 7 > as
pDJu5 1.00 2
min
DA.B 2 (fal - fiuf) DA.B+CJ
Then, provided that
1
2 2 - fint) DA. +
t t HDHS
DA,
log 8'
0D 11
we have E [1]-11]] < 28 ?|?
HD 12
Proof. Since 7 < 2( finl D+ we have that
7 (2 (Jird fial) DA.B + +C)
S 26 1012
PDHL
31

---

Published as a conference paper at ICLR 2025
p 8 If 7 2( frr DA.B+C
then, since
fir finl) DA,B+ +C
t
log 8'
1D135 11011
we obtain that
- (D-D)D S exp (exp(110 (-NADHY() S
Further, if y DA.8* then, since
DA. log t
HDHS 8"
we obtain that
(1 "yDDEy' S cxp cxp (-PTDDDD)))) S.
Thus, combining it with (35), we arrive at E [||-1]]]] - 25 11001
E.3 CONVEX ANALYSIS
Theorem 13. Let Assumptions 1, 2, 4, 15 and 16 hold. Let I be convex. Choose a stepsize
0 y S 2D Fix T 1 and let म" be chosen uniformly from the iterates x .... TT'1 of Algorithm 14. Then
1-01
E [JD(=') - pinf] -
<
+ +27 (f - -fint) PA,B+
(36)
7T
where ne def x² x E t0,..., 11.
Proof. We get
|| = 1|((x* - (66)) - =111² = =1|x² r 7961 11 = =1|r"||2- ("(r'' +² ||g²112.
Now we compute expectation of both sides of the inequality with respect to St, conditioned on x², use the fact that is continuously differentiable and use (16):
||
++ 112
||** 27 (r", + 2
[102|] |x**)
<
" [(2A (fs, (x') - ( B Vs. (x)))((())) (z')l2 +
(2))|+++()) +c]]
||'1 - (2 VfD(x*)) + 2A² (fD(x*)- - ( + 72 BE Vs (x-)
+ + 2 (C + 2A (fraf - [ 1))
We can use the convexity of f and Lemma 4 to obtain:
E
|P]
1| 2> (fo(z') - finf) (1 Ar) +22 BL max (fo(z') fiaf)
2222 DL Lmxx (flnl finn) + + (C+ 2A 2A(p - [ (la(()
< 11 27 (fo(*') - fiul) (1 (D-DDA,)
2- +22²BL, BL, L
max (rinl pinl) + + (C+ 2A ( - )
Rearranging and taking expectation, taking into account the condition on the stepsize, we have
7 (JD(2')- - - < E ||| - E [|
+ 27 (Siat - -sint) DA,B + +
32

---

Published as a conference paper at ICLR 2025
Summing over t O....,T and using telescopic cancellation gives
1E [JD(2') [b] - || EE
|-1 + 2 (Jb s +
Since |||-|1]] 2 0, dividing both sides by 7f' gives:
1
E [fo("*) fb())
+2 (fb finl) DA,B+aC.
T
T
We treat the (..... .) as it is probability vector. Indeed, using that fo is convex together with Jensen's 's inequality gives
||:013
E [fD(z') fbol]
(2 (fbe finl) DA, DC
Corollary 7. Fix > Choose the stepsize y > as
8 8/1700 12
7 min
2D
(f fint) DABBCC
2 finf
Then, provided that
2
1, fbof pint 1112011 {
T2 8
m
max
we have E [fx(*') - pintT
< 26 1/001.
1
Proof. Since 5
2 in DAB+C we have that
(2 din flal) DA.B + s 61110112
8 1
If
2 fintt D+ then, since
T> 2 ( - pinf) -fint)
DC
82 11
11 >T 5 0|00
11:9112
we obtain that
Further, if = then, since T > 2DA.9
8 we have
20
5 111 Thus, combining it with (36), arrive E [fD('' firt s 28 |11
33

---

Published as a conference paper at ICLR 2025
F DISTRIBUTED SETTING
Consider f being a finite sum over a number of machines, i.e., we consider the distributed setup:
1 11
min FER
fD(x)=
where L.D. [x)s, -B ( + S, w))]
Recall that Dinxx = max; 13, LD. LD. Lg,
max
F.1 NONCONVEX ANALYSIS: PROOF OF THEOREM 5
We solve the problem (2) with the method
x +1
7 E(s)) (S!) "
(37)
Theorem 14. Assume that each fe, E [m], differentiable, -smooth and bounded from below by finl. For every 2 O, put 8t del E [fo(x') firi] and r" def iterates {s" {z=) T-1 of Algorithm 19 satisfy
[Vfo(a')11 Fix T 2 Then the
(1 ( )
1
min OSETT
80 + Dmax
fb.
cinf
(38)
72
Proof. For [n], due L -mmottheesss Sú, we have that
fi(s (( - s)) s fi(s S (a' -)) V....
- (x'), V.ris((a))
L (s!+1) s ('is'ss) (()ss(())))
2
Taking the expectation with respect to S!+1 yields
fi, (1tt1) S fi,D (x²) ViDD (a'). vis((z)) + T. T.,
Viiss(('')??
2
Conditioned on x², take expectation with respect to S
Lg Lg..D. (2²) (
| (iis(
E [fi,, (1++2) |xy| < f.P. -
E
2
From Lemma 4 we obtain that
E fi,, (I*+)) |x'] s fi,, (+') VL.D (x')12+ LD, L.LLDLLS.. _max,? (z) - (fnt)
For every i [n], sum these inequalities, divide by n
ViLD (x9)112 + 22 Σ (x) - -f!ul)
E [f (1" |at| s fDD
1=1
= fo
ViDD (1')]1 + [o(x*) Σ
ciut
1=1
i=1
= fo (z'
Σ
Vi.D (z9)l12 +2 2.2DDmm max (fo(z")-fe
1=1
1 Σ
+ 7.2 max
pint
riuf
n
i=l
34

---

Published as a conference paper at ICLR 2025
Notice that by Jensen's inequality
1
ViDDD (d')ll²
! M V.D. (ot) -
n
Subtract f from both sides, take expectation on both sides and use the tower property:
E fD (Ilt1)
fbl S 11772 72 mmaxx
n [fo (a*) - Jb D [IV (z*)11
+ 2 max
fint
plnt
We obtain that
Yr' - (1 + 72 8 - 8²+1 + 7,2 D
cinf
1
claf
max
Notice that the iterates {x' 120 of Algorithm 14 satisfy condition (27) of Lemma 14 with M1 = D maxx
M2 = Dinax
inf JD
1
inf
1=1
Divide both sides by E w4. From ET= 1
Wt TWT-1
Tw
conclude that
14.2 max
(1 + 72 mmax)
1
min
1² <
8 + 7 Dmax
inn
0SE<T
IT
71
1=1
Corollary 8. Fix e > 0. Choose the stepsize > 0 as
62
min
Vimax 2D (Sbul
Σx=
M!n
Then, provided that
128 max T 2
max
(30", pinl
,
1
Σ rin
71 i=1
we have
min E oStTT
Proof.
Since 7 S 2Dmax (fin Ei Amr))
2
we obtain
1
8²
? D Dunxx
cin
cinf
2
i=1
Since S VDmusT we deduce that
(1 + 22 2 Dmxx)
T
< exp (Tn²² "Tn2Dmaxd Dmax)
S (1) S 3.
If 7 VDana. then, since
36 (80)² 2
Dmax
D
T
g
we have
35" 7
< c 2.
Further, if = 2ama ful 1 Di fuat) then, since
2
128 Dmaax (fint T2
1
cinf
gl
35t we have
5. Combining it with (38), we arrive at mino E [V/((>)) SE?
<
77
35

---

Published as a conference paper at ICLR 2025
F.2 STRONGLY CONVEX ANALYSIS
Theorem 15. Assume that each Jú, i E [nl,, is differentiable, L,, -smoohh and bounded from below
by ftn S is py-convex (Assumption 3 holds). Choose a stepsize O < S
Then the
max, L.g. Ss
iterations Je20 of Algorithm 19 satisfy
E |r']1]]]] S (1 (uDDu)) 1|=01 + 2y max, {Lg, LLS.L max
(fD
1 Zi=1 fint
148
where re def x² n²D
Proof. We get that
| = 1( Σ
n
x²
7
v
f1,)((z))
2
=
x²
D
2
fiss (x²)
n
i=1
i=1
= =11r'11?
27
r';
n
V Siiss((t (1+t)
+77 )
S (2+1)
n
$ 111 27
Viss(('')
) + +
1
Vi (x²)
t=1
Conditioned on x², take expectation with respect to S i [m]]
|xd S ||-1 27 re, VVo(ot)) + +,2 n
Vi.ss 100)
From Lemma 4 we obtain that
E
++1
2 x²
S ||r0|| - 27 ,-, Vfr(za))
+ 7 2 "
Σ
Tuge I.S.
maxp [[: (s + Si
-)) -5((0=)))
i=1
= =||r"||? - 27 r², (fo(zt)
+,2² 2 Σ
Lg
max (JD - statt)
1=1
$ $117911 27(++ ,², Vfo(a*))
272 max; n
{T. Lig
Imax
(fo.(z")) -fiuet
= =1|r"11 - 27 ,², V)D((*))
+ 27² max {1g max (Jo(x!)- -sinf)
1 plul
+ 27 max (Lg, LSs.x max
stul
71 iml
Since fo is convee we conclude that (r², VSD(*))) 2 fo(x') - fiaff + 1 ||,01 There- fore,
| |a|] s (1 7DDDHs) 2 -2y (1-y max [Le l.m.)) (Fo(z') - f
E
finl 1 finf
+ 2r max { Ly, Lmsas]]
Since 7
taking expectation and using the tower property we get
max:
T.y. maax
[1]] $ (1 (--DDD)) E [||-91]]] + 2y² max Lgnax
figt
1 Σ giuf
n
i=l
36

---

Published as a conference paper at ICLR 2025
Unrolling the recurrence, get
1100//2 2y max, {Lg, LS. Lmax
(for
1 Ei=1,
E
1 1
(1-NNDMA))
μDHz
Corollary 9. Fix 6 >> O. Choose the stepsize > 0 as
MD |r001
7 min
max; {L 1 2 max;
{ (.
max
(fb
1 -i=1 f
S.
Then, provided that
L L WDus
2 (f fintt
log 8
jDDDd |10011
we have E [|]]|]]]] s 26 |1111
6 61
Proof. Since Y
have that
maxx 1,,, 1. [1 Eii fnt)
27 max; f.. LS, [max) (firt 1 1 fint)
I 5 11110112
PDHL
HDM+6 1 If max, Ls, Lauax
then, since
( Σr plu))
max, L } (firr - - finf)
log 8"
131288 1101?
we obtain that
DD-DD)' S (xp((<ND)1/t) 8.
Further, i max, L,, Lg,
then, since
max;
[ max
t 2
log 8'
HDHI
we obtain that
(1 NHDH)1* < exp exp((>)DH1)) s
Thus, we arrive at [1]
< 25 110
F.3 CONVEX ANALYSIS
Theorem 16. Assume that each ft, i E [n), is differentiable, T,,, and bounded from below
by fint, f is convex. Let Assumptions I and 4 hold. Let be convex. Choose stepsize q S
Fix T 2 1 and let HT be chosen uniformly the iterates z0 .... x T- Then
max: (,, Lg.n
|
E [fD(=T) pinl" S
7T 27 max { {L. mmax
where e def se xtp; (0, (0...., 1}.
37

---

Published as a conference paper at ICLR 2025
Proof. We get that
1( x²'
2 n
|1et11p= -
Σ
v
=
r
DD
i=1
n r(
= =117018
27
Si.sy(ar)
+g2
S (a+)
n
n
$ $110911 27
Vissy (1+))
7
1
Visss
n
Conditioned on x², take expectation with respect to S i [m]:
[o0] S 1191 27 r.,, V/D(o')
+ , -
n
From Lemma and from convexity of fo we obtain that
E p +
[x0]
$ ||r-||12 - 27 (fo(x') - fb())
2 2 n
+
Σ
Lg L
S.
m [Js (s + (a((x'=s) ]
n
= =||--112- 27 (JD(x*) - sbat) +72 T. l.S.
(Jz, (a') - )
1=1
s $1/-1122 - 27 (fD(a*)- - + 2y² max, {Igg Tmax n
Σ (sD, (x²) - -Saut)
= 111-/122- 27 (VD(e") - ) + 27² max { LL, Lg Lmax (Jo(x*) - -bb0)
+ +272 "max {,g Lg.N]
/ fbo Σ finf
i=1
= -11-9112 - 27 (fp(2")))1) -bbll) (1 -1- y max { fr, max)
27²² max {Tej LS. ]
finí
inf
Rearranging and taking expectation, taking into account the condition the stepsize, we have
1 n
7 (FD(a - (b()) < E | [2]
E
[||- 11
||] + 2n2 max [L, LB maxx
)
sint
inf
n
Summing over t o.... T --1 and using telescopic cancellation gives
T-1
7 Σ t=
E [f - 1) s ||~0||- |,7
+ 2T-²2 max [L, Lmax) } ( finf
n
Σ s
n
Since R |1
2 O, dividing both sides by TT gives:
E [fo(x*) - -)bd - ||r01
n Σ
+ 27 max (L, Lgx.) ) rinf
giuf
T t=
7/
i=1
38

---

Published as a conference paper at ICLR 2025
We treat the (....... .f......f . if it is probability vector. Indeed, using that fo is convex together with Jensen's inequality gives
1 y/T
E [fo((r)) - pn S
+ 2y max { {L, I.S.x)
pint
i=1
Corollary 10. Fix 8 > 0. Choose the stepsize 7 > 0as
5 |
1
7 max; {Lg, Lg.
Fint
2i=1 iin
Then, provided that
2
S
finf fint
T 2
m
max
1,
6
11
we have E. [fx(=') - [b4
S 25 110122
8 1110112
Proof.
Since 7 S
we have that
2 max.
L,, Lg,
(1b: Ei. s(u))
n
27 max {T.g, I.S. maxx
finf
sine
S 1111.112
i=1
1|
y
then, since
max,
Ly, 8s
(fin
Σこ- f!a))
2 max, Lgn
Lmax (firt
n Er= fiml)
T2
82 11
1|:0112 7T S 8 1112 Further, i 2max. Ly, L
then, since 2 max:
L,, Lg, max
we obtain that
|
we have YT s 5 ||00
Thus, we arrive at E [fx(z") - f < 25 |1,
39

---

Published as a conference paper at ICLR 2025
Algorithm 3 Loopless Stochastic Variance Reduced Double Sketched Gradient (L-SVRDSG)
1: Parameters:
learning rate > 0; probability μ; sketches Si...., Sw; initial model and shift x²,v E IR, sketch minibatch size b; initial sketch minibatch 8 C [N]
2: Initialization: w = _ ho b X
Σ
Vss, (wo).
3: for /. 0, 1, 2.. do
4: Sample sketch: St from (Si...., SN]
5: Form gradient estimator: ht = V fs (x') (s-(w") + h:.
6: Perform a gradient-type step: x'+1 x yh
7: Sample a Bernoulli random variable 5p
8:
if Bp 1 then
9:
Sample B uniformly without replacement
10:
w(+1 = x², h++1
Σ (xx)
KEB
11: else
w++1 wt, htt1 = ht
13: end if
end for
G VARIANCE REDUCTION
In Theorem 2, established linear convergence toward a neighborhood of the solution zi for Algorithm (I). To reach the exact solution, the stepsize must decrease to zero, resulting in slower sublinear convergence. The neighborhood's size is linked the variance of gradient estimator at the solution. Various Variance Reduction (VR) techniques have been proposed to address this issue (Gower al., 2020). Consider the case when distribution is uniform and has finite support, i.e., (S,....,S. leading to a finite-sum modification of the MAST problem (2)
1 A N Σ BiS(u
fD(1)
Sx
(39)
In this situation, VR-methods eliminate the neighborhood enabling linear convergence to the solution. We utilize the L-SVRG (Kovalev et al., 2020; Hofmann et al., 2015) approach, which requires computing the full gradient with probability p. For our formulation, calculating Vss for all possible S is rarely feasible. For instance, for Rand K, there are N = dd!(( - K)!) possible operators S. Therefore, in our Algorithm 3, we employ a sketch minibatch estimator h. computed for a subset B C [N] (sampled uniformly without replacement) of sketches instead of the full gradient. Finally, we present the convergence results for the strongly convex case.
Theorem 17. Assume that T-mssmoth (2), str-slyyglyy convex (3), and S is sampled from finite
set (S S. Then, for stepsize L Limax and sketch minibatch size b E (0, NN, the iterates of Algorithm 3 satisfy
L Lmax (N b)
E [[w"] <1 P)T To pmax {1,N 1} (Jinl fint).
where p = max tYHDAA p/2} and Lyapunov function
"'r def 1|x²
x'1'"
16y²
Σ V[Vss, - V/s, (=+))
PN
Note that the achieved result demonstrates linear convergence towards the solution's neighborhood. However, this neighborhood is roughly reduced by factor of 1/b compared to Theorem 2, and scales with N b. Thus, when employing full gradient for ht with b N, the neighborhood shrinks to zero, resulting in a linear convergence rate to the exact solution.
G.1 L-SVRDSG: STRONGLY CONVEX ANALYSIS
The proof of Theorem 17 relies the following Lemma:
40

---

Published as a conference paper at ICLR 2025
Lemma 15. Let Assumptions 1 and 3 hold. Then the following inequality holds:
E ||a =51""] = (1-7WD))) ||a² - rbbl? - 2y (JD(x*) - ((xb)) + 2B [1101]]
Proof. We start from expanding squared norm:
E || zbb""
= E || - 7/1 xnn1 1|]
= E ||aa - x 11] - -2 x' xb)) + - |n'1 [-]
||Ia x 12] - (V/(V/ (x)) x² - xT)) + 7 |
-E
= (1 - ((44s) ||a2 = 27 ( f(xb)) + 2 [|]]1]²]]
Lemma 16. Let Assumptions 1 and 2 hold. Then the following inequality holds:
N
1 B [100 S 8L max (fo(z)) - fo(z))) +
(w') - (ss(((a)))
N
8(N -) b)
+ (N 1)6 Lg
max (fhr fiul)
Proof. We start the proof from the definition of h²:
|1111p] =
(s) V.ss (wo"
1 +
Σ IE
V fs: (w0)
b
(/:))(()) Vss((()))
V (10") V/s(( + Vss:(tt))
iE
(xD)
V (xb)
5 1
< 4E
V (x')
V (2 ||] + 4E
Ss (w²) - V (x>)) )1
+ 4E
1 Σ
(((s((('') - v(s(((b))
+ 4E
1 Σ v Ssp (x+))
VEB'
iEB'
(24)
+ 8 1
f (w01) - V (is 2
, (fo(z*) - ff(()))
N
+ 4E
Σ v (x})
'ES'
Using Lemma 5 we obtain
N
E [|]
< LL max (fo(z)) - fn()))) + +8- N
Σ
V (ss((w)) - Vss((at)) -V/
- b Σ
2
+ 4 1,N - 13b N
Vlss (ab)
EEB
N
S 8L, L.g xin (fo(=')- - (x +8 N
1
Σ
V (ss(((²) Vss |?
8(N - b)
I, I.g. (Staf -iint)
11, max(1,N
max
13b
41

---

Published as a conference paper at ICLR 2025
N 16 P.
Lemma 17. Let Assumptions I and hold. Let D the following inequality holds:
Sss (at) ((ss((w))
Then
E D (1-p)) 32> Ly (fo(w') -
Proof. Using the smoothness property, we obtain
16 N PN
E D (1 - D) +
Σ E
Vis (xb) Vss;(w])] -
S (1 (1- p) + 32, L (fo(w') fo(eb)).
Theorem 18. Assume that is Ly-smooth (2), pyy-strongly convex (3), and S is sampled from finite set (St.... Sy Then, for stepsize S Lg and sketch minibatch size b E (0, N). the iterates of Algorithm 3 satisfy
S (1 - p) yo L (N b) pmax (1,N- - 116 (Si -iint,
E [ r
16y def max (YNDDI P/2] and Lyapunov function T 1x
where P
(ws) (=
PN
Proof.
We combine three previous lemmas 15,16, 17:
|| -x 20/1² 2 + D
$ (1 - 7/D/U1)) 12' - x - (fo(z") - f((t)) + + [101
+ +1 Dt + L L
max (fo(w') - fo((t))
$ (1 - -/DH4s
12' - - 27 (fo(x') - f(xb)) + (1 D)D
+ 32, 2 Ls Lg (b(w")-Iold) -
max
+72 LL (fo(w') - ff((t)) + 2n²
D
8(N - -b)
max 11, 17b I, I. max (finf flul)
= (1 - (TNDDH) 1| - -x ||
- 2y (
- -f(xb)) (1 20 L L
(1- D² + 8(N b)
max - 136 72 L max (frrt -int).
Since 7 L0 get that
vl+1 < (1 (1-max max
8(N b)
2 L max (fint finl)
max(1,N 1}b
Unrolling the recursion and using max (T/DD/ }} we obtain
E [ < (1 p)² yo + L Imax (N b)
p pmax (1, N - -136 (flot finf).
42

---

Published as a conference paper at ICLR 2025
G.1.1 CONVEX ANALYSIS
Now we formulate and prove theorem for the general (non-strongly) convex regime:
Theorem 19. Assume that is -mmoohhh (2), convex, and S is sampled from finite set {S ..... (S }. Then, for stepsize y I g Igg and sketch minibatch size E (0, N) the iterates of Algorithm satisfy
4o IR [SD(z')] - ((xb) S
8(N -b) b) max m1. 1}b
I.
max (sbel
siul)
TT
where
JT
+ Et= x²
and Lyapunov function
Tre
def
1100
r 11?
16
V (w.)
Proof.
We start from the recursion in Theorem 18:
E E|| ztll" + D
- (1 (HDHD) || - xx11²
2y (fo(x') - f(xb)) - LL Lgaax max
(1 1 2 D' 8(N~ - b)
max(1,N) 1,N - -1/6" L L5 max (for fint).
Using uy = 0 and (1 (1-5) -1- [ < we have
E | rbll + D++1] s ||-2 - 2 27 (JD(x') )- s(xb)) (1 - -20y I. r. I.gax)
D2 8(N - b)
m I,N - 1}b6" L max (fbs fiul).
Since 7 S Lo we have (1 - 20 I. Lgg.x 2 and it leads to
E
| 1aatt
r 1| + D++1
$ |1 -
2 - (fo(x*) - ff(xt))
8(N - b)
+ D$ max (1,N 13b 7? LL (fi - -intt)
Using the tower property, we have
E
v 11++1
< [
yE (ff(x²)) - f(xb))]
8(N - b)
+ max 1,N -136" 72 L max (fbr - flul)
BB (ff(2+)) - -
v - [v]t1]
8(N - b)
max 1,N 1)bb 722 L,, Lg max (fr - fiut)
T-1 1
1 T-1
E ((o()))(((t))])
<
E [v]] - -E [x(-1]
T
T
E
1=
8(N - b)
+ max 1}b 72 L, max (fint - fint)
+ 8(N - b)
E [ [ < 7T
7 L, Iaaax (J - )
11, mmxx11,N - 11bb
max
43

---

Published as a conference paper at ICLR 2025
Algorithm 4 Sketched ProbAbilistic Gradient Estimator (S-PAGE)
1: Parameters: learning rate > O; probability p; sketches Si.... Sv; initial model and shift z", o E IR, sketch minibatch sizes b and d' b; initial sketch minibatch B C [N].
2: Initialization: ht"
Σ V Is, (x9).
3: for t = 0, 1, .. do
4: Perform a gradient-type step: x²+1 x² yhh
5: Sample Bernoulli random variable 3p
6: if 8, 1 then
7: Sample minibatch B with size b uniformly without replacement
Form a gradient estimator: h=+1 = Σ Vss (ss((ott)) (x++1)
8:
VEB
9: else
Sample minibatch (B') with size V/ uniformly without replacement
11: Form gradient estimator: h++1 = h² +
Σ
(Vis((=*+1) - v.ss (zl+1)
iE(BY)
12: end if end for
G.2 S-PAGE: NONCONVEX ANALYSIS
In this section, we introduce variant of the Probabilistic Gradient Estimator (PAGE) algorithm applied to the MAST formulation as defined in Equation 2 for non-convex setting. Li et al. (Li et al., 2021) showed that this method is optimal in the non-convex regime. We refer to this method as the Sketched Probabilistic Gradient Estimator (S-PAGE). Calculating the full gradient is not efficient, as the number of possible sketches when considering Rand- K is given by (n-k)lœI Consequently, we employ a minibatch estimator to achieve partial variance reduction, using a large minibatch size where b > R. For the purpose of analysis, assume that the variance of the sketch gradient is bounded.
Assumption 5 (Bounded sketch variance). The sketched gradient has bounded variance if exists oD 0, such that
E Vs (x) - V (z) S oob Ve E Rd.
Lemma 18 (Lemma 2 from (Li et al., 2021)). Suppose that function f is L-smooth and let x++1 x "ig'. Then for any g² E Rd and 7 > 0, we have
f ( - [[(00))
<s
1
-
4 2
||
2²1 ? [1 (40)
27
Lemma 19. Suppose that function o L Ly-smooth and S satisfies Assumption and let x
x² yh. Then for any ht Rd and 7 > 0, we have
lo (1+1)) S Jo (x') -
1--1 27 2
LDLD
|| 1o-t1-x'11
=01 -
(41)
Proof. Since f -mooohh and S satisfies Assumption the function fo is Lyo. Then using Lemma 19, we obtain
JD (I*+1)= S /D (x*) - 2 ||/((>>))]
(
Lif
| xe 2 2 [
27
2
Using Lemma we have Lip - LoLy. Plugging this into the inequality, we obtain
fo (1++1)= s fo (x') (x') [+++)))))) (
|| 2²11 2 + ? ||h' (f
27
2
(42)
44

---

Published as a conference paper at ICLR 2025
Lemma 20. Suppose that Assumptions 1, 2 and 5 hold. If the gradient estimator h++1 is defined according to Algorithm 4, then we have
E 1 - [1]
- (1 p) li,t fDD("))12
N V
+ (N 1)b (1 p) 1 LD | - x²I|
N - b
+ P (N- 1)b o
Proof. We start by considering two events:
27
11 = E ||1 - v/D(x*+))1 = plE
V (isy(att)) ff 111
1EB
17
+ (1 B E
h' h 1 Σ
(Vis((a*t)) (v)s(())))))))(( - Vio(o
Using Assumption 5 and Lemma 5, we obtain
= E |1 - Vfo(
S p (N - 11b
N- - ob + (1 - p) E
14 1
Σ (Vis((o-t))- - V/ss(z())
V Vss ) -
iE(ByY'
N- b
(N - 1]bb
(1 p)
x E
h'- f (xt)
1
Σ (Vis( - Viss(at)) VfD (xtt2) + Vio
6'
"E(Be)
N- 2
S (N 1)b D
+ (1 - p) - V/D (x'))1
+ EE
Σ
(viss (14 Vss (z4) - Vfr (x + Vfr (29)
VE(BLy
N b
S Sp (N - 1)6"
ob
(1
(1-p) 11h f/ (2²)
N W-B b'
1
Σ [(
( (1-t1) Vsss - -
+ #p (N 1)
E.
N
)E(B")
N -
I S (N - -1)6
0
+ (1 - p) 112² viol (22
N Aa E
1
Σ
[[ (1/ - Vss; (x')
(N 1)8
N
'E(B.)'
45

---

Published as a conference paper at ICLR 2025
Let us consider the last term:
Viss (x - - (zt) |3]
E
()) (8) E [(::)] (Vf (u+ S () 0))
E
// (u 1 (
0)))
< E Amax (S))
(u S) - -)) v/ (v S b))) |2]
< L ||( [[[f
vv + S - Vf (v + S -
-))- 0))
Using Lipschitz continuity of gradient of function f, we have
11]
Vss; (2++1)- Vss; (x')
S Lg max LB [1st 2²)
S LS L Limxmax E (S)))) I] || - -2'11
Lg L LD 1 | - 101
Plugging this inequality leads us to the final result:
E -
V (x S (1 p) |1,t {DD("')]1
+ N V
(N 158 (1 - p) 13 L LD | 1t - 291
N b
p (N 1)b 00.
Theorem 20. Assume that f is L smooth (2) and S satisfy Assumptions and 5. Then, for stepsize 1 / /1ga Lo) the iterates of Algorithum 4 satisfy
7 S V ph
E [IVID (ET))1]
2E [
(N 1)b o..
N b
7T
Proof.
E [1]
E
rinf 11(x
-
-
? 119
2p
(1 1 L
E
folx*) fine
? \fo(2"))2
1 1x - x2112+ + 2 l1 v
2
27
E
7
(1 p)
I|g \\D(>')11
(1 - p) Lmax LD
| -112 111
(N 1}b* ob }
N - b
2p
b
E
fo(a*) fing
?
(1²
2p ((1-)) - p) I|g0 (VV-Vf((>")' + + l|g Vfo
2
7 N -
1
Lip
(1- p) L LD)
2 (N - 1)b D
27
2ph'
46

---

Published as a conference paper at ICLR 2025
Using stepsize Y $
V 1 1 Lo+ L and Lemma, get
1
E [b E Th 2 2
ViD( (2 )1
7 N
2 (N - 1)b o
(43)
7
[VD(2'))
S E [
E [
+ 7 N W b
2 (N - 1)b o
(44)
7
T-1
[/r((')||
+ 7 N 5-b
2 (N - 1)b o
(45)
2
Let 3y be randomly chosen from fa") he)m we have
E [1V
2E [ol
N - b (N 1)b DD
v TT
H ADDITIONAL EXPERIMENTS AND DETAILS
First, we provide additional details on the experimental settings from Section
H.1 EXPERIMENTAL DETAILS
In Section 6 and for all further experiments, the following problem is considered:
f(x) def
log (1 exp(- x. b.)) + 2 |
(46)
where A, E (Rd, b, }-1,1] are the feature and label of i-th data point. The approximate "optimal" f* solution of optimization problem (46) is obtained by running Accelerated Gradient Descent (Nes- terov, 1983) until | f+1 < 10 -30 Our implementation is based on the public Github repository of Konstantin Mishchenko. Simulations were performed on a machine with 24 Intel(R) Xcon (R) Gold 6246 CPU @ 3.30 GHz.
Sketches. Problem (2) may not be easily solvable precisely for the most general sketches. Therefore, consider the scenario when D is uniform and has finite support similar to Section G. We use special class of diagonal permutation sparsifiers formally introduced in the following example:
def d/ K and π (Ti,..., Ta) be a random permutation Example 3. Assume! that K divides d, let = of [d]. Then for all E [K], we define Permutation sparsification (in short Perm-K) operator as
gi
S def K
Cx,y
my.
(47)
}=
1)+1
where ei....,ed E IRd are standard unit basis vectors.
In simple words Perm X sparsifiers require random shuffling and then dividing the set of indices [d] = {I....,d) d} into K non-overlapping subsets C₁ C G of equal size such that Uc,egC; = [d]. Then, every sketch S, is formed from e vectors based on indexes / C, resulting in D (Si...., } To ensure unbiasedness, sketches are sampled uniformly with probability 1/K. This class of sketches satisfies Lp μDD K and L max K²
Details for Figure (and 3, 4). Regularization parameter À in (46) is set to guarantee that the condition number of the loss function Ky equal 102. The dataset shuffled and split to train and test in 0.75 to 0.25 proportions. Initial model weights x0 Rd are generated from standard
This done for simplicity presentation and can be easily generalized (see Appendix from Szlendak al. (Szlendak et al., 2022)).
We use array_spaaa split method from NumPy (version 1.26.2) package (Harris al., 2020).
47

---

Published as a conference paper at ICLR 2025
Gaussian distribution ((0, 1). For every sparsity level Perm-K sketches D (S .Sk are generated leading to different MAST problem formulations. This process is repeated 10 times with various permutations for changing random seeds. After ERM and MAST models the accuracy of sparsified model Syz calculated for every i E [K] on the test set. This results in distributions of accuracies for every sparsity level.
obtained
0.27
2:37
66
0.35
0.94
9.35
NAST
30
72
20
25
30
90
25
20
73
20
20
Sparsity
Sparsity
Spaniiyy
(a) a5a dataset
(b) ala dataset
(c) w&a dataset
Figure 3: Accuracies distributions of sparsified solutions for the ERM (1) and MAST (2) formulations.
0.27
0,7
3.46
0.4
ERH
MAST
0.3
EMH
NAST
ESM
MAST
Sparsity
Sparsity
Spandity
(a) a5a dataset
(b) ala dataset
(c) w&a dataset
Figure 4: Test accuracies of sparsified solutions for the ERM formulation (1) and MAST problem (2).
In Figure 3, we display complete results, including the accuracies lower than 0.5 (unlike in Section 6), which occurred only for the most aggressive sparsification of ERM models. In addition, the median of accuracies (in white) is shown for every (ERM/MAST) approach and sparsification level. Additional results for ala and w&a datasets are consistent with those Section 6 MAST models' performance is higher, less variable, and more resistant to sparsification. Figure 4 also shows the swarmplots for the same experiments represent accuracy values' distributions better.
H.2 ADDITIONAL EXPERIMENTS
H.2.1 MAST LOSS TRAJECTORY
In the next experiment, we investigate the optimization efficacy of Double Sketched Gradient Descent (Algorithm 1 (II) with Perm-K X sketches (47) for MAST formulation (2) with chosen as (46) for several datasets. The model weights are divided into K = 10 groups, which allows us to solve the MAST problem, find fint and evaluate fp. VDD precisely. Moreover, unlike the previous experiment, the inexactness of the stochastic gradient estimator is introduced via uniform (single element) random subsampling of data fi as the problem enjoys finite-sum representation (17).
Figure 5 shows the trajectory of a method which averages across all sketches S which results in exact (w.r.t D) gradient estimator f... (right column on legend) and the same algorithm but with uniform sketch subsampling fi.s (left column). The methods are run with 3 different step sizes y. Subsampling of data introduces oscillations preventing the algorithms from converging to the exact optimum. Sketch subsampling leads to additional variance highlighted by the curves corresponding to the same step size. Moreover, our results clearly indicate the necessity for decreasing the learning rate y for sparse/dropout training with SGD. The method with sketch subsampling and standard SGI)
48

---

Published as a conference paper at ICLR 2025
step size (dotted blue curve) fluctuates around initialization, while full averaging across sketches (dotted cyan curve) fixes the issue partially as the error floor is lower however, there is still almost no convergence. Scaling " inversely proportionally to sparsity level LD K results in clear linear convergence to the neighborhood of the solution for V Vi, estimator. However, Vis requires decreasing the step size by L K which well agrees with conclusions of Theorems and 12.
10
MAST Loxk, fode) 18
2011
10
P/ 1-1
The 5.6
2% 1//
Fia y 7/6,
iteration number,
1.7
MAST Lowk, (622) 17
100
10
PA 114,
1000
2000
3000
0000
5000
6000
nds dataset
1.7
MAST Look. (xlP)-F -FF
10-4
10
Vie now M7 Vie
10-1
1200
2000
3030 number,
4000
3000
6000
Figure 5: Finite-sum MAST loss (17) convergence Algorithm 1 (II) with subsampling.
49

---

Published as a conference paper at ICLR 2025
H.2.2 RAND-K SKETCHES
In this experiment depicted in Fig- ure 6, we validate the claims of The- orem 2. We consider the same lo- gistic regression optimization prob- lem (46) with À set to make sure Kf =103. We use the whole dataset ala and initialization z" = O. Con- sider D as a uniform distribution over Rand-K sketches for K 1. Then, MAST stochastic optimization for- mulation (2) leads to finite-sum problem over sketches S, as defined (17). This allows us to evaluate the performance of Algorithm 1 (I), which converges linearly for the ex- act MAST loss (2). Note that apply- ing Gradient Descent requires com- puting double sketched gradient for all possible (N = d) sketches. We denote step sizes to Lp) and " = ( L according to theory.
10"
10"
MAST loss, h
10
10
G addert Desoon 50
10
Double Blositted GD with ye
Doudde Seenched on wit
Stutchied
11
5000 10000 15000 20000 25000 30000 35000 40000 Iteration number,
Figure 6: MAST loss (2) convergence for Algorithm (I) with varying step size and Rand-1 sketches.
Our findings accentuate the pivotal role of the appropriate step size fy in steering the trajectory of sparse training. Guided by the proposed theoretical framework, this step size must be adjusted in proportion to /²/d& for Rand-K. Notably, for this particular problem at hand, a larger 7 (e.g.. y 10%)) can accelerate convergence. Yet, surpassing delineated boundary of the progress (e.g.. 20%a) and, in specific scenarios, even derail the convergence altogether (c.g.. y 50%)). Such observations underscore the imperative of modulating the learning rate, especially within the realms of sparse and Dropout training.
result in stagnation
H.2.3 NEURAL NETWORKS
This section presents our deep learning experimental results. Figure 7 illustrates the loss behavior of the distributed Algorithm 2 (for M 10 clients) using Bernoulli sketches (6) for p; on the standard loss (18) (for S, = I). Our experimental setup closely follows that of Liao & Kyrillidis reiterate key details. We employ ResNet-50 model (He et al., 2016)
(2022), and for completeness, pre-trained on ImageNet as feature extractor, concatenated with two fully connected layers. This combined model is then fine-tuned the CIFAR-10 dataset (Krizhevsky al., 2009). The outputs of the re- re-trained ResNet-50 serve as input embeddings, while the logit outputs of the combined model used as labels.
The first column of Figure 7 shows how the method's performance is affected by the sparsity level (p) and step size (7). Specifically, for high sparsity p = 0.5, Figure 7(a) illustrates that an excessively large step size (y 1) may even lead to divergence of the method for ERM loss. Across all sparsity levels, we observe "sweet spot" for the step size, beyond which increasing y results in slower convergence. Furthermore, training with high sparsity (p 0.5) leads to a quick stagnation of the loss in contrast to p E {0.7, 0.9]. The second column of Figure displays the subsequent divergence (for p = 0.5), indicating that high sparsity significantly alters the minimized loss, confirming that Sparse/Dropout training indeed optimizes a different formulation than standard ERM.
In general, larger step sizes and more aggressive sparsification (lower p) result in increased loss variance, aligning with our theoretical predictions from Sections and 4. Interestingly, figs. 7(d) and 7(f) reveal that Dropout training
outperform non-sparse optimization for small step sizes (y = 0.01) or initial iterations (up to ~ 1200 for 7 = 0.1). However, the largest step size (? = 0.5) the most efficient in terms of minimizing canonical loss.
50

---

2500 Published as a conference paper at ICLR 2025
One of the key practical insights derived from theoretical analysis is that the step size y (learning rate) must be decreased for sparse optimization and Dropout training. Our neural network training results demonstrate that this insight extends to broader range of models.
0.5
2 0.5
0.001
0.5
10
0.5
-
0.9
0.01
p= 0.7
0.05
10"
10
Loss
Loss
10
10
10
10"
800
1000
1000
3000
lteration number
lleration number
(a) High sparsity
(b) Large step size
p 0.7
7 0.1
y- 0.001
7- 0.5
p 0.5
-
p- 0.9
2= 0.01
7=1
p
p=1
10
7-
10"
Loss
10
Loss
10
1
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Iteration
Iteration
(c) Mid sparsity
(d) Medium step size
0.9
0.01
p=
10
Y- 0.05
10
Loss
Loss
10"
10
500
1000
1500
2000
2500
3000
500
1000
1500
2000
3000
lleration number
lteration number
(c) Low sparsity
(f) Small step size
Figure 7: Performance of Algorithm 2 with Bernoulli sketches (6) on standard loss (18) (for S, = I).
51

---

Published as a conference paper at ICLR 2025
H.2.4 STANDARD UNBIASED DROPOUT
We supplement the results from Section H.2.3, by comparing unbiased scaled (6) and biased Dropout sketches in the same setup by running distributed Algorithm 2 on standard loss (18) (for S = I)). Figure shows the training loss curves for different sparsity levels (p 0.7, 0.9) and learning rates (y 0.05, 0.5, 1.0). The unbiased estimator includes 1/p scaling factor, while the biased version omits this scaling as the original Dropout (Hinton al., 2012).
The results demonstrate that the unbiased estimator (solid purple lines) consistently achieves lower loss values compared to the biased estimator (red lines with markers). This effect particularly pronounced at lower sparsity (p 0.7), while higher sparsity (p 0.9) the difference becomes less dramatic. Higher learning rates lead to increased variance in both estimators, though the unbiased approach maintains better overall performance, which aligns with our previous observations about the impact of sparsity and learning rates on optimization dynamics.
p=0.7
p=0.9
100
Unbiased
- 0.05
Biased
10
Loss
10
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
lteration number
Iteration number
100
y = 0.5
10-1
Loss
10-2
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
Iteration number
Iteration number
100
= 1.0
10-1
Loss
10-2
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
Iteration number
Iteration number
Figure 8: Comparison of the unbiased sketches (6) and original (biased) Dropout.
52


Reviews:
[{"review_id": "2jZikIURj0", "reviewer": "Reviewer_w3TH", "summary": "This paper introduces Model-Agnostic Sparsified Training (MAST), a novel optimization framework tailored for sparsification in model training. By defining a new objective that incorporates pre-trained models and random sketch operators, MAST allows for simultaneous model and gradient sparsification during training. The authors propose modified SGD-based algorithms to handle this objective, offering both single-node and distributed variants. Experimental results highlight the approach’s efficacy, showing MAST’s superior robustness to pruning compared to standard models.", "questions": "Could the MAST framework potentially be extended to conditional computing approaches, such as mixture-of-experts models or sparse gating mechanisms in neural networks? Specifically, do you see any pathways for adapting MAST to support dynamic routing or conditional activation, as seen in these architectures?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. **Theoretical contribution:** The framework introduces a structured approach to sparsification in training, with rigorous theoretical analysis of convergence for both convex and non-convex cases.\n\n2. **Relevance:** This paper addresses a very relevant topic in the field.\n\n3. **Scalability:** Extending the method to a distributed setting enables applications in federated learning and distributed subnetwork training.\n\n4. **Experiments:** Experiments validate MAST’s practical benefits, showing improved performance and robustness over traditional approaches.\n\n5. **Presentation:** The paper is very well-written, with clear explanations and a consistent, clean notation.", "weaknesses": "Although the assumptions for the theoretical results are standard, they are overly restrictive and may limit real-world applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1mO96wcwhI", "reviewer": "Reviewer_j9zw", "summary": "This paper introduced a new optimization formulation called MAST that considers a modified loss function $\\mathbb{E}[f_S(x)] =\\mathbb{E}[f(v+S(x-v))] $ rather than a standard loss $f(x)$. Here, $S$ is a random matrix sampled from distribution $\\mathcal{D}$, and $v$ can be viewed as a pre-trained model. The paper claims that this problem formulation can model various practical machine learning techniques including Dropout and Sparse training with a sparse matrix $S$.\n\nMain theoretical results include comparisons between the modified loss formulation and the standard loss, as well as convergence analysis in both non-convex and (strongly) convex settings. The paper also extends convergence analysis in the distributed setup. Experimental studies are also provided to test the proposed algorithms and MAST properties.", "questions": "1. I wonder if the convergence results could be extended under weaker assumptions on $f$, perhaps by assuming that $f$ only satisfies Lojasiewicz-type conditions?\n\n2. Is it true that $f_{\\mathcal{D}}^{inf}\\geq f^{\\inf}$? I don't see that in the second line of Theorem 1, but it is needed for (13), Theorem 3 and Corollary 1.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "The presentation of the paper is good, and all the proofs are well-organized and written clearly.", "weaknesses": "1. It is unclear why a linear transformation $v+S(x-v)$ in $f_S$ is good enough to model practical ML techniques that authors consider in this paper. The reason for putting such a linear transformation rather than more complicated nonlinear transformations is lacking.\n\n2. Because of this simple linear transformation, the resulted comparison and convergence analysis are quite standard and lack novelties, so is for the distributed setup. The authors also admit that the convergence results are standard for convex and non-convex analysis. \n\n3. Rigorous mathematical connections between the model and Dropout and Sparse training are missing. Since both Dropout and Sparse training are techniques in the neural network training, it would be nice to quantitatively analyze how the model $f_S$ approximate the sparsified neural networks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L9TwWsYKGe", "reviewer": "Reviewer_Lu4H", "summary": "This work presents MAST, which optimizes the expected value of a function $f$ composed with a sketching matrix $S$. The authors analyze theoretical properties and develop a solution algorithm, demonstrating its effectiveness for sparse neural network training. The authors also provide experiments to validate their method.", "questions": "- The authors formulate MAST with a general sketch matrix $S$, but in the experiments they only consider diagonal matrices as sketching matrices. Could the authors provide examples of non-diagonal sketching matrices that could be useful in practice?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is well motivated, and the formalization of the problem is interesting. The authors provide a thorough analysis of the problem properties and derive an algorithm based on SGD to solve it. The assumptions are clearly stated and the theoretical results are well proven.", "weaknesses": "While I appreciate this research, I have a few concerns and questions regarding the contribution and novelty of the work, which I explain in detail below. I hope that these remarks are helpful in improving the work and I am happy to discuss my evaluation.\n\n- MAST bears a clear relationship to standard (stochastic) gradient descent. In supervised learning, the objective is to minimize the expected loss of a model with respect to the distribution of features and labels. Similarly, MAST aims to minimize the expected value of a function $f$ composed with a sketching matrix $S$, where the expectation is taken over the distribution of the sketching matrices. This could be made more clear in the paper.\n- While Assumption 1 appears valid since it shows $S^T\\nabla f(x)$ is an unbiased estimator of $\\nabla f(x)$, there is a discrepancy in Example 1. The authors state their sketching matrix $S$ formulation models Dropout, but to the best of my knowledge, standard Dropout applies a binary mask to activations without rescaling weights. The connection between these approaches needs clarification in their problem formulation.\n- The problem properties in section 3 appear to be correct. However, apart from Theorem 1,these are well known properties of smooth and strongly convex functions. It is not clear how novel the results in this section are, given that the proofs are rather elementary and known. \n- In Theorem 1, it is assumed that Assumption 2 and Assumption 3 hold, in particular that $f$ is strongly convex, implying there is a unique minimizer $x^*$. Why is then an optimal set of solutions $\\mathcal{X}^*$ mentioned? By Lemma 3, the same should hold true for the solutions to the MAST problem. Can the theorem be simplified or made more precise? \n- The paper is well written and the proofs appear to be correct. I find the setting interesting, especially that it allows modelling sparsification. However, it is not entirely clear to me how novel the results are, as many of the results are well known properties of smooth and strongly convex functions.\n\nI appreciate the experiments, however there are a few points that could be improved:\n- The plots in the experiments should have the same scale to allow for a comparison of the methods.\n- Did the authors compare the method to standard gradient descent with Dropout? Since MAST introduces a rescaling of the gradient and that is not present in standard gradient descent with Dropout, it is not clear how the two methods compare to each other.\n- Currently, the experiments seem to be somewhat limited. Can the authors provide more experiments, potentially at a larger scale?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NWlGzVH03q", "reviewer": "Reviewer_8W9V", "summary": "This paper presents an optimization framework that incorporates model sparsification using pre-trained models and stochastic sketches for enhanced computational efficiency. The authors propose an objective function that departs from traditional approaches, aiming to improve convergence rates in stochastic gradient descent (SGD) for sparse or resource-constrained environments. They also provide theoretical analyses on the convergence of their proposed algorithms and demonstrate potential applications for both centralized and distributed machine learning scenarios.", "questions": "1. Can you give an idea of computational and memory overheads of using sketch matrices and sparsification compared to traditional methods?\n\n2. I'm curious about how MAST performs with larger, non-convex models like deep neural networks. Testing it on a complex neural network, even briefly, could demonstrate its applicability beyond logistic regression.\n\n3. How robust is MAST to variations in hyperparameters like step size and sparsity, especially in heterogeneous distributed environments?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper certainly brings a fresh perspective by introducing a sparsification-aware objective function that enables gradient sparsification during training. This idea is both innovative and timely, addressing the growing need for more efficient model training techniques.\n\nOn the theoretical side, the authors provide a strong mathematical foundation. They successfully derive convergence bounds for various setups, including challenging non-convex scenarios and methods that reduce variance, which adds significant depth to their work.\n\nThe experimental results are compelling and back up the theoretical claims. They show that the MAST method outperforms traditional approaches, especially when dealing with high levels of sparsity. This makes the findings relevant and applicable to real-world situations.", "weaknesses": "One downside is that the experimental setup doesn't include a wide range of comparisons to other sparsification or dropout methods beyond ERM. Adding these comparisons would help better highlight where MAST excels and where it might fall short. The experiments are also somewhat limited in terms of the neural network architectures tested. The study focuses on logistic regression and ResNet-50, but there's not a strong justification for this choice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)