# Nonconvex Stochastic Bregman Proximal Gradient Method with Application to
Deep Learning

Kuangyu Ding<EMAIL_ADDRESS>
Department of Mathematics
University of Singapore
10 Lower Kent Ridge Road, Singapore 119076 Jingyang Li<EMAIL_ADDRESS>
Department of Mathematics
University of Singapore
10 Lower Kent Ridge Road, Singapore 119076 Kim-Chuan Toh<EMAIL_ADDRESS>
Department of Mathematics Institute of Operations Research and Analytics
University of Singapore
10 Lower Kent Ridge Road, Singapore 119076

###### Abstract

The widely used stochastic gradient methods for minimizing nonconvex composite
objective functions require the Lipschitz smoothness of the differentiable
part. But the requirement does not hold true for problem classes including
quadratic inverse problems and training neural networks. To address this
issue, we investigate a family of stochastic Bregman proximal gradient (SBPG)
methods, which only require smooth adaptivity of the differentiable part. SBPG
replaces the upper quadratic approximation used in SGD with the Bregman
proximity measure, resulting in a better approximation model that captures the
non-Lipschitz gradients of the nonconvex objective. We formulate the vanilla
SBPG and establish its convergence properties under nonconvex setting without
finite-sum structure. Experimental results on quadratic inverse problems
testify the robustness of SBPG. Moreover, we propose a momentum-based version
of SBPG (MSBPG) and prove it has improved convergence properties. We apply
MSBPG to the training of deep neural networks with a polynomial kernel
function, which ensures the smooth adaptivity of the loss function.
Experimental results on representative benchmarks demonstrate the
effectiveness and robustness of MSBPG in training neural networks. Since the
additional computation cost of MSBPG compared with SGD is negligible in large-
scale optimization, MSBPG can potentially be employed as an universal open-
source optimizer in the future.

Keywords: Nonconvex stochastic algorithm, Bregman distance, Smooth adaptivity,
Deep neural network, Gradient explosion

##  1 Introduction

In this paper, we present and analyze a family of nonconvex stochastic Bregman
proximal gradient methods (SBPG) for solving the following generic stochastic
minimization problem:

$\min_{{\bm{x}}\in\overline{C}}\;\mathbb{E}_{{\bm{\xi}}}[f({\bm{x}},{\bm{\xi}})]+R({\bm{x}}),$
(1)

where $f(\cdot,{\bm{\xi}})$ is a nonconvex differentiable function on
$\overline{C}$, $R$ is a proper lower-semicontinuous convex function,
${\bm{\xi}}$ is a random variable, and $\overline{C}$ is the closure of $C$,
which is a nonempty convex open subset of $\mathbb{R}^{d}$. We denote
$F({\bm{x}})\coloneqq\mathbb{E}_{{\bm{\xi}}}[f({\bm{x}},{\bm{\xi}})]$, and
$\Phi({\bm{x}})\coloneqq F({\bm{x}})+R({\bm{x}})$. This stochastic
minimization problem, where an optimizer has limited access to the
distribution of ${\bm{\xi}}$ and can only draw samples from it, is prevalent
in the fields of machine learning and statistics (Hastie et al., 2009; Shapiro
et al., 2021; Zhang, 2004). In many instances, the smooth part of the
objective function $F({\bm{x}})$ can be formulated as a finite-sum structure
$F({\bm{x}})=\frac{1}{n}\sum_{i=1}^{n}f_{i}({\bm{x}})$. However, when $n$ is
extremely large, calculating the true gradient for the smooth part of the
objective function becomes extremely expensive. As a result, stochastic first-
order methods, which trace back to the work of Robbins and Monro (1951), have
become the prevailing approach for solving these large-scale optimization
problems. In particular, stochastic (proximal) gradient descent and its
numerous variants (Duchi et al., 2011; Duchi and Singer, 2009; Gu et al.,
2020; Kingma and Ba, 2014; Allen-Zhu, 2018; Wang et al., 2022) have been
widely utilized in large-scale stochastic optimization for machine learning
(LeCun et al., 2015; Shapiro et al., 2021; Zhang, 2004). From modeling
perspective, stochastic gradient descent can be viewed as minimizing a
sequence of upper quadratic approximations of the nonconvex objective
$\Phi({\bm{x}})$:

${{\bm{x}}^{k+1}}=\operatorname*{argmin}_{{\bm{x}}\in\overline{C}}\left\\{\underbrace{F({{\bm{x}}^{k}},{\bm{\Xi}}_{k})+\langle\widetilde{\nabla}_{k},\,{\bm{x}}-{{\bm{x}}^{k}}\rangle+\frac{1}{2\alpha_{k}}\|{\bm{x}}-{{\bm{x}}^{k}}\|^{2}}_{F_{{{\bm{x}}^{k}}}({\bm{x}}):\;model\;of\;F\;at\;{{\bm{x}}^{k}}}+R({\bm{x}})\right\\},$
(2)

where
$F({\bm{x}}^{k},{\bm{\Xi}}_{k})\coloneqq\frac{1}{|{\bm{\Xi}}_{k}|}\sum_{{\bm{\xi}}\in{\bm{\Xi}}_{k}}f({\bm{x}}^{k},{\bm{\xi}})$,
${\bm{\Xi}_{k}}$ is the set of samples of ${\bm{\xi}}$ at the $k$-th
iteration, and $\widetilde{\nabla}_{k}$ is an estimator of the exact gradient
$\nabla F(x^{k})$. This modeling perspective is well-known in deterministic
optimization, and has been used in methods such as Newton method, Gauss-Newton
method, bundle methods, and trust-region methods, as discussed in various
sources such as Hiriart-Urruty and Lemaréchal (1993); Nesterov (2003); Lin et
al. (2007); Paren et al. (2022).

Despite being widely used, stochastic gradient methods (2) still have several
well-known bottlenecks both in theory and practice. One of the crucial
requirements for analyzing stochastic gradient methods is the assumption of
Lipschitz continuity of the gradient of the differentiable part, which is
essential for ensuring the convergence of the algorithm. However, this
assumption does not always hold true. For example, even the seemingly simple
function $F(x)=x^{4}$ does not admit a globally Lipschitz gradient over
$\mathbb{R}$, which highlights the challenges of analyzing stochastic gradient
methods. In addition, choosing the appropriate stepsize is another challenge
in the practical usage of stochastic gradient methods. The stepsize has a
significant impact on the convergence performance of the algorithm, and
finding the optimal stepsize can be a time-consuming process. Engineers may
have to conduct multiple experiments to determine the optimal stepsize,
further compounding the complexity of the problem at hand.

To address these issues, classical approaches often resort to either line
search or more complicated inner loops, but these methods can negatively
impact the efficiency of the algorithm or even becomes intractable in a
stochastic setting. For instance, stochastic proximal point algorithm (PPA)
model the approximation of $F({\bm{x}})$ in (2) as
$F_{{{\bm{x}}^{k}}}({\bm{x}})=F({\bm{x}},{\bm{\xi}}_{k})+\frac{1}{2\alpha_{k}}\|{\bm{x}}-{{\bm{x}}^{k}}\|^{2}$
(Bertsekas, 2011; Bianchi, 2016; Patrascu and Necoara, 2017; Rockafellar,
1976), which makes the selection of stepsize $\alpha_{k}$ more robust than the
original model (2). However, the application of stochastic PPA is limited due
to the difficulty of solving the subproblems, particularly when dealing with
complicated objective functions, such as training deep neural networks. In
such cases, solving the subproblem is almost as difficult as solving the
original problem, rendering the approach impractical. Recently, Bauschke et
al. (2017); Lu et al. (2018) have proposed using Bregman proximity measures to
relax the assumption of gradient Lipschitz continuity to smooth adaptivity.
The Bregman gradient method was first introduced as the mirror descent scheme
by Nemirovskij and Yudin (1983) for minimizing convex nonsmooth functions.
From the modeling perspective, Bregman methods consider the following
subproblem at each iteration:

${{\bm{x}}^{k+1}}=\operatorname*{argmin}_{{\bm{x}}\in\overline{C}}\left\\{\underbrace{F({{\bm{x}}^{k}},{\bm{\Xi}}_{k})+\langle\widetilde{\nabla}_{k},\,{\bm{x}}-{{\bm{x}}^{k}}\rangle+\frac{1}{\alpha_{k}}\mathcal{D}_{\phi}({\bm{x}},{{\bm{x}}^{k}})}_{F_{{{\bm{x}}^{k}}}({\bm{x}}):\;model\;of\;F\;at\;{{\bm{x}}^{k}}}+R({\bm{x}})\right\\},$
(3)

where $\mathcal{D}_{\phi}$ is the Bregman distance induced by the kernel
function $\phi$. To convey the advantage of the Bregman proximity model, we
present a toy example. Consider objective function $F(x)=x^{4}$ with non-
Lipschitz gradient continuity, we compare the effectiveness of the upper
quadratic approximation model (2) and Bregman proximity model (3). As apparent
from Figure (1)(a), the Bregman proximity model (3) ($F_{2}(x)$) with the
kernel function $\phi(x)=\frac{1}{2}x^{2}+\frac{1}{4}x^{4}$ can provide a more
suitable approximation for $F(x)$ than the upper quadratic approximation model
2 ($F_{1}(x)$), as the yellow curve stays closer to the curve of the objective
function $F(x)=x^{4}$. With this improved approximation, ${x^{k+1}}$ generated
by the Bregman gradient method can make more significant progress towards
approaching the optimal solution ($x^{*}=0$) than ${x^{k+1}}$ generated by the
gradient descent method, as depicted in Figure 1(b).

(a) Two models’ approximations of $F(x)$ (b) Zoomed-in of version plot (a)

Figure 1: For function $F(x)=x^{4}$, which does not admit a globally Lipschitz
continuous gradient. We restrict the feasible set to $[-0.5,2]$. Consider the
models (2) and (3) of $F$ at ${x^{k}}=1$. The Lipschitz constant of $F$ with
respect to the kernel $\phi(x)=\frac{1}{2}x^{2}$ is 48. The Lipschitz constant
of $F$ with respect to the kernel $\phi(x)=\frac{1}{2}x^{2}+\frac{1}{4}x^{4}$
is 4. The figure in (b) is a zoomed-in version of the plot in (a) for the
range $[0.6,1]$. The unique minimum of $F(x)$ is at $x=0$.

While several stochastic extensions of Bregman methods that are based on
smooth adaptivity assumption have been developed recently, the current
literature primarily focuses on stochastic convex problems (Dragomir et al.,
2021b; Hanzely and Richtárik, 2021; Lu, 2019). The only existing convergence
analysis of Bregman methods for nonconvex problems (Latafat et al., 2022)
requires a finite-sum structure and a novel equivalent consensus
reformulation. Moreover, it is memory-intensive and requires essentially
periodic computation of the full gradient, which is expensive for large-scale
problems such as deep neural networks (Defazio and Bottou, 2019). As we can
see, stochastic Bregman methods have not been fully explored in the context of
modern large-scale nonconvex problems such as training neural networks, and
rigorous numerical evaluations of their performance are limited. Furthermore,
current literature lacks attention towards the robustness of stochastic
Bregman methods, particularly in terms of selecting stepsizes and initial
points, which can significantly impact their performance in large-scale
problems.

In this paper, we consider stochastic Bregman proximal gradient methods (SBPG)
for nonconvex problems (without finite-sum structure requirement) with
application to the training of deep neural networks. We establish the
convergence result of a vanilla SBPG without Lipschitz smoothness assumption
for nonconvex problems. Moreover, we propose a momentum-based SBPG (denoted as
MSBPG) and prove that it has improved convergence properties compared with
vanilla SBPG. We apply MSBPG to the training of deep neural networks with a
polynomial kernel function, which ensures the smooth adaptivity of the loss
function. According to our analysis, MSBPG can improve the robustness of
training neural networks by mitigating the gradient explosion phenomenon and
improve generalization performance by adopting Bregman proximity approximation
to the loss function locally. For numerical illustrations, we conduct
numerical experiments on quadratic inverse problems (QIP) and testify vanilla
SBPG’s robustness to stepsize selection and initial point’s scaling. We also
conduct extensive experiments on training CNNs for image classification and
LSTMs for language modeling by employing MSBPG, which is well-suited for
solving large-scale problems. Experimental results on representative
benchmarks show that our MSBPG has excellent generalization performance,
outperforming the most frequently used optimization algorithms, including SGD
(Robbins and Monro, 1951), Adam (Kingma and Ba, 2014), and AdamW (Loshchilov
and Hutter, 2017). Furthermore, MSBPG is demonstrated to be robustness to
large stepsize and initial point’s scaling, which are the common reasons
behind gradient explosion.

To summarize, our contributions are as follows:

  * 1.

We investigate Stochastic Bregman Proximal Gradient (SBPG) method to solve
nonconvex problem without finite-sum requirement, which employs Bregman
distance to handle the non-Lipschitz gradient continuity. We establish
convergence results for the vanilla SBPG in the sense of expectation. Further,
we propose a momentum-based SBPG (MSBPG) that is tailored for modern large-
scale applications, and prove that it has improved convergence properties
compared to the vanilla SBPG. To our best knowledge, this is the first time
that the momentum technique has been integrated into a stochastic Bregman
proximal gradient method.

  * 2.

We apply MSBPG to training deep neural networks (DNN), which leverages on a
suitable polynomial kernel function to ensure that the DNN’s loss function is
smooth adaptable with respect to the designed kernel function. MSBPG exhibits
good convergence behavior and excellent generalization performance on
extensive tasks. Moreover, MSBPG is found to be more robust than the
traditional SGD, especially when it comes to stepsize selection and
initialization. We highlight that MSBPG is a theoretically derived method that
is able to ease the difficulty of selecting stepsize, mitigate gradient
explosion, and maintains excellent generalization performance simultaneously.
This distinguishes MSBPG from many existing techniques that rely on intuition
and empirical observations.

  * 3.

We demonstrate the efficiency and robustness of SBPG or MSBPG in a range of
applications, including sparse quadratic inverse problems and large-scale deep
neural networks. In quadratic inverse problem, SBPG is more robust in terms of
both stepsize and initial point selections. In training deep neural networks,
MSBPG has been been found to achieve a superior generalization performance
compared with some of the most frequently used optimizers such as Adam and
AdamW, and also exhibits robustness to stepsize selection and initialization.
These results highlight the potential of MSBPG as a powerful tool for
optimizing complex and large-scale deep neural networks, thus offering a
promising direction for future research in this area.

The remainder of this paper is organized as follows. In Section 2, we present
notation, some related preliminaries and our problem setting. In Section 3, we
first describe SBPG and establish its convergence results in the sense of
expectation. Then, we propose a momentum-based SBPG (MSBPG) and prove its
improved convergence properties. In Section 4, we adapt MSBPG to the training
of deep neural networks and analyze its capacity in mitigating gradient
explosion and improving generalization capacity. In Section 5, we present
numerical experiments that demonstrate the efficiency and robustness of
vanilla SBPG on quadratic inverse problems and MSBPG on training deep neural
networks. Finally, we give some concluding remarks in Section 6, summarizaing
our key contributions and outlining promising topics for future research.

##  2 Preliminaries and Problem setting

In this paper, vectors are represented using boldface letters like ${\bm{v}}$,
while scalars are represented using normal font. Given a proper, lower-
semicontinuous function
$F:\mathbb{R}^{d}\rightarrow\bar{\mathbb{R}}:=[-\infty,\infty]$,
$dom\,F=\\{{\bm{x}}:F({\bm{x}})<\infty\\}$. The Fenchel conjugate function of
$F$ is defined as
$F^{*}({\bm{y}})=\sup\\{\langle{\bm{x}},\,{\bm{y}}\rangle-F({\bm{x}}):{\bm{x}}\in\mathbb{R}^{d}\\}$.
Given a set $\mathcal{S}\subset\mathbb{R}^{d}$, $\bar{\mathcal{S}}$ denotes
its closure, $int\,\mathcal{S}$ denotes the set of interior points. A function
is of class $\mathcal{C}^{k}(\mathcal{S})$ if it is $k$ times differentiable
and the $k$-th derivative is continuous on $\mathcal{S}$. We say that $F$ is
level bounded if the set $\\{x:F({\bm{x}})<\alpha\\}$ is bounded for any real
number $\alpha$. Given a matrix $A$, ${\rm Vec}(A)$ denotes the vectorization
of $A$ by column order. ${\rm Mat}(\cdot)$ is the inverse operation of ${\rm
Vec}(\cdot)$, which reshapes a vector back into its original matrix form.
Define the operator $\text{Diag}(\cdot)$ to map a vector into a diagonal
matrix with diagonal elements equal to the corresponding entries of the
vector. The Hadamard product is represented by the symbol $\circ$. If we use
the notation $\|\cdot\|$ without any additional explanation, we assume that it
refers to the Euclidean vector norm for vectors and the Frobenius matrix norm
for matrices.

