Title: PROXIMAL BACKPROPAGATION

Abstract: We propose proximal backpropagation (ProxProp) as a novel algorithm that takes implicit instead of explicit gradient steps to update the network parameters during neural network training. Our algorithm is motivated by the step size limitation of explicit gradient descent, which poses an impediment for optimization. Prox-Prop is developed from a general point of view on the backpropagation algorithm, currently the most common technique to train neural networks via stochastic gradient descent and variants thereof. Specifically, we show that backpropagation of a prediction error is equivalent to sequential gradient descent steps on a quadratic penalty energy, which comprises the network activations as variables of the optimization. We further analyze theoretical properties of ProxProp and in particular prove that the algorithm yields a descent direction in parameter space and can therefore be combined with a wide variety of convergent algorithms. Finally, we devise an efficient numerical implementation that integrates well with popular deep learning frameworks. We conclude by demonstrating promising numerical results and show that ProxProp can be effectively combined with common first order optimizers such as Adam.

1. INTRODUCTION:

In recent years neural networks have gained considerable attention in solving difficult correlation tasks such as classification in computer vision (Krizhevsky et al., 2012) or sequence learning (Sutskever et al., 2014) and as building blocks of larger learning systems (Silver et al., 2016) . Training neural networks is accomplished by optimizing a nonconvex, possibly nonsmooth, nested function of the network parameters. Since the introduction of stochastic gradient descent (SGD) (Robbins & Monro, 1951; Bottou, 1991) , several more sophisticated optimization methods have been studied. One such class is that of quasi-Newton methods, as for example the comparison of L-BFGS with SGD in (Le et al., 2011) , Hessian-free approaches (Martens, 2010) , and the Sum of Functions Optimizer in (Sohl-Dickstein et al., 2013) . Several works consider specific properties of energy landscapes of deep learning models such as frequent saddle points (Dauphin et al., 2014) and well-generalizable local optima (Chaudhari et al., 2017a) . Among the most popular optimization methods in currently used deep learning frameworks are momentum based improvements of classical SGD, notably Nesterov's Accelerated Gradient (Nesterov, 1983; Sutskever et al., 2013) , and the Adam optimizer (Kingma & Ba, 2015) , which uses estimates of first and second order moments of the gradients for parameter updates.

Nevertheless, the optimization of these models remains challenging, as learning with SGD and its variants requires careful weight initialization and a sufficiently small learning rate in order to yield a stable and convergent algorithm. Moreover, SGD often has difficulties in propagating a learning signal deeply into a network, commonly referred to as the vanishing gradient problem (Hochreiter et al., 2001) . Training neural networks can be formulated as a constrained optimization problem by explicitly introducing the network activations as variables of the optimization, which are coupled via layerwise constraints to enforce a feasible network configuration. The authors of (Carreira-Perpiñán & Wang, 2014) have tackled this problem with a quadratic penalty approach, the method of auxiliary coordinates (MAC). Closely related, (Taylor et al., 2016) introduce additional auxiliary variables to further split linear and nonlinear transfer between layers and propose a primal dual algorithm for optimization. From a different perspective, (LeCun, 1988) takes a Lagrangian approach to formulate the constrained optimization problem.

In this work, we start from a constrained optimization point of view on the classical backpropagation algorithm. We show that backpropagation can be interpreted as a method alternating between two steps. First, a forward pass of the data with the current network weights. Secondly, an ordered sequence of gradient descent steps on a quadratic penalty energy.

Using this point of view, instead of taking explicit gradient steps to update the network parameters associated with the linear transfer functions, we propose to use implicit gradient steps (also known as proximal steps, for the definition see ( 6)). We prove that such a model yields a descent direction and can therefore be used in a wide variety of (provably convergent) algorithms under weak assumptions.

Since an exact proximal step may be costly, we further consider a matrix-free conjugate gradient (CG) approximation, which can directly utilize the efficient pre-implemented forward and backward operations of any deep learning framework. We prove that this approximation still yields a descent direction and demonstrate the effectiveness of the proposed approach in PyTorch.

