Gradient flows and proximal splitting methods: A unified view on accelerated and stochastic optimization

Optimization is at the heart of machine learning, statistics, and many applied scientific disciplines. It also has a long history in physics, ranging from the minimal action principle to finding ground states of disordered systems such as spin glasses. Proximal algorithms form a class of methods that are broadly applicable and are particularly well-suited to nonsmooth, constrained, large-scale, and distributed optimization problems. There are essentially five proximal algorithms currently known, each proposed in seminal work: Forward-backward splitting, Tseng splitting, Douglas-Rachford, alternating direction method of multipliers, and the more recent Davis-Yin. These methods sit on a higher level of abstraction compared to gradient-based ones, with deep roots in nonlinear functional analysis. In this paper we show that all of these methods are actually different discretizations of a single differential equation, namely, the simple gradient flow which dates back to Cauchy (1847). An important aspect behind many of the success stories in machine learning relies on “accelerating” the convergence of first-order methods. However, accelerated methods are notoriously difficult to analyze, counterintuitive, and without an underlying guiding principle. We show that similar discretization schemes applied to Newton's equation with an additional dissipative force, which we refer to as accelerated gradient flow, allow us to obtain accelerated variants of all these proximal algorithms—the majority of which are new although some recover known cases in the literature. Furthermore, we extend these methods to stochastic settings, allowing us to make connections with Langevin and Fokker-Planck equations. Similar ideas apply to gradient descent, heavy ball, and Nesterov's method which are simpler. Our results therefore provide a unified framework from which several important optimization methods are nothing but simulations of classical dissipative systems.

Figure 1

Solutions of Eqs. (2) and (3) with $\phi =(1/2){\omega }^2{x}^2$. For constant damping (24) we choose $\eta$ slightly below the critical value, and for decaying damping (25) we choose $r=3$. (a) $\omega =1.5$. (b) $\omega =0.05$.

Figure 2

(a) Splitting approach in the flow composition (38). In one time step the error between the continuous and discrete trajectory is $\parallel x\left(t_{k+1}\right)-x_{k+1}\parallel =\mathcal{O}\left(h^2\right)$. (b) The (re)balanced splitting introduces $c_k$ which forces the discrete-time evolution to converge to a critical point of the ODE: ${\widehat{\mathrm{\Phi }}}_h^k\left(x_0\right)={\mathrm{\Phi }}_t\left(x_0\right)\to x_{\infty }$ as $k,t\to \infty$.

Figure 3

Several well-known optimization methods arise as dicretizations of the gradient flow. Moreover, “accelerated extensions” of these methods arise as discretizations of Newton's equation with a dissipative term; the gradient flow corresponds to a large friction limit of the latter. Therefore, these optimization methods consist of actual simulations of a classical dissipative physical system.

Figure 4

Different optimization algorithms arise from different discretizations (dashed lines) of the same physical system. The left column of the diagram represents stochastic processes described by a Langevin equation, while the right column represents deterministic processes from (dissipative) classical mechanics—the transition between these phases is controlled by the “temperature” $T\sim 1/S$ where $S$ is the batch size. The upper row of the diagram corresponds to an underdamped or accelerated regime, while the lower row corresponds to an overdamped regime where acceleration is negligible—the transition is controlled by the damping coefficient $\eta$. Discretizations from each colored quadrant yield (variants of) optimization algorithms in these “different phases.”

Figure 5

Comparison of the discretization provided by Algorithms 1 and 2 versus the exact solution of the second-order gradient flow (3) with Eq. (108). (a) Constant damping with $\eta =0.2$. (b) Decaying damping with $r=3$. In both cases we chose step size $h=0.01$.

Figure 6

Maximum error (109) over a range of step sizes. We consider constant damping under the setting of Fig. 5; similar results hold for decaying damping.

Figure 7

Histogram of stochastic variants of Algorithms 1 and 2 (with ${\gamma }_k=0$) for a quadratic problem which is modelled by an Ornstein-Uhlenbeck process; we considered Eq. (110) with $N=1000$, minibatch of size $S=1$, step size $h=0.1$, initial position $x_0=10,$ and 2000 Monte Carlo runs. We show the simulation for $t=2$. The Gaussian fit has mean $\mu =2.52$ and variance $\sigma =0.35$. The theoretical prediction from Eq. (111) is $\mu =2.51$, in close agreement.

Figure 8

Same simulation as in Fig. 7 for $t\in [0,10]$. The dashed line is the mean of the Ornstein-Uhlenbeck process Eq. (111), and the shaded area indicates the standard deviation with $D=4$. The markers are the means of ADMM and Davis-Yin, and the vertical red lines are standard deviations for both methods, which are very close. The inset is exactly the same plot but with markers omitted and the $y$ axis on log scale.

Figure 9

Performance of several methods on problem (112). We perform 10 Monte Carlo runs and show the mean and standard deviation (error bars) of the error $|\phi \left(x_k\right)-{\phi }^{★}|/{\phi }^{★}$ versus the iteration $k$. The accelerated variants have faster convergence.

Figure 10

Convergence of different algorithms on problem (113). We perform 10 Monte Carlo runs and indicate the mean and standard deviation (error bars) for the relative error (117) versus iteration $k$. Note the improvement of the accelerated methods.

Figure 11

Number of iterations needed to reach the termination tolerance for the problem in Fig. 10.

Figure 12

Performance of different algorithms on problem (113) with annealing on $\alpha$.

Figure 13

Number of iterations needed to reach the termination tolerance for the problem in Fig. 12.

Figure 14

(a) Image of $974\times 1194$ pixels with rank $r=70$. (b) Observed data with sampling ratio $s=0.3$. (c) Reconstructed image by running the previous algorithms a single time. The relative error is $8\times {10}^{-2}$— see Fig. 15 for convergence rates. (d) Reconstructed image with annealing. The relative error is $1.6\times {10}^{-4}$—see Fig. 16 for convergence rates. The correct rank is also recovered by all methods.

Figure 15

Convergence rates of different methods when reconstructing the observed image in Fig. 14 by solving the optimization problem (113). An example of the recovered image is in Fig. 14.

Figure 16

Convergence rates for the same problem as in Fig. 15 but with annealing on $\alpha$. An example of the recovered image with the accelerated methods is shown in Fig. 14—note that the nonaccelerated methods were unable to achieve such a small error.