Let $(\Omega,\mathcal{F},\mathbb{P})$ be a probability space. Given a random
variable ${\bm{\xi}}$ and a $\sigma$-algebra $\mathcal{F}$, we write
${\bm{\xi}}\vartriangleleft\mathcal{F}$ if ${\bm{\xi}}$ is measurable over
$\mathcal{F}$. Let $\\{{\bm{\xi}}_{k}\\}_{k\geq 0}$ be a stochastic process,
and $\\{\mathcal{F}_{k}\\}_{k\geq 0}$ be a filtration, where $\mathcal{F}_{k}$
defined by a $\sigma$-algebra
$\mathcal{F}_{k}:=\sigma(\bm{\xi}_{0},\ldots,\bm{\xi}_{k-1})$ on $\Omega$. The
conditional expectation is denoted by $\mathbb{E}[\cdot|\mathcal{F}_{k}]$. For
simplicity, we use the notation $\mathbb{E}[\cdot]$ to denote
$\mathbb{E}[\cdot|\mathcal{F}_{\infty}]$. The sequence
$\\{{\bm{x}}^{k}\\}_{k\geq 0}$ generated by our proposed method is adapted to
the filtration $\\{\mathcal{F}_{k}\\}_{k\geq 0}$, i.e.
${\bm{x}}^{k}\vartriangleleft\mathcal{F}_{k}$, for all $k\geq 0$. The notation
$\widetilde{\nabla}_{k}$ represents an estimator of the exact gradient $\nabla
F({\bm{x}}^{k})$, which satisties
$\widetilde{\nabla}_{k}\vartriangleleft\mathcal{F}_{k+1}$. This estimator is
applicable to both the vanilla and momentum cases. The stochastic error is
denoted by ${\bm{\varepsilon}}_{k}=\nabla
F({\bm{x}}^{k})-\widetilde{\nabla}_{k}$. The unbiasedness of the stochastic
error $\bm{\varepsilon}_{k}$ is assumed throughout this paper, i.e.,
$\mathbb{E}[{\bm{\varepsilon}}_{k}|\mathcal{F}_{k}]=0$. The following
supermartingale convergence theorem is a fundamental tool in the analysis of
stochastic algorithms.

###### Theorem 1

(Robbins and Monro, 1951) Let
$\left\\{y_{k}\right\\},\left\\{u_{k}\right\\},\left\\{a_{k}\right\\}$ and
$\left\\{b_{k}\right\\}$ be non-negative adapted processes with respect to the
filtration $\left\\{\mathcal{F}_{k}\right\\}$ such that
$\sum_{k}a_{k}<\infty,\sum_{k}b_{k}<\infty$, and for all $k$,
$\mathbb{E}\left[y_{k+1}\mid\mathcal{F}_{k}\right]\leq\left(1+a_{k}\right)y_{k}-u_{k}+b_{k}$
almost surely. Then, $\left\\{y_{k}\right\\}$ converges almost surely to a
non-negative finite random variable and $\sum_{k}u_{k}<\infty$ almost surely.

###  2.1 Smooth adaptable functions

In this subsection, we introduce the concept of smooth adaptivity, which was
initially proposed in Bolte et al. (2018). This concept is an extension of the
idea of relative smoothness for convex functions introduced in Bauschke et al.
(2017); Lu et al. (2018). We first give the definitions of kernel function and
Bregman distance.

###### Definition 2

(Kernel function and Bregman distance). Let $\mathcal{S}$ be a nonempty,
convex and open subset of $\mathbb{R}^{d}$. Associated with $\mathcal{S}$, a
function $\phi:\mathbb{R}^{d}\rightarrow(-\infty,+\infty]$ is called a kernel
function if it satisfies the following two conditions:

  1. 1.

$\phi$ is proper, lower-semicontinuous and convex, with
$dom\,\phi\subset\bar{\mathcal{S}}$, $dom\,\partial\phi=\mathcal{S}$.

  2. 2.

$\phi\in\mathcal{C}^{1}(\mathcal{S})$ and $int\,dom\,\phi=\mathcal{S}$.

Denote the class of kernel function associated with $\mathcal{S}$ by
$\mathcal{M}(\mathcal{S})$. Given $\phi\in\mathcal{M}(\mathcal{S})$, the
Bregman distance (Bregman, 1967) generated by $\phi$ is defined as
$\mathcal{D}_{\phi}({\bm{x}},{\bm{y}}):dom\,\phi\times
int\,dom\,\phi\rightarrow[0,+\infty)$, where

$\mathcal{D}_{\phi}({\bm{x}},{\bm{y}})=\phi({\bm{x}})-\phi({\bm{y}})-\langle\nabla\phi({\bm{y}}),\,{\bm{x}}-{\bm{y}}\rangle.$

Bregman distance measures the difference between the value of $\phi$ at
${\bm{x}}$ and its linear approximation at ${\bm{y}}$ based on the gradient of
$\phi$ at ${\bm{y}}$. Some basic properties of Bregman distance can be found
in Chen and Teboulle (1993); Teboulle (2018). Some kernel functions commonly
used in optimization are $\frac{1}{2}\|{\bm{x}}\|^{2}$,
$\frac{1}{2}\|{\bm{x}}\|^{2}+\frac{\alpha}{4}\|{\bm{x}}\|^{4}$,
$-\sum_{i=1}^{d}\log{\bm{x}}_{i}$ and
$\sum_{i=1}^{d}{\bm{x}}_{i}\log{\bm{x}}_{i}$, where
$\frac{1}{2}\|{\bm{x}}\|^{2}\in\mathcal{M}(\mathbb{R}^{d})$ recovers the
classical half squared Euclidean distance. The kernel function
$\frac{1}{2}\|{\bm{x}}\|^{2}+\frac{\alpha}{4}\|{\bm{x}}\|^{4}\in\mathcal{M}(\mathbb{R}^{d})$
has found applications in various problems, such as quadratic inverse
problems, non-negative matrix factorization, and low-rank minimization (Bolte
et al., 2018; Dragomir et al., 2021a). The entropy function
$\sum_{i=1}^{d}{\bm{x}}_{i}\log{\bm{x}}_{i}\in\mathcal{M}(\mathbb{R}^{d}_{++})$
is commonly used in applications that involve probability constraints, where
the resulting Bregman distance is known as the Kullback–Leibler (KL)
divergence. Throughout the paper we will focus on the following pair of
functions $(f,\phi)$ satisfying smooth adaptivity condition. We introduce this
concept in the following definition:

###### Definition 3

(Smooth adaptivity). Given a kernel function
$\phi\in\mathcal{M}(\mathcal{S})$, a proper lower-semicontinuous function
$f:\mathbb{R}^{d}\rightarrow(-\infty,+\infty]$ with $dom\,f\supset dom\,\phi$
that is $\mathcal{C}^{1}$ on $\mathcal{S}$. $f$ is $L$-smooth adaptable with
respect to $\phi$ if there exists $L>0$, such that $L\phi+f$ and $L\phi-f$ are
convex on $\mathcal{S}$.

Alternative definition of smooth adaptivity is the two-side descent lemma
(Bolte et al., 2018, Lemma 2.1). When both $f$ and $\phi$ belong to
$\mathcal{C}^{2}(\mathcal{S})$, we can verify their smooth adaptivity by
comparing the Hessians of $f$ and $\phi$.

###### Lemma 4

$f$ is $L$-smooth adaptable with respect to $\phi\in\mathcal{M}(\mathcal{S})$,
if and only if

$|f({\bm{x}})-f({\bm{y}})-\langle\nabla
f({\bm{y}}),\,{\bm{x}}-{\bm{y}}\rangle|\leq
L\mathcal{D}_{\phi}({\bm{x}},{\bm{y}}),\;\forall\,{\bm{x}},{\bm{y}}\in
int\,dom\,\phi.$

Moreover, when both $f$ and $\phi$ belong to
$\mathcal{C}^{2}(int\,dom\,\phi)$, then the above is equivalent to

$\exists L>0,\;L\nabla^{2}\phi({\bm{x}})-\nabla^{2}f({\bm{x}})\succeq 0,\text{
for all }{\bm{x}}\in int\,dom\phi.$

The following four-point identity is frequently employed in our proofs, and
can be easily verified.

###### Lemma 5

(Four points identity) Given points ${\bm{a}},{\bm{b}},{\bm{c}},{\bm{d}}$ and
any convex function $\phi$ which is differentiable at ${\bm{a}}$ and
${\bm{b}}$, then

$\langle\nabla\phi({\bm{a}})-\nabla\phi({\bm{b}}),\,{\bm{c}}-{\bm{d}}\rangle=\mathcal{D}_{\phi}({\bm{c}},{\bm{b}})+\mathcal{D}_{\phi}({\bm{d}},{\bm{a}})-\mathcal{D}_{\phi}({\bm{c}},{\bm{a}})-\mathcal{D}_{\phi}({\bm{d}},{\bm{b}}).$

###  2.2 Bregman Proximal Mapping

Throughout this paper, we make the following basic assumptions.

###### Assumption 1

(Basic requirements). In problem (1):

  * A1.

For every fixed $\bm{\xi}$, $f(\cdot,\bm{\xi})$ is a proper lower-
semicontinuous function with $dom\,\phi\subset dom\,f(\cdot,\bm{\xi})$, and it
is $\mathcal{C}^{1}$ on $int\,C$.

  * A2.

The Legendre kernel (Definition 6) $\phi\in\mathcal{M}(C)$ is $\mu$-strongly
convex for some $\mu>0$. For every fixed ${\bm{\xi}}$, $f(\cdot,\bm{\xi})$ is
$L_{F}$-smooth adaptable with respect to $\phi$, where $L_{F}$ is independent
of $\bm{\xi}$.

  * A3.

$R$ is is a proper, lower semicontinuous and convex function with $dom\,R\cap
int\,C\neq\emptyset$.

  * A4.

$\inf_{{\bm{x}}\in\overline{C}}\\{\Phi({\bm{x}})\\}>-\infty$.

Assumption 1 is a standard requirement for Bregman-type methods and is usually
satisfied in practice. It ensures the well-definedness of Bregman-type
methods, as shown in Bolte et al. (2018); Latafat et al. (2022). We also
recall the definition of the Legendre function in Latafat et al. (2022), which
makes additional supercoercive conditions on the concept in Rockafellar
(1997).

###### Definition 6

(Legendre kernel). Let $\phi:{\overline{C}}\rightarrow(-\infty,\infty]$ be a
proper lower-semicontinuous convex function. It is called essentially smooth
if $\text{int\;dom}\,\phi$ is nonempty and $\phi$ is differentiable on
$\text{int\;dom}\,\phi$, moreover
$\lim_{k\rightarrow\infty}\|\nabla\phi({{\bm{x}}^{k}})\|=\infty$ whenever
$\\{{{\bm{x}}^{k}}\\}_{k\in\mathbb{N}}$ converges to a boundary point of
$\text{dom}\,\phi$. The function $\phi$ is called Legendre function if it is
essentially smooth, strictly convex on $\text{int\,dom}\,\phi$ and
supercoercive, i.e.
$\lim_{\|{\bm{x}}\|\rightarrow\infty}\frac{\phi({\bm{x}})}{\|{\bm{x}}\|}=\infty$.

###### Definition 7

Given a nonempty convex open set $C$, a proper lower-semicontinuous convex
function $R$ and a Legendre kernel function $\phi\in\mathcal{M}(C)$,
${\bm{x}}\in int\,\phi$, we denote the Bregman proximal mapping by ${\rm
Prox}^{\phi}_{R}:=(\nabla\phi+\partial R)^{-1}\nabla\phi$, which is defined as

${\rm
Prox}^{\phi}_{R}({\bm{x}}):=\operatorname*{argmin}_{{\bm{u}}\in\overline{C}}\;\\{R({\bm{u}})+\mathcal{D}_{\phi}({\bm{u}},{\bm{x}})\\}.$
(4)

Note that the objective function of (4) is strictly convex on
$\text{dom}\,\phi\cap domR$, therefore (4) has at most one solution. To ensure
that (4) is well-defined, the following result claims that ${\rm
Prox}^{\phi}_{\alpha R}({\bm{x}})$ is well-defined for any $\alpha>0$, and
moreover ${\rm Prox}^{\phi}_{\alpha R}({\bm{x}})\in\text{int\;dom}\,\phi$
under standard assumptions. The proof can be found in Appendix A.

###### Lemma 8

Suppose Assumption 1 holds. Then (4) has a unique solution. Moreover, the
solution ${\rm Prox}^{\phi}_{\alpha R}({\bm{x}})\in C$.

The following proposition for Bregman proximal mapping generalizes the
nonexpansive property of the classical proximal mapping (in the case
$\phi({\bm{x}})=\frac{1}{2}\|{\bm{x}}\|^{2}$). This property is commonly used
in convergence proofs. The proof of the following proposition can be found in
Appendix A.

###### Proposition 9

Suppose Assumption 1 holds. Let ${\bm{x}}_{i}^{+}:={\rm
Prox}^{\phi}_{R}(\nabla\phi^{*}({\bm{x}}_{i}))$, $i=1,2$. Then
$\|{\bm{x}}^{+}_{1}-{\bm{x}}^{+}_{2}\|\leq\frac{1}{\mu}\|{\bm{x}}_{1}-{\bm{x}}_{2}\|$.

In this paper, we make the assumption that $R$ and $\phi$ are simple enough so
that (4) either has a closed-form solution or admits an efficient subroutine
to solve it. Using the definition of the Bregman proximal mapping, we can then
define the Bregman gradient mapping associated with (1). This mapping measures
the solution accuracy of the methods we propose. Note that $\phi$ is a
Legendre kernel, which implies that
$\phi^{*}\in\mathcal{C}^{1}(\mathbb{R}^{d})$ is strictly convex and
$(\nabla\phi)^{-1}=\nabla\phi^{*}$ (Rockafellar, 1997, Corollary 13.3.1,
Theorem 26.5). Therefore, the following concept is well-defined.

######  Definition 10 (Bregman Gradient Mapping)

Given $\alpha>0$, a nonempty convex open set $C$ and a Legendre kernel
function $\phi\in\mathcal{M}(C)$, the Bregman gradient mapping associated with
(1) is defined as follows

$\mathcal{G}_{\alpha}({\bm{x}})=\frac{{\bm{x}}-{\rm Prox}^{\phi}_{\alpha
R}\left(\nabla\phi^{*}(\nabla\phi({\bm{x}})-\alpha\nabla
F({\bm{x}}))\right)}{\alpha}.$

To simplify notation, we use $\mathcal{G}({\bm{x}})$ to denote
$\mathcal{G}_{1}({\bm{x}})$ when $\alpha=1$.

When the kernel function $\phi({\bm{x}})=\frac{1}{2}\|{\bm{x}}\|^{2}$, the
resulting Bregman Gradient Mapping becomes equivalent to the classical
Gradient Mapping (Nesterov, 2003, 2005), which measures the solution’s
accuracy for proximal gradient methods.

###### Definition 11

(Limiting subdifferential (Rockafellar and Wets, 1998, Definition 8.3))
Consider a function $f:\mathbb{R}^{d}\rightarrow\bar{\mathbb{R}}$ and a point
$x$, the regular subdifferential is defines as

$\hat{\partial}f({\bm{x}})=\\{{\bm{v}}:f({\bm{y}})\geq
f({\bm{x}})+\langle{\bm{v}},\,{\bm{y}}-{\bm{x}}\rangle+o(\|{\bm{y}}-{\bm{x}}\|)\\}.$

The limiting subdifferential is defined as

$\partial
f({\bm{x}})=\\{{\bm{v}}:{\bm{x}}_{n}\rightarrow{\bm{x}},f({\bm{x}}_{n})\rightarrow
f({\bm{x}}),{\bm{v}}_{n}\in\hat{\partial}f({\bm{x}}_{n}),\;and\;{\bm{v}}_{n}\rightarrow{\bm{v}}\\}.$

By Fermat’s rule (Rockafellar and Wets, 1998, Theorem 10.1), the set of
critical point of $\Phi$ is given by

${\rm crit}\,\Phi=\left\\{x\in\mathbb{R}^{d}:\;0\in\partial\Phi(x)\equiv\nabla
F(x)+\partial R(x)\right\\}.$

The Bregman Gradient Mapping can also be used to evaluate the solution
accuracy for Bregman methods. Let ${\bm{x}}^{+}=\text{Prox}^{\phi}_{\alpha
R}(\nabla\phi^{*}(\nabla\phi({\bm{x}})-\alpha\nabla F({\bm{x}})))$. From
Definition 10 and equation (4), it can be easily verified by definition that
$0\in\partial\Phi({\bm{x}})\Leftrightarrow 0=\mathcal{G}_{\alpha}({\bm{x}}).$
Hence, $0\in\partial\Phi({\bm{x}}^{+})$ for any $\alpha>0$. The proof of this
result is omitted for brevity. Furthermore, if $\nabla\phi$ is
$L_{\phi}$-Lipschitz continuous, then the following proposition holds,
implying that $\|\mathcal{G}_{\alpha}({\bm{x}})\|$ can be used as a reasonable
criterion to measure the accuracy of ${\bm{x}}$.

###### Proposition 12

Suppose Assumption 1 holds and that $\nabla\phi$ is $L_{\phi}$ Lipschitz
continuous. Then, we have the following inequality:

${\rm dist}\left(0,\partial\Phi({{\bm{x}}^{+}})\right)\leq(1+\alpha
L_{F})L_{\phi}\|\mathcal{G}_{\alpha}({\bm{x}})\|.$

We also define the stochastic counterpart of Definition 10, which is commonly
utilized to evaluate the accuracy of solutions for nonconvex stochastic
proximal gradient methods, as discussed in Ghadimi et al. (2016).

###### Definition 13

