Fast Stochastic Bregman Gradient Methods: Sharp Analysis and Variance Reduction

Introduction

We are interested in solving the minimization problem

Beyond simply adapting the step size, a powerful generalization of SGD consists in refining the geometry and performing instead Bregman gradient (a.k.a mirror) steps as

where the Euclidean distance has been replaced by the Bregman divergence with respect to a reference function $h$, which writes:

for all $x\in{\rm dom}\ h,y\in{\rm int}\ {\rm dom}\ h$. We make the following blanket assumptions on $h$ throughout the article, which guarantee well-posedness of the update (2).

has a unique solution, which lies in ${\rm int}\ C$.

The standard SGD algorithm corresponds to the case where $h=\frac{1}{2}\|\cdot\|^{2}$. However, a different choice of $h$ might better fit the geometry of the set $C$ and the curvature of the function, allowing the algorithm to take larger steps in directions where the objective gradient changes slowly. This choice is guided by the notion of relative smoothness and strong convexity, introduced in Bauschke et al. (2017); Lu et al. (2018). Instead of the squared Euclidean norm for standard smoothness, relative regularity is measured with respect to the reference function $h$.

The function $f$ is said to be $L$-relatively smooth and $\mu$-relatively strongly convex with respect to $h$ if it is differentiable and for all $x,y\in{\rm int}\ {\rm dom}\ h$,

where $D_{f}$ is defined similarly to (3). Note that if $\mu=0$, the left-hand side inequality reduces to assuming convexity of $f$. Similarly, if $h=\frac{1}{2}\|\cdot\|^{2}$, then $D_{h}(x,y)=\frac{1}{2}\|x-y\|^{2}$, and the usual notions of smoothness and strong convexity are recovered. If both functions are two times differentiable, Equation (4) can be turned into an equivalent condition on the Hessians: $\mu\nabla^{2}h(x,y)\preceq\nabla^{2}f(x)\preceq L\nabla^{2}h(x)$. Throughout the article, we will generally write $\mu_{f/h}$ and $L_{f/h}$ to insist on the relative aspect.

Writing the optimality conditions for the minimization problem of Equation (2), we obtain the following equivalent iteration, which is in the alternative Mirror Descent form (Nemirovsky and Yudin, 1983):

Although these updates have a closed-form solution for many choices of the reference function $h$, they may be harder to perform than standard gradient steps, since they require solving the subproblem defined in (2). Yet, this may be worth doing in some cases to reduce the overall iteration complexity, if the resulting majorization in (4) is much tighter than with the Euclidean distance. Let us list some applications of relative regularity:

Problems with unbounded curvature. Some problems have singularities at some boundary points in $C$ where the Hessian grows arbitrarily large. In this situation, smoothness with respect to the Euclidean norm does not hold globally, and standard gradient methods become inefficient as they necessit excessively small step sizes or costly line search procedures. A typical example arises in inverse problems with Poisson noise, which are used in particular for image deblurring (Bertero et al., 2009) or tomographic reconstruction (Kak and Slaney, 2001). In this case, the objective function involves the Kullback-Leibler divergence, which becomes singular as one of its arguments approaches 0. However, by choosing the reference function ${h(x)=-\sum_{i=1}^{d}\log(x^{(i)})}$, one can show that relative smoothness holds globally Bauschke et al. (2017). For more examples, see Lu et al. (2018); Bolte et al. (2018); Nesterov (2019); Mishchenko (2019).

Distributed optimization. When $h$ approximates $f$ in the sense of (4), Bregman methods can be used to speed up convergence by performing non-uniform preconditioning (Shamir et al., 2014; Reddi et al., 2016; Yuan and Li, 2020; Hendrikx et al., 2020b). Typically, $h$ is chosen as the objective function on a smaller portion of the dataset of size $n_{\rm prec}$ (e.g., the dataset of the server), which improves the conditioning by a factor of up to $n_{\rm prec}$ compared to Euclidean methods, while naturally taking advantage of an eventually small effective dimension of the dataset (Even and Massoulié, 2021). In this case, forming the gradient $g_{t}$ requires communication with the workers (where most of the data is held), and is thus expensive. Although the updates may not have a simple expression, the inner problem of Equation (2) can be solved locally at the server without additional communications. Therefore, Bregman methods allow to drastically reduce the communication cost by reducing the overall iteration complexity.

Despite these applications, there are still many gaps in our understanding of convergence guarantees of Bregman gradient methods. In particular, most existing results focus on the deterministic case $g_{t}=\nabla f(x_{t})$, or do not leverage the relative regularity assumptions.

In this work, we develop convergence theorems for Bregman SGD, for which the variance depends on the magnitude of the stochastic gradients at the optimum, and which can thus be much smaller than the one used in Hanzely et al. (2018), in particular for overparametrized models (which verify the interpolation condition that all stochastic gradients are equal to at the optimum). Our analysis relies on the Bregman generalization of a few technical lemmas such as the celebrated $\|a+b\|^{2}\leq 2(\|a\|^{2}+\|b\|^{2})$ inequality (Lemma 2) or the co-coercivity inequality (Lemma 3), which we believe to be of independent interest.

Then, we show that variance-reduction techniques, which are widely used to accelerate traditional Euclidean stochastic methods when the objective has a finite-sum structure (Schmidt et al., 2013; Johnson and Zhang, 2013; Defazio et al., 2014; Allen-Zhu, 2017), can be adapted to the Bregman setting. Although this generally requires stronger regularity assumptions (such as global smoothness of $h$ and Lipschitz continuity of $\nabla^{2}h^{*}$), we show that the asymptotical rate of convergence solely depends on relative regularity constants. The same type of results (asymptotic speedup under additional smoothness assumptions) is observed when applying Nesterov-type acceleration to Bregman gradient methods (Hanzely et al., 2018; Dragomir et al., 2019; Hendrikx et al., 2020b). We provide a summary of the rates proven in this paper in the appendix.

