Maximum Principle Based Algorithms for Deep Learning

Introduction

Supervised learning using deep neural networks has become an increasingly successful tool in modern machine learning applications (Bengio, 2009; Schmidhuber, 2015; LeCun et al., 2015; Goodfellow et al., 2016). Efficient training methods of very deep neural networks, however, remain an active area of research. The most commonly applied training method is stochastic gradient descent (Robbins and Monro, 1951; Bottou, 2010) and its variants (Duchi et al., 2011; Zeiler, 2012; Kingma and Ba, 2014; Johnson and Zhang, 2013), where incremental updates to the trainable parameters are performed using gradient information computed via back-propagation (Kelley, 1960; Bryson, 1975). While efficient to implement, the incremental updates to the parameter tend to be slow, especially in the initial stages of the training. Moreover, other than the computation of gradients through back-propagation, the specific structure of deep neural networks is not exploited. These observations point to the question of whether there exists alternative training methods tailored to deep neural networks.

In a series of papers, we introduce an alternative approach by exploring the optimal control viewpoint of deep learning (E, 2017). Our focus will be on ideas and algorithms derived from the powerful Pontryagin’s maximum principle (Boltyanskii et al., 1960; Pontryagin, 1987), which has two major components: the Hamiltonian dynamics and the condition that at each time the optimal parameters maximize the Hamiltonian. The second component suggests that optimization can be performed independently at different layers. One can also derive an explicit error control estimate based on the maximum principle (see Lemma 2 below).

In this first paper, we will consider the simplest context in which the deep neural networks are replaced by continuous (or discretized) dynamical systems, and devise numerical algorithms that are based on the optimality conditions in the Pontryagin’s maximum principle. This leads to a new approach for training deep learning models that have certain advantages, such as fast initial descent and resilience to stalling in flat landscapes. An additional advantage is that one has a good control of the error through explicit estimates.

The rest of the paper is organized as follows. In Section 2, we present a dynamical systems viewpoint of function approximation and deep learning. We then discuss the necessary optimality conditions, which is the well-known Pontryagin’s maximum principle. In Section 3 and 4, we discuss numerical methods to solve the necessary conditions and obtain error estimates and convergence guarantees. Using benchmarking examples, we then compare our method with traditional gradient-descent based methods for optimizing deep neural networks in Section 5. In Section 6, we discuss and compare our work with existing literature. Conclusion and outlook are given in Section 7.

Function Approximation by Dynamical Systems

We start with a description of the (continuous) dynamical systems approach to machine learning (see E 2017). The essential task of supervised learning is to approximate some function

Problem (2) is a special case of a class of general optimal control problem for ordinary differential equations (Bertsekas, 1995; Athans and Falb, 2013). The advantage of this formulation is that we can write down and study the optimality conditions of (2) entirely in continuous time and derive numerical algorithms that can subsequently be discretized. In other words, we optimize, then discretize, as opposed to the traditional reverse approach in deep learning.

As was suggested in E (2017), deep residual networks (He et al., 2016) can be considered as the forward Euler discretization of the continuous approach described above. In this connection, the algorithms presented in this paper can also be formulated in the context of deep residual networks. For general deep neural networks, although one can also formulate similar algorithms, it is not clear at this moment that PMP holds and these algorithms are valid (e.g. converge to the right solution) in the general setting. This issue will be studied in future work.

The optimization problem (2) can be solved by first discretizing it into a discrete problem (a feed-forward neural network) and then applying back propagation and gradient descent approaches commonly used in deep learning. However, here we will present an alternative approach. Hereafter, for simplicity of notation we shall set $K=1$ drop the scripts $i$ on all functions, noting that analogous results can be obtained in the general case since the dynamics and loss functions are decoupled across samples. Equivalently, we can think of this as effectively concatenating all $K$ sample inputs into a single input vector of dimension $d\times K$ and redefine our dynamics accordingly. Hence, all results remain valid if we perform full-batch training. The case of mini-batch training is discussed in Section 4.3.

In this section, we introduce a set of necessary conditions for optimal solutions of (2), known as the Pontryagin’s Maximum Principle (PMP) (Boltyanskii et al., 1960; Pontryagin, 1987). This shall pave way for an alternative numerical algorithm to train (2) and its discrete-time counter-part.

are satisfied. Moreover, for each $t\in[0,T]$, we have the Hamiltonian maximization condition

