Mathematics
Computer Science-2021.gz:143451
Computer Science
bbed0e1d9df24aa1b059e606e68de51ec33f566d
Accelerated Forward-Backward Optimization using Deep Learning

We propose several deep-learning accelerated optimization solvers with convergence guarantees. We use ideas from the analysis of accelerated forward-backward schemes like FISTA, but instead of the classical approach of proving convergence for a choice of parameters, such as a step-size, we show convergence whenever the update is chosen in a specific set. Rather than picking a point in this set using some predefined method, we train a deep neural network to pick the best update. Finally, we show that the method is applicable to several cases of smooth and non-smooth optimization and show superior results to established accelerated solvers.


Introduction
The tremendous growth in computing capacity of hardware and advances in numerical methods has allowed for very large-scale and detailed simulations of complex phenomena. Much focus is currently on embedding such simulation techniques into procedures for design, control, data assimilation, and inverse modeling (solving inverse problems/parameter estimation). These tasks are often formalized as continuous variable optimization problems (i.e. nonlinear programs).
An essential characteristic of the resulting optimization problems is their very large problem size, both in the number of state variables and decision variables. This is dictated by the approximation and discretization of the underlying continuous problem and the number of state variables and state equations, which in many scientific and industrial problems easily reaches the millions. Another further complicating factor is that the objective involved in the optimization may not necessarily be differentiable, like when sparsity is used in inverse modeling.
For the above stated reasons, solving large-scale non-smooth optimization problems is one of the key challenges in scientific computing. This has catalyzed a surge in research on methods that utilize specific structure in the problem and take advantage of modern multicore or distributed computing architectures. An example is decomposition algorithms that allow considerable computational speedup in parallel computing environments while addressing memory requirements. Another is operator splitting techniques, which are designed to solve optimization problems where the objective is a sum of two convex functions, one differentiable with a Lipschitz continuous gradient and the other possibly non-smooth but its proximal points are accessible. These are now very popular for addressing large-scale non-smooth optimization problems that arise in processing massive datasets. The latest line of development seeks to further improve upon computational feasibility by combining the aforementioned techniques from convex analysis with contemporary techniques from deep learning.

Mathematical setting
Let X denote a Hilbert space with inner product ·, · and corresponding norm · . Next, consider minimizing a differentiable convex function f : X → R whose gradient ∇f : X → X is Lipschitz continuous with constant β −1 > 0. This can be solved by gradient descent [16] or accelerated methods based on an overshooting idea formulated by Nesterov [33].
The objective in the aforementioned minimization is smooth, but as already mentioned, many applications involve minimizing a non-smooth objective function. To handle the latter, there has been significant effort to develop theory and algorithms suitable for non-smooth optimization. These often utilize further structure in the objective. In particular, splitting techniques have been developed to solve optimization problems where the (non-smooth) objective function is a sum of two convex functions, one smooth and the other non-smooth. Stated mathematically, this refers to optimization problems of the form min x∈X f (x) + g (x) (1) where f : X → R is a differentiable convex function and g : X → R is a proper, convex and lower semi-continuous, but possibly non-smooth function. Algorithms to solve eq. (1) include the forward-backward splitting method [22] and its accelerated variants like fast ISTA (FISTA) [10]. Much attention has also been directed towards approaches that use learning techniques to improve the speed, like in learned ISTA (LISTA) [24], see the survey [17] and references therein. However, the latter methods lack convergence guarantees. Another approach is to learn the parameters of a general optimization solver [8], but in this case there is a relatively small number of free parameters, which limits the potential of the scheme to be adapted to a particular problem class. To address this and increase the number of degrees of freedom in the learning, we will in this paper consider a deviation-based approach. Instead of learning parameters of a solver, we learn an entire updating function and then use that in an optimization solver. To ensure convergence, we constrain the updating function to sufficiently small deviations from known standard algorithms.