We start by discussing the related work in Section 2. Then, Section 3 presents the results for stochastic gradient descent, along with the main technical lemmas. Section 4 develops a Bregman version of the standard SAGA algorithm (Defazio et al., 2014). Finally, Section 5 illustrates the efficiency of the proposed methods on several applications, including Poisson inverse problems, tomographic reconstruction and distributed optimization.

Related work

The Bregman gradient method was first introduced as the Mirror Descent schemeNote that Mirror Descent and Bregman Gradient refer to the same algorithm, but that Mirror Descent is typically used when $f$ is non-smooth, or in the online optimization community, whereas Bregman Gradient is generally preferred when using the relative smoothness assumption. Yet, both names are valid and there are exceptions, for instance Hanzely and Richtárik (2018) use the Mirror Descent terminology although they assume relative smoothness. (Nemirovsky and Yudin, 1983; Beck and Teboulle, 2003) for minimizing convex nonsmooth functions, and enjoyed notable success in online learning Bubeck (2011). More recently, the introduction of relative smoothness (Bauschke et al., 2017; Lu et al., 2018; Bolte et al., 2018) has also brought interest in applying Bregman methods to differentiable objectives. This condition guides the choice of a well-suited reference function $h$ which can greatly improve efficiency over standard gradient descent. While the vanilla Bregman descent method yields the same convergence rate as the Euclidean counterpart, subsequent work has focused on obtaining better rates with acceleration schemes (Hanzely et al., 2018). However, lower bounds show that the rates for relatively smooth optimization cannot be accelerated in general (Dragomir et al., 2019), and that additional regularity assumptions are needed. Similar notions of relative regularity have also been investigated for non-differentiable functions, such as relative continuity (Lu, 2019; Antonakopoulos et al., 2019). Zhou et al. (2020) also study non-differentiable functions, but in the online setting and without relative continuity.

Stochastic optimization methods, and in particular SGD, are very efficient when the number of samples is high (Bottou, 2012) and are often referred to as “the workhorse of machine learning”. The problem with SGD is that, in general, it only converges to a neighbourhood of the optimum unless a diminishing step-size is used. Variance reduction can be used to counter this problem, and many variance-reduced methods have been developed, such as SAG (Schmidt et al., 2013), SDCA (Shalev-Shwartz and Zhang, 2013; Shalev-Shwartz, 2016), SVRG (Johnson and Zhang, 2013) or SAGA (Defazio et al., 2014).

Surprisingly, stochastic Bregman gradients algorithms have received less attention. Hanzely and Richtárik (2018); Gao et al. (2020); Hendrikx et al. (2020a) study Bregman coordinate descent methods, and Zhang and He (2018) study the non-convex non-smooth setting. Antonakopoulos et al. (2020) study stochastic algorithms for online optimization, under Riemann-Lipschitz continuity. In contrast, our work focuses on Bregman SGD for relatively-smooth objectives. Hanzely and Richtárik (2018) study the same setting and obtain comparable convergence rates, but with a much looser notion of variance, which we discuss more in details in the next section. This is problematic since their bound on the variance is thus proportional to the magnitude of the gradients along the trajectory, and may thus be very large when far from the optimum if $f$ is strongly convex. In contrast, our definition of variance leverages the stochastic gradients at the optimum, which allows us to obtain significant results without bounded gradients and in the interpolation regime (zero gradients at the optimum). In particular, our analysis can be seen as a Bregman generalization of the analysis from Gower et al. (2019). Davis et al. (2018) also analyze a similar setting, but again with more restrictive assumptions on the noise and boundedness of the gradients. Besides, to the best of our knowledge, variance reduction for Bregman stochastic methods was only studied in Shi et al. (2017) in the context of stochastic saddle-point optimization, but without leveraging relative regularity assumptions like we do in this work.

Bregman Stochastic Gradient Descent

We start by introducing a few technical lemmas, which are Bregman analogs to well-known Euclidean results, and which are at the heart of our analysis. All missing proofs can be found in Appendix A.

For $x,y\in{\rm int}\ {\rm dom}\ h$, we have $D_{h}(x,y)=D_{h^{*}}(\nabla h(y),\nabla h(x))$.

See, e.g., Bauschke and Borwein (1997, Thm 3.7.) for the proof. Using duality, we prove the following key lemma:

Let $x^{+}$ be such that $\nabla h(x^{+})=\nabla h(x)-g$, and similarly define $x^{+}_{1}$ and $x^{+}_{2}$ from $g_{1}$ and $g_{2}$. Then, if $g=\frac{g_{1}+g_{2}}{2}$, we obtain:

Lemma 2 can be adapted for any $g=(1-\alpha)g_{1}+\alpha g_{2}$ with $\alpha\in$. In the Euclidean case $h=\|\cdot\|^{2}$, we recover $\|\frac{g_{1}+g_{2}}{2}\|^{2}\leq\frac{1}{2}\left(\|g_{1}\|^{2}+\|g_{2}\|^{2}\right)$. We now generalize the cocoercivity of the gradients (Nesterov, 2003, Eq. 2.1.7) to the relatively smooth case:

If a convex function $f$ is relatively $L$-smooth w.r.t to $h$, then for any $\eta\leq\frac{1}{L}$,

2 Variance definition

for some $z_{t}\in[\nabla h(x_{t})-2\eta\nabla f_{\xi_{t}}(x^{\star}),\nabla h(x_{t})]$.

The assumption that the stochastic gradients are actual gradients of stochastic functions which are themselves smooth with respect to $h$ is rather natural, as already discussed in the introduction. It is at the heart of variance reduction in the finite sum setting (though the sum does not need to be finite in the case of Assumption 2), and is in particular verified when solving (Empirical) Risk minimization problems.

Yet, it prevents the analysis from applying to coordinate descent methods for instance, in which $g_{t}=\nabla_{i}f(x_{t})$, with $i\in\{1,\cdots,d\}$. However, in this case, the extra structure can also be leveraged to obtain similar results (Hanzely and Richtárik, 2018; Hendrikx et al., 2020a; Gao et al., 2020).

We now compare our noise assumption with (Hanzely and Richtárik, 2018, Assumption 5.1.), which writes:

for $t\geq 0$, where $g_{t}$ is the stochastic gradient estimate and $\bar{x}_{t+1}$ is the output of the (theoretical) Bregman gradient step taken with the true gradient, that is, $\nabla h(\bar{x}_{t+1})=\nabla h(x_{t})-\eta_{t}\nabla f(x_{t})$. Thus, their condition can be written:

so that $\sigma^{2}$ bounds at each step the distance (in the Bregman sense) between $x_{t+1}$ and $\bar{x}_{t+1}$, the point that would be obtained by the expected (deterministic) gradient update. To illustrate why our assumption is weaker, let us consider the case where $h$ is $\mu_{h}$-strongly convex. In this setting, a sufficient condition for (6) to hold is that

while a sufficient condition for our variance definition to hold is (using that $\nabla f(x^{\star})=0$):

3 Convergence results

We now prove the actual convergence theorems for Bregman SGD. To avoid notation clutter, we generally omit with respect to which variable expectations are taken when clear from the context.

If $f$ is $L_{f/h}$-smooth and $\mu_{f/h}$-strongly convex relative to $h$ with $\mu_{f/h}>0$, and Assumptions 1 and 2 hold, then for ${\eta\leq 1/(2L_{f/h})}$, the iterates produced by Bregman stochastic gradient (2) satisfy

By using Lemma 4 from Appendix A, we obtain:

Using Lemma 2, the last term can be bounded as $D_{h}(x_{t},x_{t+1})\leq\frac{1}{2}\left[D_{1}+D_{2}\right]$. We use Lemma 3 (Bregman co-coercivity) to write:

In the interpolation setting (when $\nabla f_{\xi_{t}}(x^{\star})=0$ for all $\xi_{t}$), we have that $\sigma^{2}=0$. Theorem 1 thus proves linear convergence in this case. For instance, when solving objectives of the form $D_{\rm KL}(Ax,b)$ (which has applications in optimal transport (Mishchenko, 2019)) or $D_{\rm KL}(b,Ax)$ (which has application in deblurring or tomographic reconstruction), then the variance as defined in Hanzely and Richtárik (2018) may be unbounded, whereas the variance as we define it is equal to if there exists $z$ such that $Az=b$.

When $f$ is convex ($\mu_{f/h}=0$), Theorem 1 can be adapted to obtain a $1/T$ decrease of the error up to a noise region.

Under the same assumptions as Theorem 1, if $\mu=0$, then

Contrary to the Euclidean case, we do not obtain a guarantee on the average iterate in general. This is because the bound is on the average of $D_{f}(x^{\star},x_{t})$ instead of $D_{f}(x_{t},x^{\star})$, and Bregman divergences are not necessarily convex in their second argument (except for the Euclidean distance and Kullback-Leibler divergence). Therefore, the final bound is obtained on $\min_{t}D_{f}(x^{\star},x_{t})$, meaning that there is at least one $x_{t}$ such that this is true. Note that the nice properties regarding interpolation still hold in this setting.

We start from Lemma 4 and bound the $D_{h}(x_{t},x_{t+1})$ in the same way as when $\mu>0$, which yields:

Averaging over $t$ and dividing by $\eta$ leads to (13). ∎

The simplicity of the proof and the generality of our technical lemmas also allow us to provide convergence results when $f$ is not convex:

If $f$ is $L_{f/h}$-smooth relatively to $h$ and Assumptions 1 and 2 hold, then for ${\eta\leq 1/(2L_{f/h})}$, the iterates produced by Bregman stochastic gradient (2) satisfy

Variance reduction

The difference with Section 3 is that we now assume that $f$ is a finite sum, which is required for variance reduction. We also assume that the minimizer $x^{\star}$ belongs to ${\rm int}\ C$, so that $\nabla f(x^{\star})=0$. The case where $x^{\star}$ lies on the border of $C$ is more delicate, as $h$ might not be differentiable there (e.g., the log-barrier); this would require an involved technical analysis which we leave for future work.

For analyzing the Bregman-SAGA scheme, we first need to introduce, in addition to relative smoothness, an assumption on the regularity of $D_{h}$.

Such structural assumptions appear to be essential for analyzing Bregman-type methods that use information provided by gradients of past iterates. The function $G$ models the fact that the Bregman divergence $D_{h^{*}}(x+v,x)$ is not homogeneous nor invariant to translation in $x$ in general (except for the Euclidean case where it is equal to $\|v\|^{2}/2$). Note that such difficulties are also encountered for obtaining accelerated rates with inertial variants of Bregman descent, where similar assumptions are needed Hanzely et al. (2018). This seems unavoidable, as suggested by the lower bound in Dragomir et al. (2019).