The proof of the PMP and its variants can be found in any optimal control theory reference, e.g. Athans and Falb (2013); Bertsekas (1995); Liberzon (2012). Some generalizations can be found in Clarke (2005) and references therein. For example, the requirement of the continuity of $f$ with respect to $t$ can be replaced by a much weaker measurability requirement if one assumes more conditions on $\nabla_{x}f$. In the statement of Theorem 1, we omitted a technicality involving an abnormal multiplier: the terminal condition for $P^{*}$ should be $P^{*}_{T}=-\lambda\nabla\Phi(X^{*}_{T})$ and the Hamiltonian should be defined as $H(t,x,p,\theta)=p\cdot f(t,x,\theta)-\lambda L(\theta)$ for some $\lambda\geq 0$ (abnormal multiplier) that we can choose. When we are forced to always take $\lambda=0$, the problem is singular and in a sense ill-posed (Athans and Falb, 2013). On the contrary, if we can take a positive $\lambda$, we can then rescale the equation for $P^{*}$ so that we can take $\lambda=1$ without loss of generality. We shall hereafter assume that this is the case.

A few remarks are in order. First, Equation 3, 4 and 5 allow us to solve for the unknowns $X^{*},P^{*},\theta^{*}$ simultaneously as a function of $t$. In this sense, the resulting optimal control $\theta^{*}$ is open-loop and is not in a feed-back form $\theta^{*}_{t}=\theta^{*}(X^{*}_{t})$. The latter is of closed-loop type and are typically obtained from dynamic programming and the Hamilton-Jacobi-Bellman formalism (Bellman, 2013). In this sense, the PMP gives a weaker control. However, open-loop solutions are sufficient for neural network applications, where the trained weights and biases are fixed and only depend on the layer number and not the inputs.

Last, we emphasize that the PMP is only a necessary condition, hence there can be cases where solutions to the PMP is not actually globally optimal for (2). Nevertheless, in practice the PMP is often strong enough to give good solution candidates, and when certain convexity assumptions are satisfied the PMP becomes sufficient (Bressan and Piccoli, 2007). In the next section, we will discuss numerical methods that can be used to solve the PMP.

Method of Successive Approximations

Now, our strategy is to devise numerical algorithms for training (2) via solving the PMP (Equation 3, 4 and 5). We derive and analyze algorithms entirely in continuous time, which allows us to characterize errors estimates and convergence in a more transparent fashion.

There are many methods for the numerical solution of the PMP, including two-point boundary value problem method (Bryson, 1975; Roberts and Shipman, 1972), and collocation methods (Betts, 1998) coupled with general non-linear programming techniques (Bertsekas, 1999; Bazaraa et al., 2013). See (Rao, 2009) for a more recent review. However, many of these methods concern small-scale problems typically encountered in control applications (e.g. trajectory optimization of spacecrafts) and do not scale well to modern machine learning problems with a large number of state and control variables. One exception is the method of successive approximations (MSA) (Chernousko and Lyubushin, 1982), which is an iterative method based on alternating propagation and optimization steps. We first introduce the simplest form of the MSA.

and is independent of the co-state $P^{*}$. Therefore, we may proceed in the following manner. First, we make an initial guess of the optimal control $\theta^{0}\in\mathcal{U}$. For each $k=0,1,2,\dots$, we first solve (3)

for $X^{\theta^{k}}$, which then allows us to solve (4)

to get $P^{\theta^{k}}$. Finally, we use the maximization condition (5) to set

for $t\in[0,T]$. The algorithm is summarized in Algorithm 1.

As is the case with the maximum principle, MSA consists of two major components: the forward-backward Hamiltonian dynamics and the maximization for the optimal parameters at each time. An important feature of MSA is that the Hamiltonian maximization step is decoupled for each $t\in[0,T]$. In the language of deep learning, the optimization step is decoupled for different network layers and only the Hamiltonian ODEs (Step 3,4 of Algorithm 1) involve propagation through the layers. This allows the parallelization of the maximization step, which is typically the most time-consuming step.

It has been shown that the basic MSA converges for a restricted class of linear quadratic regulators (Aleksandrov, 1968). However, in general it tends to diverge, especially if a bad initial $\theta^{0}$ is chosen (Aleksandrov, 1968; Chernousko and Lyubushin, 1982). Our goal now is to modify the basic MSA to control its divergent behavior. Before we do so, it is important to understand why the MSA diverges, and in particular, the relationship between the maximization step in Algorithm 1 and the optimization problem (2).

2 Error Estimate for the Basic MSA

For each $\theta\in\mathcal{U}$, let us denote

