# Simba: A Scalable Bilevel Preconditioned Gradient Method for Fast Evasion of
Flat Areas and Saddle Points

Nick Tsipinakis
Department of Mathematics
UniDistance Suisse
Brig, Switzerland
<EMAIL_ADDRESS>
&Panos Parpas
Department of Computing
Imperial College London
London, UK
<EMAIL_ADDRESS>

###### Abstract

The convergence behaviour of first-order methods can be severely slowed down
when applied to high-dimensional non-convex functions due to the presence of
saddle points. If, additionally, the saddles are surrounded by large plateaus,
it is highly likely that the first-order methods will converge to sub-optimal
solutions. In machine learning applications, sub-optimal solutions mean poor
generalization performance. They are also related to the issue of hyper-
parameter tuning, since, in the pursuit of solutions that yield lower errors,
a tremendous amount of time is required on selecting the hyper-parameters
appropriately. A natural way to tackle the limitations of first-order methods
is to employ the Hessian information. However, methods that incorporate the
Hessian do not scale or, if they do, they are very slow for modern
applications. Here, we propose Simba, a scalable preconditioned gradient
method, to address the main limitations of the first-order methods. The method
is very simple to implement. It maintains a single precondition matrix that it
is constructed as the outer product of the moving average of the gradients. To
significantly reduce the computational cost of forming and inverting the
preconditioner, we draw links with the multilevel optimization methods. These
links enables us to construct preconditioners in a randomized manner. Our
numerical experiments verify the scalability of Simba as well as its efficacy
near saddles and flat areas. Further, we demonstrate that Simba offers a
satisfactory generalization performance on standard benchmark residual
networks. We also analyze Simba and show its linear convergence rate for
strongly convex functions.

_Keywords_ saddle free optimization $\cdot$ preconditioned gradient methods
$\cdot$ coarse-grained models $\cdot$ deep learning

##  1 Introduction

We focus on solving the following optimization problem, commonly referred to
as empirical risk minimization:

$\min_{\mathbf{x}\in\mathbb{R}^{n}}f(\mathbf{x}):=\frac{1}{m}\sum_{i=1}^{m}f_{i}(\mathbf{x}).$
(1)

We assume that the model parameters are available in a decoupled form, that
is, $\mathbf{x}=\\{\mathbf{x}^{l}:l=1,\ldots L\\}$, for some positive integer
$L$. The decoupled parameter setting naturally emerges for the training of a
deep neural network (DNN) where $L$ corresponds to the number of layers. In
large-scale optimization scenarios, stochastic gradient descent method (SGD)
and its variants have become a standard tool due to their simplicity and low
per-iteration and memory costs. While easy to implement, they are nevertheless
accompanied with important shortcomings. Below, we discuss these limitations
and argue that they mainly stem from the absence of the Hessian matrix (or an
appropriate approximation that leverages its structural properties) in their
iterative scheme.

Shortfall I: Slow evasion of saddle points. The prevalence of saddle points in
high dimensional neural networks has been examined in the past and it has been
shown that the number of saddles increases exponentially with the dimension
$n$ Dauphin et al. (2014). Consequently, developing algorithms that are
capable of escaping saddle points is a primary goal when it comes to the
training of DNNs. While stochastic first-order methods have theoretical
guarantees of escaping saddle points almost always Panageas et al. (2019), it
remains unclear whether they can achieve this efficiently in practice. As a
result, given the significantly larger number of saddle points compared to
local minima in high-dimensional settings, it is likely that stochastic first-
order methods will converge to a saddle point rather than the desired local
minimum. The most natural way to ensure convergence to a local minimum in the
presence of saddles is by incorporating the Hessian matrix to scale the
gradient vector. The Cubic Newton Nesterov et al. (2018) is a key method in
non-convex optimization known to be able to escape saddles in a single
iteration resulting in rapid convergence to a local minimum. The Cubic Newton
method achieves this using the inverse of a regularized version of the Hessian
matrix to premultiply the gradient vector. However, the method has expensive
iterates and does not scale in high dimensions. Nevertheless, its theory
indicates that methods employing accurate approximations of the regularized
Hessian are expected to effectively evade saddle points and converge rapidly.

Shortfall II: Slow convergence rate near flat areas and the vanishing gradient
issue. The convergence behaviour of first-order methods can be significantly
affected near local minima that are surrounded by plateaus Boyd and
Vandenberghe (2004); Nesterov (2004). For instance, this occurs, when its
Hessian matrix has most of the eigenvalues equal to zero. Therefore, to
effectively navigate on such a flat landscape is to employ the inverse Hessian
matrix. For instance, the Newton method pre-multiplies the gradient with the
inverse Hessian to perform a local change of coordinates which significantly
accelerates the convergence rate. Intuitively, one can expect a similar
convergence behaviour by replacing the Hessian matrix with a preconditioned
matrix that retains its structure. On the other hand, the vanishing gradient
issue is commonly observed in the training of DNNs for which activation
functions that range in $[0,1]$ are required. In such cases, the gradient
value computed through the back-propagation decreases exponentially as the
number of layers $L$ of the network increase yielding no progress for the SGD.
Hence, similar to addressing the slow-convergence issue near flat areas,
employing a preconditioner is expected to mitigate the vanishing gradient
issue and thus accelerate the convergence of the SGD method.

To tackle the above shortfalls, the diagonal (first- or second-order) methods
were introduced Duchi et al. (2011); Tieleman et al. (2012); Zeiler (2012);
Kingma and Ba (2014); Zaheer et al. (2018); Yao et al. (2021); Ma (2020);
Jahani et al. (2021); Liu et al. (2023). These algorithms aim to improve the
behavior of the SGD near saddles or flat areas by preconditioning the gradient
with a diagonal vector that mimics the hessian matrix. However, the Hessian
matrices of deep neural networks need not be sparse, which means that a
diagonal approximation of the Hessian may be inadequate. This raises concerns
about the ability of diagonal methods to converge to a local minimum. A more
promising class of algorithms consists of those preconditioned gradient
methods that maintain a matrix instead of the potentially poor diagonal
approximation. As noted earlier, the Cubic Newton method can effectively
address the aforementioned shortfalls but it does not scale in large-scale
settings. Other methods that explicitly use or exploit the Hessian matrix,
such as Newton or Quasi-Newton methods Nesterov (2004); Broyden (1967), or
those based on randomization and sketching Erdogdu and Montanari (2015);
Pilanci and Wainwright (2017); Xu et al. (2016, 2020) still suffer from the
same limitations as the Cubic Newton. A potentially promising way to address
the storage and computational limitations of the Newton type methods is to
employ precondition matrices based on the first-order information Duchi et al.
(2011); Gupta et al. (2018). However, to the best of our knowledge, the
performance of such methods near saddles or flat areas is yet to be examined.
Another possible solution is to utilize the multilevel framework in
optimization which significantly reduces the cost of forming and inverting the
Hessian matrix Tsipinakis and Parpas (2021); Tsipinakis et al. (2023). But, it
remains unclear whether the multilevel Newton-type methods can be efficiently
applied to the training of modern deep neural networks.

Algorithm 1 Simba. Proposed default hyper-parameters: $\mathbf{R}_{k}$ as in
Definition 2.1, $t_{k}=10^{-2},b=0.9,r=20,n_{\ell}=0.5n,m=10^{-8}$.

$t_{k}$ (step-size), $\beta\in(0,1)$ (momentum), $\mathbf{x}_{0}$ \- (initial
point), $r$ (No of eigenvalues), $n_{\ell}$ (coarse model dimensions), $m$
(eigenvalue correction).

Set $\mathbf{G}_{0}=0$

for $k=1,2,\ldots$ do

$\mathbf{G}_{k}=\beta\mathbf{G}_{k-1}+\nabla f(\mathbf{x}_{k-1})$ (EMA of
gradients)