Although the gain function $G$ is relatively abstract at this point, it plays a key role in defining the step-size, and convergence guarantees similar those of Euclidean SAGA can be obtained provided $G$ can be chosen small enough. We first state the general Theorem 4 (convergence proof for Algorithm 1), and then detail how $G$ can be bounded in several interesting cases.

For $t\geq 0$ and step-sizes $\eta_{t}>0$, define $H_{t}=\frac{1}{n}\sum_{i=1}^{n}D_{f_{i}}(\phi_{i}^{t},x^{\star})$, and the potential $\psi_{t}$ as follows:

First note that by convexity of $h$ and of the $f_{i}$, $\psi_{t}\geq 0$ for all $t$. Our goal in this section is to show that $\{\psi_{t}\}_{t\geq 0}$ converges to at a given speed. Indeed, since $D_{h}(x^{\star},x_{t})\leq\psi_{t}$, this implies (as in Section 3) that $x_{t}$ converges to $x^{\star}$ at the same rate. To ease notations, we define

Assume that Algorithm 1 is run with a step size sequence $\{\eta_{t}\}_{t\geq 0}$ satisfying $\eta_{t}=1/(8L_{f/h}G_{t})$ for every $t\geq 0$, with $G_{t}$ decreasing in $t$ and such that for all $j\in\{1,\cdots,n\}$:

Then, under Assumptions 1 and 3, the potential $\psi_{t}$ satisfies

In the convex case ($\mu_{f/h}=0$), we obtain that

Similarly to BSGD, we apply Lemma 4 (Appendix A), which yields

Lemmas 1 and 2 yield $D_{h}(x_{t},x_{t+1})\leq(D_{1}+D_{2})/2$, with

Using Assumption 3 together with Lemma 3, we obtain:

where we used the gain function for translation and rescaling the step size. Following Hofmann et al. (2015), we write:

Therefore, we can use the $-H_{t}/n$ term to control the excess term from bounding $D_{h}(x_{t},x_{t+1})$. In the end, we obtain:

If we choose $\eta_{t}\leq 1/(8L_{f/h}G_{t})$ then the last term is positive and $1-4\eta_{t}L_{f/h}G_{t}\geq 1/2$. If $\mu_{f/h}>0$ then we use the relative strong convexity of $f$ to obtain that the right hand side is proportional to $\psi_{t}$, thus leading to a linear convergence rate. Otherwise, we obtain a telescopic sum, leading to the $1/T$ rate of Equation (19). ∎

Note that the monotonicity of $\eta_{t}$ (through $G_{t}$) is a technical condition to ensure that the Lyapunov is non-increasing. Otherwise, $\psi_{t}$ could blow up even though $x_{t+1}$ is very close to $x_{t}$, simply because $\eta_{t}$ shrinks. It could be replaced by the condition that $\eta_{t}$ does not vary too much (not more than a factor $1-O(1/n)$), which achieves the same goal. The rest of this section is devoted to shong that non-trivial $G_{t}$ can be chosen in many cases, thus leading to strong convergence guarantees. In particular, the rate recovers that of Euclidean SAGA in case $h$ is a quadratic form.

If $\nabla^{2}h$ is constant ($h$ is quadratic), then Assumption 3 is satisfied with $G=1$, so that

where $\kappa_{f/h}=L_{f/h}/\mu_{f/h}$ is the relative condition number.

If $h$ is not quadratic, but $f^{*}$ and $h^{*}$ are regular with respect to a norm, then strong guarantees can also be obtained:

If $h^{*}$ is $\mu_{h}^{-1}$-smooth and $f^{*}$ is $L_{f}^{-1}$-strongly convex with respect to a norm $\|\cdot\|^{2}$, then the stepsize can be chosen constant as $\eta_{t}=\frac{\mu_{h}}{8L_{f}}$, and

Note that following Kakade et al. (2009), having $h^{*}$ be $\mu_{h}^{-1}$-smooth is equivalent to having $h$ be $\mu_{h}$ strongly-convex.

The proof follows the same step as the proof of Theorem 4, but the translation invariance and homogeneity are obtained by comparison with the norm, instead of using Assumption 3. Thus, we pay a factor $\mu_{h}^{-1}$ when bounding $D_{h^{*}}$ by the norm, and a factor $L_{f}$ when bounding the norm by $D_{f^{*}}$. It is also possible to directly use Assumption 3, but in this case the $L_{f}$ factor is replaced by $L_{f/h}L_{h}$, which is an upper bound on $L_{f}$, and may thus be slightly looser. ∎

Note that Corollary 1 is actually a consequence of Corollary 2, since $\mu_{h}=1$ and $L_{f}=L_{f/h}$ if $D_{h}$ is a norm itself. Otherwise, the constant $G_{t}$ is chosen in a rather pessimistic way, and depends on the difference between directly bounding $D_{f}$ by $D_{h}$ (in which case we pay a factor $L_{f/h}$), or going through a norm $\|\cdot\|$ in the middle (in which we case we pay $L_{f}/\mu_{h}\geq L_{f/h}$).

As stated at the beginning of this section, one of the problems is that Bregman divergences lack translation invariance and homogeneity. However, as the algorithm converges, one can expect these conditions to hold locally, as $D_{h^{*}}(x+v,x)$ is approximated by $\frac{1}{2}\|v\|^{2}_{\nabla^{2}h^{*}(x^{*})}$ for small enough $v$, and $x$ close enough to $x^{*}$. This is indeed what happens under enough regularity assumptions on $h$.