where $X^{\theta}$ satisfies (6). Our goal is to minimize $J(\theta)$. We show in the following Lemma the relationship between the values of $J$ and the Hamiltonian maximization step. We start by making the following assumptions.

$\Phi$ is twice continuously differentiable, with $\Phi$ and $\nabla\Phi$ satisfying a Lipschitz condition, i.e. there exists $K>0$ such that

$f(t,\cdot,\theta)$ is twice continuously differentiable in $x$, with $f,\nabla_{x}f$ satisfying a Lipschitz condition in $x$ uniformly in $\theta$ and $t$, i.e. there exists $K>0$ such that

With these assumptions, we have the following estimate:

Suppose (A1)-(A2) holds. Then, there exists a constant $C>0$ such that for any $\theta,\phi\in\mathcal{U}$,

where $X^{\theta}$, $P^{\theta}$ satisfy Equations 6, 7 respectively and $\Delta H_{\phi,\theta}$ denotes the change in Hamiltonian

Proof See Appendix B for the proof and discussion on relaxing the assumptions.

In essence, Lemma 2 says that the Hamiltonian maximization step in MSA (step 5 in Algorithm 1) is in some sense the optimal descent direction for $J$. However, the last two terms on the right hand side indicates that this descent can be nullified if substituting $\phi$ for $\theta$ incurs too much error in the Hamiltonian dynamics (step 3,4 in Algorithm 1). In other words, the last two integrals measure the degree of satisfaction of the Hamiltonian dynamics (3), (4), which can be viewed as a feasibility condition, when one replaces $\theta$ by $\phi$. Hence, we shall hereafter refer to these errors as feasibility errors. The divergence of the basic MSA happens when the feasibility errors blow up. Armed with this understanding, we can then modify the basic MSA to ensure convergence.

3 Extended PMP and Extended MSA

As discussed previously in Lemma 2, the decrement of $J$ is ensured if we can control the feasibility errors in the Hamiltonian dynamics in steps 3,4 of Algorithm 1. To this end, we employ a similar idea to augmented Lagrangians (Hestenes, 1969). Fix some $\rho>0$ and introduce the augmented Hamiltonian

Then, we have the following set of alternative necessary conditions for optimality:

Suppose that $\theta^{*}$ is an essentially bounded solution to the optimal control problem (2). Then, there exists an absolutely continuous co-state process $P^{*}$ such that the tuple $(X_{t}^{*},P_{t}^{*},\theta_{t}^{*})$ satisfies the necessary conditions

Proof If $\theta^{*}$ is optimal, then by the PMP there exists a co-state process $P^{*}$ such that (3), (4) and (5) are satisfied. Then, for all $t\in[0,T]$ and $\theta\in\Theta$ we have

which implies that (9) and (10) are satisfied. Lastly, we can write

If $\mu_{k}=0$, then from the Hamiltonian maximization step (11) we must have

i.e. $(X^{\theta^{k}},P^{\theta^{k}},\theta^{k})$ satisfy the extended PMP. In other words, the quantity $\mu_{k}\geq 0$ measures the distance from a solution of the extended PMP, and if it equals 0, then we have a solution. We now prove the following result that guarantees the convergence of the extended MSA (Algorithm 2).

Let (A1)-(A2) be satisfied and $\theta^{0}\in\mathcal{U}$ be any initial measurable control with $J(\theta^{0})<+\infty$. Suppose also that $\inf_{\theta\in\mathcal{U}}J(\theta)>-\infty$. Then, for $\rho$ large enough, we have under Algorithm 2,

i.e. the extended MSA algorithm converges to the set of solutions of the extended PMP.

Proof Using Lemma 2 with $\theta\equiv\theta^{k},\phi\equiv\theta^{k+1}$, we have

From the Hamiltonian maximization step in Algorithm 2, we know that

Pick $\rho>2C$, then we indeed have $J(\theta^{k+1})-J(\theta^{k})\leq-D\mu_{k}$ with $D=(1-\frac{2C}{\rho})>0$. Moreover, we can rearrange and sum the above expression to get

and hence $\sum_{k=0}^{\infty}\mu_{k}<+\infty$, which implies $\mu_{k}\rightarrow 0$ and the extended MSA converges to a solution of the extended PMP.

Discrete-Time Formulation

In the previous section, we discussed the PMP and MSA in the continuous-time setting, where we showed that an appropriately extended version (E-MSA) converges to a solution of an extended PMP. Here, we shall discuss the discretized versions of PMP, MSA and E-MSA, as well as their connections to deep residual networks and back-propagation.

Applying Euler-discretization to Equation 1, we get