$\mathbf{G}_{\ell,k}=\mathbf{R}_{k}\mathbf{G}_{k}$ (coarse EMA of gradients)

$\mathbf{Q}_{\ell,k}=\mathbf{G}_{\ell,k}\mathbf{G}_{\ell,k}^{T}$ (EMA-based
preconditioner)

$\left[\mathbf{U}_{r+1}\mathbf{\Lambda}_{r+1}\right]=\text{T-SVD}(\mathbf{Q}_{\ell,k})$
by Halko et al. (2011)

$\left[\mathbf{\Lambda}_{r+1}\right]_{i,m}=\begin{cases}\mathbf{\Lambda}_{i},&\mathbf{\Lambda}_{i}\geq
m\\\ m,&\text{otherwise}\end{cases},\quad$ and form
$\mathbf{\Lambda}_{r+1}^{m}$

$\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}=\left[\mathbf{\Lambda}_{r+1}\right]_{r+1,m}^{-\frac{1}{2}}\mathbf{I}_{n_{\ell}}+\mathbf{U}_{r}\left(\left[\mathbf{\Lambda}_{r+1}^{m}\right]^{-\frac{1}{2}}-\left[\mathbf{\Lambda}_{r+1}\right]_{r+1,m}^{-\frac{1}{2}}\mathbf{I}_{r}\right)\mathbf{U}_{r}^{T}$

$\mathbf{x}_{k+1}=\mathbf{x}_{k}-t_{k}\mathbf{P}_{k}\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}\mathbf{R}_{k}\mathbf{G}_{k}$

end for

In this work, we introduce Simba, a Scalable Iterative Minimization Bilevel
Algorithm that addresses the above shortfalls and scales to high dimensions.
The method is very simple to implement. It maintains a precondition matrix
where its inverse square root is used to premultiply the moving average of the
gradient. In particular, the precondition matrix is defined as the outer
product of the exponential moving average (EMA) of the gradients. We explore
the link between the multilevel optimization methods and the preconditioned
gradient methods which enables us to construct preconditioners in a randomized
manner. By leveraging this connection, we significantly reduce the cost and
memory requirements of these methods to make them more efficient in time for
large-scale settings. We propose performing a Truncated Singular Value
decomposition (T-SVD) before constructing the inverse square root
preconditioner. Here, we retain only the first few, say $r$, eigenvalues and
set the remaining eigenvalues equal to the $r^{\text{th}}$ eigenvalue. This
approach aims to construct meaningful preconditioners that do not allow for
zero eigenvalues, and thus expecting a faster escape rate from saddles or flat
areas. Our algorithm is inspired by SigmaSVD Tsipinakis et al. (2023), Shampoo
Gupta et al. (2018) and SGD with momentum Bubeck et al. (2015). Simba is
presented in Algorithm 1. A Pytorch implementation of Simba is available at
https://github.com/Ntsip/simba.

###  1.1 Related Work

The works most closely related to ours is Shampoo Gupta et al. (2018) and
AdaGrad Duchi et al. (2011). However, both methods have important differences
to our approach which are discussed below. AdaGrad is an innovative adaptive
first-order method which performs elementwise division of the gradient with
the square root of the accumulated squared gradient. Adagrad is known to be
suitable for sparse settings. However, the version of the method that retains
a full preconditioner matrix is rarely used in practice due to the increased
computational cost. Moreover, Shampoo, unlike our approach, maintains multiple
preconditioners, one for each tensor dimension (i.e., two for matrices while
we always use one). In addition, it computes a full SVD for each
preconditioner. Further, all preconditioners lie in the original space
dimensions. Given the previous considerations, Shampoo becomes computationally
expensive when applied to very large DNNs. In contrast, our algorithm
addresses the computational issues of Shampoo. Another difference between
Shampoo and our approach is the way the preconditioners are defined. Shampoo
defines the preconditioners as a sum of the outer product of gradients,
whereas we employ the exponential moving average. In addition, the performance
of Shampoo around saddles and flat areas is yet to be examined. We conjecture
that Shampoo may not effectively improve the behaviour of the first-order
methods near plateaus due to the monotonic increase of the eigenvalues in the
preconditioners.

Hessian-based preconditioned methods have been also proposed to improve
convergence of the optimization methods near saddle point Reddi et al. (2018);
Dauphin et al. (2014); O’Leary-Roseberry et al. (2020). However, these methods
rely on second-order information which can be prohibitively expensive for very
large DNNs. Therefore, they are more suitable to apply in conjunction with a
fast first-order method and employ the Hessian approximation when first-order
methods slow down. However, our method is efficient from the beginning of the
training. A notable second-order optimization method is K-FAC Martens and
Grosse (2015). K-FAC approximates the Fisher information matrix of each layer
using the Kronecker product of two smaller matrices that are easier to handle.
However, K-FAC is specifically designed for the training of generative (deep)
models while our method is general and applicable to any stochastic
optimization regime.

As far as the adaptive diagonal methods are concerned, Adam and RMSprop were
introduced to alleviate the limitations of AdaGrad. RMSprop scales the
gradient using an averaging of the squared gradient. On the other hand, Adam
scales the EMA of the gradient by the EMA of the squared gradient. Adam
iterations has been shown to significantly improve the performance of first-
order diagonal methods in the optimization of DNNs. Another method, named
Yogi, has similar updates as AdaGrad but allows for the accumulated squared
root not be monotonically increasing. To provide EMA of the gradient with a
more informative scaling, several approaches have been proposed that replace
the squared gradient with an approximation of the diagonal of the Hessian
matrix Jahani et al. (2021); Yao et al. (2021); Ma (2020). AdaHessian Yao et
al. (2021) approximates the Hessian matrix by its diagonal matrix using the
Hutchinson’s method while Apollo Ma (2020) based on the variational technique
using the weak secant equation. Further, OASIS Jahani et al. (2021) is an
adaptive modification of AdaHessian that automatically adjusts learning rates.
A recent work called Sophia Liu et al. (2023) modifies AdaHessian direction by
clipping it by a scalar. The authors also provide alternatives on the diagonal
Hessian matrix approximation which they suggest to compute every $10$
iterations to reduce the computational costs. However, it is still unclear
whether the diagonal methods offer satisfactory performance near saddles
points and flat areas, or if they effectively tackle the vanishing gradient
issue in practical applications.

###  1.2 Contributions

The main contribution of this paper is the development of a scalable method
that addresses the aforementioned shortfalls by empirical evidence. The method
is scalable for training DNNs as it maintains only one preconditioned matrix
at each iteration which lies in the subspace (coarse level). Thus, we
significantly reduce the cost of forming and computing the SVD compared to
Shampoo. The fact that we use a randomized T-SVD and requiring the most
informative eigenvalues further reduces the total computational cost. The
numerical experiments demonstrate that Simba can have cheaper iterates than
AdaHessian and Shampoo in large-scale settings and deep architectures. In
particular, in these scenarios, the wall-clock time of our method can be as
much as $25$ times less than Shampoo and $2$ times less than AdaHessian. We,
in addition, illustrate that Simba offers comparable, if not better,
generalization errors against the state-of-the-art methods in standard
benchmark ResNet architectures.

Further, the numerical experiments show that our method has between two and
three times more expensive iterates than Adam. This is expected due to outer
products and the randomized T-SVD. However, this is a reasonable price to pay
to improve the escape rate from saddles and flat areas of the existing
preconditioned gradient methods. Therefore, Simba is suitable in problems
where diagonal methods suffer from at least of one of the previous shortfalls.
In such cases, we demonstrate that diagonal methods require much larger wall-
clock time to reach an error as low as our method achieves. Hence, we argue
that, besides the encouraging preliminary empirical results achieved by Simba,
this paper highlights the limitations of diagonal methods, which are likely to
get trapped near saddle points and thus return sub-optimal solution,
particularly in the presence of large plateaus. This emphasizes the need to
develop scalable algorithms that utilize more sophisticated preconditioners
instead of relying on poor diagonal approximations of the Hessian matrix.
Simba is accompanied with a simple convergence analysis assuming strongly
convex functions.