If $h$ is $L_{h}$-smooth and the Hessian $\nabla^{2}h^{*}$ is $M$-smooth, then the gain function can be chosen as:

Note that, even if the regularity conditions of Proposition 1 do not hold globally (such as for problems with unbounded curvature), they are at least valid on every bounded subset of ${\rm int}\ C$, as soon as $h$ is $C^{3}$ on ${\rm int}\ C$. We now explicit a possible explicit choice for $G_{t}$ in this setting.

Assume that $h$ is $L_{h}$-smooth, $\mu_{h}$-strongly convex and that the Hessian $\nabla^{2}h^{*}$ is $M$-smooth. Then, there exists an explicit constant $C$ such that if Algorithm 1 is run with a step size ${\eta_{t}=1/(8L_{f/h}G_{t})}$ with $G_{t}$ decreasing and satisfying

where $\lim_{t\rightarrow\infty}G_{t}=1$, or, more precisely,

The explicit expression for the constant $C$ is provided in Appendix B along with the proof. Although the result involves smoothness constants of $h$ which can be large in the relatively-smooth setting, this dependence disappears asymptotically. Hence, after some time $t$, which we can roughly estimate using Equation (26), we obtain that $G_{t}=O(1)$. Thus, we reach the same kind of convergence rate as in the ideal quadratic case, which depends only on the relative condition number $\kappa_{f/h}$, but with more general functions $h$, and thus possibly much better conditioning. Besides, the order of magnitude required for $G_{t}$ can be estimated during the optimization process using Equation (24).

2 Remarks on adaptivity

Assumption 3 highlights the fact that the key difficulty is purely geometric, and that in general we need to make up for the lack of translation invariance and homogeneity of Bregman divergences. Although Corollary 3 gives a criterion for $G_{t}$ that can be evaluated throughout training (since the constant $C$ is explicit), several approximations are required to obtain it, and it may be loose overall. Yet, for the theory to hold, it suffices to have $\eta_{t}$ small enough such that:

Experiments

In order to show the effectiveness of our method, we consider the two key settings mentioned in the introduction: problems with unbounded curvature (inverse problems with Poisson noise) and preconditioned distributed optimization. The first setting corresponds to the convex case ($\mu_{f/h}=0$), whereas the second one corresponds to the relatively strongly convex case ($\mu_{f/h}>0$). We observe that leveraging stochasticity (and, when needed, variance reduction) drastically improves the performance of Bregman methods in both cases. Additional details on the setting (such as the precise formulation of the objective or the relative smoothness constants) are given in Appendix D.

Figure 1(b) considers experiments on the tomographic reconstruction problem on the standard Shepp-Logan phantom (Kak and Slaney, 2001). Due to space limitations, the main text mainly describes the results, but the setting details can be found in Appendix D. The step-size given by theory was rather conservative in this case, so we increased it by a factor of $5$ for all Bregman algorithms (and even 10 for BGD). Figure 1(b) shows again that stochastic algorithms drastically outperform BGD. Yet, BSGD quickly reaches a plateau because of the noise. On the other hand, BSAGA enjoys variance reduction and fast convergence to the optimum. In this case, BSAGA is on par with MU, the state-of-the-art algorithm for this problem. This is because of the log barrier that allows relative smoothness to hold, but heavily slows down Bregman algorithms when coordinates are close to . Yet, these results are encouraging and one may hope for even faster convergence of BSAGA for tomographic reconstruction with a tighter reference function.

2 Statistically Preconditioned Distributed Optimization

In this section we consider the problem of solving a distributed optimization problem in which data is distributed among many workers. We closely follow the setting of Hendrikx et al. (2020a), and solve a logistic regression problem for the RCV1 dataset (Lewis et al., 2004). Function $h$ is taken as the same logistic regression objective as for the global objective $f$, but on a much smaller dataset of size $n_{\rm prec}=1000$ and with an added regularization $c_{\rm prec}=10^{-5}$. In this case, BGD corresponds to a widely used variant of DANE (Shamir et al., 2014), in which only the server performs the update. The stochastic updates in BSGD are obtained by subsampling a set of workers at each iteration, so that all the nodes do not have to participate in every iteration. Regularization is taken as $\lambda=10^{-5}$, and there are $n=100$ nodes with $N=1000$ samples each. A fixed learning rate is used, and the best one is selected selected among $[0.025,0.05,0.1,0.25,0.5,1.]$. BGD uses $\eta=0.5$ while SAGA and BSGD use $\eta=0.05$. The x-axis represents the total number of communications (or number of passes over the dataset). Note that at each epoch, BGD communicates once with all workers (one round trip for each worker) whereas BSGD and BSAGA communicate $n$ times with one worker sampled uniformly at random each time. Therefore, BSAGA requires much less gradients from the workers to reach a given precision level, yet, it is at the cost of having to solve more local iterations.

Figure 1(c) first shows that BSAGA clearly outperforms BGD. BSGD on the other hand is as fast as BSAGA at the beginning of training, until it hits a variance region at which it saturates. This is consistent with the theory, and is similar to what can be observed in the Euclidean case. An interesting feature is that although the step-size has to be selected smaller than that of gradient descent (which is also the case in the Euclidean setting since $f$ is smoother than the least smooth $f_{i}$), choosing a constant step-size is enough to ensure convergence in this case, thus hinting at the fact that the analysis is rather conservative and that $G_{t}$ does not slow down the algorithm as much as we could have feared when far from the optimum. This is consistent with the results obtained by Hendrikx et al. (2020b) on acceleration.