for $n=0,\dots,N-1$, with $\delta=T/N$ (step-size), $x_{n}:=X_{n\delta}$, $\vartheta_{n}:=\theta_{n\delta}$ and $f_{n}(\cdot):=f(n\delta,\cdot)$. Then, the discrete-time analogue of the control problem (2) is

Observe that barring the constant $\delta$, this is exactly the supervised learning problem for deep residual networksIf we pick ReLU activations (Hahnloser et al., 2000), then $\delta$ can be absorbed into $\vartheta$. Therefore, when suitably discretized, one expects that the E-MSA provides a means to train residual neural networks via the solution of the extended PMP.

We now write down formally the discretized form of the PMP. Let us use the shorthand $g_{n}(x_{n},\vartheta_{n}):=x_{n}+\delta f_{n}(x_{n},\vartheta_{n})$. Define the scaled discrete Hamiltonian

Then, a discrete-time PMP is the following set of conditions:

The issue of whether the PMP holds for discrete time dynamical systems is a delicate one and there are known counterexamples (Butkovsky, 1963; Jackson and Horn, 1965; Nahorski et al., 1984). Nevertheless, they must hold approximately for small time step size and this is the situation we will consider in the current paper. We expect Lemma 2, which implies monotonicity of the E-MSA algorithm, to hold in the discrete-time case under appropriate conditions. We leave a rigorous analysis of these statements to future work. For numerical experiments presented in the next section, we shall almost always work with residual networks that can be regarded as discretizations of continuous networks so that the PMP holds approximately at least (Halkin, 1966).

For completeness, we summarize the discrete-time version of E-MSA in Algorithm 3. Note that for residual networks ($g_{n}=x_{n}+\delta f_{n}$), this is equivalent to a forward Euler discretization on the state equation and a backward Euler discretization on the co-state equation in Algorithm 2. As before, the Hamiltonian maximization step is decoupled across layers and can be carried out in parallel.

2 Relationship to Gradient Descent with Back-propagation

We note an interesting relationship of the MSA with classical gradient descent with back-propagation (Kelley, 1960; Bryson, 1975; LeCun et al., 1998). We have shown in Lemma 2 that the divergence of MSA can be attributed to the large errors in the Hamiltonian dynamics terms caused by the maximization step, which involve drastic changes in parameter values. Assuming each $\Theta_{n}$ is a continuum and $g_{n}$, $L$ are differentiable in $\vartheta_{n}$, a simple fix is to make the maximization step “soft”: we replace step 12 in Algorithm 3 with a gradient ascent step:

for some small learning rate $\eta$. We now show that in the discrete-time setting, this is equivalent to the classical gradient descent with back-propagation.

The basic MSA in discrete-time (Algorithm 3 with $\rho=0$) with step 12 replaced by (13) is equivalent to gradient descent with back-propagation.

and the total loss function is $J(\vartheta)=\Phi(x_{N})+\delta\sum_{n=0}^{N-1}L(\vartheta_{n})$. It is easy to see that $p_{n}=-\nabla_{x_{n}}\Phi(x_{N})$ by working backwards from $n=N$ and the fact that $\nabla_{x_{n}}x_{n+1}=\nabla_{x}g_{n}(x_{n},\vartheta_{n})$. Then,

Hence, (13) is simply the gradient descent step

As the proposition shows, gradient descent with back-propagation can be seen as a modification of the basic MSA by replacing the Hamiltonian maximization step with a gradient ascent step. However, we note that the PMP (and MSA convergence) holds, at least in continuous-time, even when differentiability with respect to $\vartheta$ is not satisfied, and hence is more general than the classical back-propagation. In fact, the PMP formalism shows that the back-propagation of information through a deep network is handled by the co-state equation and there is no requirement or relationship to the gradients with respect to the trainable parameters. In other words, optimization is performed at each layer separately (with or without gradient information), and propagation is independent of optimization.

3 A Remark on Mini-batch Algorithms

So far, our discussion has focused on full-batch algorithms, where the input $x$ represents the full set of training inputs. As modern supervised learning tasks typically involve a large number of training samples, usually the optimization problem has to be solved in mini-batches, where at each iteration we sub-sample $m$ input-label pairs and optimize the parameters $\theta$ (or $\vartheta$ in discrete time) based on losses evaluated on these pairs. In the context of continuous-time PMP, we can write the batch version of the three necessary conditions as