2. MODEL AND NOTATION:

We propose a method to train a general L-layer neural network of the functional form

EQUATION

Here, J(θ; X, y) denotes the training loss as a function of the network parameters θ, the input data X and the training targets y. As the final loss function L y we choose the softmax cross-entropy for our classification experiments. φ is a linear transfer function and σ an elementwise nonlinear transfer function. As an example, for fully-connected neural networks θ = (W, b) and φ(θ, a) = W a + b1.

While we assume the nonlinearities σ to be continuously differentiable functions for analysis purposes, our numerical experiments indicate that the proposed scheme extends to rectified linear units (ReLU), σ(x) = max(0, x). Formally, the functions σ and φ map between spaces of different dimensions depending on the layer. However, to keep the presentation clean, we do not state this dependence explicitly. Figure 1 illustrates our notation for the fully-connected network architecture.

Throughout this paper, we denote the Euclidean norm for vectors and the Frobenius norm for matrices by || • ||, induced by an inner product •, • . We use the gradient symbol ∇ to denote the transpose of the Jacobian matrix, such that the chain rule applies in the form "inner derivative times outer derivative". For all computations involving matrix-valued functions and their gradient/Jacobian, we uniquely identify all involved quantities with their vectorized form by flattening matrices in a column-first order. Furthermore, we denote by A * the adjoint of a linear operator A.

Algorithm 1 -Penalty formulation of BackProp:

Input: Current parameters θ k . // Forward pass.

for l = 1 to L -2 do z k l = φ(θ k l , a k l-1 ), // a0 = X. a k l = σ(z k l )

. end for // Perform minimization steps on (3). a grad. step on E wrt. (θL-1, aL-2) for l = L -2 to 1 do b grad. step on E wrt. z l and a l-1 , c grad. step on E wrt. θ l . end for Output: New parameters θ k+1 .

Algorithm 2 -ProxProp:

Input: Current parameters θ k . // Forward pass.

for l = 1 to L -2 do z k l = φ(θ k l , a k l-1 ), // a0 = X. a k l = σ(z k l )

. end for // Perform minimization steps on (3). a grad. step on E wrt. (θL-1, aL-2), Eqs. 8, 12. for l = L -2 to 1 do b grad. step on E wrt. z l and a l-1 , Eqs. 9, 10. c prox step on E wrt. θ l , Eq. 11. end for Output: New parameters θ k+1 .

3. PENALTY FORMULATION OF BACKPROPAGATION:

The gradient descent iteration on a nested function J(θ; X, y),

EQUATION

is commonly implemented using the backpropagation algorithm (Rumelhart et al., 1986) . As the basis for our proposed optimization method, we derive a connection between the classical backpropagation algorithm and quadratic penalty functions of the form

EQUATION

The approach of (Carreira-Perpiñán & Wang, 2014) is based on the minimization of (3), as under mild conditions the limit ρ, γ → ∞ leads to the convergence of the sequence of minimizers of E to the minimizer of J (Nocedal & Wright, 2006, Theorem 17.1) . In contrast to (Carreira-Perpiñán & Wang, 2014) we do not optimize (3), but rather use a connection of (3) to the classical backpropagation algorithm to motivate a semi-implicit optimization algorithm for the original cost function J.

Indeed, the iteration shown in Algorithm 1 of forward passes followed by a sequential gradient descent on the penalty function E is equivalent to the classical gradient descent iteration. Proposition 1. Let L y , φ and σ be continuously differentiable. For ρ = γ = 1/τ and θ k as the input to Algorithm 1, its output θ k+1 satisfies (2), i.e., Algorithm 1 computes one gradient descent iteration on J.

Proof. For this and all further proofs we refer to Appendix A.

4. PROXIMAL BACKPROPAGATION:

The interpretation of Proposition 1 leads to the natural idea of replacing the explicit gradient steps a , b and c in Algorithm 1 with other -possibly more powerful -minimization steps. We propose Proximal Backpropagation (ProxProp) as one such algorithm that takes implicit instead of explicit gradient steps to update the network parameters θ in step c . This algorithm is motivated by the step size restriction of gradient descent.