Related work
The overall principle that underlies application of machine learning for solving optimization problems is to use principles from statistical decision theory to select the "best" solver by training against data. Solvers are parametrized and "best" means selecting the solver parameters that lead to the smallest average error when applied to a collection of similar optimization problems. Several authors have investigated the idea of deriving optimization schemes from statistical learning. One of the most notable schemes is LISTA introduced in [24]. This scheme was inspired by the Iterative Shrinkage-Thresholding Algorithm (ISTA), which is a popular algorithm for solving the optimization in sparse coding. Each iteration of ISTA is a combination of matrix operations followed by a thresholding function. Thus, when stopped after fixed number of iterations, the algorithm resembles a forward pass through a neural network. The idea in LISTA is to replace the handcrafted matrices in ISTA with learned ones that are trained in a supervised manner, provided that the ground truth solution to the optimization problem is available. In contrast, we follow [8] and train the parameters of our optimization scheme using an unsupervised loss that is equal to the value of the objective function attained after a fixed number of iterations.
The above idea of learning a truncated optimization scheme in a unsupervised manner is also studied in machine learning [6,11,27,29,41]. The main use case here is to solve the non-convex optimization problem that arises when training neural networks. Here the objective is to replace conventional optimizers with learned optimizers that are faster/better. This is achieved by parametrizing updates at each iteration with recurrent neural networks and training the associated hyper-parameters either by variants of stochastic gradient descent [6,41], reinforcement learning [11,27] or by using evolutionary strategies [32]. The challenge is to ensure the learned optimizer generalizes from one optimization task to another. A difficulty here is to identify a class of optimization problems that share similar structure. We instead consider (large) classes of optimization problems that naturally share similarities, like those in image reconstruction for computed tomography (CT) and regularization of other inverse problems. This makes it easier to specify the generalization requirement.
Convergence of the above mentioned learned optimization schemes is yet to be mathematically proven. Instead, their usage relies on (very strong) heuristics, so these data driven approaches can not be seen as optimization solvers in the classical sense. As we consider convex problems, our approach is not applicable for training neural networks, but we can on the other hand provide convergence guarantees.
Some authors have investigated optimal parameter choices in FISTA when iterations are truncated, see e.g. [28] and notably [18] that proves linear convergence of a related scheme under some further assumptions (sparsity, finite dimension, optimal training). Another approach is [8], which proves convergence of a learned primal-dual scheme under weaker assumptions. The analysis is however only performed in the parametric case, e.g. with an explicit form of the updating operation.
A key element in the above approaches to optimization lies in using neural networks obtained by unrolling an existing optimization scheme. This provides the opportunity for domain adaptation and such network architectures have also been successfully used to solve challenging inverse problems, but this typically involves training them in a supervised manner [2,3,25,42]. Note that even though these networks have an architecture obtained by unrolling an optimization solver, they are not themselves optimization solvers, see [7, Section 5] for further details. Another example from scientific computing is [26], which deals with solving partial differential equations. Hence, unrolling should primarily be seen as a way to assemble an appropriately domain adapted deep neural network architecture [7,Section 4.9.1].
An attempt to analyze and interpret learned optimizers by considering linear approximations of optimizers close to the stationary point is presented in [29].

Overview of the paper
Section 2 recalls the necessary notions from convex optimization. This is followed by subsection 3.1, which establishes a worst-case convergence result for a non-parametric gradient descent scheme for smooth optimization problems. Next, subsection 3.2 establishes a slightly weaker convergence theorem for a non-parametric forward-backward scheme applicable for non-smooth problems. Section 4 discusses the details of learning a non-parametric optimization scheme. Section 5 concludes with numerical experiments from an imaging application.

Background from optimization theory
The aim here is to introduce the notation and basic notions from optimization theory that are used throughout the paper. The main reference is [9], but the reader may also consult [34,36] for related results.
A function f : It is convex if f ((1 − λ)x + λy) ≤ (1 − λ)f (x) + λf (y) holds for all x, y ∈ X and 0 ≤ λ ≤ 1 and strictly convex if this inequality is strict for x = y and 0 < λ < 1. A strictly convex function has at most one minimizer [9,Corollary 11.9]. Next, the sub-differential of a convex function f : X → R is a (set-valued) mapping ∂f : X → 2 X defined as If f is Fréchet differentiable at x, then ∂f (x) consists of a single element, namely the gradient of f at x [9, Proposition 17.31], i.e. ∂f (x) = {∇f (x)}. Finally, the proximal point of a proper, convex, and lower semi-continuous function f : X → R at x ∈ X with step-size γ > 0 is defined as The proximal point can be used to characterize points in the sub-differential [9, Proposition 16.44]: The operator Prox γ f is then a firmly non-expansive single-valued mapping from X to X [9, Proposition 12.28], Furthermore, p = Prox γ f (x) if and only if the following variational inequality holds [9,Proposition 12.26]: If ∇f is Lipschitz continuous with constant β −1 , then the following inequality, which is sometimes called descent lemma (see [9,Theorem 18.15]), holds: Likewise, if f is in addition assumed to be convex [34, Theorem 2.1.5]:

Deviation-based optimization schemes
We call an optimization scheme parameter-based if convergence criteria are expressed as constraints on its parameters. An example of a parameter-based optimization scheme is the classical gradient descent method for minimizing a convex differentiable function f : X → R where ∇f : X → X is β −1 -Lipschitz. To see this, note that such a method relies on an updating scheme of the following form: For convergence to a minimizer for f , t n is a parameter that needs to satisfy −1 < t n < 1 and ∞ n=0 (1 + t n )(1 − t n ) = +∞ [9, Proposition 4.39, Corollary 5.16, and Corollary 18.17]. Thus, a convergence criterion for the above scheme is expressible in terms of a specific choice of parameters, so we are dealing with a parameter-based optimization method.
In fact, almost all methods in classical optimization theory are written as schemes of this form. Furthermore, it has been considered a virtue to have as few parameters as possible (preferably none). On the other hand, recent advancements in machine learning allow us to deal with extremely large numbers of parameters while avoiding over-fitting effects and achieve significant improvements over established methods. While optimization schemes can be formulated using neural networks, and their convergence can be studied in terms of their parameters, this is typically very cumbersome and restricts neural network architectures to, e.g., one layer networks [8].
To allow for a wider range of neural network architectures, we need to move away from convergence guarantees that involve the parameters of the scheme. The idea is to give the network as much freedom as possible while still retaining the stability properties of the underlying optimization method. Such schemes that do not impose an a-priori restriction on the number of parameters will be referred to as deviation-based.
To introduce a deviation-based scheme for smooth optimization as in the setting above, we modify the steepest descent method eq. (5) by introducing a sequence (∆x n ) n≥0 of deviations: x n+1 := x n − β(∇f (x n ) + ∆x n ) for some deviation/offset ∆x n ∈ X . (6) The classical convergence proof for the gradient descent implies convergence if ∆x n = t∇f (x n ) for −1 < t < 1. This is, however, a one-dimensional set and hardly interesting to do learning on, and in particular not deep learning. In what follows, we consider schemes of the above form and prove convergence for sets of updates ∆x n ∈ X that have full dimension, thus allowing for a much greater degree of freedom in selecting the updates. The learning itself will be discussed in section 4.

Deviation-based optimization schemes for smooth optimization
In previous literature, the sequence (∆x n ) n≥0 of deviations in eq. (6) was used to model errors in the evaluation of the gradient, see for example [13,19,40]. In our setting, however, we are interested in consciously choosing ∆x n ∈ X of eq. (5) which will preserve or, by a clever choice, might improve its convergence properties.
As we show below the condition ∆x n ≤ ∇f (x n ) guarantees that the function value f (x n ) is non-increasing, and it is the weakest condition which guarantees this property. In particular, we will show that the function values f (x n ) converge to the minimal value of the function f . Special cases of the following theorem are discussed in remark 3.1 below. Theorem 3.1. Consider minimizing f : X → R that is convex and differentiable with a gradient that is β −1 -Lipschitz continuous for some β > 0. Next, consider the iterative scheme where x 0 ∈ X is some fixed initial point and (∆x n ) n≥0 is any sequence in X such that for some sequence (α n ) n≥0 of positive numbers with α n ≤ 1 for all n ≥ 0.
1. Then (f (x n )) n≥0 is a non-increasing sequence of objective function values.
2. The following inequality holds: wherex denotes a global minimizer of f .

If we further assume that
Then f (x n ) → f (x) as n → +∞ withx denoting a global minimizer of f . Remark 3.1. Choosing α n = α for all n in theorem 3.1 simplifies eq. (9) to Hence, using Stirling's formula, Hence, we have (asymptotically) recovered the optimal rate from [20, Theorem 1] up to a constant of 2.
Note also that all the calculations work for a smaller 0 < β ′ < β since 1/β ′ will also be a Lipschitz constant of ∇f . In addition to the convergence of the function values, we are also able to show the convergence of the iteration sequence with stronger assumptions on the objective function f : Proof. The sequence (x n ) n≥0 is bounded: by item 1 of theorem 3.1, the sequence (f (x n )) n≥0 is bounded from above by f (x 0 ) ∈ R. The asserted boundedness then follows from the coercivity of f . If x ∈ X is a weak sequential cluster point, i.e., there exists a subsequence (x n k ) k≥0 of (x n ) n≥0 with x n k ⇀ x, then [9, Theorem 9.1] and item 3 of theo- The proof of item 1 now follows from the observation that We will next prove the formula in eq. (9) by proving the following stronger relation: wherex ∈ X denotes a global minimizer of f , i.e. ∇f (x) = 0. To prove the above, we proceed by induction and begin by considering the case n = 0. By eq. (4), we have This proves the claim for n = 0, i.e. the induction basis holds. Next, assume eq. (13) holds for some n ≥ 0 and we seek to prove that it holds for n + 1. Then Here, we estimate the first two terms with eq. (4) and the last one with the induction hypothesis eq. (13). This results in Using eq. (7) and reordering the terms gives On the other hand, using again eq. (7), Adding this to eq. (14), followed by a simple calculation, yields By eq. (8) and taking into account the non-negativity of norms, we get which finishes the induction. Let us now prove the convergence f (x n ) → f (x) as n → +∞. Taking the logarithm in eq. (9) gives Since log is a concave function, log Because of the assumption in eq. (10), the right-hand side converges to −∞, so