##  2 Description of the Algorithm

In this section, we begin by discussing the standard Newton-type multilevel
method and its main components, i.e., the hierarchy of coarse models and the
linear operators that will used to transfer information from coarse to fine
model and vice versa. Then, as all the necessary ingredients are in place, we
present the proposed algorithm. Even though, we assume a bilevel hierarchy,
extending Simba to several level is straightforward.

###  2.1 Background Knowledge - Multilevel Methods

We will be using the multilevel terminology to denote $f$ as the _fine_ model
and also assume that a _coarse_ model that lies in lower dimensions is
available. In particular, the coarse model is a maping
$F:\mathbb{R}^{n_{\ell}}\rightarrow\mathbb{R}$, where $n_{\ell}<n$. The
subscript $\ell$ will be used to denote quantities, i.e., vectors and scalars,
that belong to the coarse level. We assume that both fine and coarse models
are bounded from below so that a minimizer exists.

In order to reduce the computational complexity of Newton type methods, the
standard multilevel method attempts to solve (1) by minimizing the coarse
model to obtain search directions. This process is described as follows.
First, one needs to have access to linear prolongation and restriction
operators to transfer information from the fine to coarse model and vice
versa. We denote $\mathbf{P}\in\mathbb{R}^{n\times n_{\ell}}$ be the
prolongation and $\mathbf{R}\in\mathbb{R}^{n_{\ell}\times n}$ the restriction
operator and assume that they are full rank and $\mathbf{P}=\mathbf{R}^{T}$.
Given a point $\mathbf{x}_{k}$ at iteration $k$, we move to the coarse level
with initialization point $\mathbf{y}_{0}=\mathbf{R}\mathbf{x}_{k}$.
Subsequently, we minimize the coarse model $F$ to obtain $\mathbf{y}^{*}$ and
construct a search direction by:

$\mathbf{d}_{k}:=\mathbf{P}_{k}(\mathbf{y}^{*}-\mathbf{y}_{0}).$ (2)

To compute the search directions effectively, the standard multilevel method
considers the Galerkin model as a choice of the coarse model:

$\displaystyle F(\mathbf{y})$ $\displaystyle:=\langle\mathbf{R}_{k}\nabla
f(\mathbf{x}_{k}),\mathbf{y}-\mathbf{y}_{0}\rangle+\frac{1}{2}\langle\mathbf{R}_{k}\nabla^{2}f(\mathbf{x}_{k})\mathbf{P}_{k}(\mathbf{y}-\mathbf{y}_{0}),\mathbf{y}-\mathbf{y}_{0}\rangle.$
(3)

It can be shown that combining (2) and (3) we obtain a closed form solution
for the coarse direction:

$\mathbf{d}_{k}=-\mathbf{P}_{k}\left(\mathbf{R}_{k}\nabla^{2}f(\mathbf{x}_{k})\mathbf{P}_{k}\right)^{-1}\mathbf{R}_{k}\nabla
f(\mathbf{x}_{k})$

However, it may be ineffective to employ solely the above strategy when
computing the search direction if $\mathbf{R}$ is fixed or defined in a
deterministic fashion at each iteration. For instance, this can occur when
$\nabla f(\mathbf{x})\in\operatorname{null}(\mathbf{R})$ while $\|\nabla
f(\mathbf{x}_{k})\|_{2}\neq 0$, which implies no progress for the multilevel
algorithm. The following conditions were introduced to prevent the multilevel
algorithm from taking an ineffective coarse step:

$\|\mathbf{R}\nabla f(\mathbf{x})\|>\xi\|\nabla
f(\mathbf{x})\|\quad\text{and}\quad\|\mathbf{R}\nabla f(\mathbf{x})\|>e,$

where $\xi\in(0,\min(1,\|\mathbf{R}\|))$ and $e>0$ are user defined
parameters. Thus, the standard multilevel algorithm computes the coarse
direction when the above conditions are satisfied and the fine direction
(Newton) otherwise. On the other hand, it is expected that the multilevel
algorithm will construct effective coarse directions when $\mathbf{R}$ is
selected randomly at each iteration. In particular, it has been demonstrated
that the standard multilevel method can perform always coarse steps without
compromising its superlinear or composite convergence rate (see for example Ho
et al. (2019); Tsipinakis and Parpas (2021); Tsipinakis et al. (2023)).

###  2.2 Simba

The standard multilevel algorithm is well suited for strictly convex functions
since the Hessian matrix is always positive definite. Due to the absence of
positive-definiteness in general non-convex problems, Newton type methods may
converge to a saddle point or a local maximum since they cannot guarantee the
descent property of the search directions. They can even break down when the
Hessian matrix is singular. These limitations have been efficiently tackled in
the recent studies Nesterov and Polyak (2006). In multilevel literature,
SigmaSVD Tsipinakis et al. (2023) addresses these limitations by using a
truncated low-rank modification of the standard multilevel algorithm that
efficiently escapes saddle points and flat areas. While SigmaSVD has been
shown to perform well near saddles and flat areas, it requires
$\mathcal{O}(mn_{\ell}^{2})$ operations to form the reduced Hessian matrix
which is prohibitively expensive for modern deep neural network models. Here,
to alleviate this burden, we replace the expensive Hessian matrix with the
outer product of the gradients. For this purpose, we consider the following
coarse model:

$\displaystyle F(\mathbf{y})$
$\displaystyle:=\langle\mathbf{R}_{k}\mathbf{G}_{k},\mathbf{y}-\mathbf{y}_{0}\rangle+\frac{1}{2}\langle(\mathbf{R}_{k}\mathbf{H}_{k}\mathbf{P}_{k})^{\frac{1}{2}}(\mathbf{y}-\mathbf{y}_{0}),\mathbf{y}-\mathbf{y}_{0}\rangle,$
(4)

where $\quad\mathbf{H}_{k}:=\mathbf{G}_{k}\mathbf{G}_{k}^{T}$ and
$\mathbf{G}_{k}:=\nabla f(\mathbf{x}_{k})$. Thus, by defining
$\hat{\mathbf{d}}_{\ell}:=\mathbf{y}^{*}-\mathbf{y}_{0}$ and
$\mathbf{Q}_{\ell,k}:=\mathbf{R}_{k}\mathbf{H}_{k}\mathbf{P}_{k}$, the coarse
direction can be computed explicitly:

$\displaystyle\hat{\mathbf{d}}_{k}$
$\displaystyle=\mathbf{P}_{k}\left(\underset{\mathbf{d}_{\ell}\in\mathbb{R}^{n_{\ell}}}{\operatorname{arg\
min}}\langle\mathbf{R}_{k}\mathbf{G}_{k},\mathbf{d}_{\ell}\rangle+\frac{1}{2}\langle\mathbf{Q}_{\ell,k}^{\frac{1}{2}}\mathbf{d}_{\ell},\mathbf{d}_{\ell}\rangle\right)$
$\displaystyle=-\mathbf{P}_{k}\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}\mathbf{R}_{k}\mathbf{G}_{k}.$

To construct $\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}$, we perform a randomized
truncated SVD to obtain the first $r+1$ eigenvalues and eigenvectors, that is,
$\left[\mathbf{U}_{r+1},\mathbf{\Lambda}_{r+1}\right]=\text{T-SVD}(\mathbf{Q}_{\ell,k})$
Halko et al. (2011). Subsequently, we bound the diagonal matrix of eigenvalues
from below by a scalar $m>0$:

$\left[\mathbf{\Lambda}_{r+1}\right]_{i,m}:=\begin{cases}\mathbf{\Lambda}_{i},&\mathbf{\Lambda}_{i}\geq
m\\\ m,&\text{otherwise},\end{cases}$ (5)