4.1. GRADIENT DESCENT AND PROXIMAL MAPPINGS:

Explicit gradient steps pose severe restrictions on the allowed step size τ : Even for a convex, twice continuously differentiable, L -smooth function f : R n → R, the convergence of the gradient descent algorithm can only be guaranteed for step sizes 0 < τ < 2/L . The Lipschitz constant L of the gradient ∇f is in this case equal to the largest eigenvalue of the Hessian H. With the interpretation of backpropagation as in Proposition 1, gradient steps are taken on quadratic functions.

As an example for the first layer,

EQUATION

)

In this case the Hessian is H = XX , which is often ill-conditioned. For the CIFAR-10 dataset the largest eigenvalue is 6.7•10 6 , which is seven orders of magnitude larger than the smallest eigenvalue. Similar problems also arise in other layers where poorly conditioned matrices a l pose limitations for guaranteeing the energy E to decrease.

The proximal mapping (Moreau, 1965) of a function f : R n → R is defined as:

EQUATION

)

By rearranging the optimality conditions to (5) and taking y = x k , it can be interpreted as an implicit gradient step evaluated at the new point x k+1 (assuming differentiability of f ):

EQUATION

The iterative algorithm ( 6) is known as the proximal point algorithm (Martinet, 1970) . In contrast to explicit gradient descent this algorithm is unconditionally stable, i.e. the update scheme ( 6) monotonically decreases f for any τ > 0, since it holds by definition of the minimizer

x k+1 that f (x k+1 ) + 1 2τ ||x k+1 -x k || 2 ≤ f (x k

). Thus proximal mappings yield unconditionally stable subproblems in the following sense: The update in θ l provably decreases the penalty energy E(θ, a k , z k ) from (3) for fixed activations (a k , z k ) for any choice of step size. This motivates us to use proximal steps as depicted in Algorithm 2.

4.2. PROXPROP:

We propose to replace explicit gradient steps with proximal steps to update the network parameters of the linear transfer function. More precisely, after the forward pass

EQUATION

we keep the explicit gradient update equations for z l and a l . The last layer update is

EQUATION

and for all other layers,

EQUATION

EQUATION

where we use a k+ 1 /2 l and z

k+ 1 /2 l

to denote the updated variables before the forward pass of the next iteration and multiplication in (9) is componentwise. However, instead of taking explicit gradient steps to update the linear transfer parameters θ l , we take proximal steps

EQUATION

)

This update can be computed in closed form as it amounts to a linear solve (for details see Appendix B). While in principle one can take a proximal step on the final loss L y , for efficiency reasons we choose an explicit gradient step, as the proximal step does not have a closed form solution in many scenarios (e.g. the softmax cross-entropy loss in classification problems). Specifically,

EQUATION

Note that we have eliminated the step sizes in the updates for z l and a l-1 in ( 9) and ( 10), as such updates correspond to the choice of ρ = γ = 1 τ in the penalty function (3) and are natural in the sense of Proposition 1. For the proximal steps in the parameters θ in (11) we have introduced a step size τ θ which -as we will see in Proposition 2 below -changes the descent metric opposed to τ which rather rescales the magnitude of the update.

We refer to one sweep of updates according to equations ( 7) -( 12) as ProxProp, as it closely resembles the classical backpropagation (BackProp), but replaces the parameter update by a proximal mapping instead of an explicit gradient descent step. In the following subsection we analyze the convergence properties of ProxProp more closely.

4.2.1. CONVERGENCE OF PROXPROP:

ProxProp inherits all convergence-relevant properties from the classical backpropagation algorithm, despite replacing explicit gradient steps with proximal steps: It minimizes the original network energy J(θ; X, y) as its fixed-points are stationary points of J(θ; X, y), and the update direction θ k+1 -θ k is a descent direction such that it converges when combined with a suitable optimizer. In particular, it is straight forward to combine ProxProp with popular optimizers such as Nesterov's accelerated gradient descent (Nesterov, 1983) or Adam (Kingma & Ba, 2015) .