Deviation-based schemes for non-smooth optimization
The theory in subsection 3.1 only applies to optimization problems with a smooth objective. Theory for optimization with a non-smooth objective needs to be based on sub-differential calculus. Our focus is on minimizing a non-smooth objective with an additional structure as in eq. (1). The proposed deviation-based iterative scheme is inspired by the classical forward-backward (or proximal-gradient/ISTA) scheme [22], but with deviations in the spirit of the gradient method outlined in subsection 3.1. The next theorem states the scheme and analyses its convergence properties for solving the nonsmooth minimization problem in eq. (1) where g is not necessarily differentiable.
Theorem 3.2. Consider minimizing f + g where f : X → R is convex and differentiable with a gradient ∇f : X → X that is Lipschitz continuous with constant β −1 for some β > 0, and g : X → R is a proper, convex and lower semi-continuous function that is not necessarily differentiable. Next, consider the iterative scheme where x 0 ∈ X is some fixed initial point, 0 < γ n < 2β for all n ≥ 0, and holds for all n ≥ 1. Then, the following holds: Then V n ≥ f (x n+1 ) + g(x n+1 ) for all n ≥ 0 and the sequence is monotonically non-increasing if α (1) n ≤ 1 and α (2) n ≤ 1 for all n ≥ 0; 2. If f + g is bounded from below and Under the conditions of item 2, any weak sequential cluster point of (x n ) n≥0 is a minimizer of f +g. If there are no such minimizers, then x n → +∞ as n → ∞. Conversely, a weak sequential cluster point exists whenever the sequence (x n ) n≥0 is bounded. Corollary 3.2. If, in addition to the assumptions initem 3 of theorem 3.2, the objective function f + g is coercive, i.e., f (x) + g(x) → +∞ as x → +∞, and strictly convex, then the sequence (x n ) n≥0 converges weakly to the unique minimizer of f + g.
Proof. The sequence (x n ) n≥0 is bounded: by item 1 of theorem 3.2, the nonincreasing real sequence in (24) guarantees that (V n ) n≥0 is bounded from above and greater or equal to (f (x n+1 ) + g(x n+1 )) n≥0 . The asserted boundedness then follows from the coercivity of f + g.
The statement of the corollary then follows from [9, Lemma 2.46], taking into account item 3 of theorem 3.2 and that the strict convexity of f + g ensures the uniqueness of the minimizer of f + g.
Proof of theorem 3.2. We start by proving the claims in item 1. The first inequality follows from eq. (3), so we focus on proving the monotonic decrease of the sequence in eq. (17). Note first that, for n ≥ 0, eq. (15b) is equivalent to This yields By eq. (4) we get the following inequality for all n ≥ 1: From eqs. (21) and (22) and the definition of V n , we therefore obtain Using eq. (16), we then obtain the following inequality for all n ≥ 1: The claim that the sequence in eq. (17) is monotonically non-increasing now follows from ignoring the two last terms in eq. (24), which are non-negative. This concludes the proof of the claims in item 1. We next turn our attention to proving the claim in item 2. First, let M ∈ R such that f (x) + g(x) ≥ M > −∞ holds for all x ∈ X as f + g is (by assumption) bounded from below. For fixed N ≥ 1, consider the sum over n = 1, . . . , N for the terms in eq. (24): By the claim in item 1, we know Next, by eqs. (18) and (19), there exist some n 0 ≥ 0, 0 < γ − ≤ γ + < 2β, 0 ≤ α + for all n ≥ n 0 .
We therefore get Since the right-hand side is finite and independent of N in both cases, we get Using eq. (16) and the fact that the convergence of the series does not depend on the (finite) number of elements with indices below n 0 , we can immediately conclude that The last assertion of item 2 follows from the observation below, which also concludes the proof of the claims in item 2.
∞ n=0 We finally prove the claim in item 3. Let x be a weak sequential cluster point of the sequence (x n ) n≥0 with a sub-sequence (n k ) k≥0 such that x n k ⇀ x as k → ∞. By eq. (20), we have for all k ≥ 0. The assertion in item 2 combined with the Lipschitz continuity of ∇f yields the following as k → ∞: Hence, the left-hand side of eq. (25) converges to 0 as k → ∞. Thus, by [9, Proposition 20.38] we get x is a minimizer of f + g. The last statement follows from the fact that if x n → +∞, then there exists a bounded sub-sequence that has a weakly convergent sub-sequence whose weak limit is a minimizer of f +g. This concludes the proof of the claim in item 3.

Remark 3.2.
if g = 0 on X , then eq. (15) reads as This is equivalent with the algorithm in subsection 3.1, provided that γ n = β for all n ≥ 0 and On the other hand, when g = 0 the condition in eq. (16) becomes In the case that γ n = β for all n ≥ 0, this simplifies to the following inequality that is different from the one in eq. (26):

