A Stochastic Proximal Polyak Step Size
Authors: Fabian Schaipp
Robert M. Gower
Michael Ulbrich
Venue: Accepted by TMLR ()

Abstract: Recently, the stochastic Polyak step size (SPS) has emerged as a competitive adaptive step size scheme for stochastic gradient descent. Here we develop ProxSPS, a proximal variant of SPS that can handle regularization terms. Developing a proximal variant of SPS is particularly important, since SPS requires a lower bound of the objective function to work well. When the objective function is the sum of a loss and a regularizer, available estimates of a lower bound of the sum can be loose. In contrast, ProxSPS only requires a lower bound for the loss which is often readily available. As a consequence, we show that ProxSPS is easier to tune and more stable in the presence of regularization. Furthermore for image classification tasks, ProxSPS performs as well as AdamW with little to no tuning, and results in a network with smaller weight parameters. We also provide an extensive convergence analysis for ProxSPS that includes the non-smooth, smooth, weakly convex and strongly convex setting.

Published in Transactions on Machine Learning Research (05/2023)
A Stochastic Proximal Polyak Step Size
Fabian Schaipp
Sabian..shhappp@lnng
Department of Mathematics
Technical University of Munich
Robert M. Gower
rgouer@ffuaiionnnsstttuu..o.y
Center for Computational Mathematics Flatéron Institute, New York
Michael Ulbrich
m.ulbrich@tum.de
Department of Mathematics
Technical University of Munich
Reviewed OpenReview: https: opwnrevvvee. net/ forum? id= id=jWr
Abstract
Recently, the stochastic Polyak step size (SPS) has emerged as competitive adaptive step size scheme for stochastic gradient descent. Here we develop ProxSPS, a prozimal variant of SPS that can handle regularization terms. Developing proximal variant of SPS is par- ticularly important, since SPS requires lower bound of the objective function to work well. When the objective function is the sum of a loss and a regularizer, available estimates of a lower bound of the sum can be loose. In contrast, ProxSPS only requires a lower bound for the loss which is often readily available. As a consequence, we show that ProxSPS casier to tune and more stable the presence of regularization. Furthermore for image classification tasks, ProxSPS performs as well AdamW with little to no tuning, and results in network with smaller weight parameters. We also provide extensive convergence analysis for ProxSPS that includes the non-smooth, smooth, weakly convex and strongly convex setting.
1 Introduction
Consider problems of the form
min ((x), f(x) = Ep([(r; S) (x; s))P(s)
(1)
#ER"
where S a sample space (or sample set). Formally, we can see S as a random variable mapping to and Ps)) as the associated probability measure. Let us assume that for each a E S, the function f(::)) : R² + R is locally Lipschitz and hence possesses the Clarke subdifferential 0f(:;) s) (Clarke, 1983). Problems of form (1) arise in machine learning tasks where S is the space of available data. points (Bottou et al., 2018). An efficient method for such problems is stochastic (sub)gradient descent (Robbins & Monro, 1951; Botton, 2010; Davis & Drusvyatskiy, 2019), given by
k+1
Oudk 9h E 3 where Sy P.
Sc),
(SGD)
Moreover, will also consider the composite problem
min (xxx), xxx) = f(x) + p(x), #ER*
(2)

---

