Source-linked AI summary
Maximum Principle Based Algorithms for Deep Learning
Qianxiao Li, Long Chen, Cheng Tai, Weinan E
TL;DR
Deep learning training is recast as continuous-time optimal control to address limitations of gradient-based optimization. The paper derives PMP conditions and modifies successive approximations into E-MSA, which converges favorably early in training, while efficient Hamiltonian maximization remains the main bottleneck.
Problem
Gradient-based training can converge slowly initially, stall on flat landscapes, and cannot directly handle discrete trainable variables, motivating alternative training methods.
Method
The paper models deep residual networks as discretized dynamical systems, formulates learning as optimal control, and modifies successive approximations to solve a weaker sufficient condition derived from PMP.
Results
E-MSA performs favorably on a per-iteration basis, especially early in training, compared with SGD, Adagrad, and related gradient-based approaches when Hamiltonian maximization is sufficiently effective.
Takeaways & Limitations
The approach provides convergence results and explicit error estimates while offering a possible route beyond gradient descent for flat landscapes and discrete parameter spaces.
Takeaways & Limitations
For discrete-time systems, PMP and the monotonicity implications are expected only approximately under appropriate conditions, and rigorous analysis is left for future work.
Abstract
from arXiv · showhide
The continuous dynamical system approach to deep learning is explored in order to devise alternative frameworks for training algorithms. Training is recast as a control problem and this allows us to formulate necessary optimality conditions in continuous time using the Pontryagin's maximum principle (PMP). A modification of the method of successive approximations is then used to solve the PMP, giving rise to an alternative training algorithm for deep learning. This approach has the advantage that rigorous error estimates and convergence results can be established. We also show that it may avoid some pitfalls of gradient-based methods, such as slow convergence on flat landscapes near saddle points. Furthermore, we demonstrate that it obtains favorable initial convergence rate per-iteration, provided Hamiltonian maximization can be efficiently carried out - a step which is still in need of improvement. Overall, the approach opens up new avenues to attack problems associated with deep learning, such as trapping in slow manifolds and inapplicability of gradient-based methods for discrete trainable variables.
1. Introduction
Efficient training of very deep networks remains challenging because gradient-based updates can be slow and exploit limited structural information. The paper develops an optimal-control alternative based on Pontryagin’s maximum principle, with fast initial descent and explicit error control.
- Gradient-based training methods are efficient to implement but can make slow parameter updates, especially early in training.
- The optimal-control viewpoint formulates deep learning through Hamiltonian dynamics and layerwise Hamiltonian maximization.This formulation also provides an explicit error-control estimate.
- The proposed algorithms replace deep neural networks with continuous or discretized dynamical systems and use PMP optimality conditions for training.
- The approach is reported to provide fast initial descent, resilience to stalling in flat landscapes, and explicit control of error.
2. Function Approximation by Dynamical Systems
The paper casts supervised learning as an optimal-control problem for dynamical systems, deriving continuous-time PMP conditions and numerical algorithms that can be discretized. This framework connects network states, controls, losses, and co-states while retaining important scope limitations.
- Function Approximation by Dynamical Systems: Supervised learning is represented by an ODE whose initial state contains inputs, control parameters represent trainable weights, and terminal states produce predictions.
- Function Approximation by Dynamical Systems: The loss compares terminal predictions with labels, while the dynamics are decoupled across samples except through the shared control.
- Function Approximation by Dynamical Systems: Optimizing in continuous time before discretizing permits continuous-time optimality analysis and subsequent numerical algorithms, unlike the traditional reverse order.
- Function Approximation by Dynamical Systems: Deep residual networks can be viewed as forward Euler discretizations of the continuous dynamical-system formulation.
- Pontryagin’s Maximum Principle: PMP supplies Hamilton equations and a Hamiltonian maximization condition, with the co-state acting as a continuous-time analogue of a Lagrange multiplier.
- Pontryagin’s Maximum Principle: Unlike a stationarity condition, PMP requires global Hamiltonian maximization and remains applicable when parameters are discrete or dynamics are non-smooth in parameters.
- Pontryagin’s Maximum Principle: PMP is only a necessary condition, so solving it need not produce a globally optimal solution unless additional conditions make it sufficient.
3. Method of Successive Approximations
The method of successive approximations solves the PMP through alternating Hamiltonian propagation and parameter maximization, then extends this procedure to control feasibility errors and guarantee convergence under stated assumptions.
- Continuous-Time Strategy: Continuous-time algorithms solve the PMP while making error estimates and convergence behavior more transparent.The analysis is performed entirely in continuous time.
- Basic MSA: MSA alternates forward-backward Hamiltonian propagation with pointwise Hamiltonian maximization over the control.The maximization step is decoupled across time and, in deep-learning terms, across network layers.
- Basic MSA: The maximization step can be parallelized, although it is typically the most time-consuming part of the algorithm.
- Basic MSA: Basic MSA may diverge in general, particularly from a poor initial control, despite convergence for a restricted class of linear-quadratic regulators.
- Error Estimate for the Basic MSA: Lemma 2 identifies Hamiltonian maximization as an optimal descent direction while showing that feasibility errors can nullify descent and cause divergence.These errors measure violations introduced into the Hamiltonian dynamics when replacing one control by another.
- Extended PMP and Extended MSA: The extended method adds augmented-Hamiltonian penalties for feasibility errors, producing a weaker necessary condition whose maximization step targets extended PMP solutions.The Hamiltonian dynamics remain unchanged because the added terms do not affect optimal solutions.
- Extended PMP and Extended MSA: For sufficiently large ρ, extended MSA converges to the set of solutions of the extended PMP from any admissible initial control with finite objective.The convergence theorem also assumes the objective is bounded below over the control set.
4. Discrete-Time Formulation
The paper discretizes PMP, MSA, and E-MSA, connecting the resulting control formulation to residual networks and back-propagation. The discrete-time E-MSA offers a parallelizable training procedure, although rigorous discrete-time guarantees require further analysis.
- 4.1 Discrete-Time PMP and Discrete-Time MSA: The discrete formulation applies Euler discretization with step size δ = T/N to the continuous-time state dynamics.The discretized variables are x_n := X_nδ, ϑ_n := θ_nδ, and f_n(·) := f(nδ, ·).
- 4.1 Discrete-Time PMP and Discrete-Time MSA: x_n+1 = x_n + δf_n(x_n, ϑ_n) defines the discrete-time analogue of the control problem.The initial condition is x_0 = x, with 0 ≤ n ≤ N − 1.
- 4.1 Discrete-Time PMP and Discrete-Time MSA: Apart from the constant δ, the discrete dynamics exactly match the supervised-learning formulation of deep residual networks.This motivates using E-MSA to train residual neural networks through the extended PMP.
- 4.1 Discrete-Time PMP and Discrete-Time MSA: The discrete PMP is formulated using g_n(x_n, ϑ_n) := x_n + δf_n(x_n, ϑ_n) and a scaled discrete Hamiltonian.The passage introduces the Hamiltonian construction before stating the discrete-time PMP conditions.
- 4.1 Discrete-Time PMP and Discrete-Time MSA: With ReLU activations, the discretization factor δ can be absorbed into the trainable parameters ϑ.This makes the residual-network interpretation compatible with the parameterization used in the discrete formulation.
- 4.1 Discrete-Time PMP and Discrete-Time MSA: Discrete-time PMP validity is delicate: known counterexamples exist, while approximate validity and E-MSA monotonicity are expected only under suitable small-step conditions.A rigorous discrete-time analysis is left for future work.
- 4.1 Discrete-Time PMP and Discrete-Time MSA: For residual networks, discrete-time E-MSA corresponds to forward Euler state updates and backward Euler co-state updates.The Hamiltonian maximization step is decoupled across layers and can therefore be performed in parallel.
- 4.1 Discrete-Time PMP and Discrete-Time MSA: Algorithm 3 summarizes discrete-time E-MSA, beginning by initializing ϑ_n^0 ∈ Θ_n for n = 0, . . . , N − 1 and a hyper-parameter ρ.The algorithm is presented as the discrete-time counterpart of the extended method of successive approximations.
14 end
The paper relates softened Hamiltonian maximization in MSA to gradient descent with back-propagation, while retaining a broader PMP framework and layerwise optimization. It also describes how propagation and Hamiltonian maximization adapt to mini-batches, with rigorous mini-batch error estimates left open.
- Relationship to Gradient Descent with Back-propagation: Softening MSA’s Hamiltonian maximization step yields an update equivalent to gradient descent with back-propagation in discrete time.The equivalence is stated for basic MSA with ρ = 0 when step 12 is replaced by the softened update.
- Relationship to Gradient Descent with Back-propagation: Gradient descent with back-propagation is therefore a modification of basic MSA that replaces Hamiltonian maximization with gradient ascent.
- Relationship to Gradient Descent with Back-propagation: The PMP framework remains applicable without differentiability in trainable parameters, separating layerwise optimization from information propagation through the co-state equation.The paper states that propagation is independent of optimization and that layers can be optimized separately with or without gradient information.
- Mini-batch Algorithms: In mini-batch training, propagation steps decouple across samples, while the maximization step uses a mini-batch Hamiltonian aggregate.For sufficiently large independently and identically drawn batches, the mini-batch Hamiltonian mean converges uniformly to the full-batch sum.
- Mini-batch Algorithms: Rigorous error estimates for the mini-batch version remain outside the paper’s scope, so mini-batch feasibility is demonstrated numerically rather than theoretically.
5. Numerical Experiments
The experiments compare E-MSA with gradient-based methods on sine-function approximation and MNIST classification. E-MSA shows favorable early convergence per iteration and can escape some stalled regimes, but its Hamiltonian maximization makes wall-clock performance uncompetitive in the reported MNIST setup.
- Sine-function approximation: E-MSA shows favorable convergence rate per iteration on continuous sine-function approximation compared with gradient-based optimization.The experiment uses a continuous 5-dimensional dynamical system with 1000 training and 1000 test samples.
- Sine-function approximation: With all weights and biases initialized to zero, gradient-based methods tend to become stuck, whereas E-MSA better escapes slow manifolds when ρ is well chosen.Figure 1 reports ρ = 1.0 for the illustrated successful E-MSA behavior.
- MNIST: On MNIST, E-MSA has good performance per iteration, especially during the early stages of training, relative to SGD, Adagrad, and Adam.The setup uses 55,000 training samples, 10,000 test samples, and mini-batches of 100.
- MNIST: E-MSA’s wall-clock performance on MNIST is not currently competitive because Hamiltonian maximization is time consuming and outweighs its per-iteration gains.Wall-clock comparisons were performed on CPUs because the L-BFGS implementation lacked GPU support.
- MNIST: Layer decoupling makes the L-BFGS Hamiltonian maximization step tractable because it involves fewer parameters than directly minimizing the objective.
6. Discussion and Related Work
E-MSA differs from gradient-based training by using Hamiltonian maximization, yielding favorable early per-iteration behavior while incurring additional computational cost. Its broader competitiveness depends on making Hamiltonian maximization more efficient, and the paper situates the approach within optimal-control and deep-learning research.
- Distinguishing features: E-MSA does not require parameter gradients and can decouple optimization across layers through the Pontryagin maximum principle.The framework may therefore support discrete trainable parameters and potentially parallelize layerwise optimization.
- Empirical comparisons: E-MSA performs favorably per iteration, especially during early training, on MNIST and Fashion-MNIST residual CNN experiments.The reported comparisons use train and test loss and accuracy, with hyper-parameters tuned separately for the experiments.
- Computational efficiency: Efficient Hamiltonian maximization remains the key requirement for E-MSA to outperform gradient-based methods in wall-clock efficiency.The Hamiltonian problem remains non-convex, but layer decoupling reduces parameter-space size and exact maximization is unnecessary.
- Computational efficiency: Currently, E-MSA’s per-iteration gains are outweighed by additional wall-clock costs on the MNIST residual CNN.The experiments used a non-GPU L-BFGS implementation for augmented-Hamiltonian maximization, and the authors expect improved timing with more efficient implementations.
- Related work: The paper extends optimal-control ideas into deep-learning training, contrasting PMP-based optimization with gradient descent, ADMM, and related layer-decoupling approaches.It presents the work as an early translation of numerical PMP algorithms into deep-learning training beyond gradient descent.
7. Conclusion and Outlook
The paper views residual neural networks as discretized continuous-time dynamical systems and formulates supervised learning as an optimal-control problem. It develops and proves convergence for E-MSA, which shows favorable early per-iteration performance, while identifying computational efficiency and broader applicability as open issues.
- Conclusion: Deep residual neural networks can be viewed as discretizations of continuous-time dynamical systems, making supervised learning a continuous-time optimal-control problem.The paper uses this connection to motivate PMP-based training algorithms.
- Conclusion: E-MSA modifies successive approximations to solve a weaker sufficient condition derived from the extended PMP and has a proved convergence result.The resulting algorithm is tested on benchmark problems against SGD, Adagrad, and Adam.
- Conclusion: E-MSA performs favorably per iteration, particularly at early training stages, compared with SGD, Adagrad, and Adam.This conclusion is reported across the paper’s benchmark experiments rather than as a single dataset result.
- Outlook: Future work targets more efficient Hamiltonian maximization, adaptive regularization, discrete-time analysis, and applications to networks with discrete weights.The paper also notes that PMP validity and convergence for general deep neural networks remain unclear.
Appendix A. Function Space Formulation
The appendix formulates supervised learning over conditional probability densities as a function-space optimal-control problem. A controllable dynamical system evolves an initial conditional-density approximation, and PMP supplies formal state, co-state, and Hamiltonian optimality conditions.
- Function-space formulation: The function-space formulation controls a continuity equation describing the evolution of probability densities rather than directly referencing a fixed set of input-output pairs.The formulation is presented as an alternative, non-rigorous optimal-control view of supervised learning.
- Density representation: A target joint density is represented as ρ(x,y)=ρ(y|x)ρ(x), with deterministic labels expressed using a point-mass conditional density.The conditional-density formulation accommodates both deterministic and non-deterministic labels.
- Density evolution: The model begins with an initial conditional-density guess and evolves it through a controllable dynamical system to obtain ρ_t(y|x).The evolved approximation is defined using the pushforward density of the dynamical-system state.
- Optimization objective: A differentiable loss compares the evolved and target conditional densities, while regularization on trainable parameters completes the function-space optimal-control problem.Examples of the loss include L2 loss and K-L divergence.
- PMP conditions: The formal PMP introduces a co-state satisfying a time-reversed adjoint Liouville equation and requires Hamiltonian optimality for the control.The Hamiltonian is defined as a functional over time, state, co-state, and parameters.
- Scope: Practical MSA-like implementation still depends on sampled data because the co-state terminal condition depends on the target joint density.A rigorous analysis of this function-space formulation is left for future work.
Appendix B. Proof of Lemma 2
The appendix proves the loss-decrement result behind approximate Hamiltonian maximization by bounding state and co-state perturbations under regularity assumptions. The analysis relies on Gronwall-type estimates and requires conditions such as bounded derivatives or controlled parameter trajectories.
- Loss estimate: Taylor expansions and bounded second derivatives reduce the loss-difference estimate to quadratic perturbation terms controlled by a constant K′′.The proof bounds terms of the form |δZ_t · A · δZ_t| by K′′∥δZ∥².
- Perturbation bounds: Under assumptions (A1)–(A2), the state and co-state perturbations admit uniform bounds over the time interval.The proof obtains these estimates using Gronwall’s inequality.
- Proof strategy: The proof combines Hamiltonian identities, integration by parts, and perturbation estimates to establish the claimed loss-decrement relation.The argument begins from the Hamiltonian definition and rewrites boundary terms before applying the bounds.
- Assumptions: The regularity analysis may replace global Lipschitz continuity with local Lipschitz continuity when trajectories remain bounded.Bounded parameter spaces, suitable regularization, or projection steps can enforce the needed boundedness.