Deep-learning accelerated optimization
Subsections 3.1 and 3.2 introduced deviation-based iterative schemes for solving smooth and non-smooth minimization problems that contain many traditional iterative schemes as special cases. These schemes rely on a rule for selecting deviations as iterates progress. Selection rules that give rise to convergent sequences will be referred to as convergent and theorems 3.1 and 3.2 provide criteria that imply convergence.
A natural next step is to look for a selection rule that has "optimal" convergence. This requires one to formalize the notion of "optimal". Such a selection rule will most likely be problem dependent, i.e. we cannot expect to have a single selection rule that is optimal for all minimization problems. One approach is to search for the optimization method that is optimal on average when applied to a specific class of minimization problems.
Based on the above, we adopt principles from statistical decision theory to learn an optimal convergent selection rule from suitable example data. It is clearly unfeasible to search over all possible convergent selection rules, so attention is restricted to those that are encoded by a (deep) neural network.
A key issue is to select a neural network architecture that yields a convergent selection rule. The following sections provide details of this approach.

Optimization parametrized by an additional input
The first task is to specify the class of minimization problems. Here we simply consider minimizing an objective function on X that is parametrized by an additional input y ∈ Y, i.e. min x∈X F y (x) for y ∈ Y.
F y : X → R above is a proper, convex and lower semi-continuous function and Y is the set containing the additional input for the objective, i.e., the input data. For simplicity, we will assume that X is finite-dimensional and that each F y is strictly convex and coercive so that corollaries 3.1 and 3.2 are applicable and provide conditions for strong convergence of the respective iteration sequences. Such problems naturally arise in many applications, e.g., when using variational regularization to solve ill-posed inverse problems (see section 5). Next is to formalize the notion of optimal convergence for a optimization scheme that aims to solve problems of the form in eq. (27), which in turn dictates how to train the neural network. Computational feasibility was earlier stated as our main concern, so we consider a setting where the computational budget is limited in terms of a fixed number of iterations N , other possible choices are discussed in remark 4.1. Consider an iterative scheme with parameters θ ∈ Θ applied for N ∈ N iterations when the initial point x 0 ∈ X is fixed, this defines a function of the input y ∈ Y for the objective: Next, the aforementioned optimality means we look for the optimization method, which on average (as y ∈ Y varies) renders an iterate after N steps that is as close as possible to the true minimal value of the objective in eq. (27). Stated mathematically, let Y denote a Y-valued random variable generating y ∈ Y. For optimality we look for R N,θ ( · ) : Y → X wherê In practice, the distribution of Y is not available, instead one has samples We therefore consider its empirical counterpart: The unsupervised learning problem 4 eq. (29) results in a scheme that provides an output which is on average closest to minimizing the objective in eq. (27) after N iterations. Combining this with convergence means ensuring the following holds for all y ∈ Y: lim n→∞ R n,θ (y) = x * (y) (30) where x * (y) ∈ X solves eq. (27). Hence, a solver given byθ that is trained to be optimal for N iterates will converge to a solution of eq. (27) if we run more iterations, even when it is not trained for this scenario. The next two sections provide further details on how to define a deviationbased iterative solver R N,θ : Y → X in terms of a parametrized selection rule Ψ θ for θ ∈ Θ. The optimal scheme is obtained by training as in eq. (29) and convergence is ensured by an appropriate choice of neural network architecture.
Remark 4.1. The topic of how to define "optimal convergence" is of high importance to optimization, but the choice becomes especially apparent in learningbased schemes. In the classical optimization literature [4,5,23,35] the most widely used definition is upper bounds on the worst-case performance as the number of iterations go to infinity. For learning this is not an attractive definition since it is hard to evaluate computationally. In this work we chose to optimize the average function value after a fixed number of steps, but it would also be possible to optimise e.g. the average number of steps needed for convergence.

Smooth convex objective
Assume that F y : X → R in eq. (27) is differentiable with a β −1 -Lipschitz continuous gradient. Next, consider iterates (x n (y)) n≥0 ⊂ X generated by an iterative scheme for solving eq. (27). The mapping R N,θ : Y → X is defined as the N -th term in that sequence, i.e. R N,θ (y) = x N (y).
A natural approach is to consider the deviation-based scheme in eq. (7), which, when applied to the Y-parametrized minimization in eq. (27), reads as x n+1 (y) := x n (y) − β ∇F y x n (y) + ∆x n (y) for all n ≥ 0.
The selection rule for setting the deviations ∆x n (y) above is defined as ∆x n (y) := Ψ θ x n (y), ∇F y x n (y) , ∆x n−1 (y) .
where Ψ θ : X × X × X → X is an updating function that ensures that the resulting deviations satisfy the conditions in eq. (8). The above completes the definition of R N,x0 : Y × Θ → X and the optimal solver can be obtained from example data by training as in eq. (29). The updating function Ψ θ will be given by a (deep) neural network that is parametrized by θ ∈ Θ. We construct the neural network Ψ θ by composing a conventional neural networkΨ θ : X × X × X → X with a normalizing function N : X × X → X , which ensures that the outcome is a feasible deviation. More precisely, we define h n :=Ψ θ x n (y), ∇F y x n (y) , ∆x n−1 (y) .
Here α ∈ [0, 1) controls the size of the learned deviations within the feasible region. Subsection 5.6 presents performance for different values of α in ablation experiments. Regarding the choice of N , the essential property is that its range is contained in the ball that guarantees that the iterative scheme is convergent as formalized in eq. (8) with α n = α. Thus, one can essentially replace N with any transformation with this property.