for samples $i=1,\dots,M$. We omit for brevity the equivalent expressions for discrete-time. In particular, notice that the propagation steps are decoupled across samples, and hence can be carried out independently. The only difference is the maximization step, where in a mini-batch setting we would evaluate instead

If $m$ is large enough and the samples are independently and identically drawn, then uniform law of large numbers (Jennrich, 1969) holds under fairly general conditions and ensures that the mini-batch mean of Hamiltonians converges uniformly in $\theta$ to the full-batch sum. Hence, maximization performed on the mini-batch sum should be close to the actual maximization on the full Hamiltonian. Rigorous error estimates for the mini-batch version of our algorithm is out of the scope of the current work, and we use instead numerical results in Section 5 to demonstrate that the algorithm can also be carried out in a mini-batch fashion.

Numerical Experiments

In this section, we investigate the performance of E-MSA compared with the usual gradient-based approaches, namely stochastic gradient descent and its variants: Adagrad (Duchi et al., 2011) and Adam (Kingma and Ba, 2014). To illustrate key properties of E-MSA, we shall begin by investigating some synthetic examples. First, we consider a simple one-dimensional function approximation problem where we want to approximate $F(x)=\sin(x)$ for $x\in[-\pi,\pi]$ using a continuous time dynamical system. Let $T=5$ and consider

Next, we consider a familiar supervised learning test problem: the MNIST data set (LeCun, 1998) for handwritten digit recognition, with 55000 training samples and 10000 test samples. We employ a continuous dynamical system that resembles a (residual) convolution neural network (LeCun and Bengio, 1995) when discretized. More concretely, at each $t$ we consider the map $f(t,x,\theta)=\tanh(W\star x+b)$ where $W$ is a $3\times 3$ convolution filter with $32$ input and output channels. To match dimensions, we introduce two projection layers at the input (consisting of convolution, point-wise non-linearities followed by $2\times 2$ max-pooling). We also use a fully-connected classification layer as the final layer, with softmax cross-entropy loss. Note that the input projection layers and fully-connected output layers are not of residual form, but we can nevertheless apply Algorithm 3 with the appropriate $g$. We use a total of 10 layers (2 projections, 1 fully-connected and 7 residual layers with $\delta=0.5$, i.e $T=3.5$). The model is trained with mini-batch sizes of 100 using E-MSA and gradient-descent based methods, namely SGD, Adagrad, and Adam. For E-MSA, we approximately solve the Hamiltonian maximization step using either 10 iterations of L-BFGS. Note that since we have decoupled the layers through the PMP, the L-BFGS step used to maximize $H$ is tractable since it involves much fewer parameters than directly minimizing $J$. Figure 2 compares the performance of E-MSA with the other gradient-descent based methods, where we observe that E-MSA has good performance per-iteration, especially at early stages of training. However, we also show in Figure 3 that the wall-clock performance of our methods are not currently competitive, because the Hamiltonian maximization step is time consuming and the performance gains per iteration is outweighed by the running time. Note that wall-clock times are compared on the CPU for fairness since we did not use a GPU implementation of L-BFGS. As a further test, we train the same model on a different data set, the fashion MNIST data set (Xiao et al., 2017), where we again observe similar phenomena (see Figure 4). Experiments on more complex data sets such as ImageNet (Deng et al., 2009) with larger residual networks is a direction of future work. In particular, this may require further improvements to the Hamiltonian maximization step current handled by direct minimization with L-BFGS, which can be significantly slower (on a wall-clock basis) for larger networks and data sets.

Discussion and Related Work

We commence this section by highlighting the distinguishing features of E-MSA from traditional gradient-descent based training methods. First, the formulations of PMP and E-MSA do not involve gradient information with respect to the trainable parameters. In fact, Theorem 1 and Algorithm 2 remain valid even when the trainable parameters can only take values in a discrete set. Second, due to a more drastic argmax step taken at each iteration, E-MSA tends to have better convergence rates at the early steps of training, as observed in our numerical experiments (Section 5). Third, in the PMP formalism, the Hamiltonian equations for the state and co-state are the “forward and backward propagations”, whereas given the state and co-state values, the optimization step is decoupled across layers. This allows one to potentially parallelize the often time-consuming optimization step. Moreover, from Lemma 2, we show that as long as the Hamiltonian is sufficiently increased in a layer without causing too much loss in the Hamiltonian dynamics feasibility conditions, we can ensure decrement of the loss function. This is the reason why we can use a small number of iterations of L-BFGS at each step. Moreover, this suggests that the argmax updates need not happen synchronously, i.e. the optimization in each layer can be a separate thread or process that computes the argmax and updates that layer’s parameters independent of other layers. The propagation may also potentially be allowed to happen asynchronously as long as updates are sufficiently frequent. We leave a rigorous analysis of an asynchronous version of the current approach to future work. In summary, the main strength of the PMP (over e.g. solving the KKT conditions using gradient methods) is that PMP says that at the optimum, the Hamiltonian is not only stationary (KKT), but globally maximized. This hints that heuristic global optimization methods can be applied to $H$ to obtain algorithms that are very different in behavior compared with gradient-descent based approaches. Again, Lemma 2 ensures that such heuristic global maximization need only be approximate.