Conclusion

Throughout the paper, we have (i) given tight convergence guarantees for Bregman SGD that allow to accurately describe its behaviour in the interpolation setting, and (ii) introduced and analyzed Bregman analogs to the standard variance-reduced algorithm SAGA. These convergence results require stronger assumptions on the objective than relative smoothness and strong convexity, but we show that fast rates can be obtained nonetheless when $h$ is nicely behaved (quadratic or Lipschitz Hessian). We also prove that these fast rates can be obtained for more general functions $h$ after a transient regime. Besides, we show experimentally that variance reduction greatly accelerates Bregman first-order methods for several key applications, including distributed optimization and tomographic reconstruction. In particular, there does not seem to be a slow transient regime in the applications considered, despite the lack of regularity of the objectives. This need for higher order regularity assumptions but great practical performance is consistent with the results obtained for acceleration in the Bregman setting. Better understanding the transient regime (in which $G_{t}$ can be high) and finding better reference functions $h$ for the tomographic reconstruction problem are two promising extensions of our work.

Acknowledgements

RD was supported by an AMX fellowship. RD would like to acknowledge support from the Air Force Office of Scientific Research, Air Force Material Command, USAF, under grant number FA9550-19-1-7026/19IOE033 and FA9550-18-1-0226. HH was funded in part by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001(PRAIRIE 3IA Institute). HH also acknowledges support from the European Research Council (grant SEQUOIA 724063) and from the MSR-INRIA joint centre.

References

Appendix A Missing proofs for Bregman SGD (Section 3)

Let $x^{+}$ be such that $\nabla h(x^{+})=\nabla h(x)-g$, and similarly define $x^{+}_{1}$ and $x^{+}_{2}$ from $g_{1}$ and $g_{2}$. Then, if $g=\frac{g_{1}+g_{2}}{2}$, we obtain:

where the inequality step is obtained by the convexity of the Bregman divergence in its first argument. The final result is obtained by using duality back. ∎

Note that this descent lemma is an equality, and we can then use standard assumptions to bound the different terms.

We start by writing $V_{t}(x)=\eta_{t}g_{t}^{\top}x+D_{h}(x,x_{t})$. Since $x_{t+1}$ is defined as $\arg\min_{x}V_{t}(x)$ and by Assumption 1, we have $x_{t+1}\in{\rm int}\ C$ then $\nabla V_{t}(x_{t+1})=0$ and so:

since $\nabla^{2}V_{t}=\nabla^{2}h$. This writes:

Combining Equations (29), (30) and (31), we obtain:

If a convex function $f$ is relatively $L$-smooth w.r.t to $h$, then for any $\eta\leq\frac{1}{L}$,

Let $y\in{\rm int}\ {\rm dom}\ h$ and consider the function $g_{y}$ defined by

for $x\in C$. $g_{y}$ is nonnegative, convex and relatively $L$-smooth with respect to $h$, since it has the same Hessian than $f$. Therefore, for $\eta\in(0,\frac{1}{L}]$ the relative smoothness inequality (4) implies that for every $u\in{\rm int}\ {\rm dom}\ h$ we have $D_{g_{y}}(u,x)\leq\frac{1}{\eta}D_{h}(u,x)$, that is

The right-hand side $Q_{y}(u,x)$ is a convex function of $u$ and is minimized for a point $u^{+}$ such that

and the result follows from the fact that $\nabla g_{y}(x)=\nabla f(x)-\nabla f(y)$. ∎

Appendix B Missing proofs for Variance Reduced methods (Section 4)

First, we use the following Bregman counterpart of a standard variance identity (Pfau, 2013), which we prove for completeness.

B.2 Proof of Theorem 4: generic Bregman-SAGA convergence bound

In this subsection, we give a more detailed proof of Theorem 4, and include derivations that had to be skipped in the main text because of space limitations.

Similarly to BSGD, we start by applying Lemma 4 (Appendix A), which yields

Lemmas 1 and 2 yield $D_{h}(x_{t},x_{t+1})\leq(D_{1}+D_{2})/2$, with

Using the gain function with the fact that $\eta_{t}\leq 1/L_{f/h}$ and Lemma 3, we have

Note that we can pull the $G_{t}$ term out of the expectation over the choice of $i$ since $G_{t}$ holds for all $i$. For bounding $D_{2}$, Lemma 5 with $V=-2\eta_{t}(\nabla f_{i}(x^{\star})-\nabla f_{i}(\phi_{i}^{t}))$ leads to

Recall that $H_{t}=\frac{1}{n}\sum_{j=1}^{n}D_{f_{j}}(\phi_{j}^{t},x^{\star})$. Plugging the expressions for $D_{1}$ and $D_{2}$ into Equation (37), we obtain:

Following Hofmann et al. (2015), we write:

Indeed, $\phi_{j}^{t+1}=\phi_{j}^{t}$ with probability $1-1/n$, and $\phi_{i}^{t+1}=x_{t}$ with probability $1/n$. Therefore, we can use the $-H_{t}/n$ term to control the excess term from bounding $D_{h}(x_{t},x_{t+1})$. In the end, using that $G_{t}$ is decreasing and so $\eta_{t}$ is increasing, we obtain the following recursion:

If we choose $\eta_{t}\leq 1/(8L_{f/h}G_{t})$ then the last term is positive and $1-4\eta_{t}L_{f/h}G_{t}\geq 1/2$, so that using the relative strong convexity of $f$ leads to:

The result can then be obtained by chaining this inequality. If $\mu_{f/h}=0$ then we start back from Equation (B.2), use that $D_{f}(x^{\star},x_{t})\geq 0$ and the same fact that $1-4\eta_{t}L_{f/h}G_{t}\geq 1/2$ to obtain:

The result is obtained by averaging over $T$, since the right hand side yields a telescopic sum, leading to the $1/T$ rate of Equation (19). ∎

B.3 Lipschitz-Hessian setting

In this section, we add the additional assumption that $h$ is $L_{h}$-smooth, and that the Hessian $\nabla^{2}h^{*}$ is $M$-smooth in the operator norm, that is

If $h$ is $L_{h}$-smooth and the Hessian $\nabla^{2}h^{*}$ is $M$-smooth, then the gain function can be chosen as:

Using the fact that is $h$ is $L_{h}$-smooth, $h^{*}$ is $1/L_{h}$-strongly convex and hence $\|v\|^{2}\leq 2L_{h}D_{h^{*}}(y+v,y)$, leading to

Assume that $h$ is $L_{h}$-smooth and the Hessian $\nabla^{2}h^{*}$ is $M$-smooth. Then, there exists an explicit constant $C$ such that if Algorithm 1 is run with a step size $\eta_{t}=1/(8L_{f/h}G_{t})$ with $G_{t}$ decreasing in $t$ and satisfying

where $\lim_{t\rightarrow\infty}G_{t}=1$, or, more precisely,

Using the gain function from Proposition 1, to satisfy the assumptions of Theorem 4 it is sufficient to choose $G_{t}$ such that

As the quantities involving $\nabla f_{i_{t}}(x^{\star})$ are unknown, we provide an uper estimate. We can proceed in the following way, using the fact that, due to relative regularity, $f_{i}$ is also smooth with constant $L_{h}L_{f/h}$, and $f$ is strongly convex with constant $\mu_{h}\mu_{f/h}$:

And similarly, we can estimate the second term from

which leads to the following upper estimate of the RHS of Condition (46):

Now, with such choice of $G_{t}$, Theorem 4 applies and the convergence rate (44) holds. It remains to prove the estimate for the convergence rate of $G_{t}$ towards 1. To this end, we show that it is upper bounded by $\mathcal{O}(1+\psi_{t}^{1/2})$ since

Since we imposed a safeguard such that $G_{t}\geq\frac{L_{f/h}L_{h}}{\mu_{h}}$, the convergence rate of $\psi_{t}$ is bounded by

as stated by Corollary 2. Indeed, the assumptions are verified as $h^{*}$ is $1/\mu_{h}$-smooth and $f^{*}$ is $1/L_{f}$-strongly convex with $L_{f}=L_{h}L_{f/h}$. This worst-case estimate for $\psi_{t}$, along with the majorization (47), gives the resulting rate for $G_{t}$. ∎

Appendix C Bregman SVRG

We consider in this section the convergence guarantees of Bregman SVRG (BSVRG), which is presented in Algorithm 2. We consider the same variant as Hofmann et al. (2015), in which the full gradient used for variance reduction is recomputed at each step with a small probability $p$, instead of after a fixed number of steps. We study this variant of BSVRG since it is very closely related to BSAGA. The main difference is that instead of updating $\phi_{i_{t}}$ when $i_{t}$ is picked, the algorithm chooses only one common $\phi_{t}$ to perform variance reduction, and this common $\phi_{t}$ is updated with probability $p$ at the end of each iteration. Thus, the convergence Theorem for Algorithm 2 closely follows Theorem 4.

Assume that Algorithm 2 is run with a step size sequence $\{\eta_{t}\}_{t\geq 0}$ satisfying $\eta_{t}=1/(8L_{f/h}G_{t})$ for every $t\geq 0$, with $G_{t}$ decreasing in $t$ and such that for all $j\in\{1,\cdots,n\}$:

Then, under Assumptions 1 and 3, the potential $\psi_{t}=D_{h}(x^{\star},x_{t})+\frac{\eta_{t}}{2p}D_{f}(\phi_{t},x^{\star})$ satisfies

In the convex case ($\mu_{f/h}=0$), we obtain that

As explained before Theorem 5, the only thing that changes between BSAGA and BSVRG is that a global $\phi_{t}$ is used instead of separate $\phi_{i}^{t}$, and that it is update with probability $p$ at the end of each iteration (instead of updating $\phi_{i_{t}}^{t}$ at time $t$ for SAGA). Thus, all the derivations performed for BSAGA hold for BSVRG if we replace $\phi_{i}^{t}$ with $\phi_{t}$ for all $i$. The only equation that needs to be adapted is Equation (41), since it relies on the way the $\phi_{i}^{t}$ are updated. Yet, in the case of BSVRG, it writes:

which is the same as for BSAGA but with $p$ instead of $1/n$. Therefore, the conclusions are unchanged if we replace $n$ by $1/p$ whenever it appears in the bounds. Similar convergence guarantees hold when $\phi_{t}$ is updated every fixed number of steps $T$, but the proof is substantially more involved since Equation (50) does not hold in such a simple form. ∎

Appendix D Additional details for the experiments

Due to space limitations, some details of the experimental setting are missing from the main text, and we thus present them in this section. Note that all the experiments presented in this paper run in less than an hour on a standard laptop (and usually much less). Our code is also available in supplementary material.

where $x^{*}$ is the true unknown signal. Inverse problems with Poisson noise arise in various signal processing applications such as astronomy or computerized tomography, see Bertero et al. (2009) and references therein.