In the following, we give a detailed analysis of these properties. Proposition 2. For l = 1, . . . , L -2, the update direction θ k+1 -θ k computed by ProxProp meets

EQUATION

In other words, ProxProp multiplies the gradient ∇ θ l J with the inverse of the positive definite, symmetric matrix

EQUATION

which depends on the activations a k l-1 of the forward pass. Proposition 2 has some important implications: Proposition 3. For any choice of τ > 0 and τ θ > 0, fixed points θ * of ProxProp are stationary points of the original energy J(θ; X, y).

Moreover, we can conclude convergence in the following sense. Proposition 4. The ProxProp direction θ k+1 -θ k is a descent direction. Moreover, under the (weak) assumption that the activations a k l remain bounded, the angle α k between -∇ θ J(θ k ; X, y) and θ k+1 -θ k remains uniformly bounded away from π/2, i.e.

cos(α:

EQUATION

for some constant c.

Proposition 4 immediately implies convergence of a whole class of algorithms that depend only on a provided descent direction. We refer the reader to (Nocedal & Wright, 2006, Chapter 3.2) for examples and more details.

Furthermore, Proposition 4 states convergence for any minimization scheme in step c of Algorithm 2 that induces a descent direction in parameter space and thus provides the theoretical basis for a wide range of neural network optimization algorithms.

Considering the advantages of proximal steps over gradient steps, it is tempting to also update the auxiliary variables a and z in an implicit fashion. This corresponds to a proximal step in b of Algorithm 2. However, one cannot expect an analogue version of Proposition 3 to hold anymore. For example, if the update of a L-2 in ( 8) is replaced by a proximal step, the propagated error does not correspond to the gradient of the loss function L y , but to the gradient of its Moreau envelope. Consequently, one would then minimize a different energy. While in principle this could result in an optimization algorithm with, for example, favorable generalization properties, we focus on minimizing the original network energy in this work and therefore do not further pursue the idea of implicit steps on a and z.

4.2.2. INEXACT SOLUTION OF PROXIMAL STEPS:

As we can see in Proposition 2, the ProxProp updates differ from vanilla gradient descent by the variable metric induced by the matrices (M k l ) -1 with M k l defined in ( 14). Computing the ProxProp update direction v k l := 1 τ (θ k+1 l -θ k l ) therefore reduces to solving the linear equation

EQUATION

which requires an efficient implementation. We propose to use a conjugate gradient (CG) method, not only because it is one of the most efficient methods for iteratively solving linear systems in general, but also because it can be implemented matrix-free: It merely requires the application of the linear operator M k l which consists of the identity and an application of (∇φ(•, a k l-1 ))(∇φ(•, a k l-1 )) * . The latter, however, is preimplemented for many linear transfer functions φ in common deep learning frameworks, because φ(x, a k l-1 ) = (∇φ(•, a k l-1 )) * (x) is nothing but a forward-pass in φ, and φ * (z, a k l-1 ) = (∇φ(•, a k l-1 ))(z) provides the gradient with respect to the parameters θ if z is the backpropagated gradient up to that layer. Therefore, a CG solver is straight-forward to implement in any deep learning framework using the existing, highly efficient and high level implementations of φ and φ * . For a fully connected network φ is a matrix multiplication and for a convolutional network the convolution operation.

As we will analyze in more detail in Section 5.1, we approximate the solution to ( 16) with a few CG iterations, as the advantage of highly precise solutions does not justify the additional computational effort in obtaining them. Using any number of iterations provably does not harm the convergence properties of ProxProp: Proposition 5. The direction ṽk l one obtains from approximating the solution v k l of (16) with the CG method remains a descent direction for any number of iterations.

4.2.3. CONVERGENCE IN THE STOCHASTIC SETTING:

While the above analysis considers only the full batch setting, we remark that convergence of Prox-Prop can also be guaranteed in the stochastic setting under mild assumptions. Assuming that the activations a k l remain bounded (as in Proposition 4), the eigenvalues of (M k l ) -1 are uniformly contained in the interval [λ, τ θ ] for some fixed λ > 0. Therefore, our ProxProp updates fulfill Assumption 4.3 in (Bottou et al., 2016) , presuming the classic stochastic gradient fulfills them. This guarantees convergence of stochastic ProxProp updates in the sense of (Bottou et al., 2016, Theorem 4 .9), i.e. for a suitable sequence of diminishing step sizes.

5. NUMERICAL EVALUATION:

ProxProp generally fits well with the API provided by modern deep learning frameworks, since it can be implemented as a network layer with a custom backward pass for the proximal mapping. We chose PyTorch for our implementation 1 . In particular, our implementation can use the API's GPU compute capabilities; all numerical experiments reported below were conducted on an NVIDIA Titan X GPU. To directly compare the algorithms, we used our own layer for either proximal or gradient update steps (cf. step c in Algorithms 1 and 2). A ProxProp layer can be seamlessly integrated in a larger network architecture, also with other parametrized layers such as BatchNormalization.

5.1. EXACT AND APPROXIMATE SOLUTIONS TO PROXIMAL STEPS:

We study the behavior of ProxProp in comparison to classical BackProp for a supervised visual learning problem on the CIFAR-10 dataset. We train a fully connected network with architecture 3072 -4000 -1000 -4000 -10 and ReLU nonlinearities. As the loss function, we chose the crossentropy between the probability distribution obtained by a softmax nonlinearity and the ground-truth labels. We used a subset of 45000 images for training while keeping 5000 images as a validation set. We initialized the parameters θ l uniformly in [-1/ √ n l-1 , 1/ √ n l-1 ], the default initialization of PyTorch. While the run time is comparable between the methods, the proposed ProxProp updates have better generalization performance (≈ 54% for BackProp and ≈ 56% for ours on the test set).

Figure 2 shows the decay of the full batch training loss over epochs (left) and training time (middle) for a Nesterov momentum2 based optimizer using a momentum of µ = 0.95 and minibatches of size 500. We used τ θ = 0.05 for the ProxProp variants along with τ = 1. For BackProp we chose τ = 0.05 as the optimal value we found in a grid search.

As we can see in Figure 2 , using implicit steps indeed improves the optimization progress per epoch. Thanks to powerful linear algebra methods on the GPU, the exact ProxProp solution is competitive with BackProp even in terms of runtime.

The advantage of the CG-based approximations, however, is that they generalize to arbitrary linear transfer functions in a matrix-free manner, i.e. they are independent of whether the matrices M k l can be formed efficiently. Moreover, the validation accuracies (right plot in Figure 2 ) suggest that these approximations have generalization advantages in comparison to BackProp as well as the exact ProxProp method. Finally, we found the exact solution to be significantly more sensitive to changes of τ θ than its CG-based approximations. We therefore focus on the CG-based variants of ProxProp in the following. In particular, we can eliminate the hyperparameter τ θ and consistently chose τ θ = 1 for the rest of this paper, while one can in principle perform a hyperparameter search just as for the learning rate τ . Consequently, there are no additional parameters compared with BackProp.

5.2. STABILITY FOR LARGER STEP SIZES:

We compare the behavior of ProxProp and BackProp for different step sizes. Figure 3 : ProxProp as a first-order oracle in combination with the Adam optimizer. The proposed method leads to faster decrease of the full batch loss in epochs and to an overall higher accuracy on the validation set. The plots on the right hand side show data for a fixed runtime, which corresponds to a varying number of epochs for the different optimizers.

5.3. PROXPROP AS A FIRST-ORDER ORACLE:

We show that ProxProp can be used as a gradient oracle for first-order optimization algorithms. In this section, we consider Adam (Kingma & Ba, 2015) . Furthermore, to demonstrate our algorithm on a generic architecture with layers commonly used in practice, we trained on a convolutional neural network of the form:

Conv[16 × 32 × 32] → ReLU → Pool[16 × 16 × 16] → Conv[20 × 16 × 16] → ReLU → Pool[20 × 8 × 8] → Conv[20 × 8 × 8] → ReLU → Pool[20 × 4 × 4] → FC + Softmax[10 × 1 × 1]

Here, the first dimension denotes the respective number of filters with kernel size 5 × 5 and max pooling downsamples its input by a factor of two. We set the step size τ = 10 -3 for both BackProp and ProxProp.

The results are shown in Fig. 3 . Using parameter update directions induced by ProxProp within Adam leads to a significantly faster decrease of the full batch training loss in epochs. While the running time is higher than the highly optimized backpropagation method, we expect that it can be improved through further engineering efforts. We deduce from Fig. 3 that the best validation accuracy (72.9%) of the proposed method is higher than the one obtained with classical backpropagation (71.7%). Such a positive effect of proximal smoothing on the generalization capabilities of deep networks is consistent with the observations of Chaudhari et al. (2017b) . Finally, the accuracies on the test set after 50 epochs are 70.7% for ProxProp and 69.6% for BackProp which suggests that the proposed algorithm can lead to better generalization.

6. CONCLUSION:

We have proposed proximal backpropagation (ProxProp) as an effective method for training neural networks. To this end, we first showed the equivalence of the classical backpropagation algorithm with an algorithm that alternates between sequential gradient steps on a quadratic penalty function and forward passes through the network. Subsequently, we developed a generalization of Back-Prop, which replaces explicit gradient steps with implicit (proximal) steps, and proved that such a scheme yields a descent direction, even if the implicit steps are approximated by conjugate gradient iterations. Our numerical analysis demonstrates that ProxProp is stable across various choices of step sizes and shows promising results when compared with common stochastic gradient descent optimizers.

which yields the assertion.

Proof of Proposition 5. According to (Nocedal & Wright, 2006, p. 109, Thm. 5.3 ) and (Nocedal & Wright, 2006, p. 106, Thm. 5 .2) the k-th iteration x k of the CG method for solving a linear system Ax = b with starting point x 0 = 0 meets

EQUATION

i.e. is optimizing over an order-k Krylov subspace. The starting point x 0 = 0 can be chosen without loss of generality. Suppose the starting point is x0 = 0, then one can optimize the variable x = x -x0 with a starting point x 0 = 0 and b = b + Ax 0 .

We will assume that the CG iteration has not converged yet as the claim for a fully converged CG iteration immediately follow from Proposition 4. Writing the vectors b, Ab, . . . , A k-1 b as columns of a matrix K k , the condition x ∈ span(b, Ab, . . . , A k-1 b) can equivalently be expressed as x = K k α for some α ∈ R k . In terms of α our minimization problem becomes

EQUATION

leading to the optimality condition

EQUATION

)