to obtain the new modified diagonal matrix as $\mathbf{\Lambda}_{r+1}^{m}$. We
then construct $\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}$ by treating all its
eigenvalues below the $r$-th equal to $r+1$ eigenvalue. Formally, we define

$\displaystyle\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}$
$\displaystyle:=\left[\mathbf{\Lambda}_{r+1}\right]_{r+1,m}^{-\frac{1}{2}}\mathbf{I}_{n_{\ell}}+\mathbf{U}_{r}\left(\left[\mathbf{\Lambda}_{r+1}^{m}\right]^{-\frac{1}{2}}-\left[\mathbf{\Lambda}_{r+1}\right]_{r+1,m}^{-\frac{1}{2}}\mathbf{I}_{r}\right)\mathbf{U}_{r}^{T}.$
(6)

The key component when forming the precondition matrix as above is that we do
not allow for zero eigenvalues since zero eigenvalues indicate the presence of
flat areas where the convergence rate of optimization methods significantly
decays. Hence, by requiring $\mathbf{Q}_{\ell,k}\succeq
m\mathbf{I}_{n_{\ell}}$, we ensure that directions which correspond to flat
curvatures will turn into directions whose curvature is positive, anticipating
an accelerated convergence to the local minimum or a fast escape rate from
saddles. Similar techniques for constructing the preconditioners have been
employed in Tsipinakis et al. (2023); Paternain et al. (2019) where it is
demonstrated numerically that the algorithms can rapidly escape saddle points
and flat areas in practical applications. The computational cost of forming
the preconditioner matrix is $\mathcal{O}(n_{\ell}^{2})$ which is
significantly smaller than computing the Hessian matrix. Moreover, the
randomized SVD requires $\mathcal{O}(rn_{\ell}^{2})$ operation.

In addition, the method can be trivially modified to employ the EMA of the
gradients to account for the averaged history of the local information.
Moreover, the EMA of the gradient is used to construct more informative
preconditioner matrices. Specifically, to obtain the accelerated algorithm, we
set a momentum parameter $\beta\in(0,1)$ and replace $\mathbf{G}_{k}$ in (4)
with the EMA update:

$\mathbf{G}_{k}:=\beta\mathbf{G}_{k-1}+\nabla f(\mathbf{x}_{k}),$

where $\mathbf{G}_{0}=0$ and $k\geq 1$. Simba with momentum is presented in
Algorithm 1.

We emphasize that the method is scalable for solving large deep neural network
models. We describe the three components that render the iterations of Simba
efficient: (a) The operations can be performed in a decoupled manner. This
means that preconditioners can be computed independently at each layer,
leading into forming much smaller matrices that can be efficiently stored
during the training. (b) The restriction operator is constructed based on
uniform sampling without replacement from the rows of the identity matrix,
$\mathbf{I}_{n}$. This is described formally in the following definition:

######  Definition 2.1.

_Let the set $S_{n}=\\{1,2,\ldots,n\\}$. Sample uniformly without replacement
$n_{\ell}<n$ elements from $S_{n}$ to construct
$S_{n_{\ell}}=\\{s_{n_{1}},s_{n_{2}},\ldots,s_{n_{\ell}}\\}$. Then, the
$i^{\text{th}}$ row of $\mathbf{R}$ is the $s_{n_{i}}^{\text{th}}$ row of the
identity matrix $\mathbf{I}_{n}$ and $\mathbf{P}=\mathbf{R}^{T}$_.

This way of constructing the restriction operator yields efficient iterations
since in practice the coarse model in (4) can be formed by merely sampling
$n_{\ell}$ elements from $\mathbf{x}_{k}$ and $\nabla f(\mathbf{x}_{k})$ which
has negligible computational cost. (c) Taking advantage of the parameter
structure. Methods that employ full matrix preconditioners further reduce the
memory requirement by exploiting the matrix or tensor structure of the
parameters (for details see Shampoo Gupta et al. (2018)). For instance, for a
matrix setting, if the parameter $\mathbf{X}\in\mathbb{R}^{q\times d}$ and
select $\mathbf{R}\in\mathbb{R}^{n_{\ell}\times q}$, where $n_{\ell}<q$, then
we obtain
$\mathbf{G}_{\ell,k}:=\mathbf{R}\mathbf{G}_{k}\in\mathbb{R}^{n_{\ell}\times
d}$ and $\mathbf{Q}_{\ell,k}\in\mathbb{R}^{n_{\ell}\times n_{\ell}}$. This
results in $\mathcal{O}((r+d)n_{\ell}^{2})$ operations for forming the
preconditioner and applying the randomized SVD. Note that this number is much
smaller than $\mathcal{O}((d+q)(d^{2}+q^{2}))$ of Shampoo.

(a) Initialized at the origin

(b) Initialized randomly from $\mathcal{N}(0,1)$

(c) Gissete with $\mathbf{x}_{0}=0$

Figure 1: Convergence behaviour of various algorithms for the non-linear
least-squares problem.

##  3 Convergence Analysis

In this section we provide a simple convergence analysis of Simba when it
generates sequences using the coarse model in (4) (without momentum). We show
a linear convergence rate when the coarse model is constructed in both
deterministic and randomized manner. Our analysis is based on the classical
theory that assumes strongly convex functions and Lipschitz continuous
gradients. In deterministic scenarios, the method is expected to alternate
between coarse and fine steps to always reduce the value function. For this
reason, we present two convergence results: (a) when the coarse step is always
accepted, and (b) when only fine steps are taken. Hence the complete
convergence behaviour of Simba is provided. On the other hand, when the
prolongation operator is defined randomly at each iteration, multilevel
methods expected to converge using coarse steps only Ho et al. (2019);
Tsipinakis and Parpas (2021); Tsipinakis et al. (2023). Our numerical
experiments also verify this observation. Moreover, we derive the number of
steps needed for the method to reach accuracy $\epsilon>0$ for both cases. We
begin by stating our assumptions.

######  Assumption 3.1.

There are scalars $0<\mu<L<+\infty$ such that $f$ is $\mu$-strongly convex and
has $L$-Lipschitz continuous gradients if and only if:

  1. 1.

($L$-Lipschitz continuity) for all $\mathbf{x},\mathbf{y}\in\mathbb{R}^{n}$

$\|\nabla f(\mathbf{y})-\nabla f(\mathbf{x})\|_{2}\leq
L\|\mathbf{y}-\mathbf{x}\|_{2}$

  2. 2.

($\mu$-strong convexity) for all $\mathbf{x},\mathbf{y}\in\mathbb{R}^{n}$

$f(\mathbf{y})\geq f(\mathbf{x})+\langle\nabla
f(\mathbf{x}),\mathbf{y}-\mathbf{x}\rangle+\frac{\mu}{2}\|\mathbf{y}-\mathbf{x}\|_{2}$

Next we state the assumptions on the prolongation and restriction operators.

######  Assumption 3.2.

For the restriction and prolongation operators $\mathbf{R}$ and $\mathbf{P}$
it holds that $\mathbf{P}=\mathbf{R}^{T}$ and
$\operatorname{rank}(\mathbf{P})=n_{\ell}$.

The assumptions on the linear operators are not restrictive for practical
application. For instance, Definition 2.1 on $\mathbf{P}$ and $\mathbf{R}$
satisfies Assumption 3.2. The following assumption ensures that the algorithm
always selects effective coarse directions.

######  Assumption 3.3.

There exist $e>0$ and $\xi\in(0,\min(1,\|\mathbf{R}\|_{2}))$ such that if
$\|\nabla f(\mathbf{x})\|_{2}\neq 0$ it holds

$\|\mathbf{R}\nabla f(\mathbf{x})\|_{2}>\xi\|\nabla
f(\mathbf{x})\|_{2}\quad\text{and}\quad\|\mathbf{R}\nabla
f(\mathbf{x})\|_{2}>e.$

For our convergence result below we will need the following quantity: we
denote $\omega:=\max\\{\|\mathbf{R}\|_{2},\|\mathbf{P}\|_{2}\\}$.