Remark 4.2.
In eq. (31) we use a current iterate, gradient at that iterate, and previous deviation as input for the network Ψ θ that defines the selection rule for the deviations. This choice is largely heuristic and often (as in the problem discussed in subsection 5.2), the objective function naturally consists of two parts F y = f y + g y . Here one has the option to evaluate gradients of both parts separately instead of taking the gradient of the objective. This corresponds to a selection rule of the form ∆x n (y) := Ψ θ x n (y), ∇f y x n (y) , ∇g y x n (y) , ∆x n−1 (y) (34) for some neural network Ψ θ : X × X × X × X → X . Moreover, it is possible to incorporate information from previous iterations even though we did not pursue this approach in our experiments.

Non-smooth convex objective
Here F y : X → R in eq. (27) is convex, but not necessarily differentiable. We specialize to the case F y = f y + g y where f y : X → R and g y : X → R are convex but only f y is differentiable and ∇f y : X → X is β −1 -Lipschitz continuous, we seek to solve min A natural approach is to consider the deviation-based scheme in eqs. (15a) and (15b) for solving the above problem. Similarly as in subsection 4.2, the selection rules for the updates ∆x 1 n (y) and ∆x 2 n (y) are given by neural networks Ψ 1 and Ψ 2 , i.e. ∆x 1 n (y) := Ψ 1 θ 1 x n (y), ∇f y w n−1 (y) , ∆x 1 n−1 (y) (36a) ∆x 2 n (y) := Ψ 2 θ 2 x n (y), ∇f y w n−1 (y) , ∆x 2 n−1 (y), ∆x 1 n (y) .
The optimal solver is now obtained by training R N,x0 : Y ×Θ → X as in eq. (29).
To ensure convergence as in eq. (30), the resulting learned updates need to satisfy eq. (16) in theorem 3.2 according to corollary 3.2. Similar to the smooth case in subsection 4.2, we achieve this by normalizing the output of the last hidden layers h 1 n and h 2 n in the two neural networks: As before, α (1) , α (2) ∈ [0, 1) control the sizes of the learned deviations in the feasible region.

Numerical experiments
This section evaluates performance of the proposed learned optimizers for solving minimization problems arising in tomographic image reconstruction.

Variational regularization in tomographic imaging
Tomography is a collection of techniques that seek to visualize the interior structure of an object by probing it with penetrating particles/waves from different directions. The most well-known example is computed tomography (CT) in medical imaging. A patient is here scanned by X-rays from different directions and the aim is to recover a 2D/3D image of the interior anatomy. Tomographic imaging inevitably leads to an inverse problem since the interior structure x ∈ X (2D/3D image) is only indirectly observed through data y ∈ Y. This can be formalized as solving an equation where A : X → Y models how data is generated in absence of noise. A can be taken as a linear operator (ray transform) if one adopts a simplified model for how X-rays interact with tissue and properly pre-processes raw sensor data. Tomographic imaging problems are unfortunately often ill-posed. This means that a solution procedure that seeks to maximize data consistency, e.g. by minimizing x → Ax − y 2 , will be unstable. Variational regularization addresses these issues by adjusting the need for data consistency against the need to suppress unwanted features (over-fitting). Stated formally, instead of trying to solve eq. (38), one solvesx Here, S λ : X → R (regularization functional) stabilizes the recovery procedure, typically by enforcing certain regularity.
Remark 5.1. Much effort over the last three decades has been devoted to determining mathematical properties of solutions to eq. (39) for various choices of S λ . Such analysis is typically performed in a non-discretized setting where both x (2D/3D image) and y (data) resides in infinite-dimensional function spaces, see [12,14,38] for extensive surveys.