Note that A is symmetric positive definite and can therefore be written as

EQUATION

being symmetric positive definite. Hence, the matrix ((K k ) T AK k ) -1 is positive definite, too, and

EQUATION

)

Note that (K k ) T b is nonzero if b is nonzero, as b 2 is its first entry.

To translate the general analysis of the CG iteration to our specific case, using any number of CG iterations we find that an approximate solution ṽk l of

EQUATION

leads to ṽk l , -∇ θ l J(θ k ; X, y) > 0, i.e., to ṽk l being a descent direction.

B PROXIMAL OPERATOR FOR LINEAR TRANSFER FUNCTIONS:

In order to update the parameters θ l of the linear transfer function, we have to solve the problem (11),

EQUATION

)

Since we assume that φ is linear in θ for a fixed a k , there exists a matrix A k such that vec(θ k+1 ) = argmin (43)

To also make the above more specific, consider the example of φ(θ, a k ) = θa k . In this case the variable θ may remain in a matrix form and the solution of the proximal mapping becomes

EQUATION

)

Since a k ∈ R n×N for some layer size n and batch size N , the size of the linear system is independent of the batch size.

:

https://github.com/tfrerix/proxprop

PyTorch's Nesterov SGD for direction d(θ k ): m k+1 = µm k +d(θ k ), θ k+1 = θ k -τ (µm k+1 +d(θ k )).