Published in Transactions on Machine Learning Research (05/2023)
where op R + U is proper, closed, and convex regularization function. In practical situations, the expectation in the objective function f is typically approximated by sample average over N E N data points. We formalize this special case with
1 s {s1 , sN}. P(s.) - N f:=f(1s))) -- N.
(ER)
In this case, problem (1) becomes the empirical risk minimization problem
N
min
f(()).
1.1 Background and Contributions
Polyak step size. For minimizing a convex, possibly non-differentiable function f, Polyak (1987, Chapter 5.3) proposed
f(x²) min f E
x k+1 = x² - Org. O
19
This particular choice of ag, requiring the knowledge min f. has been subsequently called the Polyak step size for the subgradient method. Recently, Berrada et al. (2019); Loizou et al. (2021); Orvieto et al. (2022) adaptod the Polyak step size to the stochastic sctting: consider the (ER) case and assume that cach fi is differentiable and that lower bound C(s;) inf, fi(x) is known for all E [N]. The method proposed by (Loizou et al., 2021) is
x+1 = xh min {To. fig (x(x ((sin)
(x VV
(SPSmax)
with hyper-parameters 76 > and where in each iteration is is drawn from {1,..., N} uniformly at random. It is important to note that the initial work (Loizou al., 2021) used C(s.) inf fú later, Orvieto et al. (2022) established theory for (SPSmax) for the more general case of C(st) S inf, fe(x)) and allowing for mini-batching. Other works analyzed the Polyak step size in the convex, smooth setting (Hazan & Kakade, 2019) and in the convex, smooth and stochastic sctting (Prazeres & Oberman, 2021). Further, the stochastic Polyak step size is closely related to stochastic model-based proximal point (Asi & Duchi, 2019) as well as stochastic bundle methods (Paren et al., 2022).
Contribution. We propose proximal version of the stochastic Polyak step size, called ProxSPS, which explicitly handles regularization functions. Our proposal is based crucially on the fact that the stochastic Polyak step size be motivated with stochastic proximal point for a truncated linear model of the ob- jective function (we explain this in detail in Section 3.1). Our method has closed-form updates for squared E-regularizaaton. We provide theoretical guarantees for ProxSPS for any closed, proper, and convex regular- izntion function (including indicator functions for constraints). Our main results, Theorem and Theorem 8, also give new insights for SPSmax: in particular showing exact convergence for convex and non-convex settings.
Lower bounds and regularization. Methods such as SPS need to estimate lower bound C(s) for each loss function f(;; Though inf, (x;; can be precomputed in some restricted settings, in practice the lower bound C(s) 0 is used for non-negative loss functions.' The tightness of the choice C(s) is further reflected in the constant o² = min Ep(C(S)], which alfects the convergence guaranlees of SPS (Orvielo el al., 2022).
Contribution. For regularized problems (2) and if to is difforentiable, the current proposal of SPS would add yp to every loss lunction ((: In this case, for non-negative regularization terms, such as the squared fg-norm, the lower bound C(s) is always loose. Indeed, if y 2 0, then inf (f( + p(x)) 2 inf inf (x; and this incquality is strict in most practical sccnarios. For our proposed method ProxSPS.
See for instance httpa://github. com/
wtt.aaaa..coommaaaaaaaaaata Tssantaradj
2

---

Published in Transactions on Machine Learning Research (05/2023)
now need only estimate lower bound for the loss f(:; s) and not for the composite function f(z; +(((). Further, ProxSPS decouples the adaptive step size for the gradient of the loss from the regularization (we explain this in detail in Section 4.1 and Fig. 1).
Proximal and adaptive methods. The question on how to handle regularization terms has also been posed for other families of adaptive methods. For Adam (Kingma & Ba, 2015) with lz-regularization it has been observed that it generalizes worse and is harder tune than AdamW (Loshchilov & Hutter, 2019) which uses weight decay. Further, AdamW can be seen (Zhuang et al., 2022).² On the other hand, Loizou et al. (2021) showed that without regularization default hyperparameter settings for SPS give very encouraging results on matrix factorization and image classification tasks. This is promising since it suggests that SPSmax is an adaplive method, and can work well across varied tasks without the need for extensive hyperparameter tuning.
approximation to proximal version of Adam
Contribution. We show that by handling fz-regularization using a proximal stop, our resulting ProxSPS less sensitive to hyperparameter choice as compared to SPSmax. This becomes apparent in matrix factorization problems, where ProxSPS converges for much wider range of regularization parameters and learning rates, while SPSmax is more sensitive to these settings. We also show similar results for image classification over the CIFAR10 and Imagenet32 dataset when using ResNet model, where, compared to AdamM, our method is less sensitive with respect to the regularization parameter.
we will first recall how the stochastic Polyak step size, in the case of i 0, can be derived using the model-based approach of (Asi & Duchi, 2019; Davis & Drusvyatskiy, 2019) and how this is connected to SPS We then derive ProxSPS based on the connection to model-basod methods, and present thcoretical results, based on the proof techniques in (Davis & Drusvyatskiy, 2019).
The remainder of our paper is organized as follows:
2 Preliminaries
2.1 Notation
Throughout, we will write E instead of Ep. For any random variable X(),, we denote [[X(S)] = Ss (s)dP(ss). We denote (-)+ = ma. We write Ô when we drop logarithmic terms in the O-notation, egg. Ō O((mmm ((()()))))
2.2 General assumptions
Throughout the article, we assume the following:
Assumption 1. It is possible to generate infonitely many i.i.d. realizations S2, from S. Assumption 2. For every S, inf, fz: is finite and there erists C(s) satisfying C(s) s inf, f(z; s).
In many machine learning applications, non-negative loss functions are used and thus we can satisfy the second assumption choosing C(s) = 0 for all s E S.
2.3 Convex analysis
Let h Rn + R be convex and a > 0. The proximal operator is given by
1 lly 11/2.
prox (x)) == arg min min h(y) 2a
y
Further, the Moreau envelope is defined by (nv min, hly) 2u |1l xl12 and its derivative is (e) [z prox (z)) (Drusvyatskiy & Paquette, 2019, Lem. 2.1). Moreover, due to the opti- mality conditions of the proximal operator, if h C² then
at (I) |Vi(t)] a -1|x-- a|| = (envg(()))
For treating fg-regularization a part the loss can be seen be equivalent proximal version (ef. Appendix C).
3

---

Published in Transactions on Machine Learning Research (05/2023)
Davis & Drusvyatskly (2019) showed how to use the Moreau envelope measure of stationarity: if (z)]l is small, then a is close to á and & is an almost stationary point of h. Formally, the gradient of the Moreau envelope can be related to the gradient mapping (cf. (Drusvyatskiy & Paquette, 2019, Thm. 4.5) and Lemma 11).
We say that a function h R IR I-smooth if its gradient I. Lipschiit, that is
[Vh(z) - vh(yil s -|| - vll. Vz,y 1R..
(4)
If h is I-smooth, then
L 2 lly x|| for all v, E R".
Myy) s h(x) (Vhx),
A function h: R" R weakly convex for p 0 211.12 is convex. Any Lmsohhhh function weakly convex with parameter less than or equal to (Drusvyatskiy Paquette, 2019, Tem. 4.2). The above results on the proximal operator and Moreau envelope can immediately be extended h being p-weakly convex if a E (t), since then h + 511 -11 is convex.
If we assume that each s(; is Pe-Weakly convex for P. 0, then applying (Bertsekas, 1973, Lem. 2.1) to the convex function f():s) + 2
yields that f -12 is convex and thus f is p-wcakly convex for p== Elps]. In particular, f is convex if each ((;;8) is assumed to be convex. For a weakly convex function h, we denote with Ōh the regular subdifferential (cf. (Davis & Drusvyatskiy, 2019, section 2.2) and (Rockafellar & Wets, 1998, Def. 8.3)).
3 The unregularized case
For this section, consider problems of form (1), i.. no regularization term is added to the loss f.
3.1 A model-based view point
Many classical methods for solving (1) can be summarized by model-based stochastic proximal point: in
each iteration, a model (z(58) is constructed approximating f(;;)) locally around x. With Sq ~ P being drawn at random, this yields the update
1 Ily 1²112. 20k
arg min fzs Sb) + yER
(5)
The theoretical foundation for this family of methods has been established by Asi & Duchi (2019) and Davis & Drusvyatskiy (2019). They give the following three models as examples:
(i) Linear: fa(y:s) = f(x;s) + (g.y x) with g 0f(z;s).
(ii) Full: fo == s)y
(ii) Truncated: fa(y; s) max(s( + (g. y z), infser f(2;5)) where g E 0f(z;s).
It easy to see that update (5) for the linear model is equal to (SGD) while the full model results in the stochastic proximal point method. For the truncated model, (5) results in the update
x k+l ² min { ax. f(at: Sk) inf=ER" f(z: Sk) gk. }
9k E Of(a"
(6)
|19
More generally, one can replace the term infaeR" f(; Sk) with an arbitrary lower bound of f(: (cf. Lemma 10). The modcl-basod stochastic proximal point method for the truncated model is given in Algo- rithm 1. The connection between the truncated model and the method depicted in (6) is not a new insight and has been pointed out in several works (including (Asi & Duchi, 2019; Loizou et. al., 2021) and (Berrada
4

---

Published in Transactions on Machine Learning Research (05/2023)
al., 2019, Prop. 1)). For simplicity, refer to Algorithm 1 as SPS throughout this article. However, it should be pointed out that this acronym (and variations of it) have been used for stochastic Polyak-type methods in slightly different ways (Loizou et al., 2021; Gower et al., 2021).
Algorithm SPS
Require: σ0 R", step sizes ag 0. do 1. Sample S₃ and set C₁ == ((S)).
for k 0, 2,.,,.....
Choose Is E ; gk = 0, set xh+1 =xh Otherwise, set
nogh, Th min { (ore, f(z;; Ck
x x+1 x²
(7)
|1
return g
For instance consider again the SPSmax method
k+1 = x min {7:: far (xx) C(sie) (x*),
(SPS
ci (22
where c, % > O. Clearly, for c= 1 and as ". update (7) is identical to SPSmax- With this mind, we can interpret the hyperparameter 7t in SPSmax simply as a step size for the model-based stochastic proximal point step. For the parameter on the other hand, the model-based approach motivates the choice e = In this article, we will focus on this natural choice tuning. However, should point out that, in the strongly convex case, 1/2 gives the best rate of convergence in (Loizou el al., 2021).
which also reduces the amount of hyperparameter
4 The regularized case
Now we consider regularized problems of the form (2), i.e.
min (xx,, ((x)= = f(x) + p(x), FER
where : R + R (00) is proper, closed, >-strongly convex function for À 2 (we allow À 0). For S, denote by (13s) a stochastic model of the objective ub at z. We aim to analyze algorithms with the update
x k+1 arg min Wha (x; St) #ER"
|| x² 11,
(8)
2ax
where S₃ ~ P and ag > O. Naively, we know a lower bound Ö(s) of f(::)) + p()), the truncated model could be constructed for the function f(x; 9) + ((x), resulting in
(a s) max(f((;s) =p() + (g u. y z), ((s)), g 0f(z;)), E Op(x).
(9)
In fact, Asi & Duchi (2019) and Loizou et al. (2021) work in the setting of unregularized problems and hence their approaches would handle regularization in this way. What propose instead, is to only truncate lineaiization of the loss f(r; s), yielding the model
(x = (z + p(v),
fa(1)) max(f(z; + (g.v - C(s),, g E 0f(z;)).
(10)
Solving (8) with the model in (10) results in
++ arg min max (1 ( Sk) + (gk, x²),C((x)) + (vv) + 20
l1 - x112
(11)
VER=

---

Published in Transactions on Machine Learning Research (05/2023)
[2;;8) 1n(1 exp(-0.5 2)), 10.0, 0.1
ProxSPS
SPS
((:44) sp
ProxPFa model
SPS model
ProxfFe abjective
x
SPS objective
-3
2
-2
(a) Regularized logistic loss.
(b) Regularized squared loss with dg = 1, = 1.
Figure a) SPS refers to model (9) whereas ProxSPS refers (10). We plot the corresponding model (yr s) and the objective function of (8). x² (resp. a²) denotes the new iterate for ProxSPS (resp. SPS), x² is the minimizer of f(::s) p. b) Streamlines of the vector field V(=*)= == zk+1 x², for f(x) || bl1 and for the deterministic update, iee. (x; 8) [xx) ProxSPS refers to update (14) and SPS refers to (13). The circle marks the minimizer of f(x) + 211111
The resulting model-based stochastic proximal point method is given in Algorithm Lemma 12 shows that, if prox, is known, update (11) can be computed by minimizing a strongly convex function over a compact one-dimensional interval. The relation to the proximal operator of yp motivates the name ProxSPS. Further, the ProxSPS update (11) has closed form solution when yp is the squared fz-norm, as detail in the next section.
Algorithm ProxSPS
Require: Nº E R", step sizes ag > 0. do Sample Se and set Ck := ((Ss). Choose Ihe 0ff(="; ;x),
for = 0,1,2,....
Update k+1 according to (11).
return /
4.1
The special case of fz-rguularization
When ((x) = 2 || for some 0, ProxSPS (11) has closed form solution as we show next Lemma. For this lemma, recall that the proximal operator of vxx) ||] is given by prox (2) = 14ax for all a > O, x E R
Lemma 1. Let ppx) = 11/ and let E 0f(r; 8) and C(s) & infzer" ((2;8) hold for all 8 E S. For (y;s) = (v) ply) with fa( = m (x; 8) + (g.v - x), C(s)} consider the update
1
x A+1
arg min than (x; Sy) +
|| x 1112
20x
2ER"
Denote Cş == C(S)) and let ga E ; S.). Define
if 9h 0, else.
T
(1 A(g
min
akk
[ox 1
For yp 0, Algorithm 2 identical Algorithm 1.

---

Published in Transactions on Machine Learning Research (05/2023)
Update (11) is given by
b+1 1
(x*- -Tx gx g2) -prox (1 Th 9h).
(12)
+ and
See Lemma 9 in the appendix for an extended version of the above lemma and its proof. The update (12) can be naturally decomposed into two steps, one stochastic gradient step with an adaptive stepsize, that is gh+1 =x* xh x Ds followed by a proximal step z++1 prox This decoupling into two steps, makes easier to interpret the effect each slep, with T adjusting for the scale/curvature and the following proximal step shrinking the resulting parameters. There no clear separation of tasks if we apply the SPS method to the regularized problem, we next.
Algorithm ProxSPS for y = 11.
Require: xº R", step sizes ag 0. for h: 1,,...... .... K 1-11 do
1. Sample S₃ and set C₃ C(Ss).
2. Choose Ik ; Sk). If g = 0, set xk+1
I+o st
1
Otherwise, set.
2 min aki (1+ (f))(((ff" Sk) Ck) an (eAMgnnz") ghk sk).
1 1 +aad
|1
return
K
4.2 Comparing the model of SPS and ProxSPS
For simplicity, assume again the discrete sample space setting (ER) with differentiable loss functions fi and let yp = - 112. Clearly, the composite problem (2) can be transformed to an instance (1) by setting 4(x) = fz(z) + 1|111 and solving min, ((x) with f(x) Em 4(x). Assume that lower bound
& S inf, (x) is known. In this case (9) becomes
(y st) = max { (f((z) + ||1=| + (V(( + λx, y - x), {i).
Due to Lemma 10, if f,2 (x*) )+ + A * 0, the update (8) is given by
fu + A |1 li
- h+1 x min { (ou,
V (1*) + 1 ((Vfia (x*) + Az²)).
(13)
We refer to this method, which is using model (9), as SPS. On the other hand, using model (10) and if Vi (x*) o 0, the update of ProxSPS (12) is
ous { aass
(1I (x*) -Ce, Cla, -- (o ) ),} v Viic( IV [I=
x
k+1
1 Itook
x min
(14)
In Fig. la, we illustrate the two models (9) (denoted by SPS) and (10) (denoted by ProxSPS) for the logistic loss with squared fzrrouularization We can see that the ProxSPS model is a much better approximation of the (stochastic) objective function it still captures the quadratic behaviour of yp. Furthermore, noted in the previous section, ProxSPS decouples the step size of the gradient and of the shrinkage, and hence the update direction depends on Ogg- In contrast, the update direction of SPS does not depend on ag, and the regularization effect is intertwined with the adaptive step size. Another way to see that the model (10) on which ProxSPS is based on is a more accurate model as compared to the SPS model (9), is that the resulting vector field of ProxSPS takes a more direct route to the minimum, as illustrated in Fig. 1b.
Update (14) needs to compute the term fV (a*), x*) while (13) needs evaluate ||= Other than that, the computational costs are roughly identical. For (14), lower bound L is required. For non-negative loss

---

Published in Transactions on Machine Learning Research (05/2023)
functions, in practice both & and ((s)) often set to zero, in which case (10) will be more accurate model as compared to (9).
4.3 Convergence analysis
For the convergence analysis of Algorithm 2, we can work with the following assumption on y. Assumption 3. yp : + R [oo) proper, closed, X-strongly conver function with À 2 0.
Throughout this section consider model (10), i.c. for g E (z s), let
() = fa (v; + plv), (x; max (x; 8) x), C(s)).
Let us first state a lemma on important properties the truncated model:
Lemma 2. Consider fa(y; max (g. x), C(s), where g E (f is arbitrary and C(s) =
inlzeR" (2; Then, it holds:
(i) The mapping y + fo(y)s) convex.
(8) Por all x E R,, it holds f2(x;8) = [(;;)). If [(;;) is Ps weakly convea for all 8 E S, then fe(y;)) S f(u: + lyy x|1 for all a. E IR..
Proof-
(i) The maximum over constant and linear term is convex.
(ii) Recall that C(s) S f(r s) for all y RR. Therefore, (z; 8) max(C(ss, fx; sz; From weak convexity of ((;a) it follows f(x; s) + (g. x) S f(y;8) + flll z/|² and therefore
fuly; max(( f():)) 2 lly ==1 f(;;)) f lly 1y-xlt x|| for all y R
4.3.1 Globally bounded subgradients
In this section, we show that the results for stochastic model-based proximal point methods in Davis & Drusvyatskly (2019) can be immediately applied to specific model even though this model has not been explicitly analyzed in their article. This, however, requires assuming that the subgradients are bounded. Proposition 3. Let Assumption hold and assume that there is an open, convex set containing dom sp Let s()s) be Ps Pe-weakly conver for all S and let p H[ps). Assume that there exists G2 E R4+ for all S, such that = /E <0 and
[[)(;)))|| G, Vg(z;)) 0f(r;s), E.
(15)
Then, s2( satisfies (Davis & Drusvyatskiy, 2019, Assum. B), in particular it holds
fz(1;)) ( S GzIl - vll for all a E S and all x. E U.
(16)
Remark 1. We state all four properties (B1) (B4) of (Davis Drusvyatskig, 2019, Assum. B) explicitly in the Appendiz, see Proposition 14 which also contains the proof. The first three properties follow immediately in setting. Only the last properly (B4), stated in (16), requires the additional assumption (15).
Corollary 4 (Weakly convex case). Let the assumptions of Proposition 3 hold with Pe > 0 for all e S. Let Esps] < so and let 2 env 1/(2)) (a:) min v. Let tk=...... be generated by Algorithm 2 for constant step sizes an (2p+ 10G -1 Then, it holds
12 8pA + 16G
A K
E. Venv 1/(20)
2/(20
where x.K is uniformly drawn from f2
K-17.
Foor single element sampling, inf & can sometimes be precomputed (e.g. regularized logistic regression, see (Loizon al., 2021, Appendix D)). But in this restricted setting not clear how estimate inf when using mini-batchingg
8

---

Published in Transactions on Machine Learning Research (05/2023)
Proof. The claim follows from Proposition 3 and (Davis & Drusvyatskiy, 2019, Thm. 4.3), (4.16) setting 11 = 0, p= 2p, T = K - and Br = O**"
Corollary 5 ((Strongly) convex case). Let the assumptions of Proposition 3 hold with Ps 0 for all s E 8. Let A > 0 and x² = arg min, v(z).. Let J be generated by Algorithm 2 for step sizes ag = xTK+TJ
Then, it holds
A Σ (k ( 1x++)) -b(0)) (K + 1)2 100 =x"1l + X(K +1)
8G²
E 10
+2)
Proof- As Ps 0 and hence p = 0, have that (Davis & Drusvyatskly, 2019, Assum. B) is satisfied with
= 0 (in the notation of (Davis & Drusvyatskiy, 2019), see Proposition 14). Moreover, by Lemma 2, (i) and A-strong convexity of i, we have X-strong convexity of (=(ss)). The claim follows from Proposition 3 and (Davis & Drusvyatskiy, 2019, Thm. 4.5) setting X, T K and B =
4.3.2 Lipschitz smoothness
Assumption (15), i.e. having globally bounded subgradients, is strong: implies Lipschitz continuity of f (cf. (Davis & Drusvyatskiy, 2019, Lem. 4.1)) and simple functions such the squared loss do not satisfy this. Therefore, we provide additional guarantees for the smooth case, without the assumption of globally bounded gradients.
The following result, similar to (Davis & Drusvyatskiy, 2019, Lem. 4.2), is the basic inequality for the subsequent convergence analysis.
Lemma 6. Let Assumption 3 hold. Let xh41 be given by (11) and that be given in (10). For every x E R il holds
(1+ + (+ - x212 S 11. xz12 - ||a - 10112 + 200 (W2r Sx)) - th (a ;
(17)
Moreover, it holds
Viat k11
; Sk) 2 fx* Sk) + (gk. n+l 2² + p(aht1).
(18)
Proof. The objective of (11) is given by (v) == that (yr S&) + 20r by xz 112. Using Lemma 2, (i) and A-strong convexity of sp, (x) is (A+ as
L) strongly convex. As x++ is the minimizer of V (y),,, for all e R2 we have
(x) 2 U/ 2x (1) + + agd || - x1|2
20g
(1 + agd))lla 21|² < || x²|² ||a 111 + 2ax (12 (2; Sx) vh (a*+1: :
Moreover, by definition of (y: s) in (10) we have
(2 = fat (z (1 ;t) + p(a 2 ; + (ge, (ge, xt1 sty + ((xx
We will work in the setting of differentiable loss functions with bounded gradient noise.
Assumption The mapping ((;8) is differentiable for all 8 E S and there exists 3 2 such that
B [|)) SB for all x E R."
(19)
The assumption of bounded gradient noise (19) (in the differentiable sotting) is indeed weaker assumption than (15) since [[V](=; S)] v=))) and
E [)] 2 + [|; |V)())1² + B.
9

---

Published in Transactions on Machine Learning Research (05/2023)
Remark 2. Assumption (and the subsequent theorems) could be adopted to the case where f(;s) is weakly connex bul. non-differentiaall: for fined x E IR", due to (Bertsekas, 1973, Prop. 2.2) and (Davis B Drusvyatskiy, 2019, Lem. 2.1) it holds
Blof(z; S) = =E =(fo(f( + PS |1=11²) ps*| =f((x) + pr -Epsaa] 0f(z),
2 where we used p E[pps] Hence, for 9s 0f({;)) 8) have Elgs] E 0z(z) and (19) is replaced by
Ellgs - Eigslll 8 for all x E IR".
However, as will still require that f is Lipschitz-smooth, present our results for the differentiable setting.
The proof of the subsequent Cheorems can be found in Appendix A.2 and Appendix A.3.
Theorem 7. Let Assumption 3 and Assumption hold. Let f(;;s) be conver for all E S and let f be L-smooth (4). Let x² = arg min (xx) and let 0 > 1. Let fø Ik=o....K K be generated by Algorithm 2 for step sizes og > O such that
0k S 1 1/0
(20)
L
Then, holds
(1 + a1 - xx11| S Ella* 122 + 2cr 2(x(((((*)) ((+++11] + 0.a
(21)
Moreover, we have:
a) If > and an = t(k+ko) with ko 2 large enough such that (20) fulfilled, then
E
Σ
2b41) vb(x*)]
A60
1000
x* x=1
03(1 + In K)
(22)
2K
2AK
k=0
b) If X = 0 and ag
with
1-1/0 b
then
vK+1
K
||00 12
83a + + In) K)
R
apolt1)
- xb(*))
1 4a VK +I
(23)
+I-+
km
k=0
c) f is p-strongly convex with μ 2 0,5 and 0g = a fulfilling (20), then
03a
B| x111 (1 alp + 2\)) 1|200 - x=1 +
(24)
#1 2X
Remark 3. 1f X O, for the decaying step sizes in item a) we get rate of Õ x In the strongly conver case in ilem c)2 for conslant slep sizes, we gel a linear convergence uplo a neighborhood of the solution. Note that the constant on the right-hand side of (24) can forced to be small using a small a. Further, the rate (24) has a 2A term, instead of X. This slight improvement the rate occurs because we do not linearize yp in the ProuSPS model.
Theorem 8. Let Assumplion 3 and Assumplion hold. Let f(;8) be ps weakly conver for all and let p Eips] so. Let f be L-smooth and assume that inf ib > -0. Let Algorithm 2. For 0 > 1, under the condition
be generoted by
(0. (0)
Tp>A
1- 8-1 0g <
" E
(25)
else
I, + 17
Note that f(:s) convex, so J, and that allow here. 6 As p-weally convex, this implies p I.
10

---

Published in Transactions on Machine Learning Research (05/2023)
it holds
K-
2(env (x²) - inf v6)
30
K
agE (
+
0
(26)
1 - -n1(p- 2)
111
70
2))
k=0
Moreover, for the choice on = vk+1
a
and with a S L+m-T, we have
1-0-1
env (x0) inf
300
a(1 + + K)
min
E| Venv (2 )| a(1 - 71p x))(vK
2n(1 n(p ))) (VK +I-1)
If instead we choose org VK
and with S VK
1-0-1 Lo we have
2( (2²) - -inf 1b)
30 a
E
Venv (a )1
a(1 7(p - A)) +
171 7/pp A)) VK*
where a.K is uniformly drawn, from {ao, , xK-
4.3.3 Comparison to existing theory
Recalling that Algorithm is equivalent to SPS with c and % = Ok, we can apply Theorem 7 and Theorem for the unregularized case where yp = and hence obtain new theory for (SPSmax). We start by summarizing the main theoretical results for SPSmax given in (Loizou et al., 2021; Orvieto al., 2022): in the (ER) sctting, recall the interpolation constant 0² = [[f(x*:S) - C(S) = aें Σin fi(x)) C(s)) If fo is L--smooth and convex, (Orvieto et al., 2022, Thm. 3.1) proves convergence to a neighborhood of the solution, i.c. the iterates {r*) of SPSmax satisfy
1 12
2%02 a
aK +
(27)
, a == min booLeoe To and Lmax == maxse[N] L For the noncouvex case, if f is L.-smooth and under suitable assumptions on the gradient noise, (Loizou et al., 2021, Thm. 3.8) states that, for constants C1 and c we have
where jk == R Ek-1 wh,
1 ej K +
min BI Am L.ll... K
(28)
The main advantage of these results that 76 can be held constant; furthermore the convex setting (27), the choice of 7h requires no knowledge of the smoothness constants L,, For both results however, we can not directly conclude that the right-hand side goes to zero K + there is an additional constant. Choosing sufficiently small does not immediately solve this as c1. and 02 all go to zero as % goes to zero.
Our results complement this by showing exact convergence for the (weakly) convex case, i.e. without constants on the right-hand side. This comes at the cost of an upper bound on the step sizes ag which depends on the smoothness constant L. For exact convergence, it important to use decrcasing step sizes ag: Theorem 8 shows that the gradient of the Moreau envelope converges to zero at the rate O(1/VK) /K) for the choice of 0g = VR" Another minor difference to (Loizou et al., 2021) is that we do not need to assume Lipschitz- smoothness for all f(::s) and work instead with the (more general) assumption of weak convexity. However, we still need to assume Lipschitz smoothness of f-
Another variant of SPSmax; named DecSPS, has been proposed in (Orvieto et al., 2022): for unregularized problems (1) it is given by
( [(x Sx) Sx) Cn || Ck ck-1in--)
k+1 2 nig.. tn min
(DecSPS)
Ck
7The theorem also handles the mini-batch case but, for simplicity, we state the result for sampling single is in each itcration.
Notice that og then depends on the total number of iterations K and hence one would need to fix before starting the methor.
11

---

Published in Transactions on Machine Learning Research (05/2023)
where {co is an increasing sequence. In the (ER) setting, if all f, are Lipschttz-smooth and strongly convex, DecSPS converges with a rate of (( without knowledge of the smoothness convexity constants (cf. (Orvicto al., 2022, Thm. 5.5)). However, under these assumptions, the objective f is strongly convex and the optimal rate ()), which we achieve up to a logarithmic factor in Theorem 7, (22). Moreover, for DecSPS guarantees given for nonconvex problems.
For regularized problems, the constant in (27) is problematic if o² (computed for the regularized loss) is moderately large. We refer to Appendix D.5 where we show that this can easily happen. For ProxSPS, our theoretical rosults Theorem and Theorem 8 are not affected by this as they do not depend on the size of o². To the best of our knowledge, this is the first work to show theory for the stochastic Polyak step size setting that explicitly considers regularization. Moreover, our results also cover the case of non-smooth or non-real-valued regularization y where the theory in (Loizou et al., 2021) can not be applied.
5
Numerical experiments
Throughout we denote Algorithm 1 with SPS and Algorithm 3 with ProxSPS. For all experiments use PyTorch (Paszke et al., 2019)9.
5.1
General parameter setting
For SPS and ProxSPS we always use C(s) 0 for all E S. For on, use the following schedules:
constant: set og = ox for all &s and some og > 0.
sqrt: set Ork Vi for all iterations k during epoch j.
As we consider problems with ferreuularizatinn, for SPS we handle the regularization term by incorporating
it into all individual loss functions, as depicted in (13). With i à 112 for X 2 O, we denote by Ça the adaptive step size term of the following algorithms:
for SPS we have &A [(a* b1a*
l1o. + (cf. (13) with Liz =0),
for ProxSPS we have ÇE = (I+a)A)f(2* (s ( aan-annnnn 21) with C(S²) 0).
) and thus o minfar, (t (cf. Lemma
5.2
Regularized matrix factorization
Problem description: For A E Ri*p, consider the problem
min
E (W22 Ayl?=
min Wi W.
[2WWi - AlE
WaRer WI Woe
For the above problem, SPSmax has shown superior performance than other methods in the numerical experi- ments of (Loizou et al., 2021). The problem can can be turned into (nonconvex) empirical risk minimization problem by drawing N samples (y( (,))). Denote b(t) = (, Adding squared norm regularization with A 2 0 (cf. (Srebro et al,, 2004)), we obtain the problem
N
min WI WRIX
W2 p01 + (l 13 + 11....).
(29)
N
This fits the format of (2), where x (Wi, W2), using a finite sample space S = { 88). (x; (| Ay ProxSPS for problems of form (29) on synthetic data. For details on the experimental procedure, we refer
0y -y and x 1| 12. Clearly zero a lower bound of f(; for all i E [N]. We investigate
9The code for experiments and implementation PrexSPS can be found at pestpa:/ggithin htt.geiht....fffaiieeeeeeeeeeeee con//aabiiannee///rosst
12

---

Published in Transactions on Machine Learning Research (05/2023)
10
constant
10-1 sqrt
1322
30
180
0.7
0.50
125
0.81
9.8
0.53
0.4
0.6
0.12
10
10-2
10
10.0
10-1
10-4
10"5
10-5
0
10
20
30
-0
10
20
30
Epoch
Epoch
Figure 2: Objective function for the Matrix Factorization problem (29), with constant (left) and sqrt (right) step size schedule and several choices of initial values. Here ming (x**) is the best objective function value found over all methods and all iterations.
to Appendix D.1.
Discussion: We discuss the results for the setting matrix-faci in Table 1 in the Appendix. We first fix A = 0.001 and consider the three methods SPS, ProxSPS and SGD. Fig. 2 shows the objective function over 50 epochs, for both step size schedules sqrt and constant, and several initial values og- For the constant schedule, we observe that ProxSPS converges quickly for all initial values while SPS unstable. Note that for SGD we need to pick much smaller values for og in order to avoid divergence (SGD diverges for large oo). SPS for large ag is unstable, while for small ag we can expect similar performance to SGD (as 74 is capped by O ag). However, in the regime of small ow, convergence will be very slow. Hence, one of the main advantages of SPS, namely that its step size can be chosen constant and moderately large (compared to SGD), is not observed here. ProxSPS fixes this by admitting larger range of initial step sizes, all of which result in fast convergence, and therefore is more robust than SGD and SPS with respect to the tuming of an-
For the sqrt schedule, we observe in Fig. 2 that SPS can be stabilized by reducing the values of a over the course of the itcrations. However, for large ag we still see instability in the early iterations, whereas ProxSPS does not show this behaviour. We again observe that ProxSPS is less sensitive with respect to the choice of og compared to SGD. The empirical results also confirm our theoretical statement, showing exact convergence if og is decaying in the order of 1/Vk. From Fig. 3, we can make similar observations for the validation error, defined as Nowi Ei [|W:W,,/ the validation set (cf. Appendix D.1 for details).
ball? 2|, where (1() are the Nval = N measurements from
We now consider different values for À and only consider the sqrt schedule, we have seen that for constant step sizes, SPS would not work for large step sizes and be almost identical to SGD for small step sizes. Fig. 4 shows the objective function and validation error. Again, we can observe that SPS is unstable for large initial values On for all À 2 10-4 On the other hand, ProxSPS has a good performance for wide range of ag E [1] if À is not too large. Indeed, ProxSPS convergence only starts to deteriorate when both ag and À are very large. For ag 1, the two methods give almost identical results. Finally, in Fig. 5a we plot the validation error as function of À (taking the median over the last ten epochs). The plot shows that the best validation error is obtained for X = 10-4 and for large ag. With SPS the validation error is higher, in particular for large ag and X. Fig. 5b shows that ProxSPS leads to smaller norm of the iterates, hence a more effective regularization. Finally, we plot the actual step sizes for both methods in Fig. 6. We observe that the adaptive step size ÇA (Definition at end of Section 5.1) is lypically larger and has more variance for SPS than ProxSPS, in particular for large A. This increased variance might explain why SPS is unstable when og is large: the actual step size is the minimum between og and &o and hence both terms being large could lead to instability. On the other hand, if on = 1, the plot confirms that SPS and ProxSPS are almost identical methods as & OA for most iterations.
13

---

Published in Transactions on Machine Learning Research (05/2023)
coustant
sqrr
30
1322
180
0.5
10
10"
1225
0.81
0.50
0.
0.12
Volidat'on Eryor
10
10
10:2
10"7
10-3
10-1
10-4
10
20
30
40
10
20
30
Epoch
Epoch
Figure 3: Validation error for the Matrix Factorization problem (29), with constant (left) and sqrt (right) step size schedule and several choices of initial values.
EX001
A 2.01
1006
Chapetive
0.30
0117-
20.
0.394
1.5-
6000
1000
0.000
20
Ppcch
Rgooh
Ppoch
Rgoch
0.0025
03:002
Vailein Fate
3.0000
0.300
2.010
2.068
11:00
3.0010
0.7
0.301
2:00
0.25
0.000
2.000
1031
Ypoch
Rooch
Pioch
Epoch
Figure 4: Objective function value and validation error over the course of optimization. For the validation error, we plot rolling median over five epochs in order to avoid clutter.
We provide additional numerical results which confirm the above findings the Appendix: this includes the results for the setting matrix-fac2 of Table Appendix D.2 as well as matrix completion task on a real-world dataset of air quality sensor networks Aivvera-Muñoo al., 2022) in Appendix D.3.
5.3 Deep networks for image classification
We train ResNet56 and ResNet110 model (He et al., 2016) the CIFAR10 dataset. We use the data loading and preprocessing procedure and network implementation from https://githut cww.coo//aaaaaaa ham/akamaster/ pytorch_ resnet__tiie cifaree. We do not use batch normalization. The loss function is the cross-cntropy loss of the true image class with respect to the predicted class probabilities, being the output of the ReaNet56 network. We add || as regularization term, where z consists all learnable parameters of the model. The CIFAR10 datasct consists of 60,000 images, each of size 32 32, from ten different classes. We use the PyTorch split into 50,000 training and 10,000 test examples and use a batch size of 128. For AdamW, set the weight decay parameter to X and set all other hyperparameters to its default. We the
14

---

Published in Transactions on Machine Learning Research (05/2023)
10P
5.50
pou un up opell apt apet
1ps ret medl
72 danto
10"
5.25
5.00
10"3
Valdation
4.75
10
4.50
10-1
4.25
4.00
10
10-1
10-1
10-1
10
10
(a) Validation error
(b) Model norm VI + |Wa|l2
Figure 5: Validation error and model norm as a function of the regularization parameter x. Shaded area is one standard deviation (computed over ten independent runs). For all values, we take the median over epochs [40, 50].
proc-ipe,oo 3AA Te-
2RD-- 3J A Te----
procepa by BBC, All Ie-
124, no 52, A 16-41
2 3EAA A Ie
117
111
1000 JA A 00301
Jne-EX 50 a 0.001
(IOe-EE, 10. 3 C,001
L. = 00001
in 5. 3. 20001
13% IBA A 0.000
10P
100
persepp 14, ) 3.0
(1/s-4pe 51 > 6m01
160 A n811
0p. 1A A nell
7.0 A 0.00
1035) 6.00
10
10
10"
10. A- 0.3
50, 0.
10.0 >- Bo1
16, 1 300
801
160,
107
10.
L. 1-
1011 A
IM A
Ml
10
10
Epoch
Epoch
Epocss
Epoch
Epoch
Epoch
Figure 6: Adaptive step size selection for SPS and ProxSPS. We plot ÇA (see definition in Section 5.1) as dots for each iteration as well as their median over each epoch. For this plot, we use the results of only one of the runs.
Adan-iinplementation from https://github. com/zhenxun http://uiaaaaaaa/dddddaaaaaaaae- contrast to the Pytorch implementation multiply the weight decay paramcter with the learning rate, which leads to better comparability to SPS and ProxSPS for identical values of X. For SPS and ProxSPS we use the sqrt-schedule and og = 1. We run each method repeatedly using (the same) three different seeds for the dataset shuffling.
as it does not in
Discussion: For Resnet56, from the bottom plot in Fig. 7, observe that both SPS and ProxSPS work well
15

---

Published in Transactions on Machine Learning Research (05/2023)
a1m12
adtame, outaltal, 1p-1.00
Se
Se
0,008
1.0
1.0
1.0
Valication Accurwy
0.5
0.6
0.66
04
0.4
0.4
0.2
0.0
0.0
50
70
100
25
50
100
25
76
100
Epoch
Epoch
Epoch
120
70
80
10
100
Epoch
Epoch
Epoch
Figure 7: ResNet56: (Top): Validation accuracy and model norm for three values of the regularization parameter A. Validation accuracy is defined as the ratio of correctly labeled images on the validation set (i.e. Top-1 accuracy), plotted as five-epoch running median. (Bottom): With ||ø*]] we denote the norm of all learnable parameters at the k-th iteration. Shaded area is two standard deviations over three independent runs.
with ProxSPS leading to smaller weights. For À - 5e the progress of ProxSPS stagnates after roughly 25 epochs. This can be explained by looking the adaptive step size term S& in Fig. 9a: as it decays over time we have - G & Since every iteration of ProxSPS shrinks the weights by factor 1+axx this leads to bias towards zero. This suggests that we should choose ag roughly of the order of Çs. for example by using the values of &a from the previous epoch.
For the larger model Resmet110 however, SPS does not make progress for a long time because the adaptive step size is very small (see Fig. and Fig. 9b). ProxSPS does share this issue and performs well after a few initial epochs. For larger values of X, the training is also considerably faster than for AdamW. Generally, observe that ProxSPS (and SPS for Resnet56) performs well in comparison to AdamW. This is achieved without extensive hyperparameter tuning (in particular this suggests that sctting c = 1 in SPS leads to good results and reduces tuning effort).
Furthermore, we trained ResNet110 with batch norm on the Imagenet32 dataset. The plots and exper- imental details can be found in Appendix D.4. From Fig. 15, we conclude that SPS and ProxSPS perform equally well in this experiment. Both SPS and ProxSPS are less sensititve with respect the regularization parameter À than AdamW and the adaptive step size leads to faster learning in the initial epochs compared to SGD. We remark that with batch norm, the effect of E-reguularizatin is still unclear the output of batch norm layers is invariant to scaling and regularization becomes ineffective (Zhang et al., 2019).
6 Conclusion
We proposed and analyzed ProxSPS, proximal version of the stochastic Polyak step size. We arrived at ProxSPS by using the framework of stochastic model-based proximal point methods. We then used this framework to argue that the resulting model of ProxSPS is better approximation compared to the model
16

---

Published in Transactions on Machine Learning Research (05/2023)
a1m12
aitame, outaltal, 1p-1.00
Se 1e--00
0,003
1.0
1.0
1.0
Valicaaton Accuray
0.8
0.5
0.6
0.66
04
0.4
0.4
0.2
0.0
0.0
0.0
50
70
100
25
50
100
25
76
100
Epoch
Epoch
Epoch
100
80-
80
100
60
411
211
50
100
Epoch
Epoch
Epoch
Figure 8: ResNet110: Validation accuracy as five-epoch running median (top) and model norm (boltom) for three values of A. Shaded area is two standard deviations over three independent runs.
used by SPS when using regularization. Our thoorctical results cover a wide range optimization problems, including convex and nonconvex settings. We performed series of experiments comparing ProxSPS, SPS, SGD and AdamW when using fy-regularization. In particular, we find that SPS be very hard to tune when using ég-regularization, and in contrast, ProxSPS performs well for a wide choice of step sizes and regularization parameters. Finally, for our experiments on image classification, we find that ProxSPS is competitive to Adam,, whereas SPS can fail for larger models. At the same time ProxSPS produces smaller weights in the trained neural network. Having small weights may help reduce the memory footprint of the resulting network, and even suggests which weights can be pruned.
Acknowledgments
We thank the Simons Foundation for hosting Fabian Schaipp at the Flatiron Institute. We also thank the TUM Graduate Center for their financial support for the visit.
References
Hilal Asi and John C. Duchi. Stochastic (approximate) proximal point methods: convergence, optimality, and adaptivity. SIAM Journal on Optimization, 29(3):2257 00-2222779900 2290, 2019. ISSN 1052-6234. doi: 10.1137/
18M1230323.
Amir Beck. First-order methods in optimization, volume 25 of MOS-STAM Series on Optimization. Society for Industrial and Applied Mathematics (SIAM), Philadelphia, PA; Mathematical Optimization Society, Philadelphia, PA, 2017. ISBN 978-1-611974-98-0. doi: 10. 11.11871197816111974999
Leonard Berrada, Andrew Zisserman, and M. Pawan Kumar. Training neural networks for and by interpo- lation. June 2019.
Dimitri P. Bertsckas. Stochastic optimization problems with nondifferentiable cost functionals. Journal of Optimization Theory and Applications, 12:218 231, 1973. ISSN 0022-3239. doi: 10.1007 N..107//FF0034881
17

---

Published in Transactions on Machine Learning Research (05/2023)
propeggs op 1.0, 5e
*p6, <ep
1.0,
gro0cepps up 1.0, A= 06
1ghn, (40 1.0,
10a
modianiiaa
Step size
Step size:
10-2
10-1
10-6
10-3
10
proo sps,
10,
Sc
05
sps, 1.0, Sc
05
210X
10,
5c
05
596, 1.0,
50 median(Cu)
05
10a
10a
median(Cu)
(tadioniSa)
medien(Ca)
Step size:
10A
Step size
10-2
10-1
10-6
10-3
10-10
proe-ape, 000005
4p4,
0.0005
pnoc-ppe,
ape, 0.0005
median(Cc)
medianica
modivatte)
Step size
10a
Step size
10-2
10-1
10
10
10
100
100
30
Epoch
Epoch
Epoch
Epoch
(a) ResNet56
(b) ResNett10
Figure 9: Adaptive step sizes for SPS and ProxSPS. See definition of & in Section 5.1. For this plot, we use the results of only one of the three runs.
Léon Bottou. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMP- STAT'2010, pp. 177-186. Coyssca-Verlag/Springer Heidelberg, 2010.
Lóon Bottou, Frank E. Curtis, and Jorge Nocodal. Optimization methods for largo-scale machine learning.
SIAM Review, 60(2):223 311, 2018. ISSN 0036-1445. doi: 10.1137/16M10S00
Frank H. Clarke. Optimization and nonsmooth analysis. Canadian Mathematical Society Series of Mono- graphs and Advanced Texts. John Wiley & Sons, Inc., New York, 1983. ISBN 0-471-87504-X. A Wiley- Interscience Publication.
Damek Davis and Dmitriy Drusvyatskiy. Stochastic model-based minimization of weakly convex functions.
SIAM Journal on Optimézation, (9()):207 229,, 2019. ISSN 1052-6234. doi: 10.1137/18M11182444
Dmitriy Drusvyatskiy and Courtney Paquette. Efficiency of minimizing compositions of convex functions and smooth maps. Mathematical Programming, 178(1-2, Ser. A):503-555 558, 2019. ISSN 0025-5610. doi:
10.1007/s10100-018--
Robert Gower, Othmane Sebbouh, and Nicolas Loizou. SGD for structured nonconvex functions: Learning rates, minibatching and interpolation. Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of
The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 1315-1323. PMLR, 13 15 Apr 2021. URL https://proceeddnggs mlr. preas/v130/ge gwwer2aa. html.
Elad Hazan and Sham Kakade. Revisiting the Polyak step size, May 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016. doi: N0.119//CVPR.20166
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
18

---

Published in Transactions on Machine Learning Research (05/2023)
Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic Polyak step- size for SGD: An adaptive learning rate for fast. convergence. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference on Artificial Intellignnce and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 1306-1314. PMLR, 13-15 Apr 2021. URL https: html.
// //proceedings. www.gi............ wtt.ppesssvv//1iiiieesp
Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. httpp////ppaaereeeeeeee///eeeeee hatt/forummiiddBggfR...
Antonio Orvicto, Simon Lacoste-Julien, and Nicolas Loizon. Dynamics of SGD with stochastic Polyak stepsizes: Truly adaptive variants and convergence to exact solution. May 2022.
Alasdair Paren, Leonard Berrada, Rudra P. K. Poudel, and M. Pawan Kumar. A stochastic bundle method for interpolation. Journal of Machine Learning Research, 23(15):1 57, 2022. URL http://jmlr.org org/ 2apers/v23/20-1111 html.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019. URL http://papers: httpsssscccsppcppcpppppp//////// perpppatteeeesseeeeeeeeeerrrreee
pdf.
Boris T. Polyak. Introduction to optimization. Translations Series in Mathematics and Engineering. Opti- mization Software, Inc., Publications Division, New York, 1987. ISBN 0-911575-14-6. Translated from the Russian, With a foreword by Dimitri P. Bertsekas.
Mariana Prazeres and Adam M. Oberman. Stochastic gradient descent with Polyak's learning rate. Journal of Scientific Computing, 89(1):Paper No. 25, 16, 2021. ISSN 0885-7474. doi: 10.1007//110015500--0062233
L.M. Rivera-Muñooz A.F. Giraldo-Forero, and J.D. Martinez-Vargas Deep matrix factorization models for estimation of missing data in a low-cost sensor network to measure air quality. Ecological Informatics, 71: 101775, 2022. ISSN 1574-9541. doi: https://bbi.ogg///100000...oo
w0t../000000000o.22200000.... URL https://www.
Herbert Robbins and Sutton Monro. stochastic approximatiin method. Ann. Math. Statistics, 22:400-40 407, 1951. ISSN 0003-4851. doi: 10.1214 1aa1///ammm11777955555
R. Tyrrell Rockafellar and Roger J.-B. Wets. Variational analysis, volume 317 of Grundlehren der math- ematischen Wissenschaften (Fundamental Principles of Mathematical Sciences]. Springer-Verlag, Berlin, 1998. ISBN 3-540-62772-3. doi: 10.1007/978-33-422-22233
Nathan Srebro, Jason Rennie, and Tommi Jaakkola. L. Saul, Y. Weiss, and I. Bottou (eds.), Advances in Newral Information Processing Sys- tems, volume 17. MIT Press, 2004. URL https: http://poooeeeig.nneerree swtrripssccpappe/20//////// 3ew.339951111011110 4004 .df.
Maximum- matrix factorization. In
Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger B. Grossc. Three mechanisms of weight decay regu- larization. In 7th International Conference on Learning Represenlations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview net, 2019. URL https:
Zhenxun Zhuang, Mingrui Liu, Ashok Cutkosky, and Francesco Orabona. Understanding AdamW through proximal methods and ceale-freeness. Transactions on Machine Learning Research, 2022. URL https: / //openreview. wewiiiiwwwweeeeeeee www.teernurntteeoe terffrommmddiihhaati IKIEPEWGG
19

---

Published in Transactions on Machine Learning Research (05/2023)
Contents
1
Introduction
1
1.1
Background and Contributions
2
2
Preliminaries
3
2.1
Notation
3
2.2
General assumptions
2.3
Convex analysis
3
The unregularized case
3.1
A model-based view point
4
4
The regularized case
5
4.1
The special case of Egrregularization
6
4.2
Comparing the model of SPS and ProxSPS
4.3
Convergence analysis
4.3.1
Globally bounded subgradients
8
4.3.2
Lipschitz smoothness
9
4.3.3
Comparison to existing theory
11
5
Numerical experiments
12
5.1
General parameter setting
12
5.2
Regularized matrix factorization
5.3 Deep networks for image classification
14
6
Conclusion
16
A Missing Proofs
21
A.1 Proofs of model-based update formula
................................
21
A.2 Proof of Theorem 7
23
A.3 Proof of Theorem 8
24
B Auxiliary Lemmas
26
C Model equivalence for SGD
27
D Additional information on numerical experiments
28
D.1 Matrix Factorization
28
D.2 Plots for matrix-fac2
28
D.3 Matrix completion experiment
30
20

---

Published in Transactions on Machine Learning Research (05/2023)
D.4 Imagenet32 experiment
30
D.5 Interpolation constant
30
A
Missing Proofs
A.1 Proofs of model-based update formula
Lemma 9. For À 2 0, let pxx)) 2 1|1= and let E (f(;;) and C(s) S infzer" f(2;s) hold for all s E S.
For
(e(3is) = sy) + p()), fo(r;)) -
consider the update
1 _ || x* 112
*+1 arg min a (a; Sk) +
(30)
zE8"
20rk
Denote C₃ = C(Sk) and let gh 0f(=;; Sp). Define
9x 0, else.
Tk
(1tanA)((a"))))) C.) and Ngn
min
Olk::
11
Then, we have
k-
+ ag k
1+ax1
T
9h
1 +tad asd
(x* (a=-T1910 9x)))) prox aoo (xt T gg).
(31)
Define To = is gs and
and Tk = min { {ax,
else. Then, it holds Th S Tg
11xa
Va ( [(*
and I+axt (*))
T 1191112 sp(oht1).
(32)
Proof. Note that max max (f( (k) + (9k, x²), Ck} is convex as function of The update is therefore
unique. First, if It = O, then clearly xb+1 = prox proxaup (x = T+ and (32) holds true. Now, let # 0. (gh. y - 2²) Cs), in {ylf(** Sa) (gx. y - 2t) Ok3 (ok. y =²) Cx). We therefore solve three problems:
The solution of (30) is either in {y fa in (a* Sx)
(P1) Solve
A ||.112 +
1
y+ arg min Ct +
11 111
2
20y
Clearly, the solution is yt Ita * This yt solves (30) if 1(x Sx) + (gk, y a²) < Ck.
(P2) Solve
y+ arg min f(r ; Sk) ) + (gk, y 1²) + |lvl2 - 11 x 112
201
The optimality condition is 01 + +y+ - Thus, the solution is yt 1 (=*t-ogg))) This yt solves (30) if f(**; Sk) + (gk. g - x²) > Ct.
(P3) Solve
y- = arg min X
1lyl + 1
lly x* 11, s.t. ; + (ge, x² = Cr.
2 2ac
21

---

Published in Transactions on Machine Learning Research (05/2023)
The KKT conditions are given by
ag + - -x + Hgg = 0,
; + (gk-y y x²) Ck-
Taking the inner product of the first equation with gk: get
(1 + (gk 8) - (gk, + 1119x:11 0.
From the second KKT condition we have (gx, v) Ck f(x²; Sk) + (gx, x*), hence
(1 + [++a))((C+ - f(x;; + (gk, (9k. x²)) + pllga 1 =0
Solving for l gives μ = (I+a)A)(f(=,ss) (8) --)))))C))-- an
From the first KKT condition, we obtain
119211
y² 1+ agp (xt pgh) = 1 and (st (1+ax))(f(a" Sx) - Ck) (1))-app ax 9k).
|1
This y* solves (30) if neither (P1) nor (P2) provided a solution.
For all three cases, the solution takes the form y+ = Too 1at - tgx] =: v(t). As |1 > 0, the term
;( + (gk. u(t)-x*) is strictly monotonically decreasing in t. We know f(z"; S&) + (gk. y(t) xx) = C2
for = μ (from (P3)). Hence, ; + (ge, v(t) x²) < Ck (> Ck) if and only if t μ (t <)).
We conclude:
If f(a"; St)) + (ok, y(0) - 2²) Ck2 then the solution to (P1) is the solution to (30). This condition is equivalent to / < 0.
If ;(=*;Sa)+* + (gk. v(ax)) 2²) Cb, then the solution to (P2) is the solution (30). This condition is equivalent to a& < μ..
If neither a nor as < p hold, i.e if μ E [0, ag), then the solution to (30) comes from (P3) and hence given by v(u).
Altogether, we get that x++1 ItaxA 1xth - T gu] with of = min(ax, (x)+}.
Now, we prove (32). Note that if Sk & 0, then Th min(au, u} with a defined as in (P3). In the case of (P1), we have Vu (x Si)) = C& + ((xttl). Moreover, it holds µ < 0 and as ag. > we have Tk = μ. Plugging Tk = y into the right hand-side of (32), we obtain Ck + p(ak+1).
In the case of (P2) or (P3), we have Cx S f(a"; + (gk, xb+1 r²).. Due fla"; + (gk, y(t) - x²t) 1 Sx) ItoEX (gk, x*) + T= 11 we obtain (32) as b+1 y(ax) and μ > on in the case of (P2) and 211 = y(4) and μ I 0g in the case of (P3).
Lemma 10. Consider the model fa s) == max{f + (g. y - x), C(s)} where E 0f(z;)) 8) and C(s) < inf:ek" s() holds for all s E S. Then, update (5) is given as
Th = [o
is 9h = 0, else.
1+l = x² Tille
min fas, f(ah C(Sx)
cep
where gk E ; Moreover, it holds
La (ottt (t) max (ax(C((s),,((f;.55 on 11
(33)
and therefore fax (ak+1, St) f(at; - 11.1191112.
Proof. We apply Lemma with x == 0. As slat; S&) C(Sk), we have that - Tk 7k-
22

---

Published in Transactions on Machine Learning Research (05/2023)
A.2 Proof of Theorem 7
From now on, denote with Fk the filtration that is generated by the history of all Sy for j j= -0...,,k 1.
Proof of Theorem 7. In the proof, we will denote Sa V ; We apply Lemma 6, (17) with =x
Due to Lemma 2 (ii) and convexity of f(rss) it holds
Va Sk) S f(=;; + +((x))
Together with (18), we have
(1 + a+)ll//t xx111 S 11 x²11² - || +1 - x 1112 + 2axlp(x") (p(
(34)
+ 200 (sx"; Ska f(x"; Sp)) (gk, x - x 2²1|_
Smoothness of f yields
-s(a") S - + (V
) kt1 - s/y + -11bbt1-- - 1011
Consequently,
- (gh, (91, ) = f(*²) - -f(x)) - (g x+11 - -1")
S f(a²) - 11x + (vf(z*) 9k. k+1 - -I + 1 11 11 - 1/113
S f(x²) - f(xht1) + 8ok (f(x") - 1x 1 |1aht1 2/11" + f1 - 11.
2
20c
for any 0 > 0, where used Young's incquality in the last step. Plugging into (34) gives
(1 axA) 1²1 s 1| - 2²11 + la + 12 x² + 20x (xxx") - (2(
2ax (f(=: fa* Sx) + (x**)= - -f(((tt) + o (f(x* - 9n 11a
Applying conditional expectation, we have [lf(r*: [R] fx²) and
BI-s(x S&) f(x F2 [ g 11 < 3.
Morcover, by assumption, 0g I. + - 0. Altogether, applying total expectation yields
(1 + a1A)EE x+1 - xx 1| < E|æ x (2a+ vi + #3a
which proves (21).
Proof of a): let ag = A[o+ko]
Denote As == E||2* - 1=2 Rearranging and summing (21), we have
K-1
K-1
Σ
Elv(obtl) - (xx*)| 5
As
1 20
A A+41
+
05ac 2
Zan
k=
k=)
Plugging in 0k; we have
1+a+}
)(k+tko) 2
+ A and thus
20x
K
B) -b-1)")))))))) xxx*)]
A(k+4k))
A
A(*+1+ko) 2
A k-
+
83
X(k+
k==
k=0
Dividing by K and using convexity of 10 we have
K-
Ako
E
}(x*))
<
11:0
x=112 1
08
K
2K
2
k=0
k=0
By assumption is convex and therefore convex.
23

---

Published in Transactions on Machine Learning Research (05/2023)
Finally, as Ko 1, we estimate Ex=o kthp S
+ In K by Lemma 13 and obtain (22). Proof of b): Similar to the proof above, we rearrange and sum (21) from k = O,..., K 1, and obtain
E
k+1 S 11
K-1
120 -=11
03 Σ k= O
2
xx*x)]
We divide by Ex=o ag and use convexity of vp in order to obtain the left-hand side of (23). Morcover, by Lemma 13 we have
K-
K Σ
Σ a
1),
s (02+11l11)) ln).
k=0
Plugging the above estimates, gives
agoltt) ) b(x0)} s |100 -
03a(1+ + In K))
E
CRA
4a(v K -T Mvk+1 +1
Proof of c): If f is a strongly convex, then 1p is (A+ + (i) strongly convex and
v(x*) c1x S -
- || - =x11
From (21), with ag = a, we get
(1 a(u + 211 x=11 $ E11a** x=11 + 03a?.
Doing a recursion of the above from O,...,K K gives
Elia K 1+1 S (1 + alp + 2A)) K 11100 x=112++ 088a Σ + alu + 22))
+
Using the geometric series, Σ + a(k+ (2)) *
1+alu+2A)
-1 a(u+2Aj and thus
a(x+22)
0,3c
Elix - 1=2 $ (1 + a(u + 21)) - 1|100 -
--1
-x---111p 112
# + 2X"
A.3 Proof of Theorem 8
Proof of Theorem 8. In the proof, we will denote gb. = V Sz). By assumption f p-weakly convex and hence v is (p A)-weakly convex if p Å and convex if p S X. Hence, 1* == (x*) is well-defined for < 1/(p - 2) if p > À and for any n > clsc. Note that 4 is -o-msuuurable. We apply Lomma 6, (17)
with x tk. Due to Lemma 2 (ii) it holds
(# Sk) (a";; ; + ((**) < f(8"; S6) + P 1| -x x*112 (((*)).
Together with (18), this gives
(I + and)lla - 1x12 5(1 + aups, 1|1a* - 1-112-11111 || - 2112
+ 2a (p((at) - (l() + f(a"; Sx) - f Sx) - xkt1 2nh)
Analogous to the proof of Theorem 7, due Lipschitz smoothness, for all 0 > 0 we have
; (ge, 2 S 5-(xt + f(**)
((x + fag [VS((*)) -g 2000
7 |1 - x012
24

---

Published in Transactions on Machine Learning Research (05/2023)
Plugging in gives
(1+ + - S (1 + asps.)lla* - $/112+ + 20 (p(a)) s0(xht1))
20x (8 Sk)) ; + (axh) (aatt))
sau
2
V [f(xa)) 9k 13)
+ 11 + an - 1] 1| k+ - 112
It holds Eff 5p - f(z²; FF = f(²") f(x²) and ((b())) = (((h)) By Assumption 4, we have
gl
Vf )| Fx S B. Altogether, taking conditional expectation yields
(1 + au 1) 1 k+1 x* 1 - (1 aup))
1012 - -
2anll - ((aht1)(Fe)
+ 0 + 18 a 1|1|||= - [[a)
Next, the definition of the proxinal operator implies that almost surely
v((*) 1 21 S (baht!) + 2n +|
x²
and hence
E [(((2*) |(FF S E 1121 2x113
2n 11tt x² 1..1111 [x
Altogether, we have
(1+ anA)B - (1 + ax (p- - 11 1)1b $**112
+ af00 + [ + axL + au77 -1] E|| k+1 x [21]
From assumption (25), we can drop the last term Now, we aim for recursion in env". Using that
1+ axlp aa(p-1711) + axd and + ax(p 1 = 1+ aulp 17 - A)
S 1 + ax - -7 -A),
1 1ta
1+a
1 aaad
we get
1
Elenv" (xx $ E(v(
+
- 11 1F
2n
1
2
S (att)
+ 11
-
-
1011
+ 2n
Ja (p 17- -)]]]aa*
2²112=
03.
21
2n
=env (at)
Now using || 1011 1/l Venve (x+)] we conclude
Blenv (x + < (n*) + " [oo(p 1 1 Ven
a 03.
27
Due to (25), we have n- + - p 0. Taking expectation and unfolding the recursion by summing over k = 0, O,..., K 1, we gel.
K-1
K-1
Σ k=
2 - 111-1(p-1 A))E Venv (1*)) < env" Elenvy (x*
+
0g
03.
2n
k k=
Now using that env? (x inf zb almost surely, finally get
K
2(env") (z0) inf vb)
30
K-1 Σ
Veny
1 1-1(p X)
+
1)(
7p
)))
Ok
(35)
k=0
25

---

Published in Transactions on Machine Learning Research (05/2023)
which proves (26). Now choose ap VK+I and divide (35) by === ag- Using Lemma 13 for at K-1
and Σk=0 ai we have
envy (2)) inf 40
30
a + K)
min k=
B| "e
a(1 - 7 2))( K +T 2y n X)) (vK +1 1)
a instead, we can identify the left-hand-side of (35) as av KE Choosing VK avK and using K-1 o a² obtain
Venv (I )12 Dividing by
2(env ( 1 inf (5) +
30 a
E Venv (xx)) )
a(1 - (11-(pp- A))
n7(1 7pp - x)) VK
B
Auxiliary Lemmas
Lemma 11 (Thm. 4.5 in (Drusvyatskiy & Paquette, 2019)). Let f be L-smooth and ip be proper, closed, convex. For n > O, define (z(2) == 77 7 (x prox, (x - (())). It holds
Veny 1/2L) (x) - 1 S (1 + v2
Venv Veny
1/(2L) (x)1
Vx E TR.
Lemma 12. Let c E R, x E R and 3 > and let yp R" R U foo} be proper, closed, convex. The solution
1
y" arg min yER"
(c+ (a. + (p()
Ill =012
(36)
23
is given by
prox su (x - (a), pp (x²),
if o (a, proxxo(1' 0 - -Ba)} > 0,
y+
c (a, prox po (x0)) < 0,
(37)
prox - Bua)
else, for mrue [0,1] such that c + (a, p - (ua)) 0.
Remark 4. The first two conditions can not hold simultaneously due to uniqueness of the solution. If neither of the conditions of the first two cases are salisfied, we have to sind the root of u + c+ (a, prox -Baa) for u E [0, Due to strong converity of the objective in (36), we know that there exists a root and hence yt can be found efficiently with bisection.
Proof. The objective of (36) is strongly convex and hence there exists a unique solution. Due to (Beck, 2017, Thm. 3.63), y the solution to (36) if and only if it satisfics first-order optimality, i.c.
3u E 0(-) (a. v\) 0 ua + ap()) + 3 g(v - x²).
(38)
Now, y = pp))
+ 0 E 04()) + (v =), it holds
Eu 0(.)+ + (a, v\) 0 ap(v) 1 g -2x0 Bua))
(38)
3u E 0(-) (c (a, w)) y proxao(20 (x0 Bua).
We distinguish three cases:
Let ÿ = (x 3a) and suppose that c (a, Ø) > Then ((+)+(c+ (a, 0)) {1} and hence T satisfies (38) with 1. Hence, y* ÿ.
26

---

Published in Transactions on Machine Learning Research (05/2023)
2. Let prox3p(2") and suppose that c (a, 0. Then (())((c (a. {0} and hence y satisfies (38) with 0. Hence, y = y-
3. If neither the condition of the first nor of the second case of (37) are satisfied, then, as (38) is a necessary condition for the solution y², it must hold c + yo.y+) = 0. Hence, there cxists a u E (c+ (a, = [0, 1] such that
c+ (a, prox -B3)) = O.
Lemma 13. For any K it holds
k
K-1
K-1
Σ
1 +
Σ
11
/
ds = 1 + Ink,
x+1
K+I
841
0
k=1
K
K
V+4
2
Vk+1
k=0
The following is detailled version of Proposition 3. We refer to Scction 4.3 for context.
Proposition 14. Let Assumption 1 and Assumption hold and assume that there is on open, conver set U containing dom s. Let f(;;s) be Ps- la-weakly convex for all s S and let p = Elps]. Assume that there exists G, RR for all S, such that G / and
| G. ( Ve E U.
(39)
Then, 82(3;8) (given in (10)) satisfies the following:
(B1) It is possible generate infinitely many i.i.d. realizations S1, S2.... from S.
(B2) It holds E[fa (x; 5)] fxx and (1: S f(v) Glly x11 for all x. y R
(B3) The mapping (z(: 8) (( + p(.) is convex for all x E R" and all S.
(B4) For all 2, E U and S, it holds (x; - su < G.Ilæ all-
Proof. The properties (B1) (B4) are identical to (B1) (B4) in (Davis & Drusvyatskiy, 2019, Assum. B); setting r = y, fo(-.E) = fa(::) s), 7 = 0, t = p, L = G, and L(E) G.. (B1) is identical to Assumption 1. (B2) holds due Lemma 2, (ii), applying expectation and using the definition of f. i.c f(x) (x; S)]. (B3) holds due to Lemma 2, (i) and convexily of sp. For (B4), taking g (f;);) and x, y E U, have
fa(1;)) - fo( s f(z;s) f(x; 8) - (g. y - x) < llg/llyy - x I| s G, || al1-
C Model equivalence for SGD
In the unregularized case, the SGD update
k+1 x argki 9k E 0f(x:
can be seen solving (5) with the model
Sa(y; {1;; (g.v x), g E ()())))
27

---

Published in Transactions on Machine Learning Research (05/2023)
Now, consider again the regularized problem (2) with ppæ) |x|| and update (8) On the one hand, the model 1he (v;s)) = (2; 8) + (p(x) + (g + Az, y a)) with g E (2;8)) yields
1+ = x² - (x + Xx
(1 ax\)x** ag.9k
(40)
On the other hand, the model (y: (x; (g.y x) (vy) with g 0f(z;)) results in
x k+1 = proxapx and ( ox = - a+d
[1t akak] = (1 +aaa
Ok
A)a
1 a+ gk-
Ok
(41)
an Running (40) with step sizes ag 3g is equivalent to running (41) with step sizes
1+a,1 In this sense, standard SGD can be seen to be equivalent to proximal SGD for fg-roaularized problems. 1-BAA-
B
- ap
D Additional information on numerical experiments
D.1 Matrix Factorization
Synthetic data generation: We consider the experimental setting of the deep matrix factorization ex- periments in (Loizou al., 2021), but with an additional regularization. We generate data in the following way: first sample B E Raxx with uniform entries in the interval [0, 1]. Then choose R (which will be targeted inverse condition number) and compute A = DB where D is diagonal matrix with entries from I to v (equidistant on logarithmic scale)". In order to investigate the impact of regularization, we generate a noise matrix E with uniform entries in [-e,s] and set A= A @ (1 E). We then sample y( ~ N 7) and compute the targets 660 = (y A validation set of identical size created by the same mechanism, but computing its targets, denoted by (1)) val via the original matrix A instead of Á. The validation set contains Nval = N samples.
Name
p
9
N
matrix- fac1
6
10
1000
1e-5
4
0
fatrix-fac2
10
1000
1e-5
10
0.05
Table 1: Matrix factorization synthetic datasets.
Model and general sctup: Problem (29) can be interpreted as two-layer neural network without acti- vation functions. We train the network using the squared distance of the model output and 5(0) (averaged over a mini-batch) as the loss function. We run 50 epochs for different methods, step size schedules and values of X. For each different instance, we do ten independent runs: each run has the identical training set and initialization of W1 and W2, but different shuffling of the training set and different samples y for the validation set.. In order to allow fair comparison, all methods have identical train and validation across all runs. All metrics are averaged over the ten runs. We always use batch size of 20.
D.2 Plots for matrix-fac2
In this section, we plot additional results for Matrix Factorization, namely for the setting matrix-fac2 of Table 1, Fig. 10, Fig. 11, and Fig. 12. The results are qualitatively very similar to the setting matrix-fac1.
Note that (Loizon al., 2021) uses entries from to u on Knear scale which, in our experiments, did result in large condition numbers if is very small.
28

---

Published in Transactions on Machine Learning Research (05/2023)
conslant
10-1
10-1
10-2
10.
10-
10
10"4
10
20
30
10
30
Epoch
Ipock
Figure 10: Objective function for the Matrix Factorization problem (29), with constant (left) and sqrt (right) step size schedule and several choices of initial values. Here ming (bx*)) is the best objective function value found over all methods and all iterations.
sqgt
10-1
10a
10
20
Booch
Epoch
Figure 11: Validat ion error for the Matrix Factorization problem (29), with constant (left) and sqrt (right) step size schedule and several choices initial values.
(0000
0000
E.0
000033
3.0.2
11.000
000000
0.305-
3.0.0
0.308
00000)
2110
1.5
0.002
0.300
300
1.0
Ypoch
Epoch
Fpoch
Te L6
6.01
0.000
3.000
0.000
0.11
1.00
0.804
4.0
0.000
012
3.00
Fpoch
Fpech
Fgnch
Fyych
Figure 12: Objective function value and validation over the course of optimization. For the validation error, plot rolling median over five epochs in order to avoid clutter.
29

---

Published in Transactions on Machine Learning Research (05/2023)
D.3 Matrix completion experiment
Consider an unknown matrix of interest W E R4xxd, Factorizing W U with U E Rrxd, Ve R we can estimate the entries of matrix W as
We = uf By + by + bbj, e [di]. J E [d2],
(42)
where is the i-th column of U and vy is the j-th column of V, and dy E R2,, by E R2 are bias terms (Rivera-Muñoz et al., 2022).
We can interpret this as an empirical risk minimization problem as follows: let T be the set of indices (i,j) where Wa is known. With Wa as in (42) for trainable parameters (U, V,V ),,, the (regularized) problem is then given
X Il V.b 11 (W - W (WJE
Σ
min U.V.B BY
[T|
2
We use datasct containing air quality measurements of a sensor network over one month. This dataset has been studied Rivera-Muñoz al. (2022). The dataset contains measurements from 130 sensors over 720 timestamps, hence d1 130, d2 = 720. In total, there are 56158 nonzero measurements (the rest was missing data or removed due to being an outlier). We split the nonzero measurements into training set of size [T] 44926 0.8 56158 and the rest using mean and variance of the training set. We set 24 and use batch size 128. The validation error is defined as the root mean squared error on the elements of the validation set (which is not used for training).
validation set. We standardize training and validation set
The results are plotted in Fig. 13 and Fig. 14a. For all methods, we use a constant step size Og.. ProxSPS achieves the smallest error on the validation set for the two smaller values of X. For the largest X, ProxSPS, SPS and SGD are almost identical for ag = 5, but SGD with ag = 1 is the best method. However, over all tested values of X, Fig. 14a shows that ProxSPS obtains the smallest error. Again, from the lower plot in Fig. 13 we can observe that ProxSPS produces iterates with smaller norm.
Discussion:
D.4 Imagenet32 experiment
Imagenet32 contains 1,28 million training and 50,000 test images of size 32 x 32, from 1,000 classes. We train the same ResNet110 as described in Scction 5.3 with two differences: we exchange the output dimension of the final layer to 1,000 and activate batch norm. We use batch size 512. For this experiment we only run one repetition.
Similar to the sctup in Scction 5.3, we run all methods for three different values of X. For AdamW, we use a constant learning rate 0.001, for SGD, SPS, and ProxSPS we use the sqrt-schedule and on = 1. The validation accuracy and model norm are plotted in Fig. 15: we can observe that all methods perform similarly well in terms of accuracy. However, Adamä more sensitive with respect to the choice of and the norm of its iterates differs significantly from the other methods. Further, using an adaptive step size is advantageous: from Fig. 16, we see that the adaptive step size is active in the initial iterations, which leads to a faster learning of (Prox)SPS in the initial epochs compared to SGD.
D.5 Interpolation constant
We illustrate how the interpolation constant o² behaves if it would be computed for the regularized loss 4(x) = f((x) + A ||= (cf. also Section 4.2). We do a simple ridge regression experiment. Let A Nxn be matrix with row vectors 0; Tk,, E [N]. We set N = 80, 100 and generate f. R with entrics drawn uniformly from [o, 1]. We compute Af. In this case, we have f((x) (a, 6 and
f(x) = N ENN,,ta fi(x).
12eTe dataset downloaded from httpp:://itthub.cmm
fanntt////ttttooooeeeeessss//s/s ininados/Ventans.
El1 mes1.csv.
30

---

Published in Transactions on Machine Learning Research (05/2023)
20,0
1011 conolant opm
cam3
adtame, outaltal, 10-0.000
oy110
Se 05
0.0011
0,0005
0.85
0.85
Validation Eivor
0.80
0.75
0.75
0.70
0.70
0.71
0.65
25
50
100
25
75
100
50
100
Epoch
Epoch
Epoch
30
30
40-
15
25
50
70
100
50
100
25
50
100
Epoch
Epoch
Figure 13: Matrix Completion: Validation error (lop) and model norm (top) for three values of the regular- ization parameter X. Validation error plotted as five-epoch running median. Shaded area two standard deviations over ten independent runs.
ccontt 0001
centash de
7.2%
7x
ValaCarinn
68x 10"1
66x
64 10
62x 10-1
10
10-0
10
10-3
10P
102
(a)
(b)
Figure 14: (a) Matrix Completion: Validation error as a function of the regularization parameter X. Shaded area is one standard deviation (computed over ten independent runs). For all values, we take the median over epochs [90, 100). (b) Interpolation constant for ridge regression problem for varying regularization parameter X. See Appendix D.5 for details.
If would apply the theory of SPSmax for the regularized loss functions & with estimates & 0, the constant = (minzer= f(x) + vp(x) - En inf, (x)) determines the size the constant term in the convergence results of (Loizon al., 2021; Orvicto et al., 2022). We compute mingeR" f(x) + xp(x) by solving the ridge regression problem. Further, the minimizer of is given by (a,aī + AId) (l,)- a,bb We plot a² for varying X in Fig. 14b to verify that d² grows significantly if X becomes large (even if the loss could
31

---

Published in Transactions on Machine Learning Research (05/2023)
adaro, 0ym00003
Apd, aym11))
A-
0,000
1.0
1.0
1.0
0.5
11.8
0.8
Validatin
0.2
0.2
0.2
0.0
0.0
0.0
20
40
20
20
Epoch
Epoch
Epoch
À
0.00005
185
52-
300
130
260
125
50
45-
115
150
46
110
40
20
41
20
401
Epoch
Epoch
Figure 15: ResNet110 for Imagenet32: Validation accuracy as five-epoch running median (Lop) and model norm (bottom) for three values of A.
be interpolated perfoctly, i.. inf, f(x) 0)- We point out that the constant 0² does not appear in our convergence results Theorem and Theorem 8.
32

---

Published in Transactions on Machine Learning Research (05/2023)
prox = 1.0, À 5e 06
1.0, 5e 06
10a
median(Ca)
median(Cu)
Step size
10
10
10-3
prox-sps; oo 1.0, À
05
sp3, on 1.0,
10a
(Sdian((a))
median(Ce)
Step size
10
10-1
10-3
prox-sps, Ou 1.0, 0.0005
sps, 1.0, À 0.0005
10
median(Ca)
median(Cu)
Step size
10
10
0
10
20
30
40
50
10
20
30
40
50
Epoch
Epoch
Figure 16: ResNet110 for Imagenet32: Adaptive step sizes for SPS and ProxSPS. See definition of & in
Section 5.1.
33


Reviews:
[{"review_id": "nJGkEhUqC6", "reviewer": "Reviewer_24qn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n-----------\nThis work is well written and the authors address a task that is important for the community which could have impact on \nthe field of efficiently training deep learning methods.\n\nI appreciate the simplicity of the method and the fact that it is easy to implement. \n\nThe authors also provide a nice discussion of the theoretical properties of the method.\n\nProxSPS seems to perform better than SPS when larger networks are used, which is a nice result as it shows that the method \nis not limited to small networks and there is potential that it can be used to train large networks like LLMs.\n\nThe authors performed multiple runs in their experiments and report the mean and standard deviation of the results \nwhich is important for understanding the performance of the method.\n\nThe experiments are exhaustive and the authors provide a nice discussion of the results and how they compare to other methods in standard datasets. It is especially interesting to see ProxSPS works well with larger networks like ResNet110 vs. SPS, suggesting that it might be able to work well for larger networks like LLMs.\n\n\nWeaknesses\n--------------\nThe authors should test this method on larger datasets like ImageNet and LLM type of datasets in order to identify the limitations of the method in large scale settings and see how it compares against standard methods like AdamW.\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work extends the stochastic Polyak step size SPS with a proximal variant called ProxSPS, which can handle regularization terms efficiently.  The results show that ProxSPS is easier to tune and more stable in the presence of regularization than SPS for datasets like CIFAR10 with large networks like ResNet110.\n\nCompared to the stochastic Polyak step size SPS, ProxSPS only requires estimating a lower bound for the loss function and not for the composite function (that includes the L2 regularization term). Further, ProxSPS has closed-form updates for squared L2-regularization and provides theoretical guarantees for proper, closed, and convex regularization function.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Compare between ProxSPS, SPS, SGD and Adam on the following two useful tasks:\n\n- ImageNet classification.\n\n- Fine-tuning a pre-trained large language model like BERT for classifing a text-classification problem like CLINC150.\n\nThe reason for these experiments is to see how ProxSPS compares to other methods in large scale settings.\n", "broader_impact_concerns": "It is not clear how much tuning is required for this method to achieve good performance on some problem setups, which can lead to large unnecessary usage of compute like running large-scale GPU resources which can negatively contribute to the environment.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "7qG10vz31T", "reviewer": "Reviewer_1wix", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n## Theory\nFrom a theoretical point of view, all obtained results are a simple application of the Davis & Drusvyatskiy 2019 paper. That is, there is not much insight behind the analysis and the paper doesn't bring any new ideas on this front. Such results are already obtained for SGD and other variants, often with less restrictive assumptions. The theoretical guarantees of the algorithm are not better and actually even worse if we take into account that the algorithm is not implementable (see the comment on Corollary 4 below).\n\nThe minor addition like adding the prox function is also straightforward.\n\n### The truncation idea\nThe authors suggest to choose step size $\\gamma_k$ as \n$$ \\gamma_k = \\min \\left(\\alpha_k, {\\color{blue}\\mathrm{some\\ expression}} \\right)$$\n\nand then still require $\\alpha_k$ to be of the order $1/\\sqrt{k}$ or $1/k$, depending on the properties of $f$. I don't understand why we need this. The ${\\color{blue}\\mathrm{some\\ expression}}$ is well separated from zero (otherwise the problem is very special), so eventually $\\min$ will disappear and we end up just with the step $\\alpha_k$, exactly as in the classic SGD. Isn't it obvious that in the beginning of SGD we can use any step sizes whatsoever and then switch to a proper \"convergent regime\"?\n\n\n\n\n## Motivation\nRoughly speaking: why should we care? There are so many variations of SGD. Neither theory nor a basic intuition suggests to me that this should be an advantageous algorithm. Since it requires the same steps in the end, relying on numerics as a possible motivation doesn't make sense to me either.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies a variant of SGD for both smooth and non-smooth objective and also for the objectives with the composite structure. The only difference with the classical SGD is the choice of the step size. In this paper it is used as a combination of the Polyak step, adjusted for stochastic setting, and truncation idea.\n\nUnder either bounded subgradients or smoothness assumption of the main function $f$ and weakly-convexity of $f$ some convergence results are obtained.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Some comments:\n\n1. I don't think it is a good idea to use the same letter $f$ to denote both functions. Sometimes it is not clear to which $f$ is referred in the text.\n\n2. First $g$ denotes a (sub)-gradient, then it denotes a convex function, then again it is a subgradient. \n\n3. Assumption 2: Does it mean that you assume that $\\inf_x f(x; s)$ is finite? If this is the case, then it should be stated somehow in this form. Then $C(s)$ is just a notation.\n\n4. page 5, first paragraph: I didn't understand: which $c$ is better in the end and why?\n5. Why do we need to spend one page on special (and obvious) case for $\\ell_2$ norm?\nI would say just writing one equation should suffice.\n\n6. It is a bit strange to read: in the intro $f$ is assumed to be convex, then it becomes weakly-convex.\n\n7. I don't know what authors had in mind, but I don't see how the paper with the Proposition 3 in that form can be accepted.\n8. Assumption 4 is quite strong, there are more flexible assumptions for this.\n9. Corollary 4. It says that in order to run the Algorithm, we need to know $\\Delta$, which is uncomputable in our setting: neither $\\mathrm{env}_{\\psi(x^0)}$, nor $\\min_x \\psi(x)$ are available. If we take it too large, then the complexity becomes vacuous.\n10. Theorem 8. Haven't we already assumed that $\\inf \\psi > -\\infty$?", "broader_impact_concerns": " -", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rKxQhKRKeb", "reviewer": "Reviewer_J2D7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\nThe idea overall is clear and easy to follow. \n\nThe authors provide theoretical analysis and empirical justification.\n\n### Weaknesses\nThe proposed method is not substantially novel. In fact it only replaces the step-size in original proximal point method with the Polyak step-size.\n\nThe comparison with the existing methods is a little bit confusing.\n\nThe experiment part is not aligned with the theoretical analysis well. Specifically, the authors focused on the decaying step-size in the main text, but in the experiments the authors also showed the results with constant step-size and claimed the original version failed in this case but the proposed method can work. Also, the authors claim that the proposed method work better with larger $\\lambda$, but the theoretical analysis does not show that.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a stochastic proximal update with Polyak step size for regularized empirical risk minimization problem with a model-based viewpoint. Different from the previous methods that also construct the model for the regularizer, the authors proposed to construct the model only for the risk and then perform the update with the proximal point method. The authors demonstrated that, with a decaying step-size parameter, the proposed method can exactly converge to the stationary point, which is in contrast to the previous methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please provide a much more detailed comparison with the existing work, include but not limited to:\n\n* For Loizou et al, 2021 and Orvieto et al., 2022, if they use a decaying step-size, can they obtain exact convergence?\n* How does $\\lambda$ affect the convergence results with a decaying step-size for the original method and the proposed method in theory?\n\nI'm also interested in the reason that the proposed method still converges well with constant step-size but the original method fails. Can the authors provide some theoretical justification on that?", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)