As it currently stands, our experiments in Section 5 demonstrate that the Hamiltonian maximization step in E-MSA gives very different behavior compared with gradient-descent based methods. When the Hamiltonian is sufficiently maximized, we indeed obtain favorable performance compared with gradient descent based methods. Furthermore, we saw in Figure 1 that Hamiltonian maximization may avoid pitfalls such as a very flat landscape. Overall, the key to whether E-MSA (and other methods based on solving the PMP) will eventually constitute a replacement for gradient-descent based algorithm lies in the question of whether efficient Hamiltonian maximization can be performed at reasonable computational costs. Although this is still a non-convex optimization problem, it is much simpler than the original training problem because: (1) Optimization in the layers are decoupled and hence parameter space is greatly reduced; (2) The Hamiltonian is formally similar across different layers, loss functions and models, so specialized algorithms may be designed; (3) The Hamiltonian does not need to be maximized exactly, thus fast heuristic methods (Lee and El-Sharkawi, 2008) or learning (Andrychowicz et al., 2016; Jaderberg et al., 2016; Czarnecki et al., 2017) can potentially be used to perform this. All these are worthy of future exploration in order to make E-MSA truly competitive.

Next, we put our work in perspective by discussing related work in the optimal control, optimization and deep learning literature. First, the work on numerical algorithms for the solution of optimal control problem is abundant (see e.g. Rao 2009 for a survey). Many of the state-of-the-art techniques in the control theory literature assume a moderately small problem size, so that conventional non-linear programming techniques (Bertsekas, 1999; Bazaraa et al., 2013) as well as shooting (Roberts and Shipman, 1972) and collocation methods (Betts, 1998) produce efficient algorithms. This is usually not the case for large-scale machine learning problems, where often, the only scalable approach is to rely on iterative updates to the parameters. This is the reason for our focus on the MSA algorithms (Chernousko and Lyubushin, 1982), as they are straight-forward to implement and typically have linear scaling in computational complexity with respect to the input and parameter sizes. The basic MSA is discussed in Krylov and Chernousko (1962), and a number of improved variants are discussed in Chernousko and Lyubushin (1982) and references therein. For example, a popular improvement is based on needle-perturbations, where controls are varied on small intervals at each iteration. While convergent, the main issue with the needle-perturbation approach is the requirement of a sufficiently fine mesh (i.e. many layers in the discretized network), which impacts computational speed. A possible solution is the use of adaptive meshes, which is a future direction we plan to investigate. Our variant of the MSA presented in this work differs from classical approaches (Chernousko and Lyubushin, 1982) mainly in the sense that we solve a weaker sufficient condition (extended PMP, Proposition 3), which then allows us to control errors in the Hamiltonian dynamical equations at every iteration without going into finer mesh-sizes. The regularization terms proportional to $\rho$ is similar to the heuristic modifications suggested in Lyubushin (1982) by regularizing the distance between $\theta^{k}$ and $\theta^{k+1}$, but we do not have to assume convexity of $\Theta$ or that $f$ is Lipschitz in $\theta$.

In the optimization literature, our work shares some similarity with the recently proposed ADMM methods (Taylor et al., 2016) for training deep neural networks, where the authors also considered necessary conditions with Lagrange multipliers that can decouple optimization across layers. The main difference in our work is that the PMP gives a stronger necessary condition (Hamiltonian maximization) that also applies to general parameter spaces (e.g., discrete, or bounded with non-linear constraints). Our modification of the basic MSA in terms of the augmented Hamiltonian is inspired by the method of augmented Lagrangians often applied in constrained optimization (Hestenes, 1969). The idea of viewing an initially discrete system as the discretization of a continuous-time system has been explored in Li et al. (2017) in the form of stochastic optimization. Our current work is also in this flavor, but for neural network models.