(Stochastic Bregman Gradient Mapping). Given $\alpha>0$ a nonempty convex open
set $C$ and a Legendre kernel function $\phi\in\mathcal{M}(C)$, the stochastic
Bregman gradient mapping associated with (1) is defined as follows

$\widetilde{\mathcal{G}}_{\alpha}({{\bm{x}}}):=\frac{{{\bm{x}}}-{\rm
Prox}^{\phi}_{\alpha
R}\left(\nabla\phi^{*}\left(\nabla\phi({{\bm{x}}})-\alpha\widetilde{\nabla}\right)\right)}{\alpha},\;\text{where
}\widetilde{\nabla}\;\text{is an estimator of }\nabla F(x).$

##  3 Stochastic Bregman Proximal Gradient Method

In this section, we will study the Stochastic Bregman Proximal Gradient method
(SBPG) with the following update scheme:

${{\bm{x}}^{k+1}}=\operatorname*{argmin}_{{\bm{x}}\in\overline{C}}\;R({\bm{x}})+\langle\widetilde{\nabla}_{k},\,{\bm{x}}-{{\bm{x}}^{k}}\rangle+\frac{1}{\alpha_{k}}\mathcal{D}_{\phi}({\bm{x}},{{\bm{x}}^{k}}).$
(5)

We call the above method as ”vanilla” SBPG in this section, meaning that the
method we study is a basic version without any additional techniques such as
variance reduction, momentum, etc., except for the use of mini-batches. In
this case, we suppose the following assumptions.

###### Assumption 2

(Noise requirement). The estimator satisfies the following two conditions:

$\mathbb{E}[\widetilde{\nabla}_{k}|\mathcal{F}_{k}]=\nabla
F(x_{k})\quad\text{and}\quad\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}|\mathcal{F}_{k}]\leq\frac{\sigma^{2}}{m_{k}},$

where $m_{k}$ is the size of the mini-batch in the $k$-th iteration.

Note that we do not assume a finite-sum structure for $F({\bm{x}})$ in this
section. The solution of (5) can be written in the form of the Bregman
proximal mapping. This is stated in the following proposition.

###### Proposition 14

Suppose Assumption 1 holds. Then the solution of (5) can be written as the
following Bregman proximal mapping:

${{\bm{x}}^{k+1}}={\rm
Prox}^{\phi}_{\alpha_{k}R}\left(\nabla\phi^{*}\left(\nabla\phi({{\bm{x}}^{k}})-\alpha_{k}\widetilde{\nabla}_{k}\right)\right).$

Proof  From the optimality condition of the main subproblem (5), we have

$0\in\partial
R({{\bm{x}}^{k+1}})+\widetilde{\nabla}_{k}+\frac{1}{\alpha_{k}}\left(\nabla\phi({{\bm{x}}^{k+1}})-\nabla\phi({{\bm{x}}^{k}})\right).$

Let ${\bm{u}}^{k+1}:={\rm
Prox}^{\phi}_{\alpha_{k}R}\left(\nabla\phi^{*}\left(\nabla\phi({{\bm{x}}^{k}})-\alpha_{k}\widetilde{\nabla}_{k}\right)\right)$.
From the definition of Bregman proximal mapping, we have

${\bm{u}}^{k+1}=\arg\min_{\bm{u}}\;\left\\{\alpha_{k}R({\bm{u}})+{\mathcal{D}_{\phi}}\left({\bm{u}},\nabla\phi^{*}\left(\nabla\phi({{\bm{x}}^{k}})-\alpha_{k}\widetilde{\nabla}_{k}\right)\right)\right\\},$

which is equivalent to

$0\in\alpha_{k}\partial
R({\bm{u}}^{k+1})+\nabla\phi({\bm{u}}^{k+1})-\nabla\phi\big{(}\nabla\phi^{*}(\nabla\phi({{\bm{x}}^{k}})-\alpha_{k}\widetilde{\nabla}_{k})\big{)}.$

Note that the function $\phi^{*}$ is the Fenchel conjugate of the Legendre
kernel $\phi$, which implies that
$\nabla\phi(\nabla\phi^{*}({\bm{w}}))={\bm{w}}$ for all
${\bm{w}}\in\mathbb{R}^{d}$, as stated in (Rockafellar, 1997, Corollary
13.3.1, Theorem 26.5). Furthermore, since the objective function in (5) is
strictly convex, there exists a unique solution to the inclusion above. By
comparing the two inclusions, we can conclude that
${\bm{u}}^{k+1}={{\bm{x}}^{k+1}}$.
Based on Proposition 14 and definition of the definition of
$\widetilde{\mathcal{G}}_{\alpha}({\bm{x}})$, we can easily observe that
${\bm{x}}^{k+1}={\bm{x}}^{k}-\alpha_{k}\widetilde{\mathcal{G}}_{\alpha_{k}}({\bm{x}}^{k})$.
We can derive the following proposition, which bounds the difference between
$\mathcal{G}_{\alpha}({\bm{x}})$ and
$\widetilde{\mathcal{G}}_{\alpha}({\bm{x}})$ directly from Proposition 9. The
proof is omitted for brevity.

###### Proposition 15

Suppose Assumption 1 holds. At the $k$-th step, we have the estimation:

$\|\mathcal{G}_{\alpha_{k}}({{\bm{x}}^{k}})-\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})\|\leq\frac{1}{\mu}\|\nabla
F({{\bm{x}}^{k}})-\widetilde{\nabla}_{k}\|=\frac{\|{\bm{\varepsilon}}_{k}\|}{\mu},$

where $\bm{\varepsilon}_{k}=\nabla F({{\bm{x}}^{k}})-\widetilde{\nabla}_{k}.$

Before presenting the main convergence result, we state the following one-step
descent lemma below.

###### Lemma 16

Suppose Assumption 1 holds. The sequence generated by SBPG satisfies the
following condition:

$\Phi({{\bm{x}}^{k+1}})\leq\Phi({{\bm{x}}^{k}})-\frac{1}{\alpha_{k}}\mathcal{D}_{\phi}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})-\left(\frac{1}{\alpha_{k}}-L_{F}\right)\mathcal{D}_{\phi}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})+\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k+1}}-{\bm{x}}^{k}\rangle.$

Proof  By the optimality condition of (5), we obtain that

$0\in\partial
R({{\bm{x}}^{k+1}})+\widetilde{\nabla}_{k}+\frac{1}{\alpha_{k}}\left(\nabla\phi({{\bm{x}}^{k+1}})-\nabla\phi({{\bm{x}}^{k}})\right).$

Appealing to the convexity of $R$, we have

$R({\bm{x}})-R({{\bm{x}}^{k+1}})\geq\big{\langle}{-\widetilde{\nabla}_{k}-\frac{1}{\alpha_{k}}\left(\nabla\phi({{\bm{x}}^{k+1}})-\nabla\phi({{\bm{x}}^{k}})\right)},\,{{\bm{x}}-{{\bm{x}}^{k+1}}}\big{\rangle}.$

By the four points identity and the definition of ${\bm{\varepsilon}}_{k}$, we
get

$R({\bm{x}})-R({{\bm{x}}^{k+1}})\geq\frac{1}{\alpha_{k}}\left[{\mathcal{D}_{\phi}}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})+{\mathcal{D}_{\phi}}({\bm{x}},{{\bm{x}}^{k+1}})-{\mathcal{D}_{\phi}}({\bm{x}},{{\bm{x}}^{k}})\right]-\langle\nabla
F({{\bm{x}}^{k}}),\,{\bm{x}}-{{\bm{x}}^{k+1}}\rangle+\langle{\bm{\varepsilon}}_{k},\,{\bm{x}}-{{\bm{x}}^{k+1}}\rangle.$

Set ${\bm{x}}={{\bm{x}}^{k}}$ in the above inequality, we have the following
inequality:

$R({{\bm{x}}^{k}})-R({{\bm{x}}^{k+1}})\geq\frac{1}{\alpha_{k}}\left[{\mathcal{D}_{\phi}}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})+{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})\right]-\langle\nabla
F({{\bm{x}}^{k}}),\,{{\bm{x}}^{k}}-{{\bm{x}}^{k+1}}\rangle+\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k}}-{{\bm{x}}^{k+1}}\rangle.$

By the smooth adaptivity of $F$, we have

$F({{\bm{x}}^{k+1}})\leq F({{\bm{x}}^{k}})+\langle\nabla
F({{\bm{x}}^{k}}),\,{{\bm{x}}^{k+1}}-{{\bm{x}}^{k}}\rangle+L_{F}{\mathcal{D}_{\phi}}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}}).$

Combining the above two inequalities above, we complete the proof.

###  3.1 Convergence analysis of SBPG

In this subsection, we establish the convergence results for SBPG, which is an
extension of the convergence result in Ghadimi et al. (2016), in which the
lassical Lipschitz gradient assumption is required. In many literature, the
bounded sequence assumption is often required in the convergence analysis of
stochastic algorithms. However, in this section, we relax this assumption and
prove that under a certain condition, the sequence generated by (5) is bounded
almost surely. We need the following result to bound the stochastic error term
$\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k+1}}-{{\bm{x}}^{k}}\rangle$ in
Lemma 16.

###### Lemma 17

Suppose Assumption 1, 2 hold. We have the following estimation of the error
term:

$\mathbb{E}\left[\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k+1}}-{{\bm{x}}^{k}}\rangle\right]\leq\frac{\alpha_{k}}{\mu}\mathbb{E}[\|{\bm{\varepsilon}}_{k}\|^{2}]\leq\frac{\alpha_{k}\sigma^{2}}{\mu
m_{k}}.$

Proof  Define ${\bar{\bm{x}}}^{k+1}:={{\rm
Prox}_{\alpha_{k}R}^{\phi}}(\nabla\phi^{*}(\nabla\phi({{\bm{x}}^{k}})-\alpha_{k}\nabla
F({{\bm{x}}^{k}})))$. By Proposition 14 and the optimality condition for
${\bar{\bm{x}}}^{k+1}$, we have

$0\in\partial R({\bar{\bm{x}}}^{k+1})+\nabla
F({{\bm{x}}^{k}})+\frac{1}{\alpha_{k}}(\nabla\phi({\bar{\bm{x}}}^{k+1})-\nabla\phi({{\bm{x}}^{k}})).$

Similarly,

$0\in\partial
R({{\bm{x}}^{k+1}})+\widetilde{\nabla}_{k}+\frac{1}{\alpha_{k}}(\nabla\phi({{\bm{x}}^{k+1}})-\nabla\phi({{\bm{x}}^{k}})).$

By the monotonicity of $\partial R$ and Lemma 5, we have

$\big{\langle}{{\bar{\bm{x}}}^{k+1}-{{\bm{x}}^{k+1}}},\,{-{\bm{\varepsilon}}_{k}-\frac{1}{\alpha_{k}}(\nabla\phi({\bar{\bm{x}}}^{k+1})-\nabla\phi({{\bm{x}}^{k+1}}))}\big{\rangle}\geq
0.$

Therefore,

$\langle{{\bm{x}}^{k+1}}-{\bar{\bm{x}}}^{k+1},\,{\bm{\varepsilon}}_{k}\rangle\geq\langle{\bar{\bm{x}}}^{k+1}-{{\bm{x}}^{k+1}},\,\frac{1}{\alpha_{k}}(\nabla\phi({\bar{\bm{x}}}^{k+1})-\nabla\phi({{\bm{x}}^{k+1}}))\rangle\geq\frac{\mu}{\alpha_{k}}\|{\bar{\bm{x}}}^{k+1}-{{\bm{x}}^{k+1}}\|^{2}.$

By Cauchy-Schwarz inequality, we get
$\|{\bar{\bm{x}}}^{k+1}-{{\bm{x}}^{k+1}}\|\leq\frac{\alpha_{k}}{\mu}\|{\bm{\varepsilon}}_{k}\|$.

Now, we are ready to prove Lemma 17. From the definition, we know that
${\bar{\bm{x}}}^{k+1}\vartriangleleft\mathcal{F}_{k}$. Therefore,
$\mathbb{E}[\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k}}-{\bar{\bm{x}}}^{k+1}\rangle]=\mathbb{E}[\mathbb{E}[\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k}}-{\bar{\bm{x}}}^{k+1}\rangle|\mathcal{F}_{k}]]=\mathbb{E}[\langle\mathbb{E}[{\bm{\varepsilon}}_{k}|\mathcal{F}_{k}],\,{{\bm{x}}^{k}}-{\bar{\bm{x}}}^{k+1}\rangle]=0$,
where the first equality is from the tower rule of conditional expectation,
the second comes from the fact that
${{\bm{x}}^{k}}-{\bar{\bm{x}}}^{k+1}\vartriangleleft\mathcal{F}_{k}$. Hence,

$\mathbb{E}\left[\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k+1}}-{\bm{x}}^{k}\rangle\right]=\mathbb{E}\left[\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k+1}}-{\bar{\bm{x}}}^{k+1}\rangle\right]-\mathbb{E}\left[\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k}}-{\bar{\bm{x}}}^{k+1}\rangle\right]\leq\frac{\alpha_{k}}{\mu}\mathbb{E}[\|{\bm{\varepsilon}}_{k}\|^{2}]\leq\frac{\alpha_{k}\sigma^{2}}{\mu
m_{k}},$

which completes the proof.

######  Lemma 18 (Bounded sequence)

Suppose Assumption 1, 2 hold. If $\sum_{k}\frac{\alpha_{k}}{m_{k}}<\infty$,
$\sup_{k}\alpha_{k}\leq\bar{\alpha}<\frac{1}{L_{F}}$, then,

  * 1.

$\sum_{k=0}^{\infty}\mathbb{E}[\mathcal{D}_{\phi}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})]<\infty$.

  * 2.

If $\Phi$ is level bounded, then $\\{{{\bm{x}}^{k}}\\}_{k\geq 0}$ is bounded
almost surely.

Proof  By Cauchy-Young inequality, we have

$|\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k}}-{{\bm{x}}^{k+1}}\rangle|\leq\frac{\mu}{2\alpha_{k}}\|{{\bm{x}}^{k}}-{{\bm{x}}^{k+1}}\|^{2}+\frac{\alpha_{k}}{2\mu}\|{\bm{\varepsilon}}_{k}\|^{2}\leq\frac{1}{\alpha_{k}}{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})+\frac{\alpha_{k}}{2\mu}\|{\bm{\varepsilon}}_{k}\|^{2}.$

By Lemma 16, we have

$\left(\frac{1}{\alpha_{k}}-L_{F}\right){\mathcal{D}_{\phi}}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})\leq\Phi({{\bm{x}}^{k}})-\Phi({{\bm{x}}^{k+1}})+\frac{\alpha_{k}}{2\mu}\|{\bm{\varepsilon}}_{k}\|^{2}.$
(6)

Taking conditional expectation for both sides of (6), we get

$\mathbb{E}\left[\left(\frac{1}{\alpha_{k}}-L_{F}\right){\mathcal{D}_{\phi}}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})|\mathcal{F}_{k}\right]\leq\Phi({{\bm{x}}^{k}})-\mathbb{E}[\Phi({{\bm{x}}^{k+1}})|\mathcal{F}_{k}]+\frac{\alpha_{k}}{2\mu}\mathbb{E}[\|{\bm{\varepsilon}}_{k}\|^{2}|\mathcal{F}_{k}].$

Since $\sum_{k\geq
0}\frac{\alpha_{k}}{2\mu}\mathbb{E}[\|{\bm{\varepsilon}}_{k}\|^{2}|\mathcal{F}_{k}]\leq\sum_{k\geq
0}\frac{\alpha_{k}\sigma^{2}}{2\mu m_{k}}<\infty$, applying Theorem 1, we have
that $\Phi({{\bm{x}}^{k}})$ converges and $\sum_{k\geq
0}\mathbb{E}\left[\left(\frac{1}{\alpha_{k}}-L_{F}\right){\mathcal{D}_{\phi}}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})|\mathcal{F}_{k}\right]<\infty$
almost surely. By the tower rule of conditional expectation, we have
$\sum_{k=0}^{\infty}\mathbb{E}[{\mathcal{D}_{\phi}}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})]<\infty$.
Since $\Phi({{\bm{x}}^{k}})$ converges almost surely, thus
$\\{\Phi({{\bm{x}}^{k}})\\}_{k\geq 0}$ is bounded almost surely. By the level
boundness of $\Phi$, we deduce that $\\{{{\bm{x}}^{k}}\\}_{k\geq 0}$ is
bounded almost surely.

Now, we present our main convergence result for the vanilla SBPG, which is in
the sense of expectation.

######  Theorem 19 (Convergence result in expectation)

Suppose Assumption 1, 2 hold,
$\alpha_{k}<\frac{1}{L_{F}}\min\\{1,\frac{1}{\mu}\\}$. Define a random
variable $r$ with the distribution
$\mathbb{P}\\{r=k\\}=\frac{\alpha_{k}}{\sum_{k=0}^{N-1}\alpha_{k}}$ for
$k=0,...,N-1$. Then,

$\mathbb{E}[\|\widetilde{\mathcal{G}}_{\alpha_{r}}({\bm{x}}^{r})\|^{2}]\leq\frac{2\Phi(x^{0})-2\Phi^{*}+2\sum_{k=0}^{N-1}\frac{\alpha_{k}\sigma^{2}}{\mu
m_{k}}}{\mu\sum_{k=0}^{N-1}\alpha_{k}}.$ (7)