Figures and Tables:

Figure 1: Notation overview. For an L-layer feed-forward network we denote the explicit layer-wise activation variables as z l and a l . The transfer functions are denoted as φ and σ. Layer l is of size n l .

Figure 2: Exact and inexact solvers for ProxProp compared with BackProp. Left: A more precise solution of the proximal subproblem leads to overall faster convergence, while even a very inexact solution (only 3 CG iterations) already outperforms classical backpropagation. Center & Right:While the run time is comparable between the methods, the proposed ProxProp updates have better generalization performance (≈ 54% for BackProp and ≈ 56% for ours on the test set).

||A k vec(θ)vec(z k+ 1 /2 )|| 2 + 1 2τ θ ||vec(θ)vec(θ k )|| 2 , (41)and the optimality condition yieldsvec(θ k+1 ) = (I + τ θ (A k ) T A k ) -1 (vec(θ k ) + (A k ) T vec(z k+ 1 /2 )).(42)Published as a conference paper at ICLR 2018In the main paper we sometimes use the more abstract but also more concise notion of ∇φ(•, a k ), which represents the linear operator ∇φ(•, a k )(Y ) = vec -1 ((A k ) T vec(Y )).

Table1shows the final full batch training loss after 50 epochs with various τ . The ProxProp based approaches remain stable over a significantly larger range of τ . Even more importantly, deviating from the optimal step size τ by one order of magnitude resulted in a divergent algorithm for classical BackProp, but still provides reasonable training results for ProxProp with 3 or 5 CG iterations. These results are in accordance with our motivation developed in Section 4.1. From a practical point of view, this eases hyperparameter search over τ . Full batch loss for conjugate gradient versions of ProxProp and BackProp after training for 50 epochs, where "-" indicates that the algorithm diverged to NaN. The implicit ProxProp algorithms remain stable for a significantly wider range of step sizes.