Characteristics of eq. (39)
The minimization in eq. (39) is of the form in eq. (27) with tomographic data y ∈ Y as the additional variable. In our case, the objective in eq. (39) has a natural decomposition F y = f y + g where f y (x) := Ax − y 2 and g(x) := S λ (x).
Note here that f y : X → R is differentiable whenever A is differentiable, which in particular is the case when A is linear. Properties of g : X → R will depend on the choice of regularization functional, but it will be proper, convex and lower semi-continuous. Next, in tomographic imaging, both x and y are high-dimensional arrays after discretization. In the 2D setting that is considered in subsection 5.6, x represents 512 × 512 pixel 2D coronal image slices. These are obtained from normal-dose CT scans of the human abdomen provided by Mayo Clinic for the AAPM Low Dose CT Grand Challenge [31] and used as input for simulating X-ray parallel-beam tomographic data. The latter is from 1000 equidistributed source positions and 1000 detector elements with 5% additive Gaussian noise, i.e. data y is a 10 6 -dimensional array. In summary, even though we consider a 2D setting, eq. (39) is a large-scale optimization problem.
An additional aspect to consider is that most clinical imaging studies are performed in a time-critical setting. An informal rule of thumb for CT imaging is that reconstructed images need to be available within 5 minutes after scanning 5 . Hence, applicability of variational regularization in clinical CT rests to large extent on the ability to (approximately) solve eq. (39) within this time-frame, so one can typically only afford about 10 evaluations of A and A ⊤ .

Smoothed total-variation regularization
Total-variation regularization corresponds to choosing S λ (x) := λ ∇x 1 in eq. (39) for some λ > 0. Using such a regularization functional is popular since it preserves edges [15,37]. When x is discretized, the gradient ∇ is typically replaced by the linear operator D : X → X × X that computes differences in pixel values along vertical and horizontal directions.
A computational drawback with total variation regularization is that it results in minimizing a non-smooth objective. A common remedy is therefore to replace the 1-norm with the Huber function, i.e., choosing S λ (x) := λH δ (Dx) where δ > 0, x = (x i ) i∈I , and This results in a smooth objective as eq. (39) now reads aŝ which corresponds to eq. so One may furthermore normalize the linear operator A, i.e., A = 1. Then, ∇F y is Lipschitz continuous with constant β −1 = 2 A 2 + λ δ D 2 = 2 + 8λ/δ. The tests in subsection 5.6 compare various deviation-based schemes for solving eq. (41) with δ = 0.01 (Huber parameter) and λ = 0.0015 (regularization parameter). The performance of the learned deviation-based scheme in subsection 4.2 is evaluated for different values of the normalization parameter α in eq. (32).

Sparsity-promoting regularization
Sparsity is another powerful method for regularizing ill-posed inverse problems, and in particular those that are under-sampled [21]. A typical choice as regularizing functional in eq. (39) is to pick S λ (x) = λ W x 1 where λ > 0, so we end up with the following non-smooth minimization problem: x := arg min In the above, x → W x is a compression that preserves essential features (sparsifying transform). For the tests in subsection 5.6, we select it as an orthogonal wavelet transform given by Symlets with filter length 10 and 5 scale levels. The regularization parameter is set to λ = 0.0005. The minimization in eq. (44) corresponds to eq. (27) with a non-smooth As in subsection 5.2, A is normalized, so A = 1, and ∇f y (x) = 2A ⊤ (Ax − y) is Lipschitz continuous with constant β −1 = 2. Hence, f y and g satisfy the assumptions in subsection 3.2 and one can solve the problem by using the method proposed in subsection 4.3. To proceed, we use [9,Proposition 24.14] to compute the proximal in eq. (15b): where Prox γ · 1 (x) = sign(x) · max(|x| − γ, 0) is a proximal of l 1 norm and µ is such that µ Id = W W ⊤ .

Baseline methods
In the case of smooth optimization (i.e. solving eq. (41)), we compare our method to the steepest gradient descent and Nesterov's accelerated scheme [33]. For the steepest gradient descent, the updates are the same as in eq. (7), but with ∆x n = 0: In Nesterov's method the updates are: where t 0 := 1.
In the non-smooth case (i.e. solving eq. (44)), we compare our method against ISTA [22] and FISTA [10]. Iterates in ISTA are updated according to This is equivalent to updates in eqs. (15a) and (15b) with ∆x 1 n = ∆x 2 n = 0. In FISTA, updates are given as x n := Prox γn g (w n − γ n ∇f (w n )), (50a) where t 0 := 1. Finally, we do not provide a comparison to LISTA [24], because an input of dimension n it requires training n 2 parameters. The memory requirements for this method are prohibitively large for the examples in subsection 5.6 since n = 512 × 512.

