Boosting Data-Driven Mirror Descent with Randomization, Equivariance, and Acceleration
Authors: Hong Ye Tan
Subhadip Mukherjee
Junqi Tang
Carola-Bibiane Schönlieb
Venue: Accepted by TMLR ()

Abstract: Learning-to-optimize (L2O) is an emerging research area in large-scale optimization with applications in data science. Recently, researchers have proposed a novel L2O framework called learned mirror descent (LMD), based on the classical mirror descent (MD) algorithm with learnable mirror maps parameterized by input-convex neural networks. The LMD approach has been shown to significantly accelerate convex solvers while inheriting the convergence properties of the classical MD algorithm. This work proposes several practical extensions of the LMD algorithm, addressing its instability, scalability, and feasibility for high-dimensional problems. We first propose accelerated and stochastic variants of LMD, leveraging classical momentum-based acceleration and stochastic optimization techniques for improving the convergence rate and per-iteration computational complexity. Moreover, for the particular application of training neural networks, we derive and propose a novel and efficient parameterization for the mirror potential, exploiting the equivariant structure of the training problems to significantly reduce the dimensionality of the underlying problem. We provide theoretical convergence guarantees for our schemes under standard assumptions and demonstrate their effectiveness in various computational imaging and machine learning applications such as image inpainting, and the training of support vector machines and deep neural networks.

Published in Transactions on Machine Learning Research (05/2024)
Boosting Data-Driven Mirror Descent with Randomization, Equivariance, and Acceleration
Hong Ye Tan
hyt35
Department of Applied Mathematics and Theoretionl Physics University of Cambridge, United Kingdom
Subhadip Mukherjee
wwnkherrjeeeere.kkkkgp ac. ac.im
Department of Electronics and Electrical Commanniaation Engineering Indian Institute of Technology, Kharagpur, India
Junqi Tang
j.lang. 2iw.ng.2.bbmm.co uk
School of Mathematics University of Bärmingham, United Kingdom
Carola- Bibiane Schönlieb
chs31 @ao...commcoomm ak
Department of Applied Mathematics and Theoretical Physics University of Cambridge, United Kingdom
Reviewed on OpenReview: https: // openreview. formm? id= r2dois 1lg
Abstract
Loopting-t (L20) is an emerging research area in large-scale optimization with applications in data science. Recently, rescarchers have proposed novel L2O framework called learned mirror descent (LMD), based on the classical mirror descent (MD) algorithm with learnable mirror maps parameterized by input-convex neural networks. The LMD approach has been shown to significantly accelerate convex solvers while inheriting the convergence properties of the classical MD algorithm. This work proposes several practical extensions the LMD algorithm, addressing its instability, scalability, and feasibility for high-dimensional problems. We first propose accelerated and stochastic variants of LMD, leveraging classical momentum- based acceleration and stochastic optimization techniques for improving the convergence rate and per-iteration computational complexity. Moreover, for the particular application of training neural networks, we derive and propose novel and efficient parameterization for the mirror potential, exploiting the equivariant structure of the training problems to significantly reduce the dimensionality of the underlying problem. We provide theoretical convergence guarantees for our schemes under standard assumptions and demonstrate their effectiveness in various computational imaging and machine learning applications such image inpainting, and the training of support vector machines and deep neural networks.
1 Introduction
Large-scale optimization plays key role in modern data science applications. such applications, we typically seek infer collection of parameters x E R from m ii.d. data samples of a random variable {2, m i by solving optimization problems of the form:
x* E arg min f(x) = mER*
(f((;22) + R(x)),
(1)
m

---

Published in Transactions on Machine Learning Research (05/2024)
where each fü is fidelity function associated with the i-th data sample 2,, while R(x) is regularizer containing the prior knowledge of x. In the context of machine learning applications, would contain the parameters of some models such as SVMs and neural networks, while 2 contains pair of of training data samples and labels. For signal/image processing and inverse problems, 2, is usually referred to as a pair of measurement operators and measurement data, with the measurement operator appearing in the fidelity fa- Many classical variational image reconstruction models can also take the form of (1), with a simple example being total variation regularization (Rudin et al. 1992). In this case, we have m = 1, where fi depends only the corrupted image, and the total variation regularizer R(x) >|Vz la is the l norm of the discrete differences of x. In modern applications, the dimension of the optimization problems and the number of data samples involved are often large, leading to heavy demands on the development of computationally efficient optimization algorithuns.
Throughout the past decade, we have witnessed tremendous successes by researchers in the field of optimiza- tion in data science. In particular, efficient first-order algorithms utilizing momentum- basod acceleration NNesterov 1983 Beck & Teboulle, 2009) and stochastic approximation (Robbins & Monro 1951 Zhany
2004) have been proposed, with applications including sparse signal recovery and convex programming
Bubeeck 2015 Ghadimi & Lan 2016). Recently, stochastic gradient algorithms with theoretically optimal
convergence rates for solving generic classes of problems were developed (Lan 2012: Lan & Zhou 2015:
Defazio 2016 Allen-Zhu 2017 Chambolle et al.] 2018 Tang et al. 2018 Zhou et al. 2018b Driggs et al.
20210, matching complexity lower bounds for convex (Woodworth & Srebro 2016) or non-convex problems
(Arjevani et al. 2023).
Despite being theoretically optimal for generic optimization problem classes, stochastic gradient methods result in poor performance when applied to problems with particular structures, [Tang et al. (2020) demon- strate that while optimal stochastic gradient methods perform very well in terms of convergence rates for certain imaging inverse problems such as X-ray CT, they perform poorly for some other problems such as image deblurring and compressed sensing. One important aspect is the fact that these methods are designed for some wide generic class of problems, while in practice, each practitioner may be only interested in a very narrow subclass of problems with a specific structure. The optimal algorithms for generic problems may still perform suboptimally for the specific subclasses of interest.
1.1 Learning-to-Optiing to-Ommimize
Learning-to-optimize (L20), sometimes referred as meta-learning in the machine learning literature, aims to overcome this linitation when dealing with more specialized problem classes, such as natural image reconstruction subset of all imaging problems. An underlying idea is to utilize the intrinsic structure of the optimization problems to converge faster on tasks from the same distribution. This has been shown to be effective for problems such as linear regression, training neural networks, and tomographic reconstruction,
beating standard analytic methods (Andrychowicz et al.) 2016 Li & Malik) 2016: Li et al.] 2019 Banert et al. 2020). The aim is to generate optimizers that (I) optimize in-distribution functions at a faster rate, and (2) possibly return a higher quality solution given the same computing budget (Chen et al. 2022).
The L20 framework can be formally described follows (Chen et al. 2022 Given a learnable mapping function with parameters 0, well a set F of training sample optimization problems ming f(x) for f E F, the optimizer is learned to generate updates of the form Fx+1 = 22 -go(Te). Here, I, represents previous information, such as the iterates Tt, It-1, or their gradients (f(xz), Vs(=+-1) etc. For example, gradient descent with can be formulated as (T,) = fxzz) with a learnable step-size parameter 17. This formulation of L20 is flexible in the sense that all previous information is available, which should theoretically allow for better informed steps. The parameters are usually learned by minimizing an objective of the form
T Σ Wa soas(ce))
min E
Te+1 = Iz (It)
2
This can be interpreted as minimizing the objective value of the evolved iterates, with weights w,- Some standard choices include u, = 1, or wy 1. uT = 2016).
=... ... = w (Zhou et al. 2018a; Andrychowicz et al.
2

---

Published in Transactions on Machine Learning Research (05/2024)
1.2 Existing Theory in L20
Learning-to-oplimize frameworks broadly fall into two categories: fully learned "model-free" methods, and imodel-based" methods based on classical algorithms (Chen et al. 2022). Model-free methods directly pa- rametorize the optimizer updates go (Ie) using neural networks, formulating fast optimization as an unrolled
loss objectives of the form Equation (2) (Andrychowicz et al. 2016 Li & Malik, 2016). Recent advances for this paradigm include improving the generalization performance of both the optimizers and final opti- mization objective when used to train neural networks (Almeida et al. 2021 Yang al., 2003), more stable computation (Metz et al.) 2019), and learning to warm-start certain optimization methods (Sambharya et al.,
2023 However, major drawbacks include lack of convergence guaranlees, as well as requiring many training samples.
Mosel-based L20 methods instead inherit structure from classical optimization algorithms such as PDHG- like proximal splitting algorithms, where certain optimizer parameters such as step-sizes or momentum parameters are learned (Banert et al., 2020 Almeida et al. 2021; Banert et al., 2021 Gregor & LeCun 2010). These methods are lypically applied to convex problems rather than non-convex problems such neural network training, which also allows for theoretical convergence analysis of the learned schemes. One advantage of building L20 model from a classical optimization model is that convergence of the learned model can usually be derived by slightly modifying existing convergence results. A more in-depth overview of T.20 can be found in Chen et. al. (2022).
From a more statistical point of view, Chen Hazan (023) formulates learning-to to-optimize as a feedback control problem for smooth convex objective functions, showing an upper bound on the expected opti- mality gap in broader class of stabilizing controllers. In the case where the algorithm class has small pseudo-dimension with respect to set of problem instances, the empirical risk minimization problem ad- mits statistical PAC bounds for the number of problem instances required (Gupla & Roughgarden 2016). Sambharya et al. (2023) demonstrates PAC-Bayes generalization bounds for certain classes of fixed fixed-point optimization, where a model maps an initialization to a solution with access to an objective f (Amos et al. 2023).
optimization methods. A related notion is that of
Combining L20 with mirror descent has mainly been explored in the online setting, with theoretical con- tributions arising in terms of statistical transfer bounds and optimality gap bounds (Khodak et al., 2019: Denevi et al. 2019). Gao et al. (0222) also consider simple version of coordinate-based mirror descent, providing a statistical PAC generalization bound. These are more suitable to common meta-learning bench- marks such as few-shot training for classification where convexity is unavailable (Finn et al.) 2017 Nichol ct al. 2018). These meta-learning methods that are focused on neural networks generally take a different approach to the L,20 methods discussed above mela-learning accelerates by moving the initialization to "better points", in contrast the L20 objective of learning an optimization trajectory (from any initializa- tion). We consider a more general and less theoretically-backee framework where the meta-parameters are baked into the optimizer, rather than common instances such as initialization which are difficult to interpret in our context.
A provable cerniigglo-optimizz framework called learned mirror descent (LMD) was recently proposed in Tan et al. (2023b a) based on the classical mirror descent (MD) algorithm by Nemirovski & Yudin [1983). By using neural networks to implicitly learn and model the geometry of the underlying problem class via mirror maps, LMD can achieve significantly faster convergence on certain convex problem classes, including model- based image inpainting and training support vector machines. Moreover, the LMD framework proposes approximate convergence results based on the accuracy of the mirror map inversion.
Building upon the vanilla LMD framework by Tan et al. 2023b), we make the following contributions:
1. Learned accelcrated mirror descent and its thcorctical analysis. Our first contribution is de- veloping the learned accelerated mirror descent (LAMD) scheme, which was initially proposed in the preliminary work of Tan et al. (2023a) without convergence analysis. In this work, we complete the theo- retical analysis of the LAMD scheme under standard assumptions, demonstrating improved convergence rates over vanilla LMD. In particular, we demonstrate Theorem that we obtain convergence to the minimum, as opposed to constant above the minimum found in Tan et al. (003bb).
3

---