Table content: <table><tr><td>τ</td><td>50</td><td>10</td><td>5</td><td>1</td><td>0.5</td><td>0.1</td><td colspan="3">0.05 5 • 10 -3 5 • 10 -4</td></tr><tr><td>BackProp</td><td>-</td><td>-</td><td>-</td><td>-</td><td>-</td><td colspan="2">0.524 0.091</td><td>0.637</td><td>1.531</td></tr><tr><td colspan="4">ProxProp (cg1) 77.9 0.079 0.145</td><td>0.667</td><td>0.991</td><td colspan="2">1.481 1.593</td><td>1.881</td><td>2.184</td></tr><tr><td colspan="5">ProxProp (cg3) 94.7 0.644 0.031 2 • 10 -3</td><td>0.012</td><td colspan="2">1.029 1.334</td><td>1.814</td><td>2.175</td></tr><tr><td colspan="8">ProxProp (cg5) 66.5 0.190 0.027 3 • 10 -4 2 • 10 -3 0.399 1.049</td><td>1.765</td><td>2.175</td></tr></table>

Peer review (iclr2018):
 The paper uses a lesser-known interpretation of the gradient step of a composite function (i.e., via reverse mode automatic differentiation or backpropagation), and then replaces one of the steps with a proximal step. The proximal step requries the solution of a positive-definite linear system, so it is approximated using a few iterations of CG. The paper provides theory to show that their proximal variant (even with the CG approximations) can lead to convergent algorithms (and since practical algorithms are not necessarily globally convergent, most of the theory shows that the proximal variant has similar guarantees to a standard gradient step). On reading the abstract and knowing quite a bit about proximal methods, I was initially skeptical, but I think the authors have done a good job of making their case. It is a well-written, very clear paper, and it has a good understanding of the literature, and does not overstate the results. The experiments are serious, and done using standard state-of-the-art tools and architectures. Overall, it is an interesting idea, and due to the current focus on neural nets, it is of interest even though it is not yet providing substantial improvements. The main drawback of this paper is that there is no theory to suggest the ProxProp algorithm has better worst-case convergence guarantees, and that the experiments do not show a consistent benefit (in terms of time) of the method. On the one hand, I somewhat agree with the authors that "while the running time is higher... we expect that it can be improved through further engineering efforts", but on the other hand, the idea of nested algorithms ("matrix-free" or "truncated Newton") always has this issue. A very similar type of ideas comes up in constrained or proximal quasi-Newton methods, and I have seen many papers (or paper submissions) on this style of method (e.g., see the 2017 SIAM Review paper on FWI by Metivier et al. at https://doi.org/10.1137/16M1093239). In every case, the answer seems to be that it can work on *some problems* and for a few well chosen parameters, so I don*t yet buy that ProxProp is going to make a huge savings on a wide-range of problems. In brief: quality is high, clarity is high, originality is high, and significance is medium. Pros: interesting idea, relevant theory provided, high-quality experiments Cons: no evidence that this is a "break-through" idea Minor comments: - Theorems seemed reasonable and I have no reason to doubt their accuracy - No typos at all, which I find very unusual. Nice job! - In Algo 1, it would help to be more explicit about the updates (a), (b), (c), e.g., for (a), give a reference to eq (8), and for (b), reference equations (9,10). It*s nice to have it very clear, since "gradient step" doesn*t make it clear what the stepsize is, and if this is done in a "Jacob-like" or "Gauss-Seidel-like" fashion. (c) has no reference equation, does it? - Similarly, for Algo 2, add references. In particular, tie in the stepsizes tau and tau_theta here. - Motivation in section 4.1 was a bit iffy. A larger stepsize is not always better, and smaller is not worse. Minimizing a quadratic f(x) = .5||x||^2 will converge in one step with a step-size of 1 because this is well-conditioned; on the flip side, slow convergence comes from lack of strong convexity, or with strong convexity, ill-conditioning of the Hessian (like a stiff ODE). - The form of equation (6) was very nice, and you could also point out the connection with backward Euler for finite-difference methods. This was the initial setting of analysis for most of original results that rely on the proximal operator (e.g., Lions and Mercier 1970s). - Eq (9), this is done component-wise, i.e., Hadamard product, right? - About eq (12), even if softmax cross-entropy doesn*t have a closed-form prox (and check the tables of Combettes and Pesquet), because it is separable (if I understand correctly) then it ought to be amenable to solving with a handful of Newton iterations which would be quite cheap. Prox tables (see also the new edition of Bauschke and Combettes* book): P. L. Combettes and J.-C. Pesquet, "Proximal splitting methods in signal processing," in: Fixed-Point Algorithms for Inverse Problems in Science and Engineering (2011) http://www4.ncsu.edu/~pcombet/prox.pdf - Below prop 4, discussing why not to make step (b) proximal, this was a bit vague to me. It would be nice to expand this. - Page 6 near the top, to apply the operator, in the fully-connected case, this is just a matrix multiply, right? and in a conv net, just a convolution? It would help the reader to be more explicit here. - Section 5.1, 2nd paragraph, did you swap tau_theta and tau, or am I just confused? The wording here was confusing. - Fig 2 was not that convincing since the figure with time showed that either usual BackProp or the exact ProxProp were best, so why care about the approximate ProxProp with a few CG iterations? The argument of better generalization is based on very limited experiments and without any explanation, so I find that a weak argument (and it just seems weird that inexact CG gives better generalization). The right figure would be nice to see with time on the x-axis as well. - Section 5.2, this was nice and contributed to my favorable opinion about the work. However, any kind of standard convergence theory for usual SGD requires the stepsize to change per iteration and decrease toward zero. I*ve heard of heuristics saying that a fixed stepsize is best and then you just make sure to stop the algorithm a bit early before it diverges or behaves wildly -- is that true here? - Final section of 5.3, about the validation accuracy, and the accuracy on the test set after 50 epochs. I am confused why these are different numbers. Is it just because 50 epochs wasn*t enough to reach convergence, while 300 seconds was? And why limit to 50 epochs then? Basically, what*s the difference between the bottom two plots in Fig 3 (other than scaling the x-axis by time/epoch), and why does ProxProp achieve better accuracy only in the right figure?