Moreover, if $\Phi$ is level bounded,
$\sum_{k}\frac{\alpha_{k}}{m_{k}}<+\infty$ and $\sum_{k}\alpha_{k}=+\infty$,
then the sequence $\\{{{\bm{x}}^{k}}\\}_{k\geq 0}$ is bounded almost surely,
and the right hand side of (7) converges to zero.

Proof  Note that
${{\bm{x}}^{k+1}}={{\bm{x}}^{k}}-\alpha_{k}\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})$
and by the strongly convexity of $\phi$, Lemma 16 yields

$\displaystyle\mu(\alpha_{k}-\frac{L_{F}\alpha_{k}^{2}}{2})\|\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})\|^{2}$
$\displaystyle\leq\frac{1}{\alpha_{k}}\mathcal{D}_{\phi}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})+\left(\frac{1}{\alpha_{k}}-L_{F}\right)\mathcal{D}_{\phi}({{\bm{x}}^{k+1}},{{\bm{x}}^{k}})$
$\displaystyle\leq\Phi({{\bm{x}}^{k}})-\Phi({{\bm{x}}^{k+1}})+\langle{\bm{\varepsilon}}_{k},\,{{\bm{x}}^{k+1}}-{{\bm{x}}^{k}}\rangle.$

Taking expectations, telescoping from $k=0...N-1$, and using Lemma 17, we
obtain

$\sum_{k=0}^{N-1}\mu(\alpha_{k}-\frac{\mu
L_{F}\alpha_{k}^{2}}{2})\mathbb{E}[\|\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})\|^{2}]\leq\Phi({\bm{x}}^{0})-\Phi({\bm{x}}^{N})+\sum_{k=0}^{N-1}\frac{\alpha_{k}\sigma^{2}}{\mu
m_{k}}.$ (8)

By utilizing the inequality $\alpha_{k}-\frac{\mu
L_{F}\alpha_{k}^{2}}{2}\geq\frac{\alpha_{k}}{2}$, the condition
$\Phi({\bm{x}}^{N})\geq\Phi^{*}$, and considering the definition of the random
variable $r$, we can derive (7) from (8).

###### Remark 20

We give some remarks for Theorem 19.

  * 1.

The mini-batch setting is a crucial component for ensuring the convergence, as
it allows us to control the stochastic error term in Lemma 16 and provide a
bound for
$\mathbb{E}[\|\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})\|^{2}]$
that converges to zero as $k$ tends to infinity. If $m_{k}=1$ for all $k$,
then the upper bound for
$\mathbb{E}[\|\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})\|^{2}]$
will not converge to zero, no matter how $\\{\alpha_{k}\\}$ is selected.

  * 2.

In Ghadimi et al. (2016), a similar convergence result is established for
nonconvex stochastic proximal gradient methods, but our analysis differs in a
crucial aspect in that we do not assume the Lipschitz continuity of
$F({\bm{x}})$. Instead, we assume that $F({\bm{x}})$ is smooth adaptable,
which is a more relaxed assumption. Moreover, we provide a specific choice of
stepsizes $\\{\alpha_{k}\\}$ and mini-batch sizes $\\{m_{k}\\}$ that guarantee
the convergence of
$\mathbb{E}[\|\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})\|^{2}]$ to
$0$, as well as the almost sure boundedness of the sequence.

Based on Proposition 15 and Theorem 19, we can derive the following
convergence result using the measure
${\mathcal{G}}_{\alpha_{r}}({\bm{x}}^{r})$. This can be obtained by observing
that $\|\mathcal{G}_{\alpha_{r}}({\bm{x}}^{r})\|^{2}\leq
2\|\widetilde{\mathcal{G}}_{\alpha_{r}}({\bm{x}}^{r})\|^{2}+2\|\mathcal{G}_{\alpha_{r}}({\bm{x}}^{r})-\widetilde{\mathcal{G}}_{\alpha_{r}}({\bm{x}}^{r})\|^{2}.$

###### Corollary 21

Under the conditions in Theorem 19, we have

$\mathbb{E}[\|{\mathcal{G}}_{\alpha_{r}}({\bm{x}}^{r})\|^{2}]\;\leq\;\frac{4\Phi({\bm{x}}^{0})-4\Phi^{*}+4\sum_{k=0}^{N-1}\frac{\alpha_{k}\sigma^{2}}{\mu
m_{k}}}{\sum_{k=0}^{N-1}\mu\alpha_{k}}+\frac{2\sigma^{2}}{\mu^{2}m_{k}}.$

###  3.2 Momentum based Stochastic Bregman Gradient Descent Method

Remark 20 suggests that increasing the mini-batch size $m_{k}$ is almost
necessary for the error bound in Theorem 19 to converge to zero. However,
using a large mini-batch size in each iteration can be computationally
expensive in the modern large-scale problems, e.g. training deep neural
network. In this part, we resort to the momentum technique to address this
issue. Specifically, we consider using a stochastic moving average estimator
(SMAE) for the true gradient given by:

${\bm{v}}^{k}=(1-\beta_{k}){\bm{v}}^{k-1}+\beta_{k}\widetilde{\nabla}_{k},\quad\mbox{where}\quad\mathbb{E}[\widetilde{\nabla}_{k}|\mathcal{F}_{k}]=\nabla
F({{\bm{x}}^{k}}),$ (9)

where ${\bm{v}}^{k-1}$ can be viewed as the momentum which contain the
information of all historical stochastic gradients, and
$\mathbb{E}[\|\widetilde{\nabla}_{k}\|^{2}|\mathcal{F}_{k}]\leq\frac{\sigma^{2}}{m_{k}}$.
We expect that incorporating the SMAE technique can achieve a certain level of
variance reduction without increasing the mini-batch size. In our approach, we
utilize the gradient estimator ${\bm{v}}^{k}$ within SBPG, and we refer to the
resulting method as MSBPG. Specifically, we consider the following update
scheme:

${{\bm{x}}^{k+1}}=\operatorname*{argmin}_{{\bm{x}}\in\overline{C}}\;R({\bm{x}})+\langle{\bm{v}}^{k},\,{\bm{x}}-{{\bm{x}}^{k}}\rangle+\frac{1}{\alpha_{k}}\mathcal{D}_{\phi}({\bm{x}},{{\bm{x}}^{k}}).$
(10)

We need the following assumption that the difference of gradients of $F$ can
be bounded by the Bregman distance.

###### Assumption 3

There exists $\kappa>0$, such that $\|\nabla F({\bm{x}})-\nabla
F({\bm{y}})\|^{2}\leq\kappa\mathcal{D}_{\phi}({\bm{x}},{\bm{y}})$ for all
${\bm{x}}\in dom\,\phi$, ${\bm{y}}\in int\,dom\,\phi$.

###### Remark 22

This assumption generalizes the case of Lipschitz kernel function. Assume that
$F$ is $L_{F}$ smooth adaptable to $\phi$, it can be easily shown that if
$\phi$ has $L_{\phi}$-Lipschitz gradient, this assumption holds for
$\kappa\geq\frac{2L_{F}^{2}L_{\phi}^{2}}{\mu}$. In this paper, we are more
interested in polynomial kernel functions. For functions with polynomially
bounded growth rates, this assumption is not restrictive. For example,
consider the one-dimensional objective function $F(x)=\frac{1}{4}x^{4}$ and
the kernel function $\phi(x)=\frac{1}{2}x^{2}+\frac{1}{8}x^{8}$. Then, by (Lu
et al., 2018, Proposition 2.1), we know that $F$ is smooth adaptable with
respect to $\phi$. Simple algebra shows that
$\mathcal{D}_{\phi}(x,y)=\frac{1}{8}(x-y)^{2}(x^{6}+2x^{5}y+3x^{4}y^{2}+4x^{3}y^{3}+5x^{2}y^{4}+6xy^{5}+7y^{6}+4)$
and $(F^{\prime}(x)-F^{\prime}(y))^{2}=(x-y)^{2}(x^{2}+xy+y^{2})^{2}$.
Numerical computation shows that
$(x^{6}+2x^{5}y+3x^{4}y^{2}+4x^{3}y^{3}+5x^{2}y^{4}+6xy^{5}+7y^{6}+4)-(x^{2}+xy+y^{2})^{2}\geq
3.71$. Therefore, $(F^{\prime}(x)-F^{\prime}(y))^{2}\leq
8\mathcal{D}_{\phi}(x,y)$, which holds globally for any $x$ and $y$.

Next we present a recursion lemma that allows us to estimate the accuracy of
the SMAE. While similar lemmas have been proposed in the literature, such as
in Wang et al. (2017), their bounds are not directly applicable in the Bregman
setting. As a result, we have developed a version of the recursion lemma that
is tailored to our specific context.

###### Lemma 23

The following recursion holds

$\mathbb{E}[\|{\bm{v}}^{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}|\mathcal{F}_{k}]\leq(1-\beta_{k})\|{\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1})\|^{2}+\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}|\mathcal{F}_{k}]+\frac{\|\nabla
F({\bm{x}}^{k-1})-\nabla F({{\bm{x}}^{k}})\|^{2}}{\beta_{k}}.$

Proof  Note that ${\bm{v}}^{k}-\nabla
F({\bm{x}}^{k})=(1-\beta_{k})({\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1}))+(1-\beta_{k})(\nabla F({\bm{x}}^{k-1})-\nabla
F({\bm{x}}^{k}))+\beta_{k}(\widetilde{\nabla}_{k}-\nabla F({\bm{x}}^{k}))$,
and $\mathbb{E}[\widetilde{\nabla}_{k}-\nabla
F({\bm{x}}^{k})|\mathcal{F}_{k}]=0$. Then we have

$\displaystyle\mathbb{E}[\|{\bm{v}}^{k}-\nabla
F({\bm{x}}^{k})\|^{2}|\mathcal{F}_{k}]$ $\displaystyle={}$
$\displaystyle(1-\beta_{k})^{2}\|{\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1})\|^{2}+(1-\beta_{k})^{2}\|\nabla F({\bm{x}}^{k-1})-\nabla
F({\bm{x}}^{k})\|^{2}+$
$\displaystyle\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({\bm{x}}^{k})\|^{2}|\mathcal{F}_{k}]+2(1-\beta_{k})^{2}\langle{\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1}),\,\nabla F({\bm{x}}^{k-1})-\nabla F({\bm{x}}^{k})\rangle$
$\displaystyle\leq{}$ $\displaystyle(1-\beta_{k})^{2}\|{\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1})\|^{2}+(1-\beta_{k})^{2}\|\nabla F({\bm{x}}^{k-1})-\nabla
F({\bm{x}}^{k})\|^{2}+$
$\displaystyle\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({\bm{x}}^{k})\|^{2}|\mathcal{F}_{k}]+\beta_{k}(1-\beta_{k})\|{\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1})\|^{2}+\frac{(1-\beta_{k})^{3}}{\beta_{k}}\|\nabla
F({\bm{x}}^{k-1})-\nabla F({\bm{x}}^{k})\|^{2}$ $\displaystyle={}$
$\displaystyle(1-\beta_{k})\|{\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1})\|^{2}+\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({\bm{x}}^{k})\|^{2}|\mathcal{F}_{k}]+\frac{(1-\beta_{k})^{2}\|\nabla
F({\bm{x}}^{k-1})-\nabla F({\bm{x}}^{k})\|^{2}}{\beta_{k}}$
$\displaystyle\leq{}$ $\displaystyle(1-\beta_{k})\|{\bm{v}}^{k-1}-\nabla
F({\bm{x}}^{k-1})\|^{2}+\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({\bm{x}}^{k})\|^{2}|\mathcal{F}_{k}]+\frac{\|\nabla F({\bm{x}}^{k-1})-\nabla
F({\bm{x}}^{k})\|^{2}}{\beta_{k}}.$

This completes the proof.

Now we are ready to provide the convergence result for our momentum based
SBPG.

###### Theorem 24

Suppose Assumption 1, 2 and 3 hold. Let $\alpha_{k}=c\mu\beta_{k+1}$ for any
$c\in(0,\frac{1}{2\sqrt{\mu\kappa}}]$. Then

$\mathbb{E}[\|\widetilde{\mathcal{G}}_{\alpha_{r}}({\bm{x}}^{r})\|^{2}]\leq\frac{\Phi^{0}-\Phi^{*}+c\|{\bm{v}}_{0}-\nabla
F({\bm{x}}^{0})\|^{2}+c\sum_{k=1}^{N}\frac{\beta_{k}^{2}\sigma^{2}}{m_{k}}}{\sum_{k=0}^{N-1}\frac{\mu\alpha_{k}}{8}}=\mathcal{O}\left(\frac{1}{\sum_{k=0}^{N-1}\alpha_{k}}+\frac{\sum_{k=0}^{N-1}\frac{\alpha_{k}^{2}}{m_{k+1}}}{\sum_{k=0}^{N-1}\alpha_{k}}\right),$

where $r$ is a random variable with distribution
$\mathbb{P}\\{r=k\\}=\frac{\alpha_{k}}{\sum_{k=0}^{N-1}\alpha_{k}}$, for
$k=0,...,N-1$.

Proof  From Lemma 16 and Cauchy-Young inequality, we have

$\displaystyle\Phi({{\bm{x}}^{k+1}})$
$\displaystyle\leq\Phi({{\bm{x}}^{k}})-\frac{1}{\alpha_{k}}{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})+\frac{\mu}{4\alpha_{k}}\|{{\bm{x}}^{k}}-{{\bm{x}}^{k+1}}\|^{2}+\frac{\alpha_{k}}{\mu}\|{\bm{\varepsilon}}_{k}\|^{2}$
$\displaystyle\leq\Phi({{\bm{x}}^{k}})-\frac{1}{2\alpha_{k}}{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})+\frac{\alpha_{k}}{\mu}\|{\bm{\varepsilon}}_{k}\|^{2}.$

where we have defined ${\bm{\varepsilon}}_{k}:=\nabla
F({{\bm{x}}^{k}})-{\bm{v}}^{k}$. Summing the above inequality over
$k=0,\ldots,N-1$ and rearranging the terms, we get

$\sum_{k=0}^{N-1}\frac{1}{2\alpha_{k}}{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})\leq\Phi({\bm{x}}^{0})-\Phi^{*}+\sum_{k=0}^{N-1}\frac{\alpha_{k}}{\mu}\|{\bm{\varepsilon}}_{k}\|^{2}.$

By applying Lemma 23, we can obtain the following inequality:

$\beta_{k}\mathbb{E}[\|{\bm{\varepsilon}}_{k-1}\|^{2}]\;\leq\;\mathbb{E}[\|{\bm{\varepsilon}}_{k-1}\|^{2}]-\mathbb{E}[\|{\bm{\varepsilon}}_{k}\|^{2}]+\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}]+\mathbb{E}\left[\frac{\|\nabla
F({\bm{x}}^{k-1})-\nabla F({{\bm{x}}^{k}})\|^{2}}{\beta_{k}}\right].$

Hence

$\sum_{k=0}^{N-1}\beta_{k+1}\mathbb{E}[\|{\bm{\varepsilon}}_{k}\|^{2}]=\sum_{k=1}^{N}\beta_{k}\mathbb{E}[\|{\bm{\varepsilon}}_{k-1}\|^{2}]\leq\|{\bm{\varepsilon}}_{0}\|^{2}+\sum_{k=1}^{N}\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}]+\sum_{k=1}^{N}\mathbb{E}\left[\frac{\|\nabla
F({\bm{x}}^{k-1})-\nabla F({{\bm{x}}^{k}})\|^{2}}{\beta_{k}}\right].$

Since $\frac{\alpha_{k}}{\mu}=c\beta_{k+1}$ for some constant $c$, we get the
following inequality:

$\sum_{k=0}^{N-1}\frac{1}{2\alpha_{k}}\mathbb{E}[{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})]\leq\Phi(x^{0})-\Phi^{*}+c\left(\|{\bm{\varepsilon}}_{0}\|^{2}+\sum_{k=1}^{N}\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}]+\sum_{k=1}^{N}\mathbb{E}\left[\frac{\|\nabla
F({\bm{x}}^{k-1})-\nabla F({{\bm{x}}^{k}})\|^{2}}{\beta_{k}}\right]\right).$

By using Assumption 3, we obtain that

$\frac{\|\nabla F({{\bm{x}}^{k}})-\nabla
F({{\bm{x}}^{k+1}})\|^{2}}{\beta_{k+1}}\leq\frac{\kappa}{\beta_{k+1}}{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}}).$

Combining above two inequalities, we get

$\sum_{k=0}^{N-1}\frac{1}{2\alpha_{k}}\mathbb{E}[{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})]\leq\Phi({\bm{x}}^{0})-\Phi^{*}+c\left(\|{\bm{\varepsilon}}_{0}\|^{2}+\sum_{k=1}^{N}\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}]+\sum_{k=0}^{N-1}\frac{\kappa}{\beta_{k+1}}\mathbb{E}[{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})]\right).$

Since $c\leq\frac{1}{2\sqrt{\mu\kappa}}$ and
$\frac{\alpha_{k}}{\mu}=c\beta_{k+1}$, we can deduce that
$\frac{c\kappa}{\beta_{k+1}}\leq\frac{1}{4\alpha_{k}}$. Using this condition,
we obtain the inequality:

$\sum_{k=0}^{N-1}\frac{1}{4\alpha_{k}}\mathbb{E}[{\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})]\leq\Phi({\bm{x}}^{0})-\Phi^{*}+c\left(\|{\bm{\varepsilon}}_{0}\|^{2}+\sum_{k=1}^{N}\beta_{k}^{2}\mathbb{E}[\|\widetilde{\nabla}_{k}-\nabla
F({{\bm{x}}^{k}})\|^{2}]\right).$

Note that
${\mathcal{D}_{\phi}}({{\bm{x}}^{k}},{{\bm{x}}^{k+1}})\geq\frac{\mu}{2}\|{{\bm{x}}^{k}}-{{\bm{x}}^{k+1}}\|^{2}=\frac{\mu\alpha_{k}^{2}}{2}\|\widetilde{\mathcal{G}}_{\alpha_{k}}({{\bm{x}}^{k}})\|^{2}$
and by the definition of the random variable $a$, we get

$\mathbb{E}[\|\widetilde{\mathcal{G}}_{\alpha_{a}}({\bm{x}}^{a})\|^{2}]\leq\frac{\Phi^{0}-\Phi^{*}+{c\|\bm{\varepsilon}_{0}\|^{2}}+c\sum_{k=1}^{N}\frac{\beta_{k}^{2}\sigma^{2}}{m_{k}}}{\sum_{k=0}^{N-1}\frac{\mu\alpha_{k}}{8}}=\mathcal{O}\left(\frac{1}{\sum_{k=0}^{N-1}\alpha_{k}}+\frac{\sum_{k=0}^{N-1}\frac{\alpha_{k}^{2}}{m_{k+1}}}{\sum_{k=0}^{N-1}\alpha_{k}}\right),$

which completes the proof.

###### Remark 25

Now we give some remarks for Theorem 24.

  * 1.

When the sequence $\\{{\bm{x}}_{k}\\}$ is bounded, an alternative assumption
to Assumption 3 is that $C=\mathbb{R}^{d}$ and $\phi$ has a locally Lipschitz
gradient, as made in (Bolte et al., 2018, Theorem 4.1) and (Latafat et al.,
2022, Theorem 4.7). Under these assumptions, we can conclude that there exists
a compact set $\mathcal{U}$ containing $\\{{\bm{x}}_{k}\\}$. Therefore, there
exists $L_{\phi,\mathcal{U}}>0$ such that $\nabla\phi$ is Lipschitz continuous
over $\mathcal{U}$, and we can derive that $\|\nabla F({\bm{x}})-\nabla
F({\bm{y}})\|^{2}\leq
L_{F}^{2}L_{\phi,\mathcal{U}}^{2}\|{\bm{x}}-{\bm{y}}\|^{2}\leq\frac{2L_{F}^{2}L_{\phi,\mathcal{U}}^{2}}{\mu}\mathcal{D}_{\phi}({\bm{x}},{\bm{y}})$
holds.

  * 2.

Compared to the results presented in Theorem 19, it is worth noting that even
when keeping $m_{k}$ constant, we can still achieve convergence to zero error
bound by carefully selecting the stepsize sequence $\\{\alpha_{k}\\}$. A
typical stepsize condition is that
$\sum_{k}\alpha_{k}=\infty,\;\sum_{k}\alpha_{k}^{2}<\infty$, which coincides
with the classical stepsize condition that guarantees a sufficient but not too
fast decrease of the stepsize, as discussed in Bertsekas and Tsitsiklis
(2000). Therefore, by incorporating the momentum technique, we can achieve an
improved convergence with negligible additional computation costs. This
desirable convergence property theoretically supports the use of SBPG with
momentum for large-scale problems, such as deep neural networks, without the
need to increase the mini-batch size.

##  4 Application of MSBPG in deep neural networks

In this section, we present a detailed description of MSBPG applied to
training deep neural networks. Throughout this section, we assume that the
optimization domain $\overline{C}$ is the entire space $\mathbb{R}^{d}$, so
that $\phi\in\mathcal{M}(\mathbb{R}^{d})$ and
$F\in\mathcal{C}^{1}(\mathbb{R}^{d})$. For simplicity, we omit the explicit
mention of the feasible set $\mathbb{R}^{d}$ in this section. The optimization
problem we consider here is given by:

$\min_{{\bm{W}}}\;\underbrace{\frac{1}{N}\sum_{i=1}^{N}\mathcal{L}(\mathcal{DNN}({\bm{W}},{\bm{x}}_{i}),y_{i})}_{F({\bm{W}})}+\lambda\|{\bm{W}}\|_{1},$
(11)

where $\mathcal{DNN}({{\bm{W}}},{\bm{x}})$ is the neural network function with
training parameters ${\bm{W}}$ and input data ${\bm{x}}$, $\mathcal{L}$ is the
loss function that measures the difference between the output of the neural
network $\mathcal{DNN}({{\bm{W}}},{\bm{x}}_{i})$ and the label $y_{i}$,
$F({\bm{W}})$ is the training loss evaluated on the training dataset
$\\{({\bm{x}}_{i},y_{i})\\}_{i=1}^{N}$, and $\lambda\|{\bm{W}}\|_{1}$ is the
$L_{1}$ regularization term that is often used to avoid overfitting in
training deep neural networks (Ng, 2004). To illustrate the neural network
function $\mathcal{DNN}({{\bm{W}}},{\bm{x}})$, in the $L$-layer fully
connected neural network, we have
${\bm{W}}=[{\bm{W}}_{1},{\bm{W}}_{2},\cdots,{\bm{W}}_{L}]$ and

$\mathcal{DNN}({\bm{W}},{\bm{x}})=\sigma_{L}({\bm{W}}_{L}(\sigma_{L-1}({\bm{W}}_{L-1}(...(\sigma_{1}({\bm{W}}_{1}{\bm{x}}))...)))),$
(12)

where $\sigma_{i}$ is the nonlinear activation function. In this paper, we
focus on smooth activation functions.

At the $k$-th iteration, MSBPG has the following update scheme:

$\displaystyle{\bm{v}}_{k}$ $\displaystyle=$
$\displaystyle(1-\beta_{k}){\bm{v}}^{k-1}+\beta_{k}\widetilde{\nabla}_{k}$
(13) $\displaystyle{{\bm{W}}}^{k+1}$ $\displaystyle=$
$\displaystyle\operatorname*{argmin}_{{\bm{W}}}\;\langle{\bm{v}}^{k},\,{{\bm{W}}}-{{\bm{W}}}^{k}\rangle+\frac{1}{\alpha_{k}}\mathcal{D}_{\phi}({{\bm{W}}},{{\bm{W}}}^{k})+\lambda\|{\bm{W}}\|_{1},$
(14)

where $\widetilde{\nabla}_{k}$ is mini-batch gradient computed by automatic
differentiation (Griewank and Walther, 2008). Omitting all the constants, the
subproblem takes the form of:

${{\bm{W}}}^{k+1}=\operatorname*{argmin}_{{\bm{W}}}\;\phi({{\bm{W}}})+\langle{\bm{p}}^{k},\,{{\bm{W}}}\rangle+\alpha_{k}\lambda\|{{\bm{W}}}\|_{1},$
(15)

where ${\bm{p}}^{k}=\alpha_{k}{\bm{v}}^{k}-\nabla\phi({{\bm{W}}}^{k})$. Here
we adopt the kernel function
$\phi({{\bm{W}}})=\frac{1}{2}\|{{\bm{W}}}\|^{2}+\frac{\delta}{r}\|{{\bm{W}}}\|^{r}$
($r\geq 2$) for training neural networks, and then we have an explicit
solution for (15) in Proposition 26.

###### Proposition 26

Given ${\bm{p}}^{k}\in\mathbb{R}^{d}$, positive constant $\alpha_{k}$,
$\lambda$, and the kernel function
$\phi({{\bm{W}}})=\frac{1}{2}\|{{\bm{W}}}\|^{2}+\frac{\delta}{r}\|{{\bm{W}}}\|^{r}$
$(r\geq 2,\ \delta>0)$. The solution of the subproblem (15) is given by

${{\bm{W}}}^{k+1}=-t^{*}{\bm{p}}^{+},$

where $t^{*}$ is the unique positive real root of the equation

$\displaystyle(\delta\|{\bm{p}}^{+}\|^{r-2})t^{r-1}+t-1=0,$ (16)

and ${\bm{p}}^{+}$ is given by

${\bm{p}}^{+}=\operatorname*{argmin}_{\bm{p}}\Big{\\{}\frac{1}{2}\|{\bm{p}}-{\bm{p}}^{k}\|^{2}+\alpha_{k}\lambda\|{\bm{p}}\|_{1}\Big{\\}}$

which has an explicit expression given by
${\bm{p}}_{j}^{+}=\operatorname{sign}({\bm{p}}_{j}^{k})\max(|{\bm{p}}_{j}^{k}|-\alpha_{k}\lambda,0)$
for the $j$-th coordinate.

Proof  The optimality condition of (15) is given by

$0={\bm{W}}^{k+1}(1+\delta\|{\bm{W}}^{k+1}\|^{r-2})+{\bm{p}}^{k}+\alpha_{k}\lambda{\bm{\Gamma}}^{k},\;\text{where
}{\bm{\Gamma}}^{k}\in\partial\|\cdot\|_{1}({\bm{W}}^{k+1}).$

Let ${\bm{p}}^{+}={\bm{p}}^{k}+\alpha_{k}\lambda{\bm{\Gamma}}^{k}$. By the
optimality condition, we have ${\bm{W}}^{k+1}=-t{\bm{p}}^{+}$ for some
positive scalar $t$, and

$(-t-\delta\|{\bm{p}}^{+}\|^{r-2}t^{r-1}+1){\bm{p}}^{+}=0.$

If ${\bm{p}}^{+}\neq 0$, then $\delta\|{\bm{p}}^{+}\|^{r-2}t^{r-1}+t-1=0$. If
${\bm{p}}^{+}=0$, then ${\bm{W}}^{k+1}=-t{\bm{p}}^{+}=0$. Since $t>0$, then we
have
$\partial\|\cdot\|_{1}({\bm{W}}^{k+1})=\partial\|\cdot\|_{1}(-t{\bm{p}}^{+})=-\partial\|\cdot\|_{1}({\bm{p}}^{+})$.
Recall the definition of ${\bm{p}}^{+}$, we have

${\bm{p}}^{+}={\bm{p}}^{k}+\alpha_{k}\lambda{\bm{\Gamma}}^{k}\in{\bm{p}}^{k}-\alpha_{k}\lambda\partial\|\cdot\|_{1}({\bm{p}}^{+}),$

which is sufficient and necessary optimality condition of the convex
optimization problem:

${\bm{p}}^{+}=\operatorname*{argmin}_{\bm{p}}\left\\{\frac{1}{2}\|{\bm{p}}-{\bm{p}}^{k}\|^{2}+\alpha_{k}\lambda\|{\bm{p}}\|_{1}\right\\}.$

This completes the proof by noting the the above minimization problem is the
well-known soft threshold operator, see for example Friedman et al. (2010).

###### Example 1

In the absence of regularization, that is, when $\lambda=0$, then
$\bm{p}^{+}=\bm{p}^{k}$ and the update formula for MSBPG at the $k$-th
iteration simplifies to ${{\bm{W}}}^{k+1}=-t^{*}{\bm{p}}^{k}$, where $t^{*}$
is the positive root of the equation (16). In this case,
${\bm{W}}^{k+1}=t^{*}(\nabla\phi({\bm{W}}^{k})-\alpha_{k}{\bm{v}}^{k})$.

###### Example 2

If we set the $L_{1}$ regularization parameter $\lambda$ to zero and choose
the kernel function simply as the Euclidean distance, i.e. $\delta=0$, then
SBPG reduces to SGD with momentum. Specifically, we have $t^{*}=1$ and the
update

${{\bm{W}}}^{k+1}={{\bm{W}}}^{k}-\alpha_{k}{\bm{v}}^{k}.$

#### Determining degree of kernel function

We now turn our attention to selecting the appropriate parameter $r$ for the
kernel function. Intuitively, in order to bound the Hessian of the loss
function in (11), particularly when the number of layers $L$ in (12) is large,
$r$ should also be chosen to be larger, so that
$\nabla^{2}F\preceq\frac{1}{\alpha}\nabla^{2}\phi$ holds globally for some
$\alpha>0$. However, in this case, a significant numerical issue may arise
when computing $\|{{\bm{W}}}\|^{r-2}$. This problem can be avoided if the deep
neural network exhibits some special structure such that a moderate $r$ can
make $F({\bm{W}})$ smooth adaptable with respect to $\phi({\bm{W}})$. For
simplicity of analysis, we assume all the given label $y_{i}$ as zero and
consider a sum of squares error loss function. Then, we have a two-layer model
defined as follows:

$\min_{{\bm{W}}=({\bm{u}},{\bm{v}})}\;F({\bm{W}})=\frac{1}{2}\left\|\sigma\left({\rm
Mat}({\bm{u}})(g({\bm{v}}))\right)\right\|^{2},$ (17)

where ${\bm{v}}\in\mathbb{R}^{n}$, ${\bm{u}}\in\mathbb{R}^{km}$,
$g:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m}$,
$\sigma:\mathbb{R}\rightarrow\mathbb{R}$, ${\rm
Mat}({\bm{u}})\in\mathbb{R}^{k\times m}$ and $\sigma(\cdot)$ is a coordinate-
wise operator. Notably, any deep neural network can be recast as the two-layer
model given by (17). For instance, if we define
${\bm{v}}=({\bm{W}}_{1},...,{\bm{W}}_{L-1})$, ${\bm{u}}={\rm
Vec}({\bm{W}}_{L})$,
$g({\bm{W}}_{1},...,{\bm{W}}_{L-1})=\sigma_{L-1}({\bm{W}}_{L-1}(...(\sigma_{1}({\bm{W}}_{1}{\bm{x}}))...))$,
then model (12) can be reformulated as (17). We make the following assumptions
in this section, which guarantees that we can find a polynomial kernel
function $\phi$ with a moderate degree, such that $F$ in (17) is smooth
adaptable to $\phi$.

###### Assumption 4

$\sigma$ is twice differentiable and $\sigma^{\prime}$ and
$\sigma\cdot\sigma^{\prime\prime}$ are globally bounded.

###### Assumption 5

$g$ is twice differentiable. All partial derivatives of order zero, one, and
two of $g$ are globally bounded.

###### Remark 27

Now we give some remarks on the above assumptions.

  * 1.

Assumption 4 is typically valid for various commonly used smooth activation
functions. For example, the sigmoid function $\sigma(x)=\frac{1}{1+e^{-x}}$
satisfies global boundedness for both $\sigma$ and $\sigma^{\prime\prime}$.
Certain activation function may not have bounded function value, such as GELU
(Hendrycks and Gimpel, 2016), which takes the formulation of
$\sigma(x)=x\Phi(x)$ where $\Phi$ is the standard Gaussian cumulative
distribution function. Nonetheless, the product
$\sigma\cdot\sigma^{\prime\prime}$ is globally bounded. Another type of
activation function satisfying Assumption 4 is the smoothed ReLU function, for
example, the following smoothed ReLU function, which we will consider in our
numerical experiments:

$\sigma_{\epsilon}(x)=\left\\{\begin{array}[]{cc}0&x\leq 0\\\
x^{3}\left(\frac{1}{\epsilon^{2}}-\frac{x}{2\epsilon^{3}}\right)&0<x\leq\epsilon\\\
x-\frac{\epsilon}{2}&x>\epsilon.\end{array}\right.$

We observe that as $\epsilon$ tends to zero, $\sigma_{\epsilon}$ converges to
the ReLU function. It is straightforward to verify that
$\sigma_{\epsilon}\cdot\sigma_{\epsilon}^{\prime\prime}$ is globally bounded.
Specifically, $\frac{3}{4}$ is a uniform bound on
$\sigma_{\epsilon}\cdot\sigma_{\epsilon}^{\prime\prime}$ for
$\epsilon\in(0,\frac{1}{2})$.

  * 2.

In many popular neural network frameworks, batch normalization (BN) layers
(Ioffe and Szegedy, 2015) are often used before the fully connected layers.
For example, in the VGG (Simonyan and Zisserman, 2014) and ResNet (He et al.,
2016), BN layers are usually used before the last linear layer. In this case,
we can treat all layers except the last one as one layer, which can be modeled
as (17). It is expected that the BN layer can make the function $g$
sufficiently smooth, thereby satisfying Assumption 5.

By applying the chain rule, we can compute the Hessian of $F$ and determine a
suitable degree parameter $r$ in the kernel function, which will ensure that
$\nabla^{2}F$ is bounded by $\nabla^{2}\phi$ globally. Consequently, $F$ is
smooth adaptable with respect to $\phi$. In order to compute the Hessian of
$F$, two formulas are required, which can be verified directly.

###### Lemma 28

Let ${\bm{u}}\in\mathbb{R}^{km}$, ${\bm{g}}\in\mathbb{R}^{m}$,
${\bm{A}}\in\mathbb{R}^{n\times m}$ and ${\bm{b}}\in\mathbb{R}^{k}$. Consider
two linear maps: ${\bm{u}}\mapsto{\rm Mat}({\bm{u}}){\bm{g}}$ and
${\bm{u}}\mapsto{\bm{A}}({\rm Mat}({\bm{u}}))^{T}{\bm{b}}$, then, the Jacobian
of the two maps are given by

$J_{{\bm{u}}}\left[{\rm
Mat}({\bm{u}}){\bm{g}}\right]={\bm{g}}^{T}\otimes\mathbb{I}_{k},$
$J_{{\bm{u}}}[{\bm{A}}({\rm
Mat}({\bm{u}}))^{T}{\bm{b}}]={\bm{A}}\otimes{\bm{b}}^{T}.$

###### Proposition 29

Suppose Assumptions 4 and 5 hold. Then, for any given $\delta>0$ and any
$r\geq 4$, the function $F$ defined in (17) is smooth adaptable with respect
to
$\phi({\bm{W}})=\frac{1}{2}\|{\bm{W}}\|^{2}+\frac{\delta}{r}\|{\bm{W}}\|^{r}$.

Proof  We denote ${\rm Mat}({\bm{u}})$ by ${\bm{M}}$. The Jacobian of $g$ is
denoted by $Jg$, while its transpose is denoted by $J^{T}g$. $\mathbb{I}_{k}$
is $k\times k$ identity matrix. Using Lemma 28, we can compute the Jacobian
and Hessian of $F$ as follows:

Jacobian of $F$ :

$\displaystyle\frac{\partial F}{\partial{\bm{u}}}$
$\displaystyle=\left(g({\bm{v}})\otimes\mathbb{I}_{k}\right)\left[\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma({\bm{M}}g({\bm{v}}))\right],$
(18) $\displaystyle\frac{\partial F}{\partial{\bm{v}}}$
$\displaystyle=J^{T}g({\bm{v}}){\bm{M}}^{T}\left[\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma({\bm{M}}g({\bm{v}}))\right].$

Hessian of $F$ :

$\displaystyle\frac{\partial^{2}F}{\partial{\bm{u}}^{2}}=(1)+(2),$ (19)
$\displaystyle\text{where }(1)$
$\displaystyle=\left(g({\bm{v}})\otimes\mathbb{I}_{k}\right)\text{Diag}\big{(}\sigma({\bm{M}}g({\bm{v}}))\circ\sigma^{\prime\prime}({\bm{M}}g({\bm{v}}))\big{)}({\bm{g}}^{T}(v)\otimes\mathbb{I}_{k})$
$\displaystyle(2)$
$\displaystyle=\left(g({\bm{v}})\otimes\mathbb{I}_{k}\right)\text{Diag}\big{(}\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma^{\prime}({\bm{M}}g({\bm{v}}))\big{)}({\bm{g}}^{T}(v)\otimes\mathbb{I}_{k}).$

$\displaystyle\frac{\partial^{2}F}{\partial{\bm{u}}\partial{\bm{v}}}=(1)+(2)+(3),$
(20) $\displaystyle\text{where }(1)$
$\displaystyle=\left(J^{T}g({\bm{v}})\right)\otimes\left[\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma^{\prime}({\bm{M}}g({\bm{v}}))\right]^{T}$
$\displaystyle(2)$
$\displaystyle=J^{T}g({\bm{v}}){\bm{M}}^{T}\text{Diag}[\sigma({\bm{M}}g({\bm{v}}))\circ\sigma^{\prime\prime}({\bm{M}}g({\bm{v}}))]\left(g^{T}({\bm{v}})\otimes\mathbb{I}_{k}\right)$
$\displaystyle(3)$
$\displaystyle=J^{T}g({\bm{v}}){\bm{M}}^{T}\text{Diag}[\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma^{\prime}({\bm{M}}g({\bm{v}}))]\left(g^{T}({\bm{v}})\otimes\mathbb{I}_{k}\right).$
$\displaystyle\frac{\partial^{2}F}{\partial{\bm{v}}^{2}}=(1)+(2)+(3),$ (21)
$\displaystyle\text{where }(1)$
$\displaystyle=D^{2}g({\bm{v}})\left[{\bm{M}}^{T}[\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma({\bm{M}}g({\bm{v}}))]\right]=\sum
d_{i}\nabla^{2}g_{i}({\bm{v}})$ $\displaystyle(2)$
$\displaystyle=J^{T}g({\bm{v}}){\bm{M}}^{T}\text{Diag}[\sigma({\bm{M}}g({\bm{v}}))\circ\sigma^{\prime\prime}({\bm{M}}g({\bm{v}}))]{\bm{M}}Jg({\bm{v}})$
$\displaystyle(3)$
$\displaystyle=J^{T}g({\bm{v}}){\bm{M}}^{T}\text{Diag}[\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma^{\prime}({\bm{M}}g({\bm{v}}))]{\bm{M}}Jg({\bm{v}}),$

where
${\bm{d}}={\bm{M}}^{T}[\sigma^{\prime}({\bm{M}}g({\bm{v}}))\circ\sigma({\bm{M}}g({\bm{v}}))]$.
Now, we are ready to prove this proposition. For any
${\bm{w}}\in\mathbb{R}^{km+n}$ and
${\bm{h}}=[{\bm{h}}^{u};{\bm{h}}^{v}]\in\mathbb{R}^{km+n}$, it suffices to
prove that
$\langle\nabla^{2}F({\bm{w}}){\bm{h}},\,{\bm{h}}\rangle=\mathcal{O}(\langle\nabla^{2}\phi({\bm{w}}){\bm{h}},\,{\bm{h}}\rangle)$.
From (19)(20)(21) and Assumption 4, 5, we can easily get
$\langle\nabla^{2}F({\bm{w}}){\bm{h}},\,{\bm{h}}\rangle=\mathcal{O}((1+\|{\bm{w}}\|^{2})\|{\bm{h}}\|^{2})$.
On the other hand,
$\nabla^{2}\phi({\bm{w}})=I(1+\|{\bm{w}}\|^{r-2})+(r-2)\|{\bm{w}}\|^{r-4}{\bm{w}}{\bm{w}}^{T}$.
Hence
$\langle\nabla^{2}\phi({\bm{w}}){\bm{h}},\,{\bm{h}}\rangle\geq(1+\|{\bm{w}}\|^{r-2})\|{\bm{h}}\|^{2}$.
So, we only require $r-2\geq 2$. This completes the proof.

#### Layerwise kernel function

In Proposition 26, the kernel function
$\phi({{\bm{W}}})=\frac{1}{2}\|{{\bm{W}}}\|^{2}+\frac{\delta}{r}\|{{\bm{W}}}\|^{r}$
is used, which means we adopt the same Bregman distance for all layers of deep
neural networks. However, different layers have different optimization
property for deep neural networks (You et al., 2019), and computing
$\|{\bm{W}}\|^{r}$ with $r>2$ may result in numerical issues for neural
networks with millions of parameters, such as in VGG (Simonyan and Zisserman,
2014). To take advantage of the layerwise structure of neural networks, we
design a layerwise kernel function for a $L$-layer neural network as follows:

$\phi({{\bm{W}}})=\sum_{i=1}^{L}\phi_{i}({{\bm{W}}}_{i}),\quad\phi_{i}({{\bm{W}}}_{i})=\frac{1}{2}\|{{\bm{W}}}_{i}\|^{2}+\frac{\delta}{r}\|{{\bm{W}}}_{i}\|^{r}.$
(22)

Note that $\delta$ and $r$ can vary from layer to layer, here we take the same
$\delta$ and $r$ for different layers for simplicity. Then, we have the
Bregman distance taking the form
$\mathcal{D}_{\phi}=\sum_{i=1}^{L}\mathcal{D}_{\phi_{i}}$. By employing this
Bregman distance in subproblem (14), our MSBPG algorithm can be implemented in
a layerwise manner. See the details in Algorithm 1.

Algorithm 1 Momentum based Stochastic Bregman Proximal Gradient (MSBPG) for
training neural networks

1:Input: Total number of iterations $K$, stepsize $\alpha_{k}$, momentum
parameter $\beta_{k}$, $\delta$ and $r$ to determine the kernel function
$\phi$.

2:Initialize: Set ${\bm{W}}={\bm{W}}^{0}$, ${\bm{v}}_{0}=0$.

3:for $k=0,\cdots,K-1$ do

4: Compute mini-batch gradient $\widetilde{\nabla}_{k}$;

5: Compute SMAE:
${\bm{v}}^{k}=(1-\beta_{k}){\bm{v}}^{k-1}+\beta_{k}\widetilde{\nabla}_{k}$;

6: for $i=1,\ldots,L$ do

7: ${\bm{p}}_{i}^{k}=\alpha_{k}{\bm{v}}_{i}^{k}-\nabla\phi({\bm{W}}^{k}_{i})$;

8:
${\bm{p}}_{i}^{+}=\operatorname*{argmin}_{{\bm{p}}_{i}}\\{\frac{1}{2}\|{\bm{p}}_{i}-{\bm{p}}_{i}^{k}\|^{2}+\alpha_{k}\lambda\|{\bm{p}}_{i}\|_{1}\\}$;

9: Solve $(\delta\|{\bm{p}}^{+}_{i}\|^{r-2})t_{i}^{r-1}+t_{i}-1=0$ to get
$t_{i}^{k}$;

10: ${\bm{W}}_{i}^{k+1}=-t_{i}^{k}{\bm{p}}_{i}^{+}$;

11: end for

12:end for

13:Output: ${\bm{W}}^{1},\cdots,{\bm{W}}^{K}$.

#### Mitigating gradient explosion

In the training of deep neural networks, gradient explosion is a common
undesired phenomenon, where the gradients of the loss function grow
exponentially from layer to layer, leading to numerical instability or even
collapse of the training process (Hochreiter, 1991; Manchev and Spratling,
2020). The reasons for gradient explosion include selecting a large stepsize
and choosing an improper initialization for the model’s parameters (Pascanu et
al., 2013). In the following, we will show that MSBPG provides a novel
approach to mitigate gradient explosion. Considering MSBPG without
regularization, the update rule is given by:

${{\bm{W}}}_{i}^{k+1}=-t_{i}^{k}{\bm{p}}_{i}^{k}=t_{i}^{k}\left((1+\delta\|{{\bm{W}}_{i}}^{k}\|^{r-2}){{\bm{W}}_{i}}^{k}-\alpha_{k}{\bm{v}}_{i}^{k}\right),$
(23)

where $t_{i}^{k}\in(0,1)$ is the unique positive root of

$\left(\delta\left\|(1+\delta\|{{\bm{W}}}_{i}^{k}\|^{r-2}){{\bm{W}}}_{i}^{k}-\alpha_{k}{\bm{v}}_{i}^{k}\right\|^{r-2}\right)t^{r-1}+t-1=0.$
(24)

Combining (23) and (24), we have the following equivalent implicit update
scheme for the $i$-th layer:

${\bm{W}}_{i}^{k+1}=\frac{1+\delta\|{\bm{W}}_{i}^{k}\|^{r-2}}{1+\delta\|{\bm{W}}_{i}^{k+1}\|^{r-2}}{\bm{W}}_{i}^{k}-\frac{\alpha_{k}}{1+\delta\|{\bm{W}}_{i}^{k+1}\|^{r-2}}{\bm{v}}_{i}^{k}.$
(25)

It is observed in practice that with large stepsize or large initial point,
the gradient ${\bm{v}}_{i}^{k}$ tends to explode if no scaling is done, while
the norm of the weight $\|{\bm{W}}_{i}^{k+1}\|$ also tends to be large. In
(25), we see that by scaling the gradient with
$\frac{1}{1+\delta\|{\bm{W}}_{i}^{k+1}\|^{r-2}}$, the weight
${\bm{W}}_{i}^{k+1}$ is relieved from moving excessively in the direction of
the gradient to avoid rapid growth of its norm. At the same time, we can also
see that if the norm $\|{\bm{W}}_{i}^{k+1}\|$ does not change drastically, the
coefficient of ${\bm{W}}_{i}^{k}$ in (25) will be maintained to be
approximately 1. Thus in (25), we see an automatic scaling of the gradient to
avoid rapid growth of the weight and hence also mitigating subsequent gradient
explosion.

Experimental results in Section 5.2 indeed verify MSBPG ’s ability to mitigate
gradient explosion for training deep neural networks. An intuitive
illustration of MSBPG ’s “pull-back” ability is given in Figure 11 in Appendix
B, and this “pull-back” ability originates from the Bregman proximity model
and the polynomial kernel function we adopt.

#### Improving generalization capacity

From (25) we can see that MSBPG employs a scaling for the gradient during the
update of the parameters as a result of adopting a Bregman proximity model and
a polynomial kernel function. This scaling not only helps to mitigate the
gradient explosion phenomenon, but also can improve the generalization
capacity of MSBPG . To be specific, at the beginning of the training process,
the initial weight of each layer ${\bm{W}}_{i}$ tends to have a larger norm
$\|{\bm{W}}_{i}\|$ and therefore MSBPG takes a cautious update at the
beginning phase of training. As the training goes on, due to the effect of
regularization, either $L_{1}$ or $L_{2}$ regularization, the norm of each
layer’s weight $\|{\bm{W}}_{i}\|$ becomes smaller, and MSBPG can then take
bolder update of the parameters. This implicit training strategy of MSBPG is
in agreement with the idea of a heuristic deep learning training technique
called “learning rate warm-up” (Gotmare et al., 2018), which benefits the
training stability and generalization performance (Liu et al., 2019).
Experimental results in Section 5.2 also testify the excellent generalization
capacity of MSBPG .

##  5 Numerical experiments

In this section, we conduct numerical experiments to showcase the
effectiveness and robustness of MSBPG in comparison to modern solvers commonly
used in deep learning. We assess the impact of stepsize and initial point
selection on the performance of our method. Our experiments consist of two
parts. In the first part, we use a quadratic inverse problem as a toy example
to illustrate the capabilities of vanilla SBPG. The second part is the main
focus of this section, where we evaluate the performance of MSBPG in training
deep neural networks. The experiments for the quadratic inverse problem are
conducted using MATLAB R2021b on a Windows workstation equipped with a 12-core
Intel Xeon E5-2680 @ 2.50GHz processor and 128GB of RAM. For the deep learning
experiments, we conducted the experiments using PyTorch running on a single
RTX3090 GPU.

###  5.1 Quadratic inverse problem

The quadratic inverse problem, as formulated in Bolte et al. (2018), is given
by:

$\min\left\\{\Phi({\bm{x}}):=\underbrace{\frac{1}{4}\sum_{i=1}^{n}(\langle
A_{i}{\bm{x}},\,{\bm{x}}\rangle-b_{i})^{2}}_{F({\bm{x}})}+\lambda
R({\bm{x}}):{\bm{x}}\in\mathbb{R}^{d}\right\\},$

which has practical applications (Beck and Eldar, 2013) that includes the
phase retrieval problem as a special case (Luke, 2017). In this experiment, we
consider the $L_{1}$ regularization $R({\bm{x}})=\|{\bm{x}}\|_{1}$ with
$\lambda=1\times 10^{-3}$, and solve the quadratic inverse problem using SBPG
and stochastic (Euclidean) proximal gradient (SPG) method (Bertsekas, 2011).
Notably, SPG is a special case of SBPG, in which
$\phi({\bm{x}})=\frac{1}{2}\|{\bm{x}}\|^{2}$. Since the smooth term in the
objective function $F({\bm{x}})$ does not admit a globally Lipschitz
continuous gradient, we employ the kernel function
$\phi({\bm{x}})=\frac{1}{2}\|{\bm{x}}\|^{2}+\frac{1}{r}\|{\bm{x}}\|^{r}$ with
$r=4$. It has been shown in Lu et al. (2018) that any $r\geq 4$ guarantees
that $F$ is $\phi$-smooth adaptable globally. Moreover, according to Bolte et
al. (2018), the smooth adaptable constant $L_{F}$ can be chosen such that
$L_{F}\geq\sum_{i=1}^{n}(3\|A_{i}\|^{2}+\|A_{i}\||b_{i}|)$ for $r=4$. In this
experiment, we randomly generate the data by the following ${\tt MATLAB}$
commands:

$\displaystyle{\tt ai=randn(d,1);Ai=ai*ai^{\prime};}$ $\displaystyle{\tt
x\\_true=sprandn(d,1,density\\_x);b\\_i=x\\_true^{\prime}*(Ai*x\\_true);}$

The true solution for the quadratic inverse problem is chosen as a sparse
vector ${\bm{x}}^{*}$ that satisfies $\langle
A_{i}{\bm{x}}^{*},\,{\bm{x}}^{*}\rangle=b_{i}$ for $i=1,\dots,n$. We set the
mini-batch size for all algorithms to be $m=1$. To evaluate the effectiveness
of each algorithm, we use the following criterion that takes into account the
possibility of critical points being local minimum or saddle points:

$\epsilon_{k}=\max\left\\{\|\mathcal{G}({{\bm{x}}^{k}})\|,\;\epsilon_{\tt
obj}:=\frac{{\tt obj}_{k}-{\tt obj}_{*}}{1+{\tt obj}_{*}}\right\\},$

where ${\tt obj}_{k}=\Phi({{\bm{x}}^{k}})$ and ${\tt
obj}_{*}=\Phi({\bm{x}}^{*})$. The term $\|\mathcal{G}({\bm{x}}^{k})\|$
measures the stationarity of the solution, while a small $\epsilon_{\tt obj}$
indicates that the solution is a ”nearly” global minimum.

We conduct experiments on a problem with data size $d=100$ and ${\tt
density\\_x=0.05}$. All methods are run until they reach an accuracy of
$\epsilon_{k}\leq 0.01$ within a time limit of 30 seconds. To ensure
statistical significance, we run each algorithm 10 times and report the median
value. The results are presented in Figure 2. For Figures 2(a), we randomly
select initial points within a ball centered at the origin with radius
$1\times 10^{-2}$. We use the stepsize schedule of
$\alpha_{k}=\max\left\\{10^{-4},\frac{\alpha_{0}}{\sqrt{1+k}}\right\\}$, where
$\alpha_{0}$ is the initial stepsize. For Figure 2(b), we set constant
stepsize schedule $1\times 10^{-3}$. For Figures 2(c), we randomly select
initial points within a ball centered at the origin with radius $1\times
10^{-2}$. We use constant stepsize schedule. To prevent excessively small
stepsizes that can slow down all methods, we set a lower bound for the
stepsize.

Figure 2(a) demonstrates that SBPG has a larger range of convergent stepsizes
than SPG, indicating that SBPG is more robust in terms of stepsize selection.
The impact of the initial stepsize on the performance of the algorithms is
reported in this figure. Figure 2(b) shows that SBPG is much more robust than
SPG in terms of initial point selection. Specifically, SBPG exhibits high
resilience to initial point selection and avoids causing the training to
collapse. Figure 2(c) reveals that a larger degree $r$ in the kernel function
increases the safe stepsize threshold. These observations are partly explained
in Section 4. Since a large stepsize and a large radius of the initial point
tend to lead to gradient explosion, Bregman proximal mapping helps to pull
back the iterate and guide it towards a better solution.

(a) Stepsize robustness (b) Initial point robustness (c) Safe stepsize
threshold

Figure 2: Comparison of SBPG and SPG in terms of their robustness with respect
to stepsize and initial point selction. A method is considered non-convergent
if it fails to reach an accuracy of $\epsilon_{k}<10^{-2}$ within 30 seconds
or if it collapses. Generally, choosing large stepsize and large radius for
the initial point can cause an algorithm to collapse. The safe stepsize
threshold is the maximum stepsize (constant schedule) that a method does not
collapse. We run 10 tests for each algorithm and report the median of the
results.

###  5.2 Deep neural network

For the evaluation of MSBPG’s performance on training deep neural networks, we
consider the model with $L_{2}$ regularization here for its better
generalization capacity:

$\min_{{\bm{W}}}\;\underbrace{\frac{1}{N}\sum_{i=1}^{N}\mathcal{L}(\mathcal{DNN}({\bm{W}},{\bm{x}}_{i}),y_{i})}_{F({\bm{W}})}+\lambda\|{\bm{W}}\|_{2}^{2}.$
(26)

We employ MSBPG to solve this large-scale problem. Following AdamW (Loshchilov
and Hutter, 2017), we also conduct decoupled weight decay at the end of each
iteration for $L_{2}$ regularization and do not consider the $L_{2}$
regularization term when solving the subproblems. The detailed algorithm is
summarized in Algorithm 2. At iteration $k$, MSBPG first uses automatic
differentiation to compute the mini-batch gradient $\widetilde{\nabla}_{k}$.
Then, it maintains a bias-corrected gradient estimator $\bar{{\bm{v}}}^{k}$
(Kingma and Ba, 2014) and use it to calculate the layerwise
${\bm{p}}_{i}^{k}$. With ${\bm{p}}_{i}^{k}$, MSBPG solves a univariate
equation to get $t_{i}^{k}$ and update the weight of the $i$-th layer to
${\bm{W}}_{i}^{k}$. In the end, MSBPG conducts decoupled weight decay as
$L_{2}$ regularization.

Algorithm 2 MSBPG with $L_{2}$ regularization

1:Input: Total number of training epochs $K$, momentum coefficient $\beta$,
stepsize $\alpha_{k}$, weight decay coefficient $\gamma$, $\delta$ and $r$ to
determine the kernel function $\phi$.

2:Initialize:  Set ${\bm{W}}={\bm{W}}^{0}$, ${\bm{v}}^{0}=\bf{0}$.

3:for $k=1,\cdots,K$ do

4: Compute mini-batch gradient $\widetilde{\nabla}^{k}$;

5: ${\bm{v}}^{k}=\beta{\bm{v}}^{k-1}+(1-\beta)\widetilde{\nabla}^{k}$,
$\bar{{\bm{v}}}^{k}={\bm{v}}^{k}/(1-\beta^{k})$;

6: for $i=1,\cdots,L$ do

7:
${\bm{p}}_{i}^{k}=\alpha_{k}\bar{{\bm{v}}}_{i}^{k}-\nabla\phi({\bm{W}}^{k-1}_{i})$;

8: Solve $(\delta\|{\bm{p}}^{k}_{i}\|^{r-2})t_{i}^{r-1}+t_{i}-1=0$ to get
$t_{i}^{k}$;

9: $\widetilde{{\bm{W}}}_{i}^{k}=-t_{i}^{k}{\bm{p}}_{i}^{k}$;

10: end for

11: ${\bm{W}}^{k}=\widetilde{{\bm{W}}}^{k}-\alpha_{k}\gamma{\bm{W}}^{k-1}$;

12:end for

13:Output: ${\bm{W}}^{1},\cdots,{\bm{W}}^{K}$

We conducted extensive experiments on several representative benchmarks,
including VGG16 (Simonyan and Zisserman, 2014), ResNet34 (He et al., 2016) on
CIFAR10 dataset (Krizhevsky et al., 2009), ResNet34 (He et al., 2016),
DenseNet121 (Huang et al., 2017) on CIFAR100 dataset (Krizhevsky et al.,
2009), and LSTMs (Hochreiter and Schmidhuber, 1997) on the Penn Treebank
dataset (Marcinkiewicz, 1994). We compare MSBPG with the most popular
optimization algorithms used for training neural networks, including SGD
(Sutskever et al., 2013), Adam (Kingma and Ba, 2014), and AdamW (Loshchilov
and Hutter, 2017). Experimental results show that MSBPG has excellent
convergence performance and best generalization capacity for both the task
that SGD dominates (image classification with CNNs) and the task that Adam
dominates (language modeling with LSTMs). We also conducted experiments to
compare MSBPG with SGD on different initial stepsizes and different scales of
initial point. Our experimental results demonstrate the robustness of MSBPG in
training neural networks.

Before getting into the details of our experiments, we first make a
clarification about the activation function. The frequently used activation
function ReLU in VGG, ResNet, and DenseNet takes the form of
$\operatorname{ReLU}(x)=\max(0,x)$, which is not continuously differentiable.
Here we design a smoothing approximation of ReLU with coefficient $\epsilon$,
which is twice continuously differentiable and satisfies our Assumption 4,
namely,

$\sigma_{\epsilon}(x)=\left\\{\begin{array}[]{cc}0&x\leq 0\\\
x^{3}\left(\frac{1}{\epsilon^{2}}-\frac{x}{2\epsilon^{3}}\right)&0<x\leq\epsilon\\\
x-\frac{\epsilon}{2}&x>\epsilon.\end{array}\right.$

This activation function has gradient taking the form:

$\sigma_{\epsilon}^{\prime}(x)=\left\\{\begin{array}[]{cc}0&x\leq 0\\\
x^{2}(\frac{3}{\epsilon^{2}}-\frac{2x}{\epsilon^{3}})&0<x\leq\epsilon\\\
1&x>\epsilon.\end{array}\right.$

Note that as $\epsilon$ tends to 0, this twice continuously differentiable
activation function tends to the ReLU function. We conducted experiments with
VGG16 on the CIFAR10 dataset, where we replace all the activation functions in
VGG16 by $\sigma_{\epsilon}$ defined above. As shown in Figure 4, our
algorithm MSBPG ’s performance does not degrade as $\epsilon$ tends to 0.
Therefore, in the subsequent experiments, we use the original neural network
architectures with the ReLU activation function to evaluate our method MSBPG .

(a) Step decay. | (b) Cosine annealing.
---|---

Figure 3: Training and test accuracy (%) of VGG16 on CIFAR10 dataset under two
frequently used training settings. Here the activation function of VGG16
adopts smoothed ReLU activation function $\sigma_{\epsilon}$ with different
choices of $\epsilon$ ($\epsilon=0$ denotes adopting the original ReLU
activation function).

#### CNNs on image classification

We experimented with VGG16, ResNet34 on the CIFAR10 dataset, and ResNet34,
DenseNet121 on CIFAR100 dataset. SGD usually has better generalization
performance than adaptive gradient algorithms such as Adam and AdamW when
training CNNs on image classification tasks and therefore is the default
optimizer in these scenarios (He et al., 2016; Zhou et al., 2020). We used two
dominant experimental settings for training neural networks, including
reducing the stepsize to 0.1 times its original value near the end of training
(Zhuang et al., 2020; Chen et al., 2021; Luo et al., 2019) and adopting a
cosine annealing schedule for the stepsizes (Loshchilov and Hutter, 2016,
2017). The purpose of these two training strategies is to accelerate the
convergence of the optimization algorithms so as to give a fair comparison of
their generalization capacity. We use the default training hyperparameters of
SGD, Adam, and AdamW in these settings (He et al., 2016; Zhuang et al., 2020;
Chen et al., 2021), and set MSBPG ’s learning rate (initial stepsize) as 0.1,
momentum coefficient $\beta$ as 0.9, weight decay coefficient $\gamma$ as
$1\times 10^{-3}$. For the layerwise kernel function
$\phi_{i}({{\bm{W}}}_{i})=\frac{1}{2}\|{{\bm{W}}}_{i}\|^{2}+\frac{\delta}{r}\|{{\bm{W}}}_{i}\|^{r}$,
we set $r=4,\ \delta=1\times 10^{-2}$ for VGG16 and $r=6,\ \delta=1\times
10^{-3}$ for ResNet34 on CIFAR10 dataset, and $r=4,\ \delta=1\times 10^{-2}$
for ResNet34 and $r=4,\ \delta=1\times 10^{-3}$ for DenseNet121 on CIFAR100
dataset. From the experimental results in Figure 4, 6, 5, 7, we can see that
MSBPG attains 100% training accuracy in all the training settings, unlike Adam
which fails to fully converge with the training strategy of reducing the
learning rate to 0.1 times of original value at the 150th epoch. Furthermore,
MSBPG consistently achieves the best generalization performance for all
experimental settings, and attains at least 0.5% test accuracy improvement
compared with the second best optimization algorithm. This generalization
advantage of MSBPG can be attributed to the Bregman proximity model we adopt.

(a) VGG16 on CIFAR10 | (b) ResNet34 on CIFAR10
---|---

Figure 4: Training and test accuracy (%) of CNNs on CIFAR10 dataset with
learning rate reduced to 0.1 times of the original value at the 150th epoch.

(a) ResNet34 on CIFAR100 | (b) DenseNet121 on CIFAR100
---|---

Figure 5: Training and test accuracy (%) of CNNs on CIFAR100 dataset with
learning rate reduced to 0.1 times of the original value at the 150th epoch.

(a) VGG16 on CIFAR10 | (b) ResNet34 on CIFAR10
---|---

Figure 6: Training and test accuracy (%) of CNNs on CIFAR10 dataset with
learning rate using the cosine annealing schedule.

(a) ResNet34 on CIFAR100 | (b) DenseNet121 on CIFAR100
---|---

Figure 7: Training and test accuracy (%) of CNNs on CIFAR100 dataset with
learning rate using the cosine annealing schedule.

#### LSTMs on language modeling

To further evaluate the performance of MSBPG , we conducted experiments on
LSTM with the Penn Treebank dataset, and report the training and test
perplexity (lower is better). Adam generally have better generalization
capacity than SGD on language modeling (Fu et al., 2016; Siami-Namini et al.,
2019), and therefore is the default optimization algorithm for training LSTMs.
Here we follow the commonly used experimental setting for training LSTMs
(Zhuang et al., 2020; Chen et al., 2021), which reduces the stepsize to 0.1
times its original value two times (at 75th epoch and 150th epoch) during the
training process. We also conducted experiments with the cosine annealing
learning rate (stepsize) schedule (Loshchilov and Hutter, 2016), which is the
most frequently used learning rate schedule in practice. For training
hyperparameters, we use the default settings for SGD, Adam, and AdamW in
training 1-, 2-, 3-layer LSTMs (Zhuang et al., 2020; Chen et al., 2021). For
MSBPG , we set its learning rate as 25, 80, 80 for 1-, 2-, 3-layer LSTMs with
momentum coefficient $\beta=0.9$, weight decay coefficient $\gamma=2\times
10^{-6}$. For the layerwise kernel function
$\phi_{i}({{\bm{W}}}_{i})=\frac{1}{2}\|{{\bm{W}}}_{i}\|^{2}+\frac{\delta}{r}\|{{\bm{W}}}_{i}\|^{r}$,
we set $r=4$ and $\delta=1\times 10^{-6}$. From Figure 8 and Figure 9 we can
see that MSBPG converges well on training dataset for 1-, 2-, 3-layer LSTMs
with both two training strategies. On the other hand, SGD with the cosine
annealing learning rate schedule fails to get fully converged on the training
dataset as shown in Figure 9. Moreover, MSBPG consistently achieves the best
generalization performance for all the experiments with at least 1 unit of
test perplexity lower. This excellent generalization capacity again can be
attributed to the Bregman proxmity model we employ.

(a) 1-layer LSTM | (b) 2-layer LSTM | (c) 3-layer LSTM
---|---|---

Figure 8: Training and test perplexity (lower is better) of LSTMs on Penn
Treebank dataset with learning rate reduced to 0.1 times of the original value
at the 75th epoch and 150th epoch.

(a) 1-layer LSTM | (b) 2-layer LSTM | (c) 3-layer LSTM
---|---|---

Figure 9: Training and test perplexity (lower is better) of LSTMs on Penn
Treebank dataset with learning rate using the cosine annealing schedule.

#### Robustness to initial point scale and stepsize

As demonstrated in Section 4, MSBPG can mitigate the problem of gradient
explosion. Generally, choosing large stepsize and large initial point scale
will lead to gradient explosion. Here we conduct experiments with VGG16 on
CIFAR10 to verify MSBPG ’s robustness in training neural networks. To be
specific, we compare the performance of MSBPG and SGD with different scales of
initial point and different stepsizes, since MSBPG and SGD have the same
default learning rate ($1\times 10^{-1}$). Adaptive gradient algorithms, on
the other hand, have different scale of default learning rate ($1\times
10^{-3}$) and therefore we don’t include them in our comparison here. For
different choices of initial point scale and stepsize, we run the optimization
algorithm for 50 iterations and report the best test accuracy. As we can see
from Figure 10, MSBPG is more robust than SGD to large initial points and
large stepsizes. Training deep neural networks, which has millions or billions
of parameters, is sensitive to the scale of the initial point and stepsize
choice. It can be seen from Figure 10 that SGD fails to converge with a slight
increase of the initial point scale to 4.6 or increase the stepsize from 0.1
to 0.6. MSBPG , on the other hand, can converge with the initial point scale
as large as 20 and stepsize as large as 5. This robustness of MSBPG can ease
the tuning of hyperparameters for training neural networks, and can also make
the training process more robust to noises and errors.

(a) Initial point robustness | (b) Stepsize robustness
---|---

Figure 10: Test accuracy (%) of VGG16 on CIFAR10 dataset with different
initial point scale and stepsize choice.

##  6 Conclusion

In this paper, we consider the problem of minimizing nonconvex composite
objectives where the differentiable part does not satisfy Lipschitz
smoothness, which is a fundamental assumption made by classical stochastic
gradient methods. To overcome this limitation, we investigate a family of
stochastic Bregman proximal gradient (SBPG) methods that only require smooth
adaptivity of the differentiable part. From a modeling perspective, SBPG
replaces the upper quadratic approximation used in SGD with the Bregman
proximity measure, which captures the non-Lipschitz geometry and results in a
better approximation model. We first formulate the vanilla SBPG and establish
its convergence properties under the nonconvex setting without a finite sum
structure. We then propose a momentum-based version of SBPG (MSBPG) that
further improves its convergence properties, making it well-suited for large-
scale applications. To demonstrate the effectiveness of MSBPG, we apply it to
train deep neural networks with a polynomial kernel function that ensures the
smooth adaptivity of the loss function. We also demonstrate the ability of
MSBPG to alleviate gradient explosion during the training of deep neural
networks. We conduct numerical experiments on quadratic inverse problems and
training deep neural networks to validate the effectiveness of SBPG. The
experimental results on sparse quadratic inverse problems show that SBPG is
more robust than classical stochastic (proximal) gradient methods in terms of
stepsize selection and initial point selection. Additionally, the experimental
results on deep neural networks show that MSBPG outperforms state-of-the-art
optimizers in terms of efficiency, robustness of stepsize selection, and yet
achieves better generalization performance. In conclusion, our work
demonstrates that MSBPG can constitute a valuable addition to the existing
family of optimization methods for solving stochastic nonconvex optimization
problems. The enhanced robustness, improved convergence results, the ability
to alleviate gradient explosion, and negligible extra computational cost thus
make MSBPG a promising approach for a broad range of machine learning
applications and beyond.

##  A Proofs in Preliminaries

Proof of Lemma 8 First, we prove the uniqueness of the solution. Problem (4)
is equivalent to the following problem:

$\arg\min_{{\bm{u}}\in\overline{C}}\;\Psi({\bm{u}}):=\alpha
R({\bm{u}})+\phi({\bm{u}})-\langle\nabla\phi({\bm{x}}),\,{\bm{u}}\rangle.$

We have that

$\Psi({\bm{u}})\geq\alpha
R({\bm{u}})+\phi({\bm{u}})-\|\nabla\phi({\bm{x}})\|\|{\bm{u}}\|\geq\|{\bm{u}}\|\Big{(}\frac{\alpha
R({\bm{u}})+\phi({\bm{u}})}{\|{\bm{u}}\|}-\|\nabla\phi({\bm{x}})\|\Big{)}.$

As $\|{\bm{u}}\|\rightarrow\infty$, we have
$\Psi({\bm{u}})\geq\|{\bm{u}}\|\Big{(}\frac{\alpha
R({\bm{u}})+\phi({\bm{u}})}{\|{\bm{u}}\|}-\|\nabla\phi({\bm{x}})\|\Big{)}=\infty$,
where we use the fact that $\phi$ is supercoercive and $R$ is convex. Since
$\Psi$ is a proper lower-semicontinuous convex function, by the modern form of
Weierstrass theorem (Rockafellar, 1997, Chapter 1), we know that the solution
set of (4) is a nonempty compact set. Also note that $\Psi$ is a strictly
convex function, which implies the uniqueness of the solution. For any
Legendre function $\phi$, from (Rockafellar, 1997, Chapter 26), we have
$dom\,\partial\phi=int\,dom\phi$ with
$\partial\phi({\bm{x}})=\\{\nabla\phi({\bm{x}})\\}$ for all ${\bm{x}}\in
int\,dom\,\phi$. The optimality condition implies that $\partial\phi({\rm
Prox}^{\phi}_{\alpha R}({\bm{x}}))$ is nonempty, which automatically forces
${\rm Prox}^{\phi}_{\alpha P}({\bm{x}})\in\text{int\;dom}\phi$. This completes
the proof. $\Box$

Proof of Proposition 12 Note that
$\|\nabla\phi({\bm{x}}^{+})-\nabla\phi({\bm{x}})\|\leq
L_{\phi}\|{\bm{x}}^{+}-{\bm{x}}\|$ and $\|\nabla F({\bm{x}}^{+})-\nabla
F({\bm{x}})\|\leq L_{F}L_{\phi}\|{\bm{x}}^{+}-{\bm{x}}\|$. By the definition
of ${\bm{x}}^{+}$, we have

$\nabla F({\bm{x}}^{+})-\nabla
F({\bm{x}})+\frac{\nabla\phi({\bm{x}})-\nabla\phi({\bm{x}}^{+})}{\alpha}\in\nabla
F({\bm{x}}^{+})+\partial R({\bm{x}}^{+}).$

Thus, we obtain

${\rm dist}\left(0,\partial\Phi({\bm{x}}^{+})\right)\leq\bigg{\|}\nabla
F({\bm{x}}^{+})-\nabla
F({\bm{x}})+\frac{\nabla\phi({\bm{x}})-\nabla\phi({\bm{x}}^{+})}{\alpha}\bigg{\|}\leq\left(L_{F}L_{\phi}+\frac{L_{\phi}}{\alpha}\right)\|{\bm{x}}^{+}-{\bm{x}}\|.$

Note that
$\|{\bm{x}}^{+}-{\bm{x}}\|=\alpha\|\mathcal{G}_{\alpha}({\bm{x}})\|$, which
completes the proof. $\Box$

Proof of Proposition 9 By the definition of ${\rm Prox}_{R}^{\phi}(\cdot)$,
$x_{i}\in\partial R({\bm{x}}_{i}^{+})+\nabla\phi({\bm{x}}^{+}_{i})$, $i=1,2$.
Since $\partial R(\cdot)$ is monotone, then
$\langle{\bm{x}}_{1}-{\bm{x}}_{2}-(\nabla\phi({\bm{x}}^{+}_{1})-\nabla\phi({\bm{x}}^{+}_{2})),\,{\bm{x}}^{+}_{1}-{\bm{x}}^{+}_{2}\rangle\geq
0$. From the $\mu$-strong convexity of $\phi$, it follows that
$\langle{\bm{x}}_{1}-{\bm{x}}_{2},\,{\bm{x}}^{+}_{1}-{\bm{x}}^{+}_{2}\rangle\geq\langle\nabla\phi({\bm{x}}^{+}_{1})-\nabla\phi({\bm{x}}^{+}_{2}),\,{\bm{x}}^{+}_{1}-{\bm{x}}^{+}_{2}\rangle\geq\mu\|{\bm{x}}^{+}_{1}-{\bm{x}}^{+}_{2}\|^{2}$.
Therefore,
$\|{\bm{x}}^{+}_{1}-{\bm{x}}^{+}_{2}\|\leq\frac{1}{\mu}\|{\bm{x}}_{1}-{\bm{x}}_{2}\|$.
$\Box$

##  B Additional Figures

(a) Illustration of SBPG and SGD updates (b) Different initial points for SBPG

Figure 11: Figure (a) depicts SGD and SBPG updates. SBPG includes a ”pull
back” mechanism that prevents the point from moving excessively in any given
direction. $"P"$ and $"D"$ refer to the primal and dual spaces, respectively,
and these terms are commonly used in the mirror descent method literature
(see, e.g., Bubeck et al. (2015); Nemirovskij and Yudin (1983)). Figure (b)
illustrates the effect of choosing the initial point from a ball of radius $r$
on the SBPG when $r$ changes for the QIP example with $d=100$ and $n=5000$.
All initial step sizes are set to $1\times 10^{-3}$. As shown in Figure (b),
even for an initial point that is far from the optimal point, SBPG can pull
back the iterates to the optimal point.

## References

  * Allen-Zhu (2018) Zeyuan Allen-Zhu.  Katyusha: The first direct acceleration of stochastic gradient methods.  _Journal of Machine Learning Research_ , 18(221):1–51, 2018.
  * Bauschke et al. (2017) Heinz H Bauschke, Jérôme Bolte, and Marc Teboulle.  A descent lemma beyond Lipschitz gradient continuity: first-order methods revisited and applications.  _Mathematics of Operations Research_ , 42(2):330–348, 2017.
  * Beck and Eldar (2013) Amir Beck and Yonina C Eldar.  Sparsity constrained nonlinear optimization: Optimality conditions and algorithms.  _SIAM Journal on Optimization_ , 23(3):1480–1509, 2013.
  * Bertsekas (2011) Dimitri P Bertsekas.  Incremental proximal methods for large scale convex optimization.  _Mathematical Programming_ , 129(2):163–195, 2011\.
  * Bertsekas and Tsitsiklis (2000) Dimitri P Bertsekas and John N Tsitsiklis.  Gradient convergence in gradient methods with errors.  _SIAM Journal on Optimization_ , 10(3):627–642, 2000.
  * Bianchi (2016) Pascal Bianchi.  Ergodic convergence of a stochastic proximal point algorithm.  _SIAM Journal on Optimization_ , 26(4):2235–2260, 2016.
  * Bolte et al. (2018) Jérôme Bolte, Shoham Sabach, Marc Teboulle, and Yakov Vaisbourd.  First order methods beyond convexity and Lipschitz gradient continuity with applications to quadratic inverse problems.  _SIAM Journal on Optimization_ , 28(3):2131–2151, 2018.
  * Bregman (1967) Lev M Bregman.  The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming.  _USSR Computational Mathematics and Mathematical Physics_ , 7(3):200–217, 1967.
  * 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.
  * Chen and Teboulle (1993) Gong Chen and Marc Teboulle.  Convergence analysis of a proximal-like minimization algorithm using bregman functions.  _SIAM Journal on Optimization_ , 3(3):538–543, 1993.
  * Chen et al. (2021) Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu.  Closing the generalization gap of adaptive gradient methods in training deep neural networks.  In _Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence_ , pages 3267–3275, 2021.
  * Defazio and Bottou (2019) Aaron Defazio and Léon Bottou.  On the ineffectiveness of variance reduced optimization for deep learning.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Dragomir et al. (2021a) Radu-Alexandru Dragomir, Alexandre d’Aspremont, and Jérôme Bolte.  Quartic first-order methods for low-rank minimization.  _Journal of Optimization Theory and Applications_ , 189:341–363, 2021a.
  * Dragomir et al. (2021b) Radu Alexandru Dragomir, Mathieu Even, and Hadrien Hendrikx.  Fast stochastic bregman gradient methods: Sharp analysis and variance reduction.  In _International Conference on Machine Learning_ , pages 2815–2825. PMLR, 2021b.
  * Duchi and Singer (2009) John Duchi and Yoram Singer.  Efficient online and batch learning using forward backward splitting.  _Journal of Machine Learning Research_ , 10:2899–2934, 2009\.
  * 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.
  * Friedman et al. (2010) J. Friedman, T. Hastie, and R. Tibshirani.  Regularization paths for generalized linear models via coordinate descent.  _Journal of Statistical Software_ , 33(1):1–22, 2010.
  * Fu et al. (2016) Rui Fu, Zuo Zhang, and Li Li.  Using LSTM and GRU neural network methods for traffic flow prediction.  In _2016 31st Youth Academic Annual Conference of Chinese Association of Automation (YAC)_ , pages 324–328. IEEE, 2016.
  * Ghadimi et al. (2016) Saeed Ghadimi, Guanghui Lan, and Hongchao Zhang.  Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization.  _Mathematical Programming_ , 155(1-2):267–305, 2016.
  * Gotmare et al. (2018) Akhilesh Gotmare, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher.  A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation.  _arXiv preprint arXiv:1810.13243_ , 2018.
  * Griewank and Walther (2008) Andreas Griewank and Andrea Walther.  _Evaluating derivatives: principles and techniques of algorithmic differentiation_.  SIAM, 2008.
  * Gu et al. (2020) Bin Gu, Wenhan Xian, Zhouyuan Huo, Cheng Deng, and Heng Huang.  A unified q-memorization framework for asynchronous stochastic optimization.  _The Journal of Machine Learning Research_ , 21(1):7761–7813, 2020.
  * Hanzely and Richtárik (2021) Filip Hanzely and Peter Richtárik.  Fastest rates for stochastic mirror descent methods.  _Computational Optimization and Applications_ , 79:717–766, 2021.
  * Hastie et al. (2009) Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman.  _The elements of statistical learning: data mining, inference, and prediction_ , volume 2.  Springer, 2009.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on Computer Vision and Pattern Recognition_ , pages 770–778, 2016.
  * Hendrycks and Gimpel (2016) Dan Hendrycks and Kevin Gimpel.  Gaussian error linear units (gelus).  _arXiv preprint arXiv:1606.08415_ , 2016.
  * Hiriart-Urruty and Lemaréchal (1993) Jean-Baptiste Hiriart-Urruty and Claude Lemaréchal.  _Convex analysis and minimization algorithms I: Fundamentals_ , volume 305.  Springer science & business media, 1993.
  * Hochreiter (1991) Sepp Hochreiter.  Untersuchungen zu dynamischen neuronalen netzen.  _Diploma, Technische Universität München_ , 91(1), 1991.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural Computation_ , 9(8):1735–1780, 1997.
  * Huang et al. (2017) Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 4700–4708, 2017.
  * Ioffe and Szegedy (2015) Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International Conference on Machine Learning_ , pages 448–456. pmlr, 2015.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Latafat et al. (2022) Puya Latafat, Andreas Themelis, Masoud Ahookhosh, and Panagiotis Patrinos.  Bregman finito/miso for nonconvex regularized finite sum minimization without lipschitz gradient continuity.  _SIAM Journal on Optimization_ , 32(3):2230–2262, 2022.
  * LeCun et al. (2015) Yann LeCun, Yoshua Bengio, and Geoffrey Hinton.  Deep learning.  _Nature_ , 521(7553):436–444, 2015.
  * Lin et al. (2007) Chih-Jen Lin, Ruby C Weng, and S Sathiya Keerthi.  Trust region newton methods for large-scale logistic regression.  In _Proceedings of the 24th international conference on Machine learning_ , pages 561–568, 2007.
  * Liu et al. (2019) Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han.  On the variance of the adaptive learning rate and beyond.  _arXiv preprint arXiv:1908.03265_ , 2019.
  * Loshchilov and Hutter (2016) Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  _arXiv preprint arXiv:1608.03983_ , 2016.
  * Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Lu (2019) Haihao Lu.  “Relative continuity” for non-Lipschitz nonsmooth convex optimization using stochastic (or deterministic) mirror descent.  _INFORMS Journal on Optimization_ , 1(4):288–303, 2019.
  * Lu et al. (2018) Haihao Lu, Robert M Freund, and Yurii Nesterov.  Relatively smooth convex optimization by first-order methods, and applications.  _SIAM Journal on Optimization_ , 28(1):333–354, 2018.
  * Luke (2017) D Russell Luke.  Phase retrieval, what’s new.  _SIAG/OPT Views and News_ , 25(1):1–5, 2017.
  * Luo et al. (2019) Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun.  Adaptive gradient methods with dynamic bound of learning rate.  _arXiv preprint arXiv:1902.09843_ , 2019.
  * Manchev and Spratling (2020) Nikolay Manchev and Michael Spratling.  Target propagation in recurrent neural networks.  _The Journal of Machine Learning Research_ , 21(1):250–282, 2020.
  * Marcinkiewicz (1994) Mary Ann Marcinkiewicz.  Building a large annotated corpus of english: The penn treebank.  _Using Large Corpora_ , 273, 1994.
  * Nemirovskij and Yudin (1983) Arkadij Semenovič Nemirovskij and David Borisovich Yudin.  Problem complexity and method efficiency in optimization.  _Wiley Interscience_ , 1983.
  * Nesterov (2005) Yu Nesterov.  Smooth minimization of non-smooth functions.  _Mathematical Programming_ , 103:127–152, 2005.
  * Nesterov (2003) Yurii Nesterov.  _Introductory lectures on convex optimization: A basic course_ , volume 87.  Springer Science & Business Media, 2003.
  * Ng (2004) Andrew Y Ng.  Feature selection, L1 vs. L2 regularization, and rotational invariance.  In _Proceedings of the twenty-first International Conference on Machine Learning_ , page 78, 2004.
  * Paren et al. (2022) Alasdair Paren, Leonard Berrada, Rudra PK Poudel, and M Pawan Kumar.  A stochastic bundle method for interpolating networks.  _Journal of Machine Learning Research_ , 23:1–57, 2022\.
  * Pascanu et al. (2013) Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio.  On the difficulty of training recurrent neural networks.  In _International Conference on Machine Learning_ , pages 1310–1318. Pmlr, 2013.
  * Patrascu and Necoara (2017) Andrei Patrascu and Ion Necoara.  Nonasymptotic convergence of stochastic proximal point methods for constrained convex optimization.  _The Journal of Machine Learning Research_ , 18(1):7204–7245, 2017.
  * Robbins and Monro (1951) Herbert Robbins and Sutton Monro.  A stochastic approximation method.  _The Annals of Mathematical Statistics_ , pages 400–407, 1951.
  * Rockafellar (1976) R Tyrrell Rockafellar.  Monotone operators and the proximal point algorithm.  _SIAM Journal on Control and Optimization_ , 14(5):877–898, 1976.
  * Rockafellar (1997) R Tyrrell Rockafellar.  _Convex analysis_ , volume 11.  Princeton university press, 1997.
  * Rockafellar and Wets (1998) R. Tyrrell Rockafellar and Roger J.-B. Wets.  _Variational Analysis_.  Springer Verlag, Heidelberg, Berlin, New York, 1998.
  * Shapiro et al. (2021) Alexander Shapiro, Darinka Dentcheva, and Andrzej Ruszczynski.  _Lectures on stochastic programming: modeling and theory_.  SIAM, 2021.
  * Siami-Namini et al. (2019) Sima Siami-Namini, Neda Tavakoli, and Akbar Siami Namin.  The performance of LSTM and BiLSTM in forecasting time series.  In _2019 IEEE International Conference on Big Data_ , pages 3285–3292. IEEE, 2019.
  * Simonyan and Zisserman (2014) Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Sutskever et al. (2013) Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton.  On the importance of initialization and momentum in deep learning.  In _International Conference on Machine Learning_ , pages 1139–1147. PMLR, 2013.
  * Teboulle (2018) Marc Teboulle.  A simplified view of first order methods for optimization.  _Mathematical Programming_ , 170(1):67–96, 2018\.
  * Wang et al. (2022) Bokun Wang, Shiqian Ma, and Lingzhou Xue.  Riemannian stochastic proximal gradient methods for nonsmooth optimization over the stiefel manifold.  _The Journal of Machine Learning Research_ , 23(1):4599–4631, 2022.
  * Wang et al. (2017) Mengdi Wang, Ethan X Fang, and Han Liu.  Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions.  _Mathematical Programming_ , 161:419–449, 2017.
  * You et al. (2019) Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: Training bert in 76 minutes.  _arXiv preprint arXiv:1904.00962_ , 2019.
  * Zhang (2004) Tong Zhang.  Solving large scale linear prediction problems using stochastic gradient descent algorithms.  In _Proceedings of the twenty-first International Conference on Machine Learning_ , page 116, 2004.
  * Zhou et al. (2020) Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al.  Towards theoretically understanding why sgd generalizes better than adam in deep learning.  _Advances in Neural Information Processing Systems_ , 33:21285–21296, 2020.
  * Zhuang et al. (2020) Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan.  Adabelief optimizer: Adapting stepsizes by the belief in observed gradients.  _Advances in Neural Information Processing Systems_ , 33:18795–18806, 2020.