In deep learning, there are a few works that share our perspective of deep neural networks as a discretization of a dynamical system. We note that the connection between the PMP and back-propagation has been pointed out qualitatively in LeCun (1988) and in the development of back-propagation (Bryson, 1975; Baydin et al., 2015), although to the best of our knowledge, this work is the first attempt to translate numerical algorithms for the PMP into training algorithms for deep learning that goes beyond gradient descent. The treatment of machine learning as function approximation via a dynamical system has been presented in E (2017). The recent work of Haber and Ruthotto (2017); Chang et al. (2017) also propose the dynamical systems viewpoint, and the authors used continuous-time tools to address stability issues. In contrast, our work focuses on the optimization aspects centered around the PMP. We also mention other recent approaches to decouple optimization in deep neural networks, such as synthetic gradients (Jaderberg et al., 2016; Czarnecki et al., 2017) and proximal back-propagation (Frerix et al., 2017).

Conclusion and Outlook

In this paper, we discuss the viewpoint that deep residual neural networks can be viewed as discretization of a continuous-time dynamical system, and hence supervised deep learning can be regarded as solving an optimal control problem in continuous time. We explore a concrete consequence of this connection, by modifying the classical method of successive approximations for solving optimal control problems (in particular the PMP) into a method for solving a weaker sufficient condition (extended PMP). We prove the convergence of the resulting algorithm (E-MSA) and test it on various benchmark problems, where we observe that the E-MSA algorithm performs favorably on a per-iteration basis, especially at early stages of training, compared with gradient-based approaches such as SGD, Adagrad and Adam.

There are many avenues of future research. On the algorithmic side, it is necessary to further improve the computational efficiency of the E-MSA, in particular the Hamiltonian maximization step. Moreover, adaptive selection of $\rho$ depending on iteration number and/or layer can be explored, e.g. by designing adaptive tuning schemes using control theoretic tools (Li et al., 2017). Also, it is desirable to formulate and analyze the PMP and E-MSA from a discrete-time perspective in order to broaden the method’s application. From a modeling perspective, viewing deep neural networks as continuous-time dynamical systems is useful in the sense that it allows one to think of neural network architectures as dynamical objects. Indeed, at each training iteration of the E-MSA, we do not have to use the same discretization scheme to compute the Hamiltonian dynamical equations. Also, as the PMP and E-MSA assume little structure on the parameter space $\Theta$, it will also be interesting to apply the E-MSA to train neural networks that have discrete weights (e.g. those that can only take on binary values). Such networks have the advantage of fast inference speed and small memory requirement. However, training such networks is a challenge and most existing techniques rely on approximating or thresholding the derivatives (Courbariaux et al., 2015, 2016). With the PMP and MSA, we may be able to directly train discrete networks in a principled way.

The work of W. E is supported in part by Major Program of NNSFC under grant 91130005, ONR grant N00014-13-1-0338, DOE grants DE-SC0008626 and DE-SC0009248 Q. Li is supported by the Agency for Science, Technology and Research, Singapore.

A Function Space Formulation

In this section, we give an alternative, non-rigorous formulation of the supervised learning problem as an optimal control problem on function spaces. This provides an alternative formulation of (continuous-time) deep learning that does not make reference to a specific set of input-outputs, but rather their conditional distributions. The idea is to consider the control of a continuity equation that describes the evolution of probability densities. Hereafter, we proceed formally by assuming all differentiability and integrability conditions are satisfied.

As before, the idea is to consider passing the inputs through a dynamical system

We begin with a guess of a conditional density $\rho_{0}(y|x)$ of $y$ given $x$. In the deterministic case, we may set $\rho_{0}(y|x)=\delta({y-F_{0}(x)})$ for some $F_{0}:\mathcal{X}\rightarrow\mathcal{Y}$ (this is like the last layer of the neural network, be it a regressor or a classifier). Note that $F_{0}$ is potentially very different from $F$, so that $\rho_{0}(\cdot|x)$ is far from our target $\rho(\cdot|x)$.

To improve this approximation, we drive the initial condition by the controllable dynamical system (14). That is, we define the approximation at time $t$ of $\rho(y|x)$ to be $\rho_{t}(y|x):=\langle\rho_{0}(y|\cdot),u_{t}\rangle$, with $u_{t}$ denoting the probability density of $X_{t}$ at time $t$ (push-forward distribution of $X_{t}$ according to (14)). It is well-known that $u_{t}$ follows the continuity equation, or Liouville equation (Gibbs, 2014); or forward Kolmogorov equation in stochastic processes, but with zero noise (Risken, 1996),