Implementation
The architecture for the neural networks in eqs. (31) and (36) is similar to the networks representing the learned updates in [2]. Thus, it consists of three convolutional layers (two hidden and one output layer) with 3 × 3 convolutional kernels. The hidden layers have 32 filters that are followed by instance normalization and leaky ReLU as an activation function. Instead of performing explicit pre-processing of the input data, we use additional instance normalization before the hidden layers. This results in a network with a total of 1 · 10 · 32 + 32 · 10 · 32 + 32 · 10 · 1 = 10 880 free parameters that need to be set during training.
To train the networks, we use the unsupervised loss function F (x N ), where N is the total number of iterations. However, during the evaluation, we run much more iterations to verify the convergence, despite the fact that the network was not trained for this scenario. During the training, we sample N uniformly between 10 and 20. Randomizing the number of iterations is necessary to ensure a smooth transition between the scenario for which the model is trained and for which it is not. Since the loss function F (x n ) is not guaranteed to decrease monotonically, we observed that fixing the number of iterations to 10 during the training results in a slight increase of the loss after 10 iterations, before it starts to decrease again.
The neural networks are implemented in Tensorflow [30] and trained for 10 5 iterations using the Adam optimizer with learning rate 10 −3 and batch size 1. This takes about three days on a workstation with a single GeForce RTX 2080 Ti GPU. However, for most presented models 3·10 4 training iterations are sufficient. Tomography-related operations necessary for evaluating the data consistency along with bindings to TensorFlow are implemented in ODL [1] with ASTRA [39] as the computational back-end. We make the code publicly available 6 .

Results
All the solution algorithms initialize iterations to zero image x 0 = 0. Next, all deviation-based schemes with learned selection rules have been trained on data from nine of ten patients. The 210 images corresponding to the last patient are used for testing.

Smoothed total variation regularization (smooth optimization)
In figure 1 we visualize the performance of different methods averaged over the set of test images. We identify the minimal objective value F * achieved among all the methods and plot the difference F (x n ) − F * to the minimal value in logarithmic scale depending on the iteration.
We present results for the following methods: the gradient descent, Nesterov's accelerated method, and the learned scheme described in subsection 4.2 with the normalization parameter α = 0.9. For comparison, we also include a variant of the learned scheme without the normalization given by eq. (32).
We use 32-bit single-precision floating-point numbers in our experiments (that represents about 7 decimal digits) and we average over 210 images. The final loss for each image is in most cases a 3 digit number. Therefore we expect the results be precise only up to 5 digits, which is up to and including 10 −2 .
As we can see, learned schemes perform better than the baseline methods during the first 10 iterations (marked with a grey vertical line). This result is expected since it corresponds to the objective set during the training. We also evaluate a variant of the learned scheme that uses a network without the normalization in eq. (32) and, hence, does not have convergence guarantees. Not  only does it achieve a worse performance during the first 10 iterations, but it also quickly diverges afterwards.
In figure 2, we evaluate the performance of the learned scheme for different choices of the normalization parameter α, which controls the size of the learned deviation ∆x n . The largest value α = 0.999 leads to the best performance during the first 10 iterations, however, α = 0.9 performs almost as good and converges much faster. Decreasing α further makes the method approach the steepest gradient descent, by enforcing convergence and sacrificing performance during the first 10 iterations. Figure 3 visualizes the performance of different methods averaged over the set of test images. We identified the minimal objective value F * achieved among all the methods and plot the difference F (x n ) − F * to the minimal value in   [22] and the learned scheme in subsection 4.3 with convergent updates by eq. (37) for different values of α.

Sparsity-promoting regularization (non-smooth optimization)
logarithmic scale depending on the iteration. We can see that, during the first 10 iterations (marked with a grey vertical line), the learned methods perform significantly better than the two baseline methods: ISTA [22] and FISTA [10]. Moreover, the normalization step eq. (37) not only ensures convergence, it also improves the performance after 10 iterations.
Next, we show the effect of choosing different normalization parameters α in figure 4. As in smooth optimization, decreasing α improves convergence. Moreover, in this case α = 0.5 leads to the best performance among the tested methods during the first 10 iterations as well.

Conclusion
Our goal was to construct deep-learning accelerated algorithms for convex optimization, and we began by noting that these heavily over-parametrized al-gorithms do not fit into classical optimization schemes, which only use a few parameters.
To do this, we proved convergence of two convex optimization algorithms, one for smooth optimization and one for non-smooth. Instead of showing their convergence under some specific parameter choice, we showed convergence whenever the updates are chosen in particular sets of full dimension. Given that these sets are rather large, there is hence an enormous number of options to choose from. To make the best choice in any situation we used a neural network and trained it to pick the updates for optimal convergence speed.
We applied our algorithms to the problem of regularized image reconstruction in CT and showed that our algorithm out-performs classical algorithms, especially when only a few iterations are applied. Furthermore, our algorithm converged to the exact solution, as expected.
We hope that this work will open new venues in combining deep learning with provable convergence, and we see several interesting open problems arising from it. For one, our algorithms are only applicable to a subset of all convex problems and it would be particularly interesting to find a primal-dual scheme which can be combined with deep learning. Another interesting research direction is to investigate other definitions of the best optimization algorithm, for example one highly sought-after feature is the ability to train our network to converge to a given accuracy as quickly as possible. Finally, one could combine the deterministic acceleration in the spirit of Nesterov [33] or [10] with the deep learning techniques introduced in this paper and find an optimal trade-off between the two.