######  Theorem 3.4.

Let $f:\mathbb{R}^{n}\rightarrow\mathbb{R}$ be a function such that Assumption
3.1 holds. Suppose also that Assumptions 3.2 and 3.3 hold. Moreover, given
$\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}$ in (6), define

$\displaystyle\hat{\mathbf{d}}_{k}$
$\displaystyle:=-\mathbf{P}_{k}\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}\mathbf{R}_{k}\nabla
f(\mathbf{x}_{k}),$

and suppose that the sequence $(\mathbf{x}_{k})_{k\in\mathbb{N}}$ is generated
by
$\mathbf{x}_{k+1}=x_{k}+\frac{\xi^{2}m}{L\sqrt{M}\omega^{4}}\hat{\mathbf{d}}_{k}$.
Then, there exists $\hat{c}\in(0,1)$ such that

$f(\mathbf{x}_{k+1})-f(\mathbf{x}^{*})\leq\hat{c}(f(\mathbf{x}_{k})-f(\mathbf{x}^{*})).$

Moreover, at most

$\hat{K}=\frac{\log\left(f(\mathbf{x}_{0})-f(\mathbf{x}^{*})/\epsilon\right)}{\log(1/\hat{c})}$

iterations are required for this process to reach accuracy $\epsilon$.

###### Proof.

Recall from the statement of the theorem that the sequence is generated by
$\mathbf{x}_{k+1}=\mathbf{x}_{k}+t_{k}\hat{\mathbf{d}}_{k}$, where
$\hat{\mathbf{d}}_{k}=-\mathbf{P}_{k}\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}\mathbf{R}_{k}\nabla
f(\mathbf{x}_{k})$, $t_{k}=\frac{\xi^{2}m}{\omega^{4}L\sqrt{M}}$ and
$\mathbf{Q}_{\ell}^{-1/2}$ is defined in (6). We also define the following
quantity

$\hat{\lambda}({\mathbf{x}}):=\sqrt{\nabla
f(\mathbf{x})^{T}\mathbf{P}\mathbf{Q}_{\ell}^{-\frac{1}{2}}\mathbf{R}\nabla
f(\mathbf{x})}.$

It holds $\hat{\lambda}({\mathbf{x}})\geq 0$. Below, we collect two general
results that will be useful later in the proof. Given the Lipschitz continuity
in Assumption 3.1 one can prove the following inequality (Nesterov et al.
(2018)):

$f(\mathbf{y})\leq f(\mathbf{x})+\nabla
f(\mathbf{x})^{T}(\mathbf{y}-\mathbf{x})+\frac{L}{2}\|\mathbf{y}-\mathbf{x}\|_{2}$
(7)

Similarly, from the strong convexity we can obtain

$f(\mathbf{y})\leq f(\mathbf{x})+\nabla
f(\mathbf{x})^{T}(\mathbf{y}-\mathbf{x})+\frac{1}{2\mu}\|\nabla
f(\mathbf{x})-\nabla f(\mathbf{y})\|_{2}$

Replacing $\mathbf{y}$ and $\mathbf{x}$ with $\mathbf{x}_{k}$ and
$\mathbf{x}^{*}$, respectively, we have that

$f(\mathbf{x}_{k})-f(\mathbf{x}^{*})\leq\frac{1}{2\mu}\|\nabla
f(\mathbf{x}_{k})\|_{2}$ (8)

Furthermore, by construction, $\mathbf{Q}_{\ell,k}$ is bounded. It is bounded
from below by $m$ from (5) and it is bounded from above by the Lipschitz
continuity of the gradients. Then, there exists $M\geq m>0$ such that for all
$k\geq 1$ we have that

$m\mathbf{I}\preceq\mathbf{Q}_{\ell,k}\preceq M\mathbf{I}$

The above inequality implies

$\frac{1}{\sqrt{M}}\mathbf{P}_{k}\mathbf{R}_{k}\preceq\mathbf{P}_{k}\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}\mathbf{R}_{k}\preceq\frac{1}{\sqrt{m}}\mathbf{P}_{k}\mathbf{R}_{k}$

Using the above inequality we can obtain the following bound

$\displaystyle\|\hat{\mathbf{d}}_{k}\|_{2}=\|\mathbf{P}_{k}\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}\mathbf{R}_{k}\nabla
f(\mathbf{x}_{k})\|_{2}\leq\|\mathbf{P}_{k}\mathbf{Q}_{\ell,k}^{-\frac{1}{2}}\mathbf{R}_{k}\|_{2}\|\nabla
f(\mathbf{x}_{k})\|_{2}\leq\frac{1}{\sqrt{m}}\|\mathbf{P}_{k}\mathbf{R}_{k}\|_{2}\|\nabla
f(\mathbf{x}_{k})\|_{2}\leq\frac{\omega^{2}}{\sqrt{m}}\|\nabla
f(\mathbf{x}_{k})\|_{2}.$ (9)

Similarly we can show a lower bound on $\hat{\lambda}({\mathbf{x}})$

$\hat{\lambda}({\mathbf{x}})^{2}=\nabla
f(\mathbf{x})^{T}\mathbf{P}\mathbf{Q}_{\ell}^{-\frac{1}{2}}\mathbf{R}\nabla
f(\mathbf{x})\geq\frac{1}{\sqrt{M}}\|\mathbf{R}\nabla
f(\mathbf{x})\|_{2}^{2}\geq\frac{\xi^{2}}{\sqrt{M}}\|\nabla
f(\mathbf{x}_{k})\|_{2}^{2},$ (10)

where the last inequality follows from Assumption 3.3. Using now (7) and the
fact that $\hat{\lambda}({\mathbf{x}})^{2}=-\nabla
f(\mathbf{x})^{T}\hat{\mathbf{d}}_{k}$ we take

$f(\mathbf{x}_{k+1})\leq
f(\mathbf{x}_{k})-t_{k}\hat{\lambda}(\mathbf{x}_{k})^{2}+t_{k}^{2}\frac{L}{2}\|\hat{\mathbf{d}}_{k}\|^{2}_{2}.$

Combining the above inequality with (10) and (9) we have that

$\displaystyle f(\mathbf{x}_{k+1})$ $\displaystyle\leq
f(\mathbf{x}_{k})-t_{k}\frac{\xi^{2}}{\sqrt{M}}\|\nabla
f(\mathbf{x})\|_{2}^{2}+t_{k}^{2}\frac{L\omega^{4}}{2m}\|\nabla
f(\mathbf{x}_{k})\|_{2}^{2}$
$\displaystyle=f(\mathbf{x}_{k})-\frac{\xi^{4}m}{2\omega^{4}ML}\|\nabla
f(\mathbf{x}_{k})\|_{2}^{2},$

where the last inequality follows by the definition of $t_{k}$. Adding and
subtracting $f(\mathbf{x}^{*})$ on the above relationship and incorporating
inequality (8) we obtain

$f(\mathbf{x}_{k+1})-f(\mathbf{x}^{*})\leq\hat{c}(f(\mathbf{x}_{k})-f(\mathbf{x}^{*})),$

where $\hat{c}:=1-\frac{\xi^{4}m\mu}{\omega^{4}ML}$. Since $\xi<\omega,m<M$
and $\mu<L$ we take $\hat{c}\in(0,1)$. Unravelling the last inequality we get

$f(\mathbf{x}_{k})-f(\mathbf{x}^{*})\leq\hat{c}^{k}(f(\mathbf{x}_{0})-f(\mathbf{x}^{*})),$

and thus $\lim_{k\rightarrow\infty}f(\mathbf{x}_{k})=f(\mathbf{x}^{*})$.
Finally, solving for $k$ the inequality
$\hat{c}^{k}(f(\mathbf{x}_{0})-f(\mathbf{x}^{*}))\leq\epsilon$, we conclude
that at least