The goal now is to adjust $\theta\in\mathcal{U}$ so that $\rho_{t}(\cdot|x)$ is close to $\rho(\cdot|x)$. To this end, we define a differentiable loss function $\Phi(\rho_{1},\rho_{2})$ that measures distances between two conditional densities $\rho_{1},\rho_{2}$ (e.g., $L^{2}$ loss, K-L divergence). Then, the learning problem can be formulated as the following optimal control problem:

As before, $L$ is a regularizer on the trainable parameters. Now, (16) is an optimal control problem on the function space $\mathcal{H}$.

Then, the Pontryagin’s maximum principle for this system is expected to take the form: let $\theta^{*}\in\mathcal{U}$ be an optimal control, then there exists a co-state process $v_{t}\in\mathcal{H}$ such that

where $D$ denotes the usual Fréchet derivative. Note that by definition, we have $D_{v}H=-\text{div}(fu)$ and $D_{u}H=f\cdot\nabla_{x}v$. Observe that the co-state $v^{*}$ satisfies the (time-reversed) adjoint Liouville’s equation with a specified terminal condition. The PMP for similar functional optimal control problems has been studied in, among others, Pogodaev (2016); Roy and Borzì (2017), albeit without the expectation over initial density.

In summary, the advantage of this formulation is that we make no explicit reference to the training data or target functions and formulate the entire problem as a control problem on probability densities. Of course, in practice, to implement an MSA-like algorithm, the terminal condition of the co-state will depend on the target joint density, which we can only access through the sampled data. A rigorous analysis of this function space control formulation and its consequences will be explored in future work.

B Proof of Lemma 2

First, observe that assumptions (A1)-(A2) in the main text implies that the second derivatives of $f$ and $\Phi$ are bounded by $K$. Provided that $P^{\theta}_{t}$ is bounded, they also imply that the second derivatives of $H$ with respect to $x$ and $p$ are bounded when evaluated on $X^{\theta}_{t},P^{\theta}_{t},\theta_{t}$. We first establish the boundedness of $P^{\theta}_{t}$.

Assume that (A1)-(A2) hold. Then, there exists a constant $K^{\prime}>0$ such that for any $\theta$,

Using (A1)-(A2), we have $\|P^{\theta}_{T}\|=\|\nabla_{x}\Phi(X^{\theta}_{T})\|\leq K$ and $\|\nabla_{x}f(t,X^{\theta}_{t},\theta_{t})\|_{2}\leq K$. Hence,

This proves the claim since it holds for any $\tau$.

We now prove Lemma 2. The approach here is similar to that employed in Rozonoer (1959).

Proof [Proof of Lemma 2] From (6) and the definition of the Hamiltonian, we have for any $\theta\in\mathcal{U}$,

Denote $\delta X_{t}=X^{\phi}_{t}-X^{\theta}_{t}$ and $\delta P_{t}=P^{\phi}_{t}-P^{\theta}_{t}$, then we have

where in the last line we defined $Z^{\theta}:=(X^{\theta},P^{\theta})$. Similarly, from (19) we get

where we have used Taylor’s theorem in the last step with $r_{1}(t)\in$. We now rewrite the boundary terms. Since $\delta X_{0}=0$, we have

for some $r_{2},r_{3}\in$. Lastly, for each $t\in[0,T]$ we have

Substituting (20), (21), (22), (23) into (17), we obtain

The left hand side is simply $J(\phi)-J(\theta)$, and so it remains to estimate the right hand side terms. First, let us estimate $\delta X$ and $\delta P$. By definition,

and hence using Lemma 6 and assumptions (A1)-(A2),

Now, we substitute estimates (26) and (28) into (24) and rename constants for simplicity. Note that by assumptions (A1)-(A2) and Lemma 6, all the second derivative terms are bounded element-wise by some constant $K^{\prime\prime}$. Hence, we have $|\delta Z_{t}\cdot A\cdot\delta Z_{t}|\leq K^{\prime\prime}\|\delta Z\|^{2}$ for each $A$ being a second derivative matrix. Thus we obtain

For applications, the global Lipschitz condition (A2) w.r.t. $x$ on $f$ may be restrictive. Note that this can be replaced by a local Lipschitz condition if we can show that $X_{t}$, $t\in[0,T]$ is bounded for all $\theta\in\mathcal{U}$. This is true if the parameter space $\Theta$ is bounded, which we can safely assume in practice, as long as a suitable regularization is used that prevents the parameters from getting arbitrarily large. Alternatively, a projection step can be used to restrict the parameters to a bounded set. In either cases, this should not negatively affect the performance of the model.

References