As a motivating application of relative smoothness, Bauschke et al. (2017) prove that the Poisson objective $f$ is relatively smooth with respect to the log-barrier reference function

with constant $\sum_{j=1}^{n}b_{j}/n$. This constant can be quite conservative when $A$ is a sparse matrix, and so we prove a better estimate by leveraging this structure. For $j\in\{1\dots n\}$, we denote $S_{j}$ the support of the $j$-th column of $A$, that is

The Poisson objective function defined in (51) is relatively $L$-smooth w.r.t the log-barrier for

Applying the Jensen inequality to the function $t\mapsto t^{2}$ and weights $w_{ij}=A_{ij}x_{j}/(A_{i}^{\top}x)$ yields

where we used the fact that $w_{ij}\in$ if $i\in S_{j}$, and $w_{ij}=0$ otherwise. ∎

The relative Lipschitz constant provided by Proposition 2 can be considerably smaller than $\sum_{j=1}^{n}b_{j}/n$ when $A$ is sparse, which is the case in practical applications.

For our numerical experiments, we compare full-batch Bregman gradient descent (BGD), Bregman stochastic gradient descent (BSGD), and the Bregman SAGA scheme described in Algorithm 1. We also implement the Multiplicative Update (MU), also known as Lucy-Richardson or Expectation-Maximization (Shepp and Vardi, 1982), which is the standard baseline for Poisson inverse problems.

Tomographic reconstruction problem.

Computerized tomography (Kak and Slaney, 2001) is the task of reconstructing an object from cross-sectional projections, with fundamental applications to medical imaging. We study a classical synthetic toy problem for this task: the Shepp-Logan phantom (Figure 3(a)). In this setting, the observation matrix $A$ corresponds to the discrete Radon transform, which is the cross-sectional projection of the original image $x$ along different projection angles $\theta_{1},\dots,\theta_{n}$ (Figure 3(b)). That is, the objective writes

where $b_{\theta_{i}},A_{\theta_{i}}$ correspond to the observation and projection matrix along the angle $\theta_{i}$. For stochastic algorithms, the formulation (53) naturally yields a finite-sum structure: we thus take $f_{i}(x)=D_{\rm KL}(b_{\theta_{i}},A_{\theta_{i}}x)$ for $i=1\dots n$.

We corrupt the sinogram with Poisson inverse noise, and apply our algorithms. We use $n=360$ projection angles, and the image dimension is $d=100^{2}$. As the matrix $A$ has a sparse structure, we use the relative smoothness constant provided by Proposition 2 for a better estimate. The step-size given by theory was rather conservative in this case, so we increased it by a factor of $5$ for all Bregman algorithms (and even 10 for BGD).

D.2 Statistically Preconditioned Distributed Optimization

We detail in this section the setting that was used to obtain Figure 1(c). In particular, we use the following logistic regression objective with quadratic regularization, meaning that the function at node $i$ is:

where $y_{i,j}\in\{-1,1\}$ is the label associated with $a_{j}^{(i)}$, the $j$-th sample of node $i$. We use a regularization parameter of $\lambda=10^{-5}$, and the size of the local datasets is equal to $N=1000$. The local dataset is constructed by shuffling the RCV1 dataset, downloaded from LibSVM, and then assigning a fixed portion to each worker. Then, one node (without loss of generality, node 0) uses its local dataset to construct the preconditioning dataset, so that:

where $c_{\rm prec}=10^{-5}$. Tuning $c_{\rm prec}$ in order to obtain the fastest algorithms is hard in general, as detailed in Hendrikx et al. (2020b) (in which it is denoted as $\mu$). One strategy is to choose $c_{\rm prec}$ of order $1/n_{\rm prec}$ (in our case $n_{\rm prec}=N=1000$), and then decrease it as long as BGD is stable. Our chosen value ($10^{-5}$) is smaller than that of Hendrikx et al. (2020b) for this problem ($10^{-4}$), in which they used a rougher $c_{\rm prec}=c/n_{\rm prec}$ criterion with varying $n_{\rm prec}$, and a larger step-size $\eta=1$ for BGD (which is the same as DANE). Besides, we see that SPAG is slightly unstable in our example, and increasing $c_{\rm prec}$ would help with that. In this case, theory gives that $L_{f/h}\approx 1$. Yet, when $c_{\rm prec}\approx\lambda$, this step-size usually has to be chosen a bit smaller. Therefore, we choose in our case $\eta=0.5$ for BGD and SPAG, and $\eta=0.05$ for BSGD and BGD. Note that there is always a constant factor between the maximum step-size for SAGA and that of BGD, and the difference could further be explained by the difference between the batch condition number (relative smoothness of $f$) versus the stochastic one (max relative smoothness of the $f_{i}$).

We compute the minimum error as the smallest error over all iterations for all algorithms. Then, we subtract it to the running error of an algorithm to get the suboptimality at each step. Following Hendrikx et al. (2020b), local problems are solved using a sparse implementation of SDCA (Shalev-Shwartz, 2016). We warm-start the local problems (initializing on the solution of the previous one), and perform 10 passes over the preconditioning dataset at each step, or until the norm of the gradient of the inner problem is small enough ($10^{-6}$). The number of inner passes could be reduced further, but then the algorithms started to converge slightly more slowly. This results in an overall computational overhead for the server, since BSAGA and BSGD require to solve many more inner problems, which are not so cheap to compute. Yet, this overhead only affects the server, and the iteration complexity is much lower, meaning that BSAGA is indeed very efficient to reduce the communication complexity of solving distributed empirical risk minimization problems.