$\hat{K}=\frac{\log\left(f(\mathbf{x}_{0})-f(\mathbf{x}^{*})/\epsilon\right)}{\log(1/\hat{c})}$

steps are required for this process to achieve accuracy $\epsilon$. ∎

A direct consequence of the above theorem is convergence in expectation when
$(\mathbf{x}_{k})_{k\geq 1}$ is generated randomly via a random prolongation
matrix, e.g., see Definition 2.1. In this case we can guarantee that

$\mathbb{E}[f(\mathbf{x}_{k})]-f(\mathbf{x}^{*})\leq\hat{c}^{k}(f(\mathbf{x}_{0})-f(\mathbf{x}^{*})),$

which implies that
$\lim_{k\to\infty}\mathbb{E}[f(\mathbf{x}_{k})]=f(\mathbf{x}^{*})$.

Table 1: Mean training error and standard deviation for the Non-Linear Least-Squares problem. The results were obtained over 5 runs for random initialization from $\mathcal{N}(0,1)$. Algorithm | Mean $\pm$ Std
---|---
Adam | $0.0158\pm 0.0016$
AdaHessian | $0.0325\pm 0.0029$
Simba | $0.0097\pm 0.001$
Apollo | $0.0556\pm 0.0018$

Theorem 3.4 effectively shows the number of steps required for the method to
reach the desired accuracy when the coarse direction is always effective.
However, this may not be always true, which, in deterministic settings,
implies no progress for the method. As discussed previously $\xi$ and $e$ can
be viewed as user-defined parameters that prevent the method from taking the
ineffective coarse steps. Given fixed $\xi$ and $e$, the method performs
iterations in the fine level if one of the two conditions in Assumption 3.3 is
violated. In the fine level, the method constructs the preconditioner as
follows

$\mathbf{Q}_{k}:=\nabla f(\mathbf{x}_{k})\nabla f(\mathbf{x}_{k})^{T},$

and then $\mathbf{Q}_{k}^{-\frac{1}{2}}$ is constructed exactly as in
$\mathbf{Q}_{\ell,k}$ in (6). The next theorem shows a linear rate and derives
the number of steps required when the method takes only fine steps.

######  Theorem 3.5.

Let $f:\mathbb{R}^{n}\rightarrow\mathbb{R}$ be a function such that Assumption
3.1 holds. Suppose also that Assumption 3.2 holds. Moreover, define

$\displaystyle\mathbf{d}_{k}:=-\mathbf{Q}_{k}^{-\frac{1}{2}}\nabla
f(\mathbf{x}_{k}),$

and suppose that the sequence $(\mathbf{x}_{k})_{k\in\mathbb{N}}$ is generated
by $\mathbf{x}_{k+1}=x_{k}+\frac{m}{L\sqrt{M}}\hat{\mathbf{d}}_{k}$. Then,
there exists $c\in(0,1)$ such that

$f(\mathbf{x}_{k+1})-f(\mathbf{x}^{*})\leq
c(f(\mathbf{x}_{k})-f(\mathbf{x}^{*})).$

Moreover, at most

$K=\frac{\log\left(f(\mathbf{x}_{0})-f(\mathbf{x}^{*})/\epsilon\right)}{\log(1/c)}$

iterations are required for this process to reach accuracy $\epsilon$.

###### Proof.

The proof of theorem as it follows analogously to Theorem 3.4. The difference
is the term that controls the linear rate which is now given by
$c:=1-\frac{m\mu}{ML}$. It holds $0<\hat{c}\leq c<1$. ∎

As expected, theorem 3.5 shows a faster linear rate since the entire local
information is employed during the training. Combining theorems 3.4 and 3.5,
we provide the complete picture of the linear convergence rate of the proposed
method.

Figure 2: Convergence behaviour of various algorithms for the CURVES and MNIST autoencoders.  Table 2: Mean training error and standard deviation over $20$ and $5$ runs for the CURVES and MNIST autoencoders, respectively The last column reports the wall-clock time for a single run of the two autoencoders problem. Algorithm | CURVES | MNIST | Seconds CURVES/MNIST
---|---|---|---
Adam | $0.0057\pm 0.0017$ | $0.0431\pm 0.0007$ | $22/176$
AdaHessian | $0.0165\pm 0.0049$ | $0.0604\pm 0.0044$ | $48/224$
Shampoo | $-$ | $-$ | $606/12,442$
Apollo | $0.02309\pm 0.00159$ | $0.0672\pm 6.3\times 10^{-7}$ | $31/197$
Simba | $0.00250\pm 0.0005$ | $0.0043\pm 8.1\times 10^{-5}$ | $54/467$

##  4 Numerical Experiments

In this section we validate the efficiency of Simba to a number machine
learning problems. Our goal is to illustrate that Simba outperforms the state-
of-the-art diagonal optimization methods for problems with saddle points and
flat areas or when the gradients are vanishing. For this purpose, we consider
a non-linear least squares problems and two deep autoencoders where
optimization methods often converge to suboptimal solutions. In addition, we
demonstrate that our method is efficient and offers comparable, if not better,
generalization errors compared to the state-of-the-art optimization methods on
standard benchmark ResNets using CIFAR10 and CIFAR100 datasets.

Algorithms and set up: We compare Simba with momentum (algorithm 1) against
Adam Kingma and Ba (2014), AdaHessian Yao et al. (2021), Apollo Ma (2020) and
Shampoo Gupta et al. (2018) on a Tesla T4 GPU with a 16G RAM. The GPU time for
Shampoo is not comparable to that of the other algorithm and hence its
behaviour is reported only for the autoencoder problems. For all algorithms,
the learning rate was selected using a grid search for $t_{k}\in${1e-4, 5e-4,
1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1}. For all the algorithms we set the
momentum parameters to their default values. For Simba we set $r=20$ in all
experiments. Through all the experiments the batch size is set to $128$. For
all algorithms we comprehensively tuned $\operatorname{eps}$; the default
value is selected when others do not yield an improvement. In our case we
denote $m\equiv\operatorname{eps}$.

###  4.1 Non-linear least-squares

Given a training dataset $\\{\mathbf{a}_{i},b_{i}\\}_{i=1}^{m}$,
$\mathbf{a}_{i}\in\mathbb{R}^{n}$ and $b_{i}\in\mathbb{R}$, we consider
solving the following non-linear least-squares problem

$\min_{\mathbf{x}\in\mathbb{R}^{n}}\frac{1}{m}\sum_{i=1}^{m}\left(b_{i}-g(\mathbf{a}_{i}^{T}\mathbf{x})\right)^{2},\quad
g(\omega):=\frac{1}{1+\exp(\omega)},$

which is a non-convex optimization problem. Here, we consider the Gisette
dataset 111Dataset available at:
https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/binary.html for which
$m=6000$ and $n=5000$. Furthermore, for Adam we select $t_{k}=0.001$ and
$\operatorname{eps}=10^{-8}$ while for Apollo we set $t_{k}=0.01$ and
$\operatorname{eps}=10^{-4}$. For AdaHessian, $t_{k}=0.1$ and the hessian
power parameter is set to to $0.5$. For Simba we select
$n_{\ell}=250,t_{k}=0.05$ and $\operatorname{eps}=10^{-12}$. The performance
of optimization methods appear in Figure 1. Observe that this problems has
several flat areas or saddle points which slow down the convergence of all
algorithms. Nevertheless, the method with the best behaviour is Simba which
enjoys a much faster escape rate from the saddle points and thus always
returns lower training errors. This can be also observed for different
initialization points. In Table 1 we report the average training error and the
standard deviation over $5$ runs. The method that comes closer to Simba is
Adam. Further, from Figure 1 we see that, although the wall-clock time of
Simba is two to three times increased compared to its competitors, the total
time of our method is much better due to its fast escape rate near saddles and
flat areas. Further, in Figure 4 in the supplementary material we compare the
performance of Simba for different sizes in the coarse model. We observed that
even Simba with $n_{\ell}=25$ (i.e., updating only $25$ parameters at each
iteration) enjoys better escape rate than its competitors. This indicates that
diagonal methods perform poorly near saddles and flat areas and highlights the
importance of constructing meaningful preconditioners. Moreover, observe that
the best wall-clock time is achieved with only $5\%$ of the dimensions. This
indicates that Simba significantly reduces the computational cost of
preconditioned methods without compromising the convergence rate.