Published in Transactions on Machine Learning Research (05/2024)
2. Learned stochastic mirror descent and its acceleration. We propose the learned stochastic mirror descent (LSMD) and the accelerated version (LASMD) in the spirit of stochastic optimization, which is crucial for the scalability of the LMD framework in large-scale problems prevalent in modern data science applications. We also provide theoretical convergence analysis of these schemes and demonstrate their elfectiveness using numerical experiments including image processing, training of SVMs, and training neural networks for classification.
3. Efficient mirror map parameterization by utilizing the equivariant structure. We propose equivariant theory for T,20 in the case where the larget class of objectives satisfies some group sym- metries. For the training of deep neural networks, we utilize the equivariant structure of the training problems to significantly simplify the parameterization of the mirror map, further improving the effi- ciency and practicality of the LMD framework, and allowing for applications in a non-convex setting, We provide a derivation of this scheme, showing theoretically that learned optimizers trained under certain group invariance and equivariance assumptions will also satisfy a group invariance property, removing redundant parameters and reducing the problem dimensionality. We additionally utilize this to extend the LMD framework the task of training deep and convolutional neural networks, achieving competitive performance with widely used and empirically powerful optimizers such as Adam.
1.3 Mirror Descent and Learned Mirror Descent
Mirror descent (MD) was originally proposed by Nemirovski & Yudin (1883) as method of generalizing gradient descent to general infinite-dimensionall Banach spaces. The lack of Hilbert space structure and isomorphisms between X' and their duals X prevents the use of gradient descent, which usually identifies the formal derivative Df E X with the corresponding Riesz clement f E X, satisfying (Df,z)x+.x = (Vf,x)x,xx MD proposes to address this by directly updating elements in the dual space X", and 'tying' the dual space updates in X to corresponding primal space updates in X' using a "mirror map", satisfying various conditions.
Generalizing the classical Euclidean structure used for gradient descent to more general Bregman divergences, such as those induced by L norms, can improve the dimensionality scaling of the Lipschitz constant. Having smaller Lipschitz constant with respect to another norm compared to the Euclidean norm directly impacts the step size for which the methods converge, as well as the convergence rates. For example, for minimizing a convex function f on the n-simplex A" = far E IR" | x; 2 O, Σ 2;), MD is able to achieve rates of |' /vk,, opposed to (sub-)gradient descent which has rate O(IS")la///) (Beck & Teboulle, 2003 Ben- Tal et al. 2001). The ratio ||/l can be as large as vn, leading to (n/ (t/log))'/ asymptotic
O O(ylog n
speedup by using MD.
Gunasekar et al. ( demonstrate further that mirror descent corresponds to a Riemanmian gradient flow in the limiting case, and further generalize mirror descent to general Riemannian geometries. Moreover, the mirror descent framework is amenable to modifications similar to gradient descent, with extensions such ergodicity (Duchi et al., 2012), composite optimization (Duchi et al. 2010). stochasticity (Lan et all 2012 Zhou et al. 2017) and acceleration (Krichene et al.] 2015 Hovhannisyan et al. 2016). In summary, the mirror descent algorithm allows us to utilize non-Euclidean geometry for optimization, which has been shown to accelerate convergence in applications including online learning and tomography (Srebro et al.
201T Raskutti & Mukherjee 2015 Allen-Zhu Orecchia 2014
Ben-Tal et al., 2001 Orabona et al. 2015
Zimmert & Lattimore 2019
We present the MD algorithm given in Beck & Teboulle 2003), , as well as the data-driven version of which this work is based, LMD (Tan et al. 2023b). The LMD framework serves stochastic extensions that will be presented in this work. LMD aims to speed up convergence on a class of "similar" functions, which are taken to be qualitatively similar, such as image denoising on natural images or CT imaging [Banert et al. 2020). By using data to generate functions from the target function class, the geometry of the class can be learned using these sample functions. Mirror descent then exploits this learned geometry for faster convergence. Follow-up works additionally suggest that the learned geometry can be transferred to other mirror-descent--ypp algorithms without the need for retraining. Moreover, the learned
base for the accelerated and
4

---

Published in Transactions on Machine Learning Research (05/2024)
mirror maps are robust to small perturbations in the forward operator such from the identity to blur convolution, well as change of data sets (Tan et al. 2023a).
We begin with some notation and definitions that will be used for mirror descent, as well as more standard assumptions for the MD framework.
Let X' be the finite dimensional normed space (R²", || I))- Denote by X = ((R")",lI Il II.)) the dual space, where the dual norm defined as ||p||- sup (p, ||x||| S 1}, and the bracket notation denotes evaluation (p. 22 p(æ) R. Let R == R J (+ denote the extended real line. Let f : X' - R be proper, convex, and lower-semicontimoor function with well-defined subgradients. For a differentiable convex function h X + R, define the Bregman divergence as B₁ (x, y) h(x)) h(y) - (N'(v), - y). Let g X + X be a function such that g(x) E 0f(x) for all x E X. We denote a (possibly noisy) stochastic approximation to g by
G E) = g(x) +
(x)),
where Eg [c[A(x,E] = o, and (xx X + X' represents dual noise. One standard assumption that we will use is that G has bounded second moments.
While we have defined X' to be vector space, this is not necessary. Indeed, suppose instead that X were a (proper) closed convex subset of R". Assuming that the inverse mirror function maps into X, i.e. Va²": (R")* + X, such as when Vib(z) + oa as x - dx, then the proposed algorithms will still hold. We will restrict our exposition in this work to the simpler case where X = R" to avoid these technicalities, but the results can be extended.
Definition 1. We define a mirror potential to be an a-strongly convex C function &b ;:x- - R with a > 0. We define a mirror map be the gradlent of mirror potential V X X*
We further denote the convex conjugate of &b by v0" :x + R. Note that under these assumptions, we have that Vi** (Vi))* on dom w. Moreover, th* is a -smoooh with respect to the dual norm, or equivalently, Vyl* is Lipschitz (Azé & Penot 1995
For mirror potential & and step-sizes (tx) the mirror descent iterations of Beck & Teboulle (0003 applied to an initialization x(()) E X are
x (1-1) Vel (V(((()) - (z))
(3)
The roles of VΦ X' X and V4 X X' are mirror between the primal and dual spaces such that the gradient step is taken in the dual space. A special case is when (xx)) 1 | whereby V and V
are the canonical isomorphisms between X' and X", and the MD iterations (3) reduce to gradient descent. In the case where f is L-Lipschitz and va has strong convexity parameter
0, mirror descent is able to
achieve optimality gap bounds of the following form (Beck & Teboulle 2003 h:
f(x(0)) -(=)) Bylz" x + (on-'Eiea & ( Iskss
min
Ex=1th tx
LMD arises when we replace so and u with neural networks, which can then learn from data. We denote Iearned variants of v and v/* by Ma X' - R and Mī X - IR respectively. Note that Mo and Mj are not necessarily convex conjugates of each other due to the lack of a closed-form convex conjugate for general neural networks - we change the subscript to remind the reader of this subtlety. Equipped with this new notation for learned mirror maps, the learned mirror descent algorithm arises by directly replacing the mirror maps Ve and Vy* with learned variants Mo and M3 respectively, where VMJ is enforced to be approximately (VM))-1 during training (Tan el. al. 2023b):
a (k+1) V MM -
(4)
The mismatch between VMJ and (VM,) - also called the forward- inconsistency, necessarily arises as the convex conjugate of neural network is generally intractable. Tan et al. (2023b) consider analyzing LMD as approximate mirror descent scheme, where the error bounds depend on the distance wetween
5

---

Published in Transactions on Machine Learning Research (05/2024)
the computed (approximate) iterates 4 and the true mirror descent iterates (3) Under certain conditions on the target function f and the mirror potential s Ma, they show convergence in function value up to a constant over the minimum (Tan et al. 2023b. Thm. 3.1, 3.6). A sufficient condition for this convergence is that the LMD iterations 2; with inexact mirror maps given by Equation (4 is close to the exact MD iterations & (k+11) = (V - (xx(()))))))))---))))))))))))))) V)((()))),
in the sense that (Vi((Tt)) - - x(0) and (Vf()),, 24) are uniformly bounded for given minimizer x.
This work proposes to address two current restrictions of LMD. The first restriction is the approximate convergence and instability that arises due to the forward- backward inconsistency, such that the function values are only minimized up to constant above the minimum. This is demonstrated in Tan et al. (2023a, Thm. 3.1); where LMD is also shown to diverge in later iterations as the forward-backwarrd inconsistencies accumulate. Section presents an algorithm through which the constant can vanish if the forward- backward inconsistency uniformly bounded. The second drawback is the computational cost of training LMD, which is prohibitive for expensive forward operators such as CT or neural network paramcters. Sections 3 and present stochastic and accelerated stochastic extensions of LMD that show convergence in expectation even with possibly unbounded stochastic approximation errors. We demonstrate these extensions with function classes considered in previous works in Section 5 utilizing pro-trained neural networks on tasks such as image denoising, image inpainting and training support vector machines (Tan et al. 2023b al.
Section G introduces a method reducing the dimensionality of the mirror maps without affecting the ex- pressivity, by exploiting the symmetrrie in the considered functions, with applications to training mirror maps on neural networks. Moreover, the lower dimensionality allows us to consider parameterizations of mirror maps with exact inverses, bypassing the forward- backwrd inconsistency problem. We utilize equiv- ariance to extend simple version of LMD to train deep neural networks in Section 7 where we demonstrate competitive performance with Adam on non-convex problem classes.
2 Learned Accelerated MD
In this section, we first present a mirror descent analog to the Nesterov accelerated gradient descent scheme (Nestervv] 1983). The resulting scheme, named accelerated mirror descent (AMD), can be considered as a discretization of coupled ODE. Krichene et al. (2015) show accelerated convergence rates of the ODE, which translate to discrete-time algorithm after particular discretization. The analysis and convergence rales of AMD are similar to the ODE formulation for Nesterov accelerated gradient descent given in Su el al. 2014hh where both papers show O(1/k²) convergence rate of function value.
By considering the discrete AMD iterations with incxact mirror maps, we then generalize this algorithm to the approximate case, given in Algorithm [ and provide the corresponding convergence rate bound. This can then be applied in the case where the forward and backward mirror maps are given by neural networks, leading to convergence results for learned accelerated mirror descent (LAMD). Our main contribution in this section is Theorem [ which shows convergence of the function value to the minimum, rather than to the minimum plus a constant.
Algorithm Accelerated mirror descent (AMD) with approximate mirror updates
Require: [0) = z(0) = (0), step-sizes l.. parameter 2 3
1: for &e E N do
2: x (x+1)) = (kk) + (1 - (k) with d ===a = F+
g(k+1)- VMA V Ma (5 (x)) kis - (((((()))))
4: a(k+1)- = arg min 7tu (V)((( + R(a ((+1))
end for
Step 3 of Algorithm 1 is the mirror descent step of AMD, where {(k) represents a (computable) approximate mirror descent iteration. The additional step in Step arises as a correction term allow for convergence analysis (Krichene et al. 2015). Here, R is a regularization function where there exists 0 < la S I.n such
6

---

Published in Transactions on Machine Learning Research (05/2024)
that for all x, E X, we have
RR | - d'12 < R(x, x') S 2
LR 2
-
For practical purposes, R can be taken to be the Euclidean distance R(2,z')= - z'1l? in which case in Ln = 1. Taking R of this form reduces Step to the following gradient descent step:
((k+1) = (k+ 7th V (2
Similarly to Tan et al. (2033b we additionally define a sequence g(k) that comprises the true mirror descent updates, applied the intermediate itcrates g(A):
2(k+1) = (VM))- (VMo(2()) kt& v (colht11))
(5)
This additional variable will allow to quantify the approximation error, with which will show the approximate convergence. For ease of notation, will denote the forward mirror potential by v Ma. With this notation, 1² is the convex conjugate, satisfying V = (Vú)-1, The true mirror descent aux- iliary sequence can be written as 2(k+1) = Vip" (Vo(s(()) dي ( We begin with the following lemma, which bounds the diffcrence between consecutive energies to be the sum of a negative term and an approximation error lerm.
Lemma 1. Consider the appronimate AMD ilerations from Algorithm and let g(k) be the enact MD iterates given by Equation 5 Assume v is -smoohh with respect to reference norm || 1| on the dual space, i.e. By (2,3) < L 1 ylla. or equivalently that Vip* is -iischhtt Assume further that there erists < la S LR such that for all x, x, E ex, lR || z'11 s R(x, x') S LR l| r'11. Define the energy E(K) for k 2 0 as follows (where 1-1
2
k2ta-1 (f(S(h))) -5+ + rr V(((z))
E(K
(6)
Assume the slep-size conditions satisfies the following:
2 Libbr and la S La Esg.
Then the difference belween conseculive energies
-(k) (2k + 1- rrkjtx (f(E(ktl))) - (r)
2 tt) (f(E( - -5a) (ib(z(++)))- ( -x²)
r.
Proof. Deferred to Appendix A.1
Lemma gives us a way of telescoping the energy E. To turn this into a bound on the objective values (((()) f", we need the initial energy. The following proposition gives a bound on the energy BU), which when combined with telescoping with Lemma [ will be used to derive bounds on the objective.
Proposition 1. Assume the conditions as in Lemma 1 Then
E(0 (V V+((")) to (f(z!()) ft) +r (VD(E())- +r)
Vo( z(1) x²)
(7)
Proof. The proof relies bounding f(x()) fr which is done identically to that in Krichene et al. 2015
Deferred to Appendix A.2
Putting together Lemma and Proposition we get the following theorem that bounds the objective value of the approximate AMD iterates.

---

Published in Transactions on Machine Learning Research (05/2024)
Theorem 1. Assume the conditions as in Lemma 1 increasing, and that the approzimation error term, given by
Assume also that the step-sizes (tx) are non-
( - -V g( -x²).
(8)
is uniformly bounded by a constant M >20.0 0. If our step-sizes are chosen as tx = O(k for C [0, 1), then we get (hk-11 convergence objective value.
Proof. Summing the expression in Lemmn [ noting the first two terms are non-positive for 2 since la is non-increasing and r 2 3, and noting that the final term bounded by rM,
E(+t1) S E + kr
By definition of É and non-negativity of B also have
[( - -5 S (k: (k+
E(++1)
(50) + kr ().-1).
(9)
(I: + 1)2t
Remark 1. The condition that Equation S is bounded con be interpreted in terms of more natural func- tional bound on the mårror potentials Mo. Mj. Indeed, the LHS is simply (V VM - -VV (ktz/r)V f(zb)))) Therefore, if 1| Ma o VM - Ill is uniformly bounded (over the iterates in the dual space) and the iterates are bounded, by Cauchy-Schwarz, the approzimation error term Equation will also bounded.
This theorem shows that we are actually able to obtain convergence to the function objective minimum by using acceleration, as opposed to the minimum plus a constant given in Tan et al. (2023b). There are two main ideas at play here. One that the acceleration makes the function value decrease fast enough, as seen by the energy given in Equation 6 growing as Okk), while having A2tx factor in front of the objective difference term. The second idea is that the Å factor in Step 2 of Algorithm [ decreases the importance of inconsistency.
the mirror descent step, which thus reduces the effect of the
3 Learned Stochastic MD
Stochastic mirror descent (SMD) arises when we are able to sample gradients G(a .*) such that the expecta- tion of G is equal to subgradient of our convex objective f. For example, in large-scale imaging applications such as CT, computing the forward operator may be expensive, and stochastic gradient may consist of computing subset of the forward operator. Another example would be neural network training over large data sets, where there is insufficient memory to keep track of the gradients of the network parameters over all of the given data. Therefore, we wish to extend the analysis for approximate mirror descent to the case where we only have access to stochastic gradients.
While may be tempting to put the stochastic error into the approximation error terms of the previous analyses, stochastic errors may be unbounded, violating the theorem assumptions. This section will extend approximate MD the case where the approxination error is now the sum of bounded term and zero- mean term with finite second moments. In particular, Theorem 2 shows expected optimality gap bound depending on the stochastic noise based on bounded variance and bounded error assumptions.
We base the analysis of this section on the robust stochastic approximation approach (Nemirovski et al. 2009). We require two additional assumptions in this setting as follows: Assumption 1. We can draw i.i.d. samples (())) (1)),... of random vector E Assumption 2. An oracle G exists for which, given an input (x. E X E, returns a vector G(æ, {) such that g(x) = Ee(G(ar E)] is well defined and satisfies g(x) az(x).
8

---

Published in Transactions on Machine Learning Research (05/2024)
If f can be written expectation f()) = [2 E),, where F(::E) is convex with having finite values in neighborhood of point z, then we can interchange the subgradient with the expectation (ttrassen 1965
0f(z) = F(0,
Stochastic MD thus generalizes MD by replacing the subgradient ff at cach step with this gradient oracle G, which can be writlen as follows
(k+1) Vi2" (Vi(a!*)) t2 G(((())) g(1))
(10)
Observe that under this formulation, we can additionally encode noise component that arises as a result of inexact mirror descent computation. Therefore, we may redefine G(x, as an inexact stochastic oracle as in the introduction, having two components
G( ( g(x) (xx U(x),
where A(: signifies stochastic noise satisfying [+) 0, and (xx) is deterministic error to model the approximation error of computing MD steps. We will use the notation A A(x(k), UR U (xx(k) to signify these values at each iteration.
Algorithm 2 Stochastic mirror descent (SMD) [Nemirovski et. al. 2009)
Require: x0()) X, step-sizes t > 0, stochastic oracle G, random i.i.d. g(k)
for k N do
(k+1)) V1) 1. (x))) g(()))
end for
We first reformulate the MD itcrates into " "generalizod proximal" form. In particular, small modification of the argument in Beck & Teboulle (2003) shows that the SMD iterations given in Equation can be written follows:
(k+1) arg min zex
{ (G(olt) (o0()))(())))))))))))) g(())) + B& (a', 20(1))
xt)
1 4
(11)
This can be written in terms of the proz-mapping P. (B=)) + X, defined as follows:
(v(v) = arg min (1, x²> + (x',
(12x)
x'EX'
(k++)) = Pa( (1
(12b)
The following result gives optimality gap bounds of approximate SMD in the deterministic- deterministi-ppls-stoooaa noise setting. In particular, the expected ergodic average able to attain loss that is constant over the minimum.
Thcorem 2.
Consider the approzimate SMD iterations al generated by Equation 10 Suppose that the stochastic oracle satisfies BlI|G(H, () < o² for all X for some 2 0. Let x² be some minimizer of f.
(U,,xx()) x²) is uniformly bounded by C, then the eapected loss salisfies
()) m4| s Bo( x
σ²
Σ
k t
(13)
2a
f is p-strongly conver and ||:|1 uniformly bounded by C² the expected loss satisfies
0² k
k
()) - sr] < Bz .0)) +
t
(14)
20
2μ
1=0
1=0
9

---

Published in Transactions on Machine Learning Research (05/2024)
In porticular, the ergodic average defined by
2 Σ
nall,
7
2i
x
satisfies respectively
E [5(0)) - |4] s ( x000)
Σ
+
+C,
(15)
2a
1=0
0²
C
E [f())) - 14] < 20(=",2(01)
t
(16)
2a
2p
i=U
Proof. Deferred to Appendix
B
Remark 2. A stronger assumption to condition (1) in Theorem is to assume uniformly bounded ilerales, as well as uniformly bounded deterministic errors U,. In LMD, deterministic errors arise due to mismatches in the learned mirror maps VM (VMa)-+, which can be controlled using a soft penalty.
We observe in Equations 15 and 16 that f2 4 summability (where Σ < +oo but Σ4=+00 = is a sufficient condition to remove the contribution of the first term to the ergodic average error. This is consistent with the empirical observation in Taan et al. (2023a) that extending the learned step-sizes using reciprocal rule ls e/k gives the best convergence results.
4 Learned Accelerated Stochastic MD
In the previous sections, we considered accelerated and stochastic variants of mirror descent and presented results pertaining to the convergence of these algorithms with noisy mirror maps. These two variants improve different parts MD, with acceleration improving the convergence rates, while stochasticity improves the computational dependency on the gradient. Indeed, one of the drawbacks of SMD is the slow convergence rate of O(L/Vk) in expectation, where the constant depends on the Lipschitz constant of f. Acceleration tool to counteract this decrease in convergence rate for SMD has recently been explored, with convergence results such as convergence in high probability (Ene & Nguyen 2022 Lan 2020 and in expected function value (uu et al. 2018 Lan 2012). These approaches decouple the gradient and stochastic noise, resulting in
O(L/k² +o/v convergence rate, where C depends on the Lipschitz constant of V,, and o² is the variance of the stochastic gradient.
We note that is possible to extend Algorithm [ (approximate AMD) to the stochastic case by directly replacing the gradient V with stochastie oracle G, albeit currently without convergence guarantees. In this section, consider another version of approximate accelerated SMD that comes with convergence guarantecs, using a slightly different proof technique.
For our analysis, we follow the accelerated SMD setup in Xu et al. 2018 in particular of Algorithm 2, replicated as below. Suppose that f is differentiable and has Ly- Lipschitz gradient, and let G(z, = ((,( be stochastic gradient satisfying
E (() c(0))] v)(((()), E [IGI? (+) <
For simplicity, we let A; = ((r) ((K)) (x)())) denote the zero mean stochastic component.
An inexact convex conjugate can be modeled using noise term in Step 3. With a noise term U, the step would instcad read
(k-1))
Th
(b) + (((lt)) Th +
(())
(17)
Th +
10

---

Published in Transactions on Machine Learning Research (05/2024)
Algorithm 3 Accelcrated stochastic mirror descent (ASMD) (Xu et al. 2018
Alg. 2)
Require: Inputs (0 E EX (0) = Vib(x((0) E E Ao = 1/2,
1: for & & N do
2:
AR+1 = (}: + 1)(k 2)/2, To = (Ax+1 - As)/Ak; 36 3x= = (k + 1)3/2
(b+1) = r+tl Vi (y + r++1" s!(8)
Th
y ((++1) y(t) de+1-Ag G(zlk+1), ( {((tt))
end for
The corresponding optimality conditions become
A
(B()) = =x xx(+1)=
(x (k+ - xx())) -U(k)
(18)
A&+1 A&
A++1 A ((z(ktt1)) ((((())))),((t1)))))
((++ yy(
(19)
8k
We can perform a similar convergence analysis using the energy function given in Xu al. (2018). Let g()) E[Ax(f(a -[))((( sx Vy (v(*))) We require a bound on the diameter of the optimization domain, stated in the following assumption.
Assumption 3. There exists a constant My such that M. sup,,x'ex (z(x,z').
Theorem 3. Suppose I has T.p-hipschitz gradient, the mirror mop is o--ttrongyy convex, the approzimation error ((k) is bounded, and that the stochastic oracle is otherwise unbiased with bounded second moments. Assume Assumption 3 holds, and suppose further that there exists a constont such that for every iterate (K)), have
(k)),, ),U()) K.
Then the convergence rate of approcimate ASMD is,
(k + 1)3/2
ML + 2ao + 4/ (xx)))
B -(x))((x)] K ( /Ax
Mp.
(20)
As = K + OKK -2 -- k-11/2).
a
Proof. Deferred to Appendix
Theorem thus extends the classical O(k-2 descent, to the approximate case, up to a constant depending on the approximation error. We note that f having bounded gradients is sufficient to satisfy the second assumption that (V)((())) (((()) is bounded, as we have also assumed that (A) is bounded.
2 + k1/) convergence rate for accelerated stochastic mirror
5 Experiments
In this soction, we demonstrate the performance of the proposed algorithms on image reconstruction and various training tasks. As we will be using neural networks model the mirror potentials, we refer to the three proposed algorithms as learned AMD (LAMD), learned SMD (LSMD), and learned ASMD (LASMD) for the algorithms proposed in Sections 2 to respectively with the learned mirror potentials. We will primarily consider GD, Nesterov accelerated GD, and the Adam optimizer as basclines for these tasks, with SGD where stochasticity is applicable. A grid search was used to find the baseline optimizer step-sizes, chosen to minimize the loss at iteration 100. Specific details for the choice of baseline step-sizes can be found in Appendix E. In the image reconstruction experiments, proximal operators are casily available, which allows us compare with learned primal-dual (LPD) method, given by Algorithm 4.4 in Banert et al. (2020). This LPD scheme be interpreted generalization of the PDHG and Douglas Rachford algorithms.
11

---

Published in Transactions on Machine Learning Research (05/2024)
For fairness, we pre-trained mirror maps given in Tan et al. (2023b a), and replicate selected exper- iments from these works to review the convergence behavior as compared to LMD. The mirror maps were trained to minimize a loss function of the form
I(0,0) = Σ [( + sxll(VM3 VM - (=
(21)
where ( are the approximate MD iterations as in Equation and the mirror maps are tied across each iteration. The loss T. is averaged over a family of training objectives J, such as TV-regularized image reconstruction or SVM hinge losses, and the regularization parameter sy is increased as training progresses. The maximum unrolled iterations is set to N 10, and the LMD step-sizes are also learned for these iterations (Tan et al. 2023b
All implementations were done in PyTorch, and training was done on Quadro RTX 6000 GPUs with 24GB of memory Paszke et al. 2019).
5.1 Numerical Considerations and Implementation
For numerical stability, additionally store the dual variables V.)( similarly to ASMD in Algo- rithm 3 This avoids computing V. VMJ, which would be the identity in the exact MD case where M is the convex conjugate of Ma, and appears to be a source of instability. For example, the current LMD scheme (with step-size t) computes
-V x (k+1) = VM -VMo(ll")-t
(22)
We propose to replace this with updates in the dual
(k+1) (k) (V (())
(o)
(23)
In the case where VM (VM))- both schemes correspond exactly to mirror descent. The key dif- ference is that the inconsistency is now within the Vʃ which would heuristically be close to 0 around the minimum, instead of on the entire dual term. To put this in terms of forward-backward error, let ((++1) = VM, (VMaz(*) (VMz - (z)(()))) be the exact mirror update on (*) with mirror potential Mg. The
heward-backwart inconsistencies (.Mo(((tt)) V (a(++))))) for Equation 22 and Equation then be given respectively as
23 would
(o( - (o(+t()) - (Mg(a(++1)) (E(++1)))) (V)(((()) - +VF(V (M)
- (VM3-((VM)) () (vMaa! (V)(())))) for the primal update, (24)
for the dual update.
(25)
By pulling the difference between VMj and (VM into the dual term, we empirically observe much lower borward-backward error, which would correspond to smaller constants and tighter convergence bounds for our presented theorems. We formalize the proposed learned AMD (LAMD) with R(x,,') -|l I/I|?, learned SMD (LSMD), and learned ASMD (LASMD) algorithms with this dual update modification in the following Algorithms to
Since the original pre-trained models only have step-sizes available up to N = 10, we need to choose step-sizes to continue optimizing. We consider running LAMD, LSMD and LASMD for 2000 iterations with various choices of step-size. In particular, we consider three constant step-sizes t2 = ē, well as step-size regimes of the form and 44 = c/k and 4 e//V/ for LAMD and LSMD. These step-sizes were derived from step-sizes given from the pre-trained models, given for the image denoising, image inpainting and SVM experiments in Tan et al. (2023bla a). In particular, for the constant step-size extensions, we consider taking ē to be the mean, minimum and final learned step-size, i.e. Ei 4/10, minigesio 4, 410 respectively. For
12

---

Published in Transactions on Machine Learning Research (05/2024)
the reciprocal step-size extensions c/k, compute e by taking the average E it, For root- reciprocal extensions 1ę = d/jk, we similarly take Di= iy. These extensions act as best-fit constants for the given learned step-sizes.
10
Recall that the conditions of convergence for AMD and SMD desire non-increasing step-size condition. Morcover, convergence of the ergodic average for SMD requires bz-41 summability condition, which is satisfied by the 4 = c/k extension. We will demonstrate the behavior of these algorithms under both constant, reciprocal and root-reciprocal step-size regimes.
We note that while ASMD does have a step-size in its definition, we can artificially introduce step-size by adding step-size in the dual update step in Step of Algorithm 3
y (k++)) = y - 41 A++1 - A ((x
((((()))) g(k+1)).
In the case where tg = t is constant, this can be interpreted instead running ASMD without this step-size modification on the scaled function if. Therefore, we still get convergence guarantees, up to a constant multiple factor.
Algorithm Learned AMD (LAMD)
Require: Input f:()) = =(0) x((0) X, parameter r 2 step-sizes tx, number of iterations K (9) (2(0)
for k = 0, .... K do
3: (b++)) = λ v MM (2 (1) + +1 - -x) with d
((++1) (k) ktg (c((tt))
5: g(++1) = (b+1) - 7/e V (c(((t1)
6: end for
7: return (K+1) AK (2) + (1- (K)
Algorithm 5 Learned SMD (LSMD)
Require: Input ((()) E X, step-sizes tk, number of iterations K, stochastic gradient oracle 1: y (0) (o) for k= 0, .... K do
(k+1) ((h) (v)
end for
5: return (K+1)= (M((K())1))
Algorithm 6 Learned ASMD (LASMD)
Require: Input (0) E X, Ag = 1/2, step-sizes tk, number of iterations K
1: (0) (o)
2: for h: = O,... K do
3:
AR+1 = (A: + 1)(k + 22, 2, To = (Ax41 - Ax)/As, 35 3;= = (k + 1)3/2
(k++)) xx+1 VM (y (k) + r++1 x(K)
((++1) -y(b) to
A (kx(1++1) e(k)+))))))
2k
end for
return (K+1) тx+1 VM + xk)
5.2 Image Denoising
We consider image denoising on an ellipse data set, with images of size 128 128, where we apply LAMD, LSMD and LASMD with pre-trained mirror maps given in Tan et al. (023a)). The mirror maps V
13

---

Published in Transactions on Machine Learning Research (05/2024)
(a) Clean phantom
(b) Noisy phantom
(c) True TV reconstruction
Figure 1: Example of the ellipse phantom data, well the reconstruction gained by using gradient descent
the TV function.
X + X², VMG X + X' are modelled using input-convex neural networks (ICNNs) (Amos et al. 2017 and are trained to minimize the function values while also minimizing the borward-backung inconsislency, which is defined to be (VM3 VM - (=))))) The training data used for the pre-trained mirror map are noisy ellipse phantoms in X-ray CT, generated using the Deep Inversion Validation Library (DIVal) (Wang & Zhou 2006 Leuschner et al. 2019). The phantoms were generated in the same manner as in Tan et al. (023a)). The target functions to optimize were given by TV model-based denoising,
F [=((()) +
where y are noisy phantoms, and Vr is the pixcl-wisc wise discrete gradient, with the regularization parameter X 0.15.
Figure ! presents some of the noisy ellipse data that we aim to reconstruct using TV model. We observe that the noise artifacts are non- non-Gaussian and have some ray-like artifacts. We model the stochasticity by artificially adding Gaussian noise to the gradient G(v, - g(x) = (xx, ~ NN0, Inxa), with = 0.05.
Figure 2 demonstrates the effect of extending the learned step-sizes for LMD, LAMD, LSMD and LASMD, compared the baseline optimization algorithms GD, Nesterov accelerated GD, and Adam, and LPD. We observe that step-sizes following 4 c/k perform well for L.MD and LSMD, which is consistent with observations made in Tan et al. (023a),, Theorem 2 For LASMD, observe that a constant step- size extension choice is optimal. We also observe the poor performance of Adam, for which various non- convergence results are known for convex settings (Reddi et al., 2018 Zou et al.) 2011)). The plateauing phenomenon with the baselines can not be avoided in this setting, with smaller step-sizes delaying the "decrease phase" and lowering the eventual plateau. Moreover, LPD also plateaus, possibly due to non- generalization of the learned method higher iteration counts. We will use the aforementioned optimal step-size extensions as found in Figure 2 to extend step-sizes for the following experiments.
5.3 Image Inpainting
We additionally consider the image inpainting setting in Tan al. (003b).. The images reconstruct are STL-10 images, which are x 96 color images. The images are corrupted with a fixed 20% missing pixel mask Z, then 5% Gaussian noise is added to all color channels to get. the noisy image. The function class optimize over is the class of TV-reullarized variational forms, with regularization parameter À = 0.15 chosen in Tan et al. (2023b), given by
F (f(x) 12 (o v)113 + \||V>)11
where y images corrupted with missing pixels. We use mirror maps that are pre-trained with the training regime given in 'Tan et al. (2023b), and mcasure their performance when applied with LAMD, LSMD and LASMD as compared to MD.
We first note that this setting does not admit a stochastic interpretation, and thus we set the artificial noise in LSMD and LASMD to zero. Notice that this casc, LSMD updates revert to mirror descent updates. The main difference between the LSMD that we will present here with the LMD used in Tan 2023b is that in this work's computations, the dual iterates saved, described at the start of Section
14

---

Published in Transactions on Machine Learning Research (05/2024)
Reconstruction
Reconstruction
104
10a
100
10f
10
Maan
101
Mean
Mi
Final
2e1k
2ck
62K
10-4
10-1
10
10
10
10a
107
Iteration
Thee
(a) LAMD
(b) LSMD
10a
10a
Rnal
Adam
2X
LPD
LAND
LSMD
107
60k
10
LASMO
100
10-3
10-1
10P
Itzration
Iteration
(c) LASMD.
(d) Comparison with baselines.
Figure 2: Evolution of losses for the proposed LAMD, LSMD and LASMD methods with various step-size extensions (a-c), as well as baselines, evaluated over 2000 iterations for TV model-based denoising of noisy ellipse phantoms. We observe that reciprocal stop sizes generally work well for LAMD and LSMD, with larger step-sizes (given by the "mean" plot) resulting in non-convergence minima. Constant step-sizes work better for LASMD, which can be attributed the already decaying step-sizes in LASMD, We observe that the learned MD methods exhibit similar convergence behavior, which may be due to the relative strong convexity of the underlying optimization problem.
15

---

Published in Transactions on Machine Learning Research (05/2024)
Reconstruction Loss
GD
Nesterov
10a
LAMD
LSMD
LASMD
10a
LMD
fxx
10a
10t
10a
10a
Iteration
Figure Plot of the function optimality gap [z(()) - fr for 2000 iterations various MD methods for image inpainting. The step-size regime is reciprocal for LAMD, LSMD and LMD, and constant for LASMD. We observe that each of LAMD, and LASMD both exhibit approximately O(k-2) convergnce. LMD and LSMD exhibit approximately O(k-1) convergence, with instability around iteration 100 for LMD causing the increase in function value. As are not applying added stochastic noise in this setting, the only difference between LSMD and LMD is that the dual variables are stored as described in Section which allows the LSMD iterates to stay bounded.
For comparison with basclines, we apply the proposed methods to optimize function that arises from TV- based variational model for a single noisy masked image from the STL-10 test data set. The step-sizes were chosen by considering the best step-size extension out of those described at the start of Section 5 which were reciprocal step-size extensions for LAMD, LSMD and LMD, and constant step-size for LASMD given by the mcan of the learned step-sizes. To compute the global minimum of the optimization problem, we run gradient descent for 15000 iterations with step-size of 5 x 10-4, followed by another 5000 iterations with step-size of 1 10-4
Figure 3 demonstrates approximately that LAMD and LASMD are able to achieve O(k convergence. Moreover, we see the effect of storing the dual iterates instead of the primal iterates, as the LSMD con- verges without the divergont bchavior of LMD. Morcover, the convergence rate of LAMD is faster than is suggested by Theorem for step-size 4 = O(k This suggests that the approximation error term given in Equation
[g(++1) x-, ( Ve((((+t))).
is decaying well. This could be attributed to the iterate approaching the global minimum such that the first. term g(k+1) xx is decaying. We additionally observe the speedup of LAMD compared to Nesterov accelerated GD and Adam in the earlier iterates, coming from the learned mirror maps, as well as better generalization performance to higher iterations as compared to LPD.
5.4 SVM Training
To demonstrate an application of stochasticity, we consider training a support vector machine (SVM) over a dataset with many samples. Similarly to Tan et al. (202bb), the task classify the digits 4 and from the MNIST data set. 5-layer neural network was used compute feature vectors Φ; e I0 from the MNIST images for classification. The SVM training problem in this to find weights R50 and bias b such
16

---

Published in Transactions on Machine Learning Research (05/2024)
Hinge Loss
classification enror
104
107
10
10:14
GD
60
Restarov
hastarev
Adam
10
LASND
UMD
10
LND
LASMD
560
SGD
10a4
Reration
Barmatio
(a) SVM hinge loss
(b) Test prediction error
Figure 4: Plots the SVM hinge loss and test accuracy for the varions optimization methods tested. We observe that LAMD able to outperform the baselines and SGD terms of hinge loss, without the plateauing bchavior of Nesterov accelerated GD. We additionally observe that the end fluctuations of LSMD are smaller due the reciprocal step-size regime chosen. However, SGD able obtain solutions that have lower classification error.
that for given set of features and targels (øi, M) R {+1}, the prediction sign(w" di + b) matches Ys for most samples. This can be reformulated into a variational problem as follows, where C' >> 0 is some positive constant,
min
w
w
m0, 10,1- Vi(w 6, +b).
(26)
w.b 2
iE
Using the pre-trained neural network, train SVMs using the training fold of MNIST using the 4 and 9 classes, which consists of 11791 images. For lesting the generalization accuracy of these methods, use the 4 and 9 classes in the testing fold, which consists of 1991 images. For the stochastic methods, we consider using a batch-size of 500 random samples from the testing fold, and the same for stochastic gradient descent. To train the SVMs, full-batch gradient is used for LAMD, LMD, GD and Nesterov accelerated GD. Batched gradients with batch-size 500 were used for LSMD, LASMD and SGD. We consider the same number of optimization for both the full-batch methods and the stochastic variants. This is slightly different from typical learned schemes where if the data is batched, then more training iterations have to be taken to keep the number of "epochs" the same.
The optimization problem and initialization was taken to be the same as in Tan et al. (002bb where 100 SVMs are initialized using (w, ~ No 15041). Each SVM has its parameters optimized with respect to the hinge loss (26 using the various optimization methods. As the previous section, we consider running each of the methods with all the step-size choices, and present the best choice for comparison. We note that the LPD scheme is applicable in this setting as the sum in 26 does not admit an easily computable proximal term.
Figure compares the losses on the entire training set, well as the generalization error of the computed SVMs on the test set. The step-size regimes chosen are reciprocal for LAMD, LSMD and LMD, and constant for LASMD and SGD. LSMD is able to perform on par with full-batch GD despite only having access stochastic gradients, until eventually platcauing around the same loss as SGD. We observe significant acceleration of LAMD compared to all the compared methods. LAMD also has almost identical final test classification error to the baseline methods.
17

---

Published in Transactions on Machine Learning Research (05/2024)
6 Equivariant MD for NNs
In the previous sections, we demonstrated some applications of LMD for convex optimization. However, the mirror map models used still require hundreds of thousands of parameters. Compared to classical optimizers such as SGD or Adam with less than 10 hyperparameters, this is a significant limitation in scaling LMD to complex problems such as training neural networks. In this section, we propose to remedy this by reducing the number of learnable parameters for LMD while retaining expressiveness. This is done by exploiting symmetries of the problems at hand, sometimes known as equivariance in the litcrature.
Equivariance is a group-theoretic concept, where a group action commutes with a function application. In practical terms, symmetries can be exploited to reduce problem complexity. For example, translations and rotations of medical histopathology or electron microscopy images can also be considered as images from the same distribution. [Bekkers et. al. (2018) exploit this to construet convolutional neural networks exploiting this symmetry to achieve improved performance various medical imaging segmentation tasks.
In the context of L20, we aim to exploit symmetrics to construct simpler or faster optimizers. Equivariance for neural nel works have recently been explored, from both an external viewpoint of optimization or training losses, as well as the internal viewpoint of learned filters. Lenc & Vedaldi (2015) consider using equivariance to interpret the filters learned in a convolutional neural network. CCen et al. (2021) consider using equivariance in the unsupervised learning setting, using known equivariance properties to construct virtual operators and drastically increase the reconstruction quality. A major field is the usage of equivariant neural networks that enforce some type of group symmetry by using specific network architectures and activations, used in inverse problems and graph learning problems (Celledoni et al. 2021 Cohen & Welling 2016 Worrall et al., 2017). By strictly enforcing symmetry, the resulting networks are more measurement seeasurement-coonisstene and have better reconstruction quality. We focus on the use of equivariance for dimensionality reduction, provably demonstrating "weight-tying" when using LMD for neural networks.
To simplify the process of training LMD for neural networks, we consider exploiting the architecture of neural network. In particular, the weights of neural network typically carry permutation symmctries, which are exploited in works such as the Neural Tangent Kernel or Neural Network Gaussian Processes (Jacot et al. 2018 Lee et al. 2018). For example, neural network with a dense layer and layer-wise activations between two feature layers will be invariant to permutations of the weights and feature indices. Another example is convolutional networks, where permuting the kernels will result in a permuted feature layer. In this section, we aim to formalize equivariance in the context. of neural networks, and eventually show that we can MD mothods that are equivariant stay equivariant along their evolution. This means that we can replace parameterizations of MD with equivariant parameterizations, redueing the number of learnable parameters while retaining the expressivity.
Let us first formalize the notion of symmetry that we aim for equivariance with. Let G be group of symmetries acting linearly on a real Hilbert parameter space (Z, (s))). In other words, G acts on Z via the
binary operator : G x Z + Z, and moreover, for any z E Z, x E R. we have
2. 21, 9 (Az) = >g 2.
Assune further that G has an adjoint in the inner product, Lee. (g 2, ) = (=.g- One example is (g-2,=))
when Z = =R" equipped with the standard inner product, and G is a subset of the orthogonal group O(n) (the group of n x n orthonormal matrices) with the natural action (matrix multiplication).
Let L 2 + IR be loss function for our parameters. Suppose (0)) is initialized with some distribution (0()) ~ (Z,p) that is stable under G, i.., p()("))
(0)) Vg E G, and further that L is stable under the permutations L(g - I)) [xx). For example, for feature vector x E R", the loss function L(x)= = ||æ|| and initialization distribution (0, I,) are invariant under the orthogonal group G = 0(n). The loss functions
[(x)= = ||0||, || and any component-wise i.i.d. distribution are invariant under the permutation group G = Sym(n..
Recall that the objective of LMD is to minimize the loss after some number of optimization iterations Tan
et al.
2023b). To this end, let Ω Z - Z be optimization iteration. For example, gradient descent on L.
18

---

Published in Transactions on Machine Learning Research (05/2024)
can be written R(2)) (2)...... The objective of LMD can be written, for example,
min E 12
((z))
Suppose that we optimize Q using some sort of gradient descent (in, say, the Sobolev space (2) of 1?² functions with L derivative), and further that fg satisfies
Rolg 2) - g [Ro(=)], GgEG,
This reads that Go G-invariant under conjugation, where the conjugation action of G' on functions F 2 - 2 is (g (=)= = 9 (F(g-1.2)). We equivalently denote this G-equivariance, as g commutes with Gg. For sanity check, consider the case where Go is gradient descent with on the objective function L. The following proposition says that Sg is indeed G-equivariant under these conditions.
Proposition 2. no is gradient descent the G-invuriant objective function L, then it is G-equivarianl, i.e. Go(g z) = [=)( for all E G. Z.
Proof. Relies on the Riesz representation theorem and adjoint property. Deferred to Appendix D.1
Suppose first that N = that we want to minimize (initialized at G-invariant Ω = Qo)
(()) = Exo [Mz1)] = Ezo~p [L(2(z))))
(27)
We wish to show that after optimization, Ω commutes with y (i.e. Ω stays invariant under the conjugation
action of G on functions Z + 2). To do this, we observe that the discrete differences computed by
autograd are invariant under the group action of g. Assume further No E ((2,,,p) where is the initialization probability measure. The corresponding inner product
(0, (0, = E. ((())) (()))
G naturally acts on (2, p.p) by the conjugation action
(g ())2)) g- (
To check that this conjugation is an action, first note that the action by is isotropic since G S O(n):
/ 110e (d)))))))()
/ [1(g (dp)(()))(2
/ llg (2(( (2))
/ [dd=))))pp==).
Here, used invariance of g under p, and also that [1g y1 (g v.g v) = (v. |lyl from the adjoint. The required composition and identity properties of the action are clear. Moreover, the action is linear, and we have the adjoint property
(g,
2))
/ (100 - 2),9 -(z))dp(z)
(((=)) 'B(g d))
(0,g-1 .)..
19

---

Published in Transactions on Machine Learning Research (05/2024)
Assume that fg is G-equivariant. Then the next iterate satisfies the following
dC
Ω S2 - -1 =m (Ro).
(28)
But we also have for any and 52,
(()) Eza [[[((z)))
Eo [[g (2())))
since L(z) Lig 2) since p(2) = p(g 2)
Ezompl [L[g E.lg.900 22g 20))
Ezoppl [b[(g (2)
Lig 9)..
Hence if Qo 920 for all g, then
d ((9x) - d(g- 22)
d
((sg)
(29)
dn
Since the variation of the Ω and g is the same at Ωo, we have that g 22 = Q1 i.. Qn is G-equivariant. This argument can be repeated to get that On G-equivariant (under conjugation) for all n E N. This is summarized in the following proposition.
Proposition 3. Let (2,()))) be Hillert parameter space. Suppose the following hold:
G a group acting linearly on 2 with an adjoint;
2. The loss function L Z + R and are stoble under G, that is, Lig 2) L(z) for any g G and
E Z;
The initialization distribution 20) ~ (2,p) is stable under G, that is, the laaus p(z(0)) and plg 20)) coincide for any g E G.
Let Qg Z > Z be an optimization iteration that G-equivariant, such that Dolg 2) g -2(2) for all 2,9 E G. So is evolved using gradient step (with slep-size 11) to minimize
L(l2) (zz~p[[(Zz))
then the evolved oplimizer Ω1 = So (ddf/d((Ss))) is also G-equivurianl.
6.1 Application: Deep Neural Networks
For the sake of exposition with a concrete example, let us consider a two-hidden-layer dense neural network of the form
N(; Asoz(A2o + b2) + by) + by
where Rdo di coordinate-wise or otherwise permutation-invariae softmax, and A, - b, E Idd: are the weights and biases. In this case, the parameter space will be
activation functions such
3
Z = Rdd
Σ
(d. + d:).
Let us consider lease-squares regression problem with data (2,, j 1, 1,... k which can be written minimizing the loss
L(x) Σ [N(2)22) 211
Suppose that the weights are initialized entry-wise iid. for cach A,, b,, such as in common initializa- tions where A NN of Ta. 1d 1dxx,-1d A Uniform(-d,t) . -1/2). Consider the symmetry group
20

---

Published in Transactions on Machine Learning Research (05/2024)
G Sym(d1) x Sym(d22, consisting of permutations of the intermediate feature maps with appropriate permutations on the weights and biases, i.e. of permutations Pr. P2 on [da] respectively with the group action as
(z) (z)[A1lpsg - (A1)
9ps p2
9p P2 (z)[A2lps = (A2)p1(0)oaal))
9p ps (=)
-(A3)pa()).o
9ps P2 (z) (2)[b1l (b1)pr(a)
pp (z) (b2ja y (b2))2(9))
P2 (z) (bagg = (bz)q.
We can check directly that this group action satisfies the requirements of Proposition 3
(Linearily) Permutation of elements linear.
(Initialization Invariance) Follows from ii.d. property of entries of A,, b,.
(Loss Equivariance) Follows from o; acting coordinatewise and invariance of suns under permula- tion.
Consider the simple approach of scaling each component of the gradient by learned constant, i.e. MD with mirror map of the form &b x Dz for some learned positive definite diagonal D. Notice that starting with D = I yields gradient descent, which is equivariant under G. Using Proposition 3 we find that D will continue to be G-equivariant. In other words, each of the components of D must be fixed under permutation actions of the form (p1, G.
We can explicitly demonstrate this experimentally on a small-scale neural network. In particular, we consider training a classifier using three-layer neural network with fully connected layers on the 2D moons data set. We arbitrarily choose 50 dimensions for the middle feature layer. The neural network takes the form
N(x; A, b.) o2 (Ai bi) b2),
where A; Rd.xd,-,, bi Rd,, do 2, 50, d2 1. In this case, we choose ai to be ReLU activation and 02 to be log-softmax. We train using the binary cross-entropy loss. Note that in the case of a log-softmax, which takes the form
exp(xi) Σn exp(xx)
LogSoftmax(x): = log
the activation still equivariant under permutations G, as the sum of the exponentials stays the same. Therefore, loss equivariance holds and the above theory given by Proposition 3 appllies.
Using the diagonal quadratic LMD introduced above, we thus expect from the equivariance theory that the diagonal weights in 1 will be invariant under the effect of any permutations (pi) on [di]. Using superscripts to denote the weights corresponding to each learnable neural network component, we thus have that for any
indices i E [di] = j. E [d2] = [50], k E [dz] = [1]:
(D^L) )is = (D^1)Es
(D'1); (D'1)
(D (DA2)
Therefore, from G-oquivariance of LMD with diagonal weighting D. we have that the weights constant across the orbits of Pl, and thus fully determined by the weights
21

---

Published in Transactions on Machine Learning Research (05/2024)
Componentwise weighting
A1
25
b1
A2
20
15
10
5
Figure Componentwise gradient weightings learned for two layer neural network on 2D moons data. Each band consists of 50 datapoints, corresponding to elements of the weight matrix D for each neural network parameter group. We observe strong grouping of the weightings within each parameter group. In particular, we find that the weightings of A1 are split into two groups, which correspond to the two input dimensions.
(DA)),1, (DAI) (D (DAt)) ),1, (D This reduces the number of variables that need to be learned from dim(D) = 201 to only 5. The equivariance effect can be clearly seen in Figure 5 which empirically demonstrates what happens if we train all the variables of D without directly enforcing equivariance. By initializing D to be the identily, the initial mapping Qg Z + Z induced by D is gradient descent, which we have shown before to be equivariant under permutations of the intermediate feature layer. The theory presented thus suggests that the optimization path of Ω (or in this case, of D), will continue to be equivariant under such permutations. Indeed, at the end of the optimization for D, we observe that the weights do indeed satisfy approximate equivariance, of the weights. We observe that the weights D become two groups, and Db, D becoming one group each, which supports the theory that if the initial map Op is equivariant, then so are subsequent maps after training.
demonstrated by the tight grouping phenomenon
We note that this analysis can be used to partially explain the existing houristic method of applying mirror maps to neural networks, such as using p-norms vlww) || (Axizan et al. 2021]; ; squared p-norms (w)) [|ll| (D'Orazio et al. 2021 Gunasekar et al. 2018), or element-wise maps induced by projections such hyperbolic tangent softmax function (Ajanthan et al. 2021). For these maps, when considering the weights as vectors, permuting components does not change the value of the mirror map (though general orthogonal transformations will). This permutation symmetry over all parameters can be thought of as a special case of the presented theory which suggests only equivariance within a layer.
7 Experiments: Neural Network Training using Equivariance
We consider the problem of training LMD on neural networks, based on the group equivariance demonstrated in Section 6 This is a challenging non-convex optinization lask, where the approximate theory for LMD (and its extensions) does not apply, due to significant errors that may occur from small perturbations in the parameters. This motivates us to consider simpler yet exact mirror maps to learn, which can be done layer-wise by the analysis in Section 0 We demonstrate competitive performance with Adam, one of the most widely used optimizers for neural network training (Kingma & Ba 2015
22

---

Published in Transactions on Machine Learning Research (05/2024)
7.1 Deep Fully Connected Neural Network
We first consider training classifier neural network on the MNIST image dataset. Neural networks are usually trained using minibatched gradients, as computing gradients over the whole dataset is infeasible. This is a prime use case for the stochastic variants of LMD. The neural network to train as well as the LSMD parameterization are detailed as follows.
Algorithm 7 Learned AMD with minibatched gradients (LAMD*)
Require: Input x(0) = g(0) = x0) E X, parameter r 2 3, step-sizes tj, number of iterations K
1: (0)) = V (z(0))
2: for k = 0, K do
3: x (b+1) = dR V M3 + (1 Xu Ax) with λu = r+E
z(++1) = 2(k) He V (((
x(k+1) ((++1) ttp f (2((tt)))
end for return ((++1) AK (Mz(z(k)) + (1 AK)
Neural network architecture. The number of features in each layer of the neural network is 784 500500000 40- 30-20-10, with dense lincar layers with bias between the layers and ReLU activations, for total of 43350 parameters. The MNIST images are first flattened (giving 28 28 784 features), before being fed into the network. The function objective corresponding to a neural network is the cross-entropy loss for classification over the MNIST dataset. For LSMD, we associate each layer with two splines - one for the weight matrix and one for the bias vector. This results in a total of 10 learned splines, a total of 400 learnable paramoters for T.SMD.
LMD parameterization. Utilizing the reduced parameterization as suggested by Section 6 we can con- sider only layer-wise mirror maps V. This allows to use more complicated mirror maps from R + R. One parameterization is to use monotonic splines as mirror maps, which are derivatives of convex mirror potentials Vτ = a. We directly model the mirror map using a knot-based piecewise-linear spline param- eterization as Goujon et al. (2022), where knots are defined at pre-determined locations, and the spline is (uniquely) determined by the value of a at each knot. We note that since the monotonic spline o R - R is continuous, the corresponding mirror potential given by the integral (xx) = fo att) dt is C.. For these experiments, we use 41 knots equispaced on the interval [-1, with gap 0.05, which can be parameterized using 40 parameters. More details the construction of a can be found in Appendix G
The goal of LSMD is to train neural networks with architecture faster, given MNIST data. The function class that we wish to optimize over is thus
F C(N)) neural networks No
(30)
where L is the cross-entropy of neural network (on an image mini-batch), and Ne is the neural network after evolving the parameters through LSMD for k iterations. Note that the number of iterations that we can unroll is significantly higher than in the image processing examples, due to the dimensionality reduction from the equivariant theory. The minibatching process means that we are training using SMD instead of the MD framework, similarly to how SGD is used instcad of GD to train neural networks on large datasets. We are allowed use three epochs of MNIST to train, which we find to be sufficient for these experiments. The MNIST batch-size is taken to be 500, resulting in unrolling N 360 steps. The neural networks No are initialized using the default PyTorch initialization, where dense layer from Cin -Cut features has the entries of its weight matrix A and bias matrix b initialized from Uniform(-C,,, C. We use 20 such neural network initializations as a batch for each LMD training iteration, and average the cross-entropy over each neural network as our loss function.
Figure 0 shows the performance of LSMD on this task, as well as minibatched LAMD* Algorithm 7 applied with the spline-based mirror maps learned using the LSMD framework. LMD and LAMD* are compared
23

---

Published in Transactions on Machine Learning Research (05/2024)
Training negative log-lkelhoad
accuracy
140
120
3.6-
120-
SGD
2.5
20
Acam
Aetmm SpITHALAND
SARNLLAND
SpiinoLM
3.2-
15:0
230
200
300
(a) Cross-entropy loss
(b) Train accuracy
Figure 6: Spline-based LMD applied to training a 4-hidden-layer dense neural network for MNIST classifi- cation. We observe significant acceleration of LAMD as compared to LMD. There is a slight performance gap between LAMD and Adam on this task.
against: SGD with learning rate 10-1, and Adam with learning rate x 10-2, where other parameters are kept as default. The values are averaged over 20 neural network initializations. We observe that Adam with this learning rate is able to train this network architecture very rapidly. LAMD* is able to almost match the performance of Adam at the end of the 3 training epochs, while both LAMD* and LMD significantly outperform SGD.
7.2 Convolutional Neural Network
We additionally consider LMD to train a 2-layer convolutional neural network. We impose a diagonal quadratic LMD prior on the weights, similarly to Section 6.1 The CNN is taken to have a single convolutional layer with 8 channels and kernel size 3, a max pooling layer with window size 2, followed by a dense layer from 288 features 10 output features. The inputs of the CNN are MNIST images that have been downscaled to have dimension 14 x 14. The target is to minimize the cross-entropy loss of the CNN after 2 epochs of MNIST, with mini-batch size 1000, resulting in total unrolling count of N = 120.
The equivariance here is due to the permutation invariance of each of the eight 3 x 3 kernels. Therefore, we can reduce the number of parameters of LMD by tying the diagonal weights across the channels, reducing the number of parameters for this layer from x (9+1) = 72 to only 10, with 9 from the sliding window weights and from the bias term. While not strictly induced by the training problemt we additionally impose equivariance on the dense layer by tying the 288 features together, reducing the number of parameters from 288 10 10 to 1 10 10. In total, learn 30 parameters.
Figure 7 shows the performance of LMD and LAMD* when averaged over training 20 randomly initialized CNN problems, trained for 2 epochs of MNIST. We observe competitive performance with Adam for both LAMD and LMD. Interestingly, we obscrve that the testing accuracy of the training result of LAMD* is marginally higher than that of Adam. This distinctly highlights the potential of learning the geometry for gradient- based methods, as we achieve significant acceleration with only 30 parameters to the point of being competitive with Adam, while retaining principled approach to learning the model.
We note that the size of the CNN is limited due to the memory constraints of LMD. To train one iteration of LMD, all the intermediate activations of the neural network need to be stored. In the case of convolutional layer, the intermediate layer is also mike,-like, which results in very high memory usage.
7.3 Few-Shot Learning
Since the layer before dense layer convolutional layer (followed by max-poo pooling)) the features this intermediate layer still contain local information. Therefore, do have permutation invariance of the features.
24

---

Published in Transactions on Machine Learning Research (05/2024)
Training regative lngllkelhoad
accuracy
3.6-
2.5
563
AdNm
Adam
3.2-
(a) Cross-entropy loss
(b) Test accuracy
Figure Diagonal LMD and LAMD applied to training one-hidden-laye layer convolutional neural networks with 2 epochs of MNIST, averaged over 20 runs. We observe that LAMD can achieve higher test accuracy and lower training loss by the end of the training epochs. Both LMD and LAMD are competitive with Adam in this example.
To further compare the proposed LMD methods non-convex problems, consider the problem of few-shot learning, where classifier is trained on limited training data. We utilize the Omniglot dataset, consisting of 20 instances each of 1623 characters (Lake et al. 2011). The characters are split into training and testing sets of 1200 and 423 characters respectively as in Finn et al. (2017). We consider a 5-shot 10-way problem to utilize the small CNN architecture as detailed in the previous section. In other words, a problem training instance consists of choosing 10 characters from the respective training or testing set, training the network to classify instances of each character, and testing on the other 15 instances of the character.
As baseline meta-learning approach, consider the Reptile method (Nichol et al. 2018), which finds good neural network initialization such that a standard optimization algorithm such as SGD can minimize the loss quickly. Further details on how the Reptile baseline is trained can be found in Appendix H
The L20 approach of LMD instead considers a different approach of fast optimization by changing the optimization trajectory, allowing for random neural network initializations. We consider training LMD on the training character set in two settings: one with standard random initialization, and one from the initialization learned using Reptile. These methods then compared with SGD with learning rate 10-1 and Adam 10-2 as before, with random initializations and Reptile initializations. We additionally consider LAMD with the same learned parameters from both random initialization and from the Reptile initialization. The methods are evaluated on 10 randomly sampled 5-shot 10-way problems from characters drawn from the test split, with gradients caleulated in a full-batch manner.
Figure contains comparison of the equivariant LMD and LAMD methods with SGD and Adam, with ran- dom intializations and from the Reptile initialization. We note that LMD outperforms its SGD counterpart in both scenarios. Morcover, LAMD is able to achieve a significantly lower log-likelihood, outperforming Adain with random network initialization well as with Reptile initializations. Nonetheless, there is still gap between the LMD methods and Adam in terms of generalization performance, suggesting that the trajectory than LMD takes trades off generalization performance for faster training. Interestingly, LAMD with random initializations outperforms LAMD with the Reptile initialization, suggesting that random ini- tializations plays a role in learning the geometry over a wider set of trajectories.
8 Conclusions
In this work, we present multiple extensions of the learned mirror descent method introduced a recent work on learning optimize (Tan et al.) (2023b/a). In particular, we introduce accelerated and stochastic versions, well training paradigm that reduces the number of learnable parameters while maintaining
25

---

Published in Transactions on Machine Learning Research (05/2024)
Training negative lngllkelhoad
accuracy
2.3
2.5
20-2
5G0 MdN
5SD
Adam
20-1
0.3
LM3
9.2
Seprile-AAN
Rezti m4AM
103
(a) Cross-entropy loss
(b) Test accuracy
Figure 8: Training loss and test accuracy for 5-shot 10-way learning on the Omniglot dataset. We observe that LMD with Reptile initialization able to achieve similar test accuracy to the compared methods. Moreover, LAMD is able to outperform Adam in terms of training loss, but seems to converge to a solution that generalizes worse.
the expressivity of the mirror maps by exploiting symmetries. Numerical experiments demonstrate that the learned accelerated and stochastic MD methods are able to outperform their corresponding GD counterparts as well as previous LMD with the same mirror maps, and we empirically improve on the constant optimization gap by utilizing the dual domain during mirror steps. These experiments additionally demonstrate that mirror maps trained using LMD as in [Tan et al.] (023b are able to successfully generalize to different LMD-type optimization schemes. We empirically show that mirror maps respect group invariances under training to support our equivariant thoory, and exploit this to train LMD the non-convex problems of training deep and convolutional neural networks.
This work presented results as well as techniques for developing L20 schemes for large-scale applications. Given the various modifications to LMD presented, the mirror maps were pretrained using a basic scheme that penalizes the objective function as well as the forward--backward inconsistency for fixed number of unrolled iterates (Tan al., 2023b). One consideration is the modification of the training scheme instead use the accelerated or stochastic mirror descent rather than full-batch mirror descent, which could help in large-scale applications. Another possible direction would be to consider convergence properties of the approximate MD methods for non-convex functions, similarly to the non-mirrored cases presented in Berahas et al. 2017;; Jin et al. (017)).
This work additionally presented an application of equivariance for dimensionality reduction of LMD, reduc- ing the total number of parameters required to train. We demonstrated that the reduced parameterization arises naturally from common training scenarios, and that the resulting LMD is able to perform compet- itively with Adam. While limit our equivariance directions could consider equivariance under the scope of other optimization algorithms utilizing techniques such scaling, momentum, or using proximal operators. Another direction to continue pushing LMD is to address the current. limitations caused by needing to store intermediate activations, allowing for this method to used for larger and more practical neural networks.
simple gradient-based methods, interesting future
References
Thalaiyasingam Ajanthan, Kartik Gupta, Philip Torr, Richad Hartley, and Puneet Dokania. Mirror descent view for neural network quantization. In International conference artificial intelligence and statistics, pp. 2809 2817. PMLR, 2021.
Zeyuan Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. The Journal of Machine Learning Research, 18(1):8194-822 2017.
26

---

Published in Transactions on Machine Learning Research (05/2024)
Zeyuan Allen-Zhu and Lorenzo Orecchia. Linear coupling: An ultimate unification of gradient and mirror descent. arXie preprint arXiv:1407. 1537, 2014.
Diogo Almeida, Clemens Winter, Jie Tang, and Wojciech Zaremba. A generalizable approach to learning optimizers. arXiv preprint arXiv:2106.000 009588 2021.
Brandon Amos, Lei Xu, and J. Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 146 155. PMLR, 2017.
Brandon Amos et al. Tutorial on amortized optimization. Foundations and TrendsS in Machine Learning, 16(5):502-732, 2023.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. In Advances in neural information processing systems, pp. 3981 3989, 2016.
Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199(1-2) 214, 2023.
Dominique Azé and Jcan-Paul Penot. Uniformly convex and uniformly smooth convex functions. In Annales de la Facullé des sciences de Toulouse: Malhémaliques, volume 4, pp.
1995.
Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent overparameterized nonlinear models. IEEE Transactions on Newral Nelworks und Learning Systems, 33(12)77717-777 2021.
Sebastian Banert, Axel Ringh, Jonas Adler, Johan Karlsson, and Ozan Oktem. Data-driven nonsmooth optimization. SIAM Journal Optimization, 30(1):102- 131, 2020.
Sebastian Banert, Jevgenija Rudzusika, Ozan Öktem, and Jonas Adler. Accelerated forward-backer optimization using deep learning. arXiv preprint arXiv:2105. 05210, 2021.
A. Beck and M. Teboulle. Fast gradient- based algorithms for constrained total variation image denoising and deblurring problems. IEFF Transactions on Image Processing, 18(11):2414) 2434, 2009.
Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-117 175, 2003.
Erik J Bekkers, Maxime W Lafarge, Mitko Veta, Koen AJ Eppenhof, Josien PW Pluim, and Remco Duits, Roto-translation covariant convolutional networks for medical image analysis. In Medical Image Comput- ing and Computer Assisted Intervention-MICC.. 2018: 21st International Conference, Granada, Spain, Seplember 16-20, 2018, Proceedings, Part I, pp. 440 448. Springer, 2018.
Aharon Ben-Tal, Tamar Margalit, and Arkadi Nemirovski. The ordered subsets mirror descent optimization method with applications to tomography. SIAM Journal on Oplamization, 12(1):79- 1089 2001.
Albert S Berahas, Raghu Bollapragada, and Jorge Nocedal. An investigation of newlon-sketch and subsam- pled newton methods. orXiv preprint vrXii:1705. 062111, 2017.
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 8(3-4)2231-35) 2015.
Elena Celledoni, Matthias J Ehrhardt, Christian Etmann, Brynjull Owren, BabolaeBibiane Schönlieb, and Ferdia Sherry. Equivariant neural networks for inverse problems. Inverse Problems, 37(8):085006, 2021.
Antonin Chambolle, Matthias J Ehrhardt, Peter Richtarík, and Carola-Bibiane Schonlieb. Stochastic primal- dual hybrid gradient algorithm with arbitrary sampling and imaging applications. SIAM Journal on Optimization, 28(4):2783 -80888 2018.
Dongdong Chen, Julián Tachella, and Mike E Davies. Equivariant imaging: Learning beyond the range space. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4389,4388, 2021.
27

---

Published in Transactions on Machine Learning Research (05/2024)
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Zhangyang Wang, Howard Heaton, Jialin Liu, and Wotao Yin. Learning to optimize: A primer and a benchmark. The Journal of Machine Learning Research, 23 (1):8562 8620, 2022.
Xinyi Chen and Elad Hazan. A nonstochastic control approach to optimization. arXiv preprint arXiv:2301.07000 2023.
Edward Chlebus. An approximate formula for partial sum of the divergent p-series. Applied Mathematics Letters, 22(5):732-737, 2009.
Тaco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pp. 2990 2999. PMLR, 2016.
Aaron Defazio. A simple practical accelerated method for finite sums. D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems 29, pp- 684. Curran Associates, Inc., 2016.
Giulia Denevi, Dimitris Stamos, Carlo Ciliberto, and Massimiliano Pontil. Online-within-online meta-
learning. Advances in Neural Information Processing Systems, 32, 2019.
Ryan D'Orazio, Nicolas Loizou, Issam Laradji, and Toannis Mitliagkas. Stochastic mirror descent: Con- vergence analysis and adaptive variants via the mirror stochastic polyak stepsize. arXiv preprint arXiv:2110.1544 15412 2021.
Derek Driggs, Junqi Tang, Jingwei Liang, Mike Davies, and Carola-Bibiann Schönlieb. A stochastic proximal alternating minimization for nonsmooth and nonconvex optimization. SIAM Journal on Imaging Sciences, 14(4):1932 1970, 2021.
John C Duchi, Shai Shalev-Shwarde Yoram Singer, and Ambuj Tewari. Composite objective mirror descent. In COLT, volume 10, pp. 14-26. Citeseer, 2010.
John C Duchi, Alekh Agarwal, Mikael Johansson, and Michael I Jordan. Ergodic mirror descent. SIAM Journal on Optimization, (2(4):1549-157) 2012.
Alina Enc and Huy L Nguyen. High probability convergence for accelerated stochastic mirror descent. arXiv preprint arXi::2210.006 0067,,, 2022.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. Tn International conference on machine learning, pp. 1126-1135. PMLR, 2017.
Boyan Gao, Henry Gouk, Hae Beom Lee, and Timothy M Hospedales. Meta mirror descent: Optimiser
learning for fast convergence. arXiv preprint arXim:2203.021) 2022.
Saeed Ghadimi and Guanghui Lan. Accelerated gradient methods for nonconvex nonlinear and stochastic programming. Mathematical Programming, 156(1-2):59 99, 2016.
Alexis Goujon, Sebastian Neumayer, Pakshal Bohra, Stanislas Ducotterd, and Michael Unser. A neural- network-based convex regularizer for image reconstruction. arXiv preprint arXiv:2211.1440 2022.
Karol Gregor and Yann TeCun. Learning fast approximations of sparse coding. Tn Proceedings of the 27th international conference on international conference on machine learning, pp. 399-406, 2010.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pp. 1832 1841. PMLR, 2018.
Suriya Gunasekar, Blake Woodworth, and Nathan Srebro. Mirrorless mirror descent: A natural derivation of mirror descent. In International Conference on Artificial Intelligence and Statistics, pp. 2305-2313. PMLR, 2021.
Rishi Gupta and Tim Roughgarden. A pac approach application-specifs algorithm selection. In Proceed- ings of the 2016 ACM Conference on Innovations in Theoretical Computer Science, pp. 134, 2016.
28

---

Published in Transactions on Machine Learning Research (05/2024)
Vahan Hovhannisyan, Panos Parpas, and Stefanos Zafeiriou. MAGMA: Multilevel accelerated gradient mirror descent. algorithm for large-scale convex composite minimization. SIAM Journal on Imaging Sciences, 9 (4):1829- 1857, 2016.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems, 31, 2018.
Kyong Hwan Jin, Michael T McCann, Emmanuel Froustey, and Michael Unser. Deep convolutional neural network for inverse problems in imaging. IEEE Transactions Image Processing, 26(9):4509-455 2017.
Mikhail Khodak, Maria-Florina Balcan, and Amect S Talwalkar. Adaptive gradient-basol meta-learning methods. Advances in Neural Information Processing Systems, 32, 2019.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. International Conference on Learning Representations, 2015.
Walid Krichene, Alexandre Bayen, and Peter L Bartlett. Accelerated mirror descent in continuous and discrete time. Advances in Neural Information Processing Systems, 28, 2015.
Brenden Lake, Ruslan Salakhutdinov, Jason Gross, and Joshua Tenenbaum. One shot learning of simple visual concepts. In Proceedings of the annual meeting of the cognilive science sociely, volume 33, 2011.
Guanghui Lan. An optimal method for stochastic composite optimization. Mathematical Programming, 133 (1-2):365-3399 2012.
Guanghui Lan. Pirst-order and stochastic oplimvization methods for machine learning, volume Springer, 2020.
Guanghui Lan and Yi Zhou. An optimal randomized incremental gradient method. arXiv preprint arXiv:1507.020 0,0000, 2015.
Guanghui Lan, Arkadi Nemirovski, and Alexander Shapiro. Validation analysis of mirror descent stochastic approximation method. Mathematical programming, 134(2):425-455, 2012.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl- Dickstein. Deep neural networks as Gaussian processes. In International Conference on Learning Repre- sentations, 2018
Karel Lenc and Andrea Vedaldi. Understanding image representations by measuring their equivariance and equivalence. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 991 991-999, 2015.
Johannes Leuschner, Maximilian Schmidt, and David Erzmann. Deep inversion validation library. Software available from Attps://gitthubbcom http:/lww...m//ll/eeeeeee comm/jeuschn//eeee /dival, 2019.
Ke Li and Jitendra Malik. Learning to optimize. In International Conference on Learning Represenlations, 2016.
Yanjun Li, Kai Zhang, Jun Wang, and Sanjiv Kumar. Learning adaptive random features. In Proceedings of the AAAT Conference on Artificial Intelligence, volume 33, pp. 1229-42366 2019.
Luke Metz, Niru Maheswaranathann Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Under- standing and correcting pathologies in the training of learned optimizers. In International Conference Machine Learning, pp. 4556 4565. PMLR, 2019.
Arkadi Nemirovski and David Berkovich Yudin. Problem Complezity and Method Efficiency in Optimization / translated by E.R. Dawson. Interccnne scries in discrete mathematics. Wilcy, Chichester, 1983.
ISBN 0471103454.
Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approxima- tion approach to stochastic programming. SIAM Journal Optimizalion, 19(4):1574-1600 1609, 2009.
29

---

Published in Transactions on Machine Learning Research (05/2024)
Yurii Evgen'en oegen'evich Nesterov. A method of solving convex programming problem with convergence rate O(1/k²) In Doklady Akademäi Nauk, volume 269, pp. 143-547. Russian Academy of Sciences, 1983.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint urXi::1803002090 2018.
Francesco Orabona, Koby Crammer, and Nicolo Cesa-Bianchi. A generalized online mirror descent with applications to classification and regression. Machine Learning, 99:411 435, 2015.
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-performanc deep learning library. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d' D'Alché-Bue Buc, E. Fox, and R. Garnett (cds.), Advances in Neural Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019.
Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent. IEEE Transactions on Information Theory, 61(33:1451-1457 2015.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, 22(3):400-400 407, 1951.
Leonid T Rudin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: nonlinear phenomena, (0(14))2 268, 1992.
Rajiv Sambharya, Georgina Hall, Brandon Amos, and Bartolomeo Stellato. Learning to warm-start fixed- point optimization algorithms. arXie preprint arXiv:2309.002 07835, 2023.
Nati Srebro, Karthik Sridharan, and Ambuj Tewari. On the universality of online mirror descent. Advances in Neural Information Processing Systems, 24, 2011.
Volker Strassen. The existence of probability measures with given marginals. The Annals of Mathematical Slatisties, 36(2):423 439, 1965.
Weijie Su, Stephen Boyd, and Emmanuel Candes. A differential equation for modeling Nesterov's accelerated gradient method: theory and insights. Advances in neural information processing systems, 27, 2014.
Hong Ye Tan, Subhadip Mukherjee, Junqi Tang, Andreas Hauptmann, and Carola-Bibiane Schönlieb. Robust data-driven accelerated mirror descent. In 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023a.
Hong Ye Tan, Subhadip Mukherjoe, Junqi Tang, and Carola-Bibianc Schönlicb. Data-driven mirror descent with input-convex neural networks. SIAM Journal on Mathemalics of Dula Science, 5(2):558 587, 2023b.
Junqi Tang, Mohammad Golbabacc, Francis Bach, and Mike E Davics. Rost-Katyusha: Exploiting the solutions' structure via scheduled restart schemes. In Advances in Neural Information Processing Systems, pp. 427-438. Curran Associates, Inc., 2018.
Junqi Tang, Karen Egiazarian, Mohammad Golbabaee, and Mike Davies. The practicality of stochastic optimization in imaging inverse problems. IEEE Transactions on Computational Imaging, 6:1471-1485, 2020.
Yang Wang and Haomin Zhou. Total variation wavelet-based medical image denoising, International Journal of Biomedical Imaging, 2006, 2006.
Blake E Woodworth and Nati Srebro. Tight complexity bounds for optimizing compasite objectives. In Advances in Neural Information Processing Systems, pp. 2016.
30

---

Published in Transactions on Machine Learning Research (05/2024)
Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetor, and Gabriel J Brostow. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the TEEE conference on computer vision and pattern recognition, pp. 5028 5037, 2017.
Pan Xu, Tianhao Wang, and Quanquan Gu. Accelerated stochastic mirror descent: From contimuous- time dynamics to discrete-time algorithms. In International Conference Artificial Intelligence and Statistics, pp. 1087-1096. PMLR, 2018.
Junjic Yang, Tianlong Chen, Mingkang Zhu, Fengxiang He, Dacheng Tao, Yingbin Liang, and Zhangyang Wang. Learning to generalize provably in learning to optimize. In International Conference on Arlificial Intelligence and Statistics, pp. 9807-9825. PMLR, 2023.
Tong Zhang. Solving large scale linear prediction problems using stochastic gradient descent algorithms. In International Conference on Machine learning, pp. 116. ACM, 2004.
Joey Tianyi Zhou, Kai Di, Jiawei Du, Xi Peng, Hao Yang, Sinno Jialin Pan, Ivor Tsang, Yong Liu, Zheng Qin, and Rick Siow Mong Goh. Sc2net: Sparse Istms for sparse coding. In Proceedings of the AAAI Conference Artificial Intelligence, volume 32, 2018a.
Kaiwen Zhou, Fanhua Shang, and James Chong. simple stochastic variance reduced algorithm with fast convergence rates. In International Conference on Machine Learning, pp. 5975 5984, 2018b.
Zhengyuan Zhou, Panayotis Mertikopoulos, Nicholas Bambos, Stephen Boyd, and Peter W Glynn. Stochastic mirror descent in variationally coherent optimization problems. Advances Newral Information Processing Systems, 30, 2017.
Julian Zimmert and Tor Lattimore. Connections between mirror descent, Thompson sampling and the information ratio. Advances in Neural Information Processing Systems, 32, 2019.
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of Adam and RMSprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 11127-11135, 2019.
A Proofs in Section 2
A.1 Proof of Lemma
Lemma 2. Consider the approzimate AMD iterations from Algorithm and let g(k) be the enact MD ilerates given by Equation 5 Assume v/* is Ly- -smooth with respect to reference norm || ll- on the dual space, i.e. (2,3) < 1z ylli, or equivalently that Vip -iipschttt Assume further that there exists 0 < lR S LR such that for all x, E X, fallz 2'12 2 < R(xxx²) S LR - x'12. Define the energy ((k) for k 2 0 as follows (where t- = 0):
B(00 k24x-1 (I(Z(M) - (r + rr (Vi(E())), vrb(x)))
(31)
Assume the step-size conditions y LRL and to S L.p. Then the difference between consecutive energies satisfies the following:
B(.) (2k + 1 ( (r)
2-p -te) te) (S(Z(N)) -
5r)
r
((((t()))))- -V g(K+1) 1 -x²)
Proof. The proof closely follows that of Lemma. 2 in Krichene et al. (015 We will use the following lommas, which will be stated without proof. The proofs of these results can be found in the supplementary material of Krichene et al. (2015).
31

---

Published in Transactions on Machine Learning Research (05/2024)
Lemma 3. Let f be convez with Ly-Lipschtte gradient with respect to ||. Then for all x, x', xxx
L
S f(=)) + (Vf(x), - x²) +
|| - xl1?.
2
Lemma 4. For any differentiable convex x²" and u, U, W E dom v"
By (u, u) - By-(w0 v) = Bw- (w, u) + (Vi)(()) - Vib -w)
Lemma 5. If has bug- Lipschitz gradient, then for all v E (RR"),
1 Vel (u)
(w)) BB- (u, v) S L 1u 01
2
We begin our analysis with these lemmas in place. Let x² = x¹ be some minimizer. We begin by bounding the difference in Bregman divergence.
B (Ve(E((t))), ve(x")) - B. (Vo(2(t)) va((z))
BBy (V VO( + (VW( VV(E(A)), g(k+1) - x²)
(b)
| =200)/2 + ( (VN(z - V z(6+1) - -x²)
2L
(c)
1 2
- 1100 + (- -kt (- (((l((t2) s(K+1) xz + (V - Vib(((0t2)), 2(k+2) - x²) |a(k+1]
ktx
(32)
In Step (a) above, we used V a Vi and 2- = x' together with Lemma In Step (b), we use Lemma to bound the first term. In Step (c), we decomposed the right-hand side of the inner product using
5
kls V (z((1)))) + Ve(e(k+))) (v(E(())= (o( (e(e()))))
- -
r
which follows directly from Equation 5 Recall Step of Algorithm
(k+1) arg min (zx .x) ))+ R(ã, 2(0,2((+1)) (k++1)), 8ER
1
ytv
where ln |z v11 < R(x S La | v12 From the definition of z((+1), have that for any R".
2
R(x, (k+t)) 2 ((+(++1)) ((())) + 7ha (V (z)()()-1) (((+1) - =x).
(33)
Recalling the step for (k+1) in Step 2 of Algorithm
we can write
- p((+1)) g(++1) - g(k) 1 (deg()+1) + (1 - -x0))())=2((t++)
1 (a((+t)) (((ttt))((ttt)) (((t1))
do
where we define (++1)= = >1 + +(1 (x)) We then compute
||( - =200)12= = |a((+11) -20041))=(001))) 11 2 (((()+11) (l+t))), since R(x, x') S La || - 2/122.
2
x²
LR
2
- [|[++]]] ==(10112 2 In>i (( (k++))) 7/4 (v) (z((((++))) (((+1)) dlett))).
2
by setting x = d(b+1) in Equation
33 Now, using R( ,((+1)) > اا | - x(x+2)122 we have:
[|(()+]] - 2(h))1
2
Lx
[++++ (k+1) + ztx (Vfla( (( adutt)))
LR>
2
32

---

Published in Transactions on Machine Learning Research (05/2024)
By replacing d(k+1) -x2((++1) + (1 - (x))(()) and multiplying both sides by Au 2ny
kag
we get
KLR | As 2m7
kl |z(((1]) - y1
- =2(01122 2
2r 2x
ktk V f (o((t))) 1 -((tt)] - g(++1) +
1 1- de g000
(34)
d
An
Subtracting Equation 34 from Equation 32
klR
By (Vb(z((++1)), Vi((x*)) - BB- (Vb()(()), Vab(x*) S -axe|=((+11) - =/00/p2-
2r |
-00+
ktx (f(((tt))) -x" + d
(k+t)) de 2000) ) + ( (v 11) - (V((((tt)), 2 x²)
d
where Ork
MhLn 2r
Defining Dit (k+1) ||=((++1) 1.( and Dba+t) || =0,) we rewrite
2L
the last incquality as
By- (Vb()((+2)), Vv(z²)) By- (VW(E(1)), Vxb(z')
un ((++1))
kl
D (1+1) kl ( ((x((()11) [(+++)) - -x²)
2r>
r
1 - du ktk (-Vf(r(itl)), a((t)))) - =(t)) + (vu(s(+t)) - Fi(EEktl)), - 2²)
+
Using Lemma 3 with 2, xtxx od as follows, we bound the inner products:
( g(++2)) x(6)) S f( f(1(kt1)) + Li ((++1)), (x = ((++1), xt g(k+1), and x' ((());
2
(V)(((ttl)) ((++l) x²)) s -((((+)))
L p(tt)), (x (x+1), xt = g(k+1), and x' = x²).
+
2
Combining the inequalities and using 1- Au
8, we obtain
BB (3-((V5(((tt)) Vv(x*)) - -B (V)(()))))) Ve((x"))
S (haD + x²
12tu ((()()) - (( + Lis D((+2)
+ ktv (r- ((- [(E()+11))
+ Lg D 1(++1))
2
2
klR (1+t2) (VitE)HH))) V ((,()))),,+))))))) g(k+1) - xx))
2rd47
k2
- (o00)))((((+))))))
kts (fr - -s(x(ht1)) az DD (+t1)) ((++1)
22
& + (VV((((AA1)) (el(2((t))), g(k++)) - x²)
klg
11b21x
Lykil
where 3x = 2r\
Re-introducing the energy function for k: 2 1,
2r
2r
E(K) = k2tx
(f(E(N)) - ff") + B ( ve(x"))
33

---

Published in Transactions on Machine Learning Research (05/2024)
can compute the difference E(++1) - -E(h);
((++1) - F.(K) (k + 1) (s(olut))) - h-le-1 (S(Z(h)) -r)
+ +r(Boo (V))((()(()()) vxb(x*)) - Bi- (Vi)(((A)), vab("'))
(f(E(ht))) - -5t) (2k: + +>1x (f(z( - (r)
tttx (I(Z(0))- (000)-99 -3) 12tx-1- tk
k2t (s(E(E)) (-)
r
k2tx
(F(E /(E(t)))) + ktx (f- - f(a(kt)))) 12 r.2 (k+1)
r
+ (V()t)))))- ( ,, s(k+1) -x*)
- 1x) (2k: rki)lk (F(ECAtD)) - -1) 1.2 (1x1-1
(f(E(H))) -8 -5t) -raaDfat)) - r. (R+1)
r
+ tr +VV V z(1+1) x²)
For the desired inequality to hold, it suffices that ax. Bu 2 0. Recalling their definitions (and Aa), we want
kLR 2 0 and k(r + k)ln
L2khtx
L L
2 0.
2
2(r + k)y
2r2y
27
2r
which are equivalent to Under these conditions, we get the bound as required
kr
lR Li7
7 kr n22
LRLØ and tu S
It thus suffices to have
In 7 2 Unber to S
Ln
E(++1) - (k()) <
(2k + 1 - -kk)t
(S(E(Atll) - (r)
( -tx) ( - 5*)
r.
tr (W(Et)+)))- VW((s(+1)), - =²))
A.2 Proof of Proposition
1
Proposition 4. Assume the conditions in Lemma 1 Then
E0 s (V Vi(x")) to (f(((0)) f³) + +V -V V z(1) +00)=(000 (()=x')) -xr)
(35)
r
Proof. From Lemma applied with x: 0, we have
E0) < E(0 + to ()( - )²) + r (V)()()) - Vib(a(1), z(1) -
=00)==))) -a*)
rB... (2)) vib(*)) + -/2 (F( - )m) + (V)()(()) Vob(((1). 2(1) - x²)
By definition, ((1) arg min yto (oo (V)()()),,) + R(i, ,()),, thus (since R(x, x) = 0)
7o ()(((()))())) + 11 (1) S 5/40 (V)(z(1)), .0.1))
(36)
34

---

Published in Transactions on Machine Learning Research (05/2024)
Therefore, we get
L |=(0)
- V)((( .0(1) -x²)
1(0-=(00))) 11
by Lemma
T:
$ ()(( ),((2) x²) R(E 2000)
by assumption on R
LR
$ (V((E()) ) ),((1) x²)
+ R(a xx(1))
1
L1 R(i x(11)
using
2L
710
lR
7/0
$ (V)(E()) x,x() x²)
Li R( 2011)
using Equation
36
I
S((a0) -5 + Tef |a((1) x²
T.
R( x0())
using Lemma
2
lR
$((00)) - fr
by assumption on
Now recalling that (1) X0 + (1 - Ao)=(() and A = 1, we have x(1) = a(0) = (0) Therefore, we have f(x(1)) - f S fxx(0)) - f Further using 2(0) (0) shows the desired inoquality.
B Proof of Section (approximate SMD)
The following lemma characterizes the effect of a mirror step on the Bregman divergence. We use this lemma show an optimality gap bound on SMD with the inexact stochastic oracle G(z, E) similarly to Nemirovski et al. 2009),, where the bound depends the variance of the stochastic component (xx, and the norm of the inexactness (xx).
Lemma 6 (Nemirovski et al. 2009 Lem. 2.1). For any u, X and y E (RR"), have (where sb is a-strongly convez),
Bylu, Pa())) S (B + (y, u - x) + lvv
(37)
2a
Applying Equation 37 with x = x()), t ((x g(()), given any point u E X, we have
th
4. (k)) c(k)), x - un S (u, - (u + (
(38)
2a
Using the definition of G, we can expand
to ())) (0)) - u) S Bvlu, x(1) Bv, x (1+1))
t
[|( ((3))
2a
(39)
tx (A kr. (()) u) - to (Uk, (b) - u).
Observe that since f is convex and g(x) E 0f(z), we have (0(x(()), x(()) x,(() u) 2 f(x( - f(u). Using this,
summing Equation 39 from 1 to k, and noting that Br 2 O, we have
Σ ()((((0)- - f(u)] s (u, + 1=0
-0) E4(A,, -w) Σ ti(U,, - -u). (40)
2
1=0
Observing that EEA ..., for any i,
6--1] = 0 and that (0 lives in the o-algebra defined by Ei., we have that
E((A,, x0 -w)] = E. E (A,, x(0) u) So. ...,
:......; 6o [[]]]= =0.
We can take expectations over to get
k
12
k Σ
E Σ ()) - s(w)] < Belu, (11) +
Σ
E [[G((((0,(0))))
t (U., 1 mn)
2a
1=0
35

---

Published in Transactions on Machine Learning Research (05/2024)
Assuming that the stochasticity is bounded
(z,, & o²² Vz E X
this gives an optimality gap bound, which can be extended to convergence of the ergodic average.
If f were j-strongly convex, then we can remove the uniform boundedness assumption since this allows us to control ||x - w|l, using the fact that (g(x(()), xx()) - u) f(x(1) f(u) + =nlx /2/2//2. Equation 40 instead becomes
Σ [[S(x!9) s(a)] (u,
Σ
G( {(0))
i=0
& Σ4
4(A,,, - u)
[U - u) + = (aa0))_ll]]]] (41)
i=U
Taking expectations before and using Young's inequality on the final term, we get
k
t
E
Σ
(alf f(w]] s Bp (0)
+
{0)
2p |u:
2xx
1=0
C Proof of Theorem 3 (approximate ASMD)
Theorem 4. Suppose I has TTppsshct gradient, the mirror mop is o-strongly convex, the approzimation error (Rk is bounded, and that the stochastic oracle is otherwise unbiased with bounded second moments. Assume the diameler of X finile, that is, there erists a constant Ma such that
sup (1(
xx''E
Suppose further that there erists a constant K such that for every iterate x(8), we have
(V)(a! U(h)) K.
Then the convergence rale of approcimate ASMD is (where the right hand side is also given by Equation 44
S Okk k-12))
Proof. Recall the definition of the energy,
- f(x")) + Ss Ba V (1))))) (()) ([A)((()(())
(42)
Using the following identity (which can be shown by expanding the Bregman distances using their definitions and using Val" = 11)
Bx" Vub² (b+) Bx(w", V (v())) By (v Vib"" (y(kt2)) (x) y ((+ x0 V
compute the difference in energy
g(++1) {()) = E[AA+ ([ - f(x))) - (n( -
+ [s 8k Bp V "y (1())))
B V (y Bulx" (+)))
= Ef - f(x')) (x( - ff(x") (){ E[ - 8/ Bx)B
V (3 (1(++)))))
+ V (1 (1)), Vi (1(+1)))) + (g) - y x²' - V (1(()))
= =f( f(x') Ax (f( - ff(x")) E - 8/ By (x V (g
+ESs (4), V (b(+1))) + (Ax-1 Az) ((+++))), ), V (1())))
36

---

Published in Transactions on Machine Learning Research (05/2024)
where the second equality comes from the above identity, and the third equality from the definition of y(u) We further compute using the optimality conditions in Equations 18 and 19 and expanding g(k+1))= = (f(x(tt)) ((ott)))
+ AR+1:
[8 - 8K Be Vab" (y (1)++)) g(b+1) g(h) = E[[A+ ( ((tt))tt))))))) - f(x))) - (a( - (x=')))
+ B(ss B V (V ),) V V
+ (A - A&) (V I (z x'' (1(1)) A Ax (Vf(x(())) x. x (8)))
+ BA (
VV (p)()(1)))
+ B[(A+t1 - (G (k+1), (((tt)),((()))) [0((1))]
2 [ - f((")) - (A(((((())) f(x"))) [[((++1 sk) Bu Vv (b((((tt)) ((((+)))))
+ B. Be (V Vi ((((t((t)))
+ B[(AA+I (zx)) ((x( + (n( (cx((t2)))111
+ B[(A++1 - (A (v (1)))
+ B[(A++1 - (G (1+1) ,5 g(K+1), [(k))
E[(sk+1 - Bx)) (x", V (y (++11))] + (Vb Ve (V'((((tt))))
+ BAA+++ A&) f ( 1). ((0))
where the second inequality follows from convexity of f, and the final equality from zero mean of A441 and independence with ( (7()).
Recalling that since v a-strongly convex, we have that of is 1/a-strongly smooth. Using this, we have the following bound on the Bregman divergence:
B (VV (y(tt))) By (g((+1), ,,)t))),(()) (1( - (1y - -1101))
2a
ML - 2002 + 4 (12) (A Ax)2
a
20s
where the final iterate comes from the fact that for any X, [V)s( < L V s (x* Mee, which is Vo
derived from the smoothness of f, strong convexity of sb and bounded domain assumption. Furthermore, we have Big (x"; Vit" (y(++)) I Mg. Further assume that (VI(x((++1), Ud0y is bounded, say by K. Substituting back we have
g(++1) - g(k) < Mo(sk+1 - -k))+ + ML..Mo + ao² + +2[ (Ak+1 - Ax)² +1 Ag|K.
a²
Sk
Assuming that As is monotonically increasing (which will be justified later), we can sum from to l and get
g(k) < (0) M 4L + ao + +2
(Aj+1 A,)²
a²
$j
j=0
+ +A K.
Taking A₂ أ0ح + 1)/2 (which is monotonically increasing) and 8, (j + 1)3/2 get
g()) < g() Mb 1)3/2 4L3 + ao + 21 V V 11 k-1
Vj+I+ I ArK
a²
<9(() MM. (k 1)3/2 8L + 200 + {+))x()))
(k- + 1)3/2 + AcK.
a²
where we use the following lemma for divergent series for the final inequality.
37

---

Published in Transactions on Machine Learning Research (05/2024)
Lemma 7
CChebbu 2009). For p 0,
kt-P
(k + 1+ -1
1 1
(43)
1-p
Dividing throughout by A&, we get
(k (
8L 200
4 V
- f(x=)) /A4 +
Mg +
+ K
(44)
Ak
a²
=K OK -k
D Proofs in Section 6
D.1 Proof of Proposition 2
Proposition 2. If S2 is gradient descent on the G-invariant objective function L, then it is G-equivariant,
i.e. Goig z) = [=)())] for all E G. Z.
Proof. We first check that V commules with actions. For any E 2, the derivalive operator DL: 2 - Z" satisfies
L(g + 1y) Iig 2)
DL(g z)()) lim t-0
t L(2 tg v) - L(z)
lim
where the second incquality holds since L(g 2) L(2), applied with g-1 and lincarity of actions. Let denote the canonical isomorphism s: 2 - Z. For any y E 2,
(1 Vi 2), v) DL(g (v)
-DL(2)(g y)
= =(1-VL(()) yy
- v).
Hence, G commutes with LL DL. We now have that
Dolg 2) = g. 2 (V (g 2)
= -g.2 2 - 7g (((2)
= g
L()))
g Do(=),
where the first equality comes from the fact that g trivially commutes with (and thus&-! using the induced group action, and the second and third equality comes from the linearity of actions. Since this is true for all y E 2, by the Riesz representation theorem, we have VL(g =) = VL(=).
E
Step-size comparison for baseline methods
Here we detail the step-sizes used to compare. In general, we choose the best performing method at 100 iterations,with preference given to the one converging to smaller value at higher iterations in case of ambiguity. Note later iteration cutoffs usually mean insignificant performance earlier iterations. The grid used for the step-sizes are of the form (1,2,5) 10-k The endpoints
38

---

Published in Transactions on Machine Learning Research (05/2024)
GD
[2 10"4 ,, 10 1]
Adam
[1 x 10 1 x 1
Nesterov [1 x 10 10- 5,5 x 30
Figure | shows the loss evolution for the proposed step-sizes up to 5000 iterations. We note that the grid is sufficiently wide such that smaller or larger step-sizes are suboptimal, lead to slow convergence the low iteration count setting. The specific choices are as below.
Image denoising. GD: 1x 1 Adam: 2 10-3, Nesterov: X 10
Image inpainting. GD: 1x 10
Adam: 1 x 10-?.
Nesterov: 2 x 10 10-3,
SVM training. GD: 1x 2 Adam: 10-2.
Nesterov: 5 x 10-3,
facesutration Jasa
30'
105
24-4
10'
24
10
56-2
10a
21-1
10
20F
201
122
201
(a) GD (denoise)
(b) GD (inpaint)
(c) GD (SVM)
16
30
10"
10"
2e-2
107
(d) Adam (denoise)
(e) Adam (inpaint)
(f) Adam (SVM)
Receustationaaa Lass
Necoranvdiee Loas
Facenutraion La1a
17
302
10'
2e-4
10-1
300
107
Denellan
(g) Nesterov (denoise)
(h) Nesterov (impaint)
(i) Nesterov (SVM)
Figure 9: Comparison of the baseline methods with various step-sizes for the ellipse denoising and image inpainting tasks. Vertical gray line represents iteration 2000, which is the maximum iteration considered in the main text figures.
F Comparison of loss against time
We consider the image inpainting experiment and plot against the wall-clock time instead of the iteration count. In Figure 10 we consider various different step-sizes for the baselines GD, Nesterov accelerated GD, and Adam, chosen in Appendix P Each solid line corresponds different choice of step-size. We observe
39

---

Published in Transactions on Machine Learning Research (05/2024)
that the baseline methods plateau after period of fast optimization, while the proposed LAMD method is able to continue decreasing.
Reconstruction Loss
10a
10a
10
GD
Nesterov
Adam
LASMD
LMD
10
15
20
25
Time (s)
Figure 10: Plot of reconstruction loss against wall-clock time for image inpainting experiment.
In Figure 11 plot the training loss of CNN training MNIST as in Section 7.2. Due to the compact parameterization, we observe that the diagonal LMD methods introduce minor overheads compared to the baseline SGD and Adam methods, which have been optimized in the PyTorch framework, retaining its competitive nature when plotting against iteration count.
Training negative log-likelihood
40
Cross-entropy Loss
SGD
Adam
LAMD
20
LMD
10
0.00
0.05
0.10
0.15
0.20
0.25
Time (s)
Figure 11: Plot training cross-entropy loss against wall-clock time when training CNN on MNIST. The wall-clock time includes only the optimization phases and excludes overhead resulting from testing at each iteration.
40

---

Published in Transactions on Machine Learning Research (05/2024)
G Monotonic spline parameterization
In Section 7 we utilize an clement-wise splinc-based parameterization (shared across parameters of the same layer) of the mirror map. Here we detail the implementation of the spline.
Recall that we define the mirror map (acting element- wise) to be
: R IR, (xx) / ott)t..
Here, learnable piecewise-line -ineer monotonic spline, with o(0) 0. Thus (x(x) is piecewise quadratic monotonic spline. be parameterized in terms of knot locations t_K < ... < the values G at,,, with linear interpolation between the intervals (and extrapolation (lx, +oo)). The monotonic condition can be enforced during training by forcing c- < c-K++ < < We note that shifting up or down by a constant does not change the mirror map dynamics, so without loss of generality ep can be taken to be 0.
... kk well (-o0, (-k and
+ Z- (c++1 ct)
[t,4::]]:
b++1 -
(2))
K+1
L- K+1 x
K-1 K-1-1p bx
(C-K+1 c-K) < L-K;
CK-1
1x-tk-1
2- IK-1 (CK CK-1)
lg-
In our implementation, we fix the knots as in Goujon al. (2022) be t, = 0.05i, ......,, 20, which reduces the number of learnable parameters to only the values of at the knots, which 2K 40. One useful point of this characterization is that the inverse of a can be easily computed by computing the values at the knots (t;) and lincarly interpolating.
H Reptile Initialization for Section 7.3
In Section 7.3
compared LMD with baseline methods from mela-learning objective. Informally, Reptile can be treated as an algorithm that finds solution that is close to each task's manifold of optimal solutions. Here we detail how the Reptile initialization is computed for our purpose. The algorithm is as follows:
Algorithm 8 Reptile (Nichol et al. 2018 Alg. 1)
1: Initialize parameters 0
2: for == K do
3: Sample task f E
4: Perform n steps of SGD/Adam on f from 0 to obtain parameters /
Update initialization f c(θ
end for
return initialization 0
We consider training on the 5-shot 10-way Omniglot dataset for K 10 meta-epochs. The inner optimization loop in Step 4 is done using full-batch SGD with learning rate 10- for n 10 iterations. The meta-stepsize e is decreased lincarly from = 0.1 to 0 over the entire training process. This returns the Reptile initialization 0x-
41


Reviews:
[{"review_id": "ULPMYQlsnv", "reviewer": "Reviewer_Czv7", "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\n* Nice and extensive set of follow-up contributions to a specific method (LMD) in the L2O literature.\n* Relevant to the TMLR audience.\n\n\n**Weaknesses**\n\n* The paper may benefit from restructuring. Right now, the algorithmic treatment is somewhat repetitive with little qualitative or analysis insights separating the various accelerated versions. Moving large chunks of Sections 2,3,4, to an appendix can make the paper much more readable and will have the added benefit of cutting down its length (currently, 31 full pages of text). \n* The part about equivariance is a bit orthogonal to the central problem addressed by this paper (mirror descent and its accelerated variants), and feels tacked on rather than being germane to the overall discussion. In any case it seems that any benefits to acceleration are low, since the authors could only demonstrate results on very simple toy neural networks. \n* I am not very convinced by the experiments on equivariance, and would appreciate it if the authors provided further quantitative evidence (beyond trying it out for MNIST). \n* Related to the above two points: it is not entirely clear to me whether reducing the number of parameters via equivariance is a good idea. We do now know that overparameterization can be a benefit in various ways in learning deep neural networks -- either by making the optimization landscape more benign, or by adding additional implicit regularization to improve generalization. Weight tying may not always be a good thing. I would encourage the authors to think a bit more deeper into whether this section really makes the paper stronger.", "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": "There is a growing body of literature on \"learning-to-optimize\" (L2O) methods. Broadly speaking, these methods use learning-based strategies to accelerate the solution of specific instances of optimization algorithms, given a training set of solution outputs from previously seen instances. This paper focuses on learning-based extensions of mirror descent (MD), which enables optimization over convex functions over non-Euclidean spaces (such as the simplex). \n\nIn previous work, Tan-Mukherjee-Tang-Hauptmann-Schonlieb (2023a,b) proposed and studied learned mirror descent (LMD). At a high level this takes the standard mirror descent algorithm of Nemirovski and Yudin (1983) and replaces the mirror potential (and its convex conjugate) with learned functions parameterized by a pair of neural networks. This work extends this earlier study from 2023 in a few different ways: fixing certain algorithmic inconsistencies in LMD; achieving further acceleration; analyzing stochastic variants; and reducing the sample complexity of learning.\n\nWithin this overall context, the authors make the following specific contributions: \n* They propose and analyze learned accelerated mirror descent (LAMD), which fixes the *forward-backward* inconsistency that arises in LMD due to the fact that the \"mirror map\" neural network and the \"conjugate\" neural network are in fact not guaranteed to be convex conjugates of each other. In the previous work on LMD, the analysis swept this discrepancy \"under the rug\" using an inexact gradient oracle; here, the authors provide a more salient algorithmic fix. Theoretical analysis shows that under certain assumptions, LAMD gives a true minimum, as opposed to giving an additive error factor above the minimum.\n* They propose and analyze stochastic versions of LMD and LAMD, where the gradients are supplied by noisy/stochastic oracles. \n* They propose using *equivariance* to reduce the number of parameters required to learn the mirror maps. This arises due to the property that MD methods remain permutation invariant throughout training,", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Consider restructuring the paper where the first 10-12 pages summarize the most important parts of the paper, and relegate many of the more repetitive parts to an appendix.\n* The paragraph \"We note that this analysis can be used to partially explain the existing heuristic...\" is confusing. Can you explain more clearly what these earlier methods do, and how your analysis partially explains existing heuristics?\n* More clarity on the monotonic splines experiment can be helpful -- why this is a meaningful choice of mirror map, what the spline order is, what the variables are etc. Maybe a figure, or clear equations, can help.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3tMsE4uHbG", "reviewer": "Reviewer_QA8S", "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:\n1. This paper is well-written. The proof looks correct but I didn't have time to check all the details.\n2. The assumptions are all reasonable and many are standard in the optimization literature.\n3. The proposed methods are validated in various different settings to justify their effectiveness.\n\nWeaknesses:\n1. In the comparison with the GD, Nesterov and Adam baselines, the authors did not discuss how they picked the learning rates and whether they are thoroughly tuned with a grid search. If the learning rates are not tuned, then the current experiments may underestimate the optimization speed of the baselines.\n2. The experiments only measure the number of iterations, but applying the learned mirror maps can lead to extra computation costs. It would be valuable if the authors could report the wall-clock time as well.\n3. Experiments on neural net training illustrate the effectiveness of the proposed methods beyond convex optimization, but these experiments are limited to rather shallow neural nets trained with only a couple of epochs.", "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 aims to improve the optimization speed of the learned mirror descent method. This method is similar to classical mirror descent but uses a learnable mirror map (and its convex conjugate) parameterized by input-convex neural nets.\n\nThis work proposes the Nesterov-accelerated version, the stochastic version, and the stochastic accelerated version of learned mirror descent. And then it provides theoretical analyses showing that these methods converge to optimal values on convex functions, assuming certain error bounds on how well the learned maps are convex conjugate to each other.\n\nExperiments are conducted on several convex optimization problems and neural network training.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The theory part is good overall, but for the experiment part I hope the authors can resolve Weaknesses 1 and 2.", "broader_impact_concerns": "/", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "p85mTm8teb", "reviewer": "Reviewer_QBq1", "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:\n1. While the idea of analyzing approximate mirror descent in order to justify learned mirror maps is not new, the result showing that acceleration suffices to converge to the true optimum is an interesting advance upon past work.\n2. The authors take seriously the question of scaling up the algorithms using stochastic optimization rather than full-batch optimization, as well as for reducing the number of parameters on larger-scale (neural network) problems.\n3. Empirically, the problem of amortizing optimization cost by learning an optimizer across many optimization tasks seems well-motivated, and the authors demonstrate improved performance of their approach over baselines/past unaccelerated methods on some of the convex problems.\n\n## Weaknesses:\n1. Theoretically, it is unclear to me how reasonable the Theorem 1 assumption given by equation (12) is. While some bound on the error of the approximate mirror map is necessary, the specific bound used here assumes the algorithmic trajectories of the approximate and non-approximate versions are close across all iterations of the algorithm. Conditions under which this holds are not given, and it is unclear to me why the assumption is somehow more natural or simple than the (un-interesting) approach of just assuming the result of the theorem directly. A more natural assumption seems to me some type of functional bound on the potentials $M_\\varphi$ and $M_\\theta$ themselves. \n2. There is a large amount of missing related work in the theory community. Learned mirror descent, including the Bregman divergence, has been studied in the online setting by Denevi et al. (2019) and Khodak et al. (2019), as well as in the statistical setting by Gao et al. (2022). There is a variety of other related theoretical work—although not regarding mirror descent explicitly—under the names “amortized optimization” (Amos, 2023), “meta-optimization” (Chen & Hazan, 2023), and “data-driven algorithm design” (Gupta & Roughgarden, 2017). While the theory presented here is distinct, comparison (both theoretical and possibly empirical) is warranted.\n3. The empirical results are somewhat underwhelming overall, especially in the neural network case. The meta-learned optimizer does not significantly outperform Adam—despite presumably much greater time spent due to meta-optimization—and there is no comparison to standard benchmarks in the learning-to-optimize literature such as Omniglot (Lake et al., 2011) and Mini-ImageNet (Vinyals, 2016). Would the proposed (fairly-involved) approaches outperform basic initialization-learning methods like MAML (Finn et al., 2017) or Reptile (Nichol et al., 2018) on such benchmarks? If such few-shot learning benchmarks are not the target application, what deep learning application is, beyond MNIST?\n4. The paper is extremely lengthy and poorly organized. While it is the authors’ decision, it seems to me the theoretical arguments could largely be moved to the appendix, as they do not seem to significantly factor into the main message of the paper.\n\n## References:\n1. Amos. *Tutorial on amortized optimization*. FTML 2023.\n2. Chen, Hazan. *A nonstochastic control approach to optimization*. NeurIPS 2023.\n3. Denevi, Stamos, Ciliberto, Pontil. *Online-within-online meta-learning*. NeurIPS 2019.\n4. Finn, Abbeel, Levine. *Model-agnostic meta-learning for fast adaptation of deep networks*. ICML 2017.\n5. Gao, Gouk, Lee, Hospedales. *Meta mirror descent: Optimiser learning for fast convergence*. 2022.\n6. Gupta, Roughgarden. *A PAC approach to application-specific algorithm selection*. SIAM J. Computing, 2017.\n7. Khodak, Balcan, Talwalkar. *Adaptive gradient-based meta-learning methods*. NeurIPS 2019.\n8. Lake, Salakhutdinov, Gross, Tenenbaum. *One-shot learning of simple visual concepts*. CogSci, 2011.\n9. Nichol, Achiam, Schulman. *On first-order meta-learning methods*. 2018.\n10. Vinyals, Blundell, Lillicrap, Kavukcuoglu, Wierstra. *Matching networks for one shot learning*. NeurIPS 2016.", "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 the problem of mirror descent optimization in the case of the distance generating functions in the Bregman divergences being learned. Building upon past work that analyzes the convergence of full batch mirror descent in the case of approximate mirror maps, they show that under a good-enough approximation an accelerated variant will converge to the optimal value. They further extend these results to the stochastic setting with unbounded errors, and conduct several empirical studies in convex and non-convex settings.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. [critical] Substantively addressing the first two weaknesses above, namely (1) explaining why the assumption in Equation 12 is reasonable or why a more natural assumption would not work and (2) addressing related theoretical work, especially that related to learned mirror descent. I believe the first point is critical for establishing convincing evidence of the authors' claims and the second point is critical for clarifying the claims in comparison with past work.\n2. While not critical due to the TMLR guidelines, more substantive empirical evaluations, in terms of benchmarks and baselines, in the neural network setting would significantly strengthen the work.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a correct and useful result as identified by the reviewers. The reviewers all voted leaning accept to accept. The paper meets all the relevant criteria for TMLR.
Decision: Accept (poster)