Figure 3: Convergence behaviour of various algorithms on CIFAR10 using ResNet110 and CIFAR100 using ResNet9. Table 3: Mean training error and accuracy $\pm$ standard deviation of various optimization algorithms over $5$ on CIFAR10 and CIFAR100, respectively. | ResNet-110 - CIFAR10 | ResNet-9 - CIFAR100
---|---|---
Algorithm | Training Error | Accuracy | Training Error | Accuracy
Adam | $0.0284\pm 0.0060$ | $91.03\pm 0.16$ | $0.0237\pm 0.0005$ | $73.91\pm 0.18$
AdaHessian | $0.0560\pm 0.0427$ | $91.34\pm 1.04$ | $0.0081\pm 0.0036$ | $72.93\pm 0.66$
Apollo | $0.0315\pm 0.0066$ | $91.18\pm 0.25$ | $0.0076\pm 0.0003$ | $73.70\pm 0.19$
Simba | $0.0097\pm 0.0023$ | $92.07\pm 0.09$ | $0.0133\pm 0.0001$ | $75.41\pm 0.22$

###  4.2 Deep autoencoders

In this section we investigate the performance of Simba on two deep
autoencoder optimization problems that arise from (Hinton and Salakhutdinov
(2006)), named CURVES and MNIST. Both optimization problems are considered
difficult to solve and have become a standard benchmark for optimization
methods due to the presence of saddle points and the issue of the vanishing
gradients. The CURVES222Dataset available at:
www.cs.toronto.edu/~jmartens/digs3pts_1.mat autoencoder consists of an encoder
with layer size $28\times 28,400,200,100,50,25,6$, a symmetric decoder which
totals to $0.8$M parameters. For all layers the $\operatorname{relu}$
activation function is applied. The network was trained on $20,000$ images and
tested on $10,000$. We set the learning rate equal to $0.001,1,0.05,0.5$ and
$0.01$ for Adam, Apollo, AdaHessian, Shampoo and Simba, respectively. The
$\operatorname{eps}$ parameter is selected $10^{-15}$ and $10^{-8}$ for Apollo
and Simba, respectively, and set to its default value for the rest algorithms.
The hessian power parameter for AdaHessian is selected $1$. Further, the
MNIST333Dataset available at: https://pytorch.org/vision/main/datasets.html
autoencoder consists of an encoder with layer size $28\times
28,1000,500,250,30$ and a symmetric decoder which totals to $2.8$M parameters.
In this network we use the $\operatorname{sigmoid}$ activation function to all
layers. The training set for the MNIST dataset consists of $60,000$ images
while the test set has $10,000$ images. Here, $t_{k}$ is set equal to
$0.001,0.9,0.1,0.05$ and $0.05$ for Adam, Apollo, AdaHessian, Shampoo and
Simba, respectively. The $\operatorname{eps}$ parameter is set equal to
$10^{-8}$ for Simba and to its default value for all other algorithms; the
hessian power of AdaHessian is set equal to $1$. For both autoencoder problems
the coarse model size parameter is set $n_{\ell}=0.5N$.

The comparison between the optimization algorithms for the two deep
autoencoders appears in Table 2 and Figure 2. Clearly, Simba performs the best
among the optimization algorithms resulting in the lowest training error in
both CURVES and MNIST autoencoders. As a result, Simba performs better than
its competitors on the test set too. It is evident that all the other methods
get stuck in stationary points with large errors. In these points we observed
that the gradients become almost zero which results in a very slow progress of
Adam, AdaHessian, Apollo and Shampoo. On the other, Simba enjoys a very fast
escape rate from such points and thus it yields a very fast decrease in the
training error. In addition, Table 2 reports the wall-clock time for each
optimization algorithm. We see that Simba has an increased per-iteration costs
(by a factor between two and three compared to Adam), nevertheless this is a
reasonable price to pay for achieving a good decrease in the value function.
The averaged behaviour of Shampoo is missing due to its expensive iterations.
For this reason we omit Shampoo from the following experiments.

###  4.3 Residual Neural Networks

In this section we report comparisons on the convergence and generalization
performance between the optimization algorithms using ResNet110444Network
implementation available at:
https://github.com/akamaster/pytorch_resnet_cifar10 and ResNet9555Network
implementation available at:
https://www.kaggle.com/code/yiweiwangau/cifar-100-resnet-
pytorch-75-17-accuracy with CIFAR-10 and CIFAR-100 datasets, respectively.
ResNet110 consists of 110 layers and a total of $1.7$M parameters while
ResNet9 has 9 layers and $30$M parameters. Hence, our goal is to investigate
the performance of Simba on benchmark problems and different architectures,
i.e., deep and wide networks.

For ResNet110 we set an initial learning rate parameter equal to
$0.001,0.1,0.1,0.01$ for Adam, Apollo, AdaHessian and Simba, respectively. We
train the network for a total number of $60$ epochs and decrease the learning
rate by a factor of $5$ at epochs $30$ and $40$. The hessian power parameter
for AdaHessian is selected $1$ and for Simba the hyper-parameter
$\operatorname{eps}$ is set to $10^{-8}$ while $n_{\ell}=0.5N$. The
convergence and generalization performance between the optimization algorithms
is illustrated in Figure 3 and Table 3. We see that Simba is able to achieve
results that are comparable, if not better, with that of Adam, AdaHessian and
Apollo in both training and generalization errors. The total GPU time of Simba
is about three times larger than that of Adam, which is the fastest algorithm,
but it is considerably smaller than that of AdaHessian (see Figure 5 in the
supplement).

For ResNet9, we set an initial value on $t_{k}$ as follows:
$0.0005,0.001,0.05,0.005$ for Adam, Apollo, AdaHessian and Simba,
respectively. We train the network for $60$ epochs and use cosine annealing to
determine the learning rate at each epoch and set minimum $t_{k}$ value equal
to $0.01t_{k},0.1t_{k},0.01t_{k},0.02t_{k}$ for Adam, Apollo, AdaHessian and
Simba, respectively. The hessian power, $\operatorname{eps}$ and $n_{\ell}$
parameters are set as above. Since classifying CIFAR100 is a much more
difficult problem than CIFAR10, to improve the accuracy we use weight decay
and gradient clipping with values fixed at $0.001$ and $0.01$, respectively,
for all algorithms. Figure 3 and Table 3 indicate that Simba is able to offer
very good generalization results on a difficult task. Further, we see that
Simba is less than two times slower in wall-clock time than Adam and Apollo
but again considerably faster than AdaHessian (see Figure 6 in the
supplement). Figure 6 also shows that the total GPU time of Simba for reaching
the desired classification accuracy is comparable to that of Adam and Apollo
which indicates the efficiency of our method on large-scale optimization
problems.

##  5 Conclusions

We present Simba, a scalable bilevel optimization algorithm to address the
limitations of first-order methods in non-convex settings. We demonstrate the
fast escape rate from saddles of our method through empirical evidence.
Numerical results also indicate that Simba achieves good generalization errors
on modern machine learning applications. Convergence guarantees of Simba when
locally assume convex functions are also established. As a future work, we aim
to apply Simba for the training of LLMs such as GPT.

## References

  * Dauphin et al. [2014] Yann N Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio.  Identifying and attacking the saddle point problem in high-dimensional non-convex optimization.  _Advances in neural information processing systems_ , 27, 2014.
  * Panageas et al. [2019] Ioannis Panageas, Georgios Piliouras, and Xiao Wang.  First-order methods almost always avoid saddle points: The case of vanishing step-sizes.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Nesterov et al. [2018] Yurii Nesterov et al.  _Lectures on convex optimization_ , volume 137.  Springer, 2018.
  * Boyd and Vandenberghe [2004] Stephen Boyd and Lieven Vandenberghe.  _Convex optimization_.  Cambridge University Press, Cambridge, 2004.  ISBN 0-521-83378-7.  doi:10.1017/CBO9780511804441.  URL https://doi.org/10.1017/CBO9780511804441.
  * Nesterov [2004] Yurii Nesterov.  _Introductory lectures on convex optimization_ , volume 87 of _Applied Optimization_.  Kluwer Academic Publishers, Boston, MA, 2004.  ISBN 1-4020-7553-7.  doi:10.1007/978-1-4419-8853-9.  URL https://doi.org/10.1007/978-1-4419-8853-9.  A basic course.
  * Duchi et al. [2011] John Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  _Journal of machine learning research_ , 12(7), 2011.
  * Tieleman et al. [2012] Tijmen Tieleman, Geoffrey Hinton, et al.  Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude.  _COURSERA: Neural networks for machine learning_ , 4(2):26–31, 2012.
  * Zeiler [2012] Matthew D Zeiler.  Adadelta: an adaptive learning rate method.  _arXiv preprint arXiv:1212.5701_ , 2012.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Zaheer et al. [2018] Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar.  Adaptive methods for nonconvex optimization.  _Advances in neural information processing systems_ , 31, 2018.
  * Yao et al. [2021] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney.  Adahessian: An adaptive second order optimizer for machine learning.  In _proceedings of the AAAI conference on artificial intelligence_ , volume 35, pages 10665–10673, 2021.
  * Ma [2020] Xuezhe Ma.  Apollo: An adaptive parameter-wise diagonal quasi-newton method for nonconvex stochastic optimization.  _arXiv preprint arXiv:2009.13586_ , 2020.
  * Jahani et al. [2021] Majid Jahani, Sergey Rusakov, Zheng Shi, Peter Richtárik, Michael W Mahoney, and Martin Takáč.  Doubly adaptive scaled algorithm for machine learning using second-order information.  _arXiv preprint arXiv:2109.05198_ , 2021.
  * Liu et al. [2023] Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma.  Sophia: A scalable stochastic second-order optimizer for language model pre-training.  _arXiv preprint arXiv:2305.14342_ , 2023.
  * Broyden [1967] Charles G Broyden.  Quasi-newton methods and their application to function minimisation.  _Mathematics of Computation_ , 21(99):368–381, 1967.
  * Erdogdu and Montanari [2015] Murat A Erdogdu and Andrea Montanari.  Convergence rates of sub-sampled newton methods.  In _Proceedings of the 28th International Conference on Neural Information Processing Systems-Volume 2_ , pages 3052–3060. MIT Press, 2015.
  * Pilanci and Wainwright [2017] Mert Pilanci and Martin J Wainwright.  Newton sketch: A near linear-time optimization algorithm with linear-quadratic convergence.  _SIAM Journal on Optimization_ , 27(1):205–245, 2017.
  * Xu et al. [2016] Peng Xu, Jiyan Yang, Farbod Roosta-Khorasani, Christopher Ré, and Michael W Mahoney.  Sub-sampled newton methods with non-uniform sampling.  In _Advances in Neural Information Processing Systems_ , pages 3000–3008, 2016.
  * Xu et al. [2020] Peng Xu, Fred Roosta, and Michael W Mahoney.  Second-order optimization for non-convex machine learning: An empirical study.  In _Proceedings of the 2020 SIAM International Conference on Data Mining_ , pages 199–207. SIAM, 2020.
  * Gupta et al. [2018] Vineet Gupta, Tomer Koren, and Yoram Singer.  Shampoo: Preconditioned stochastic tensor optimization.  In _International Conference on Machine Learning_ , pages 1842–1850. PMLR, 2018.
  * Tsipinakis and Parpas [2021] Nick Tsipinakis and Panos Parpas.  A multilevel method for self-concordant minimization.  _arXiv preprint arXiv:2106.13690_ , 2021.
  * Tsipinakis et al. [2023] Nick Tsipinakis, Panagiotis Tigkas, and Panos Parpas.  A multilevel low-rank newton method with super-linear convergence rate and its application to non-convex problems.  _arXiv preprint arXiv:2305.08742_ , 2023.
  * Halko et al. [2011] N. Halko, P. G. Martinsson, and J. A. Tropp.  Finding structure with randomness: probabilistic algorithms for constructing approximate matrix decompositions.  _SIAM Rev._ , 53(2):217–288, 2011.  ISSN 0036-1445.  doi:10.1137/090771806.  URL https://doi.org/10.1137/090771806.
  * Bubeck et al. [2015] Sébastien Bubeck et al.  Convex optimization: Algorithms and complexity.  _Foundations and Trends® in Machine Learning_ , 8(3-4):231–357, 2015.
  * Reddi et al. [2018] Sashank Reddi, Manzil Zaheer, Suvrit Sra, Barnabas Poczos, Francis Bach, Ruslan Salakhutdinov, and Alex Smola.  A generic approach for escaping saddle points.  In _International conference on artificial intelligence and statistics_ , pages 1233–1242. PMLR, 2018.
  * O’Leary-Roseberry et al. [2020] Thomas O’Leary-Roseberry, Nick Alger, and Omar Ghattas.  Low rank saddle free newton: A scalable method for stochastic nonconvex optimization.  _arXiv preprint arXiv:2002.02881_ , 2020.
  * Martens and Grosse [2015] James Martens and Roger Grosse.  Optimizing neural networks with kronecker-factored approximate curvature.  In _International conference on machine learning_ , pages 2408–2417. PMLR, 2015.
  * Ho et al. [2019] Chin Pang Ho, Michal Kočvara, and Panos Parpas.  Newton-type multilevel optimization method.  _Optimization Methods and Software_ , pages 1–34, 2019.
  * Nesterov and Polyak [2006] Yurii Nesterov and Boris T Polyak.  Cubic regularization of newton method and its global performance.  _Mathematical Programming_ , 108(1):177–205, 2006\.
  * Paternain et al. [2019] Santiago Paternain, Aryan Mokhtari, and Alejandro Ribeiro.  A newton-based method for nonconvex optimization with fast evasion of saddle points.  _SIAM Journal on Optimization_ , 29(1):343–368, 2019.
  * Hinton and Salakhutdinov [2006] Geoffrey E Hinton and Ruslan R Salakhutdinov.  Reducing the dimensionality of data with neural networks.  _science_ , 313(5786):504–507, 2006.

##  Appendix A Extra Numerical Results

Here, we report some extra numerical results. Figure 4 shows how the size of
the coarse model affects the convergence behaviour of Simba on the non-linear
least squares problem. Additionally, Figures 5 and 6 show the wall-clock time
for training the residual neural networks for 60 epochs.

Figure 4: Convergence behaviour of Simba for different coarse models. In all
cases Simba is initialized at the origin. The parameter $r$ is set to $10$ for
$n_{\ell}=0.005N$ and $20$ otherwise. Further, $t_{k}$ is selected $0.01$ for
$n=0.5N$ or $n=0.1N$ and $0.05$ otherwise. Interestingly, we see that Simba
with $n=25$ matches the performance of Simba that use much larger coarse
models. Note that, by Definition 2.1, Simba only updates $25$ parameters at
each iteration which indicates that near saddles and flat areas one should
employ more informative directions than those of the first-order methods with
second-order momentum. As expected, we see that a better convergence and
escape rate is expected for larger coarse models but this comes at the cost of
more expensive iterations. The figure shows that Simba with $n=250$ has the
best total complexity.

Figure 5: Convergence behaviour and wall-clock time of various algorithms on
CIFAR10 using ResNet110.

Figure 6: Convergence behaviour and wall-clock time of various algorithms on
CIFAR100 using ResNet9.