PaperlayerView exact revision on arXiv

Source-linked AI summary

A Mean-Field Optimal Control Formulation of Deep Learning

Weinan E, Jiequn Han, Qianxiao Li

arXiv:1807.01083v1math.OCcs.LG

TL;DR

The paper addresses the need for a systematic mathematical analysis of deep learning viewed as an optimal control problem. It formulates population risk minimization as mean-field optimal control, develops HJB and PMP characterizations, and relates population and empirical problems through quantitative error estimates.

  • Problem

    Prior dynamical-systems work on deep learning focused mainly on algorithms and network structures, leaving the formulation's fundamental mathematical aspects insufficiently explored.

  • Method

    The paper formulates continuous-time population risk minimization as mean-field optimal control and analyzes it through dynamic programming and Pontryagin maximum-principle approaches.

  • Results

    The paper derives an infinite-dimensional HJB characterization, proves mean-field PMP conditions, and establishes that stable mean-field PMP solutions have nearby sampled solutions with explicit error estimates under appropriate conditions.

  • Takeaways & Limitations

    The framework provides a mathematical foundation for studying algorithmic and theoretical connections between optimal control and deep learning, including population-to-empirical relationships.

  • Takeaways & Limitations

    The HJB approach is generally difficult to solve numerically, and its global characterization requires solving over the entire space of input-target distributions.

Abstract

from arXiv · show

Recent work linking deep neural networks and dynamical systems opened up new avenues to analyze deep learning. In particular, it is observed that new insights can be obtained by recasting deep learning as an optimal control problem on difference or differential equations. However, the mathematical aspects of such a formulation have not been systematically explored. This paper introduces the mathematical formulation of the population risk minimization problem in deep learning as a mean-field optimal control problem. Mirroring the development of classical optimal control, we state and prove optimality conditions of both the Hamilton-Jacobi-Bellman type and the Pontryagin type. These mean-field results reflect the probabilistic nature of the learning problem. In addition, by appealing to the mean-field Pontryagin's maximum principle, we establish some quantitative relationships between population and empirical learning problems. This serves to establish a mathematical foundation for investigating the algorithmic and theoretical connections between optimal control and deep learning.

1 Introduction

This paper develops a mathematical study of deep learning by formulating population risk minimization as a mean-field optimal control problem, analyzed through dynamic programming and maximum-principle methods.

  • Deep learning is widely used in tasks such as image classification and segmentation, motivating a solid mathematical framework for analyzing deep neural networks.
  • ResNet can be viewed as a time-discretization of a continuous-time dynamical system, allowing learning to be recast as an optimal control problem.
  • The paper studies population risk minimization as mean-field optimal control, where trainable weights depend on the population distribution of input-target pairs.
  • The analysis derives an infinite-dimensional HJB characterization and develops a mean-field Pontryagin maximum principle for optimal controls.
  • Under appropriate conditions, stable mean-field PMP solutions have nearby sampled-PMP solutions with high probability and explicit error estimates.
  • The paper uses convenient reasonable assumptions rather than pursuing the sharpest estimates under the most general conditions.

2 Related work

Related work connects deep learning with dynamical systems, optimal control, and mean-field control, while distinguishing this paper’s population-coupled learning formulation from prior control models.

  • Earlier work connected back-propagation and optimal control, while later studies used dynamical systems for training algorithms, continuum limits, and network architecture design.
  • Mean-field control research includes mean-field games, McKean-Vlasov systems, Cucker-Smale systems, and stochastic maximum principles.
  • In this paper, randomness comes from coupled input-target pairs, while the dynamics conditional on each pair are otherwise deterministic.
  • Unlike McKean-Vlasov control and mean-field games, each realization’s dynamics do not directly depend on the population distribution; realizations are coupled through shared controls.

3 From ResNets to mean-field optimal control

The paper formulates continuous-time residual-network learning as a stochastic ODE control problem whose objective is population risk over random input-target pairs.

  • A T-layer residual network is represented by difference equations, with each layer updating the state using trainable parameters.
  • The input is propagated to a final output, which is compared with its target through a loss function while trainable parameters are tuned.
  • Continuous-time residual learning replaces the difference equation with a differential equation, yielding the idealized model studied in the paper.
  • The initial input-target pair is random with joint law µ0, while admissible training weights belong to a possibly constrained control set Θ.
  • The resulting population risk minimization problem is posed as a mean-field optimal control problem.
  • The mean-field designation reflects shared trainable controls whose optimization depends on the law of input-target variables, despite the forward dynamics not explicitly depending on that law.
  • Empirical risk minimization instead samples independent input-target pairs, producing sampled control problems whose solutions are typically random variables.

Additional Notation

The notation introduces extended state, loss, probability-law, path-space, and dynamical-flow conventions used to define the mean-field control problem.

  • The variable w denotes the concatenated input-target variable (x,y), with corresponding extended feed-forward and regularization functions.
  • The paper uses Fréchet derivative notation for mappings between Banach spaces and negative-semidefinite notation for matrices.
  • L∞([0,T],E) denotes essentially bounded measurable paths, including the control paths used throughout the analysis.
  • L2 random variables correspond to square-integrable probability measures in P2(Rd+l), with norms defined through expected squared Euclidean magnitude.
  • The notation W^{t,ξ,θ} describes the dynamical system evolving from a random variable ξ under a control process θ.
  • The flow notation expresses the semigroup property of the dynamical evolution of probability distributions.
  • Generic constants may depend implicitly on the time horizon and ambient dimensions, even when this dependence is omitted.

4 Mean-field dynamic programming principle and HJB equation

The paper formulates deep-learning population risk minimization as mean-field optimal control with probability measures as states. Dynamic programming yields a value function and HJB equation that characterize optimal loss and feedback controls, while the value function is Lipschitz under the stated assumptions.

  • Mean-field formulation: The value function is defined on time and probability measures, extending the deep-learning control problem to a Wasserstein-space state.The formulation treats µ ∈ P2(Rd+l) as the state and J(t, µ, θ) as the objective.
  • Value-function properties: Under boundedness and Lipschitz assumptions, both the objective functional and value function are Lipschitz continuous in time and the Wasserstein state.The assumptions include bounded f, L, and Φ, with Lipschitz continuity in x and an initial distribution in P2(Rd+l).
  • Dynamic programming principle: The dynamic programming principle states that the remaining trajectory from any intermediate time and state is again optimal.This recursive relation is the basis for deriving the HJB equation.
  • Mean-field HJB equation: The probability-measure derivative is constructed by lifting functions to L2(Ω, Rd+l), enabling a chain rule and an infinite-dimensional HJB equation.The lifted derivative depends on the law of the random variable rather than its particular representation.
  • Optimal control characterization: A smooth HJB solution attaining the Hamiltonian infimum equals the value function and supplies an optimal feedback control policy.An open-loop solution can then be generated sequentially from the feedback policy along the feed-forward ODE.
  • Scope and computational boundary: The feedback characterization depends on solving the HJB globally over the space of input-target distributions, making numerical solution difficult for nontrivial learning problems.Local solutions around trajectories may nevertheless yield policies applicable to nearby input-label distributions.

5 Viscosity solution of HJB equation

Because smooth HJB solutions are generally unavailable, the paper uses viscosity solutions in a lifted Hilbert-space formulation. It proves that the value function exists as a viscosity solution and is uniquely identified by comparison.

  • Motivation: Smooth solutions to the HJB equation cannot generally be expected, motivating a viscosity-solution formulation.The approach extends the classical PDE notion to the Wasserstein space of probability measures.
  • Definition: The paper defines viscosity solutions on probability measures by lifting them to functions on L2(Ω, Rd+l) and applying the lifted Bellman equation.Subsolution and supersolution conditions are imposed using smooth test functions and terminal inequalities.
  • Existence: The value function v∗(t, µ) is a viscosity solution of the mean-field HJB equation.The proof establishes both subsolution and supersolution properties.
  • Analytical ingredients: The Hamiltonian’s continuity properties support the viscosity analysis in the lifted space.The proof uses continuity of H together with dynamic programming and local test-function arguments.
  • Uniqueness: Any viscosity subsolution lies below any viscosity supersolution, so the value function is the unique viscosity solution.This comparison result establishes well-posedness in the viscosity sense.
  • Implications: The HJB characterization can identify an optimal control by solving the Hamiltonian infimum, while a unique Hamiltonian minimizer yields a unique open-loop control.The latter property is illustrated for a residual-network example with quadratic parameter regularization.

6 Mean-field Pontryagin’s maximum principle

The paper develops a mean-field Pontryagin maximum principle for population-risk deep-learning control problems, yielding forward-backward optimality conditions with an expectation over the input-target distribution. It relates this local characterization to HJB theory and shows that the mean-field principle also covers sampled control problems.

  • Mean-field PMP: Because one control is shared across input-target pairs, the PMP imposes Hamiltonian maximization in expectation over the population distribution.This expectation distinguishes the mean-field condition from the classical PMP formulation.
  • Mean-field PMP: The mean-field PMP gives necessary conditions for optimal controls through forward state dynamics, backward co-state dynamics, and Hamiltonian maximization.Under assumptions (A1′) and (A2′), an optimal control admits absolutely continuous stochastic state and co-state processes.
  • Interpretation: The co-state acts analogously to a Lagrange multiplier for the forward ODE constraint, while the Hamiltonian condition globally maximizes over training weights.This formulation remains applicable when dynamics are not differentiable in controls or optimal weights lie on the boundary of Θ.
  • Scope: The PMP is weaker than HJB because it supplies necessary conditions for an open-loop control, without explicitly characterizing an optimal control policy.Uniqueness can make the PMP sufficient under additional conditions discussed in the following section.
  • Population and sampled problems: The mean-field PMP includes sampled optimal-control necessary conditions as a special case by substituting the empirical measure for the population measure.The resulting equations are random because the empirical measure is random.
  • Connection to HJB: Hamilton’s equations from the PMP are characteristics of the HJB equation, so the PMP specifies the necessary condition along a characteristic originating from the initial distribution.This explains why PMP is a local optimality characterization compared with the HJB equation’s global characterization.

7 Small-time uniqueness

The paper studies when mean-field PMP solutions are unique and therefore sufficient for optimality. It proves uniqueness for sufficiently short time horizons under strong concavity of the Hamiltonian in the control, while relating short horizons to limited reachable dynamics.

  • Problem setting: The PMP system is a nonlinear two-point boundary-value problem coupled through the laws of the state and co-state, so uniqueness is not automatic.This remains true even when the Hamiltonian maximization has a unique solution for each state-co-state law.
  • Uniqueness result: Under stronger smoothness assumptions and strong concavity of H(x, p, θ) in θ, the paper establishes uniqueness of PMP solutions for sufficiently small T.The proof combines flow-map difference estimates with a contraction argument.
  • Proof strategy: The uniqueness proof compares trajectories and co-states generated by two controls, bounds their differences, and concludes equality when the horizon is short enough.The final step uses KLC(T) < λ0 to force the control difference to vanish.
  • Deep-learning interpretation: In deep learning, small T roughly corresponds to a small reachable set and can be interpreted as low model capacity or expressive power.The number of parameters remains infinite because controls are only required to be essentially bounded and measurable in time.
  • Deep-learning interpretation: Strong concavity of the Hamiltonian does not imply that the loss J is convex, because nonlinear state transformations can make J highly non-convex.Thus, PMP uniqueness is obtained under a short-time condition rather than from convexity of the learning objective.

8 From mean-field PMP to sampled PMP

The paper connects mean-field and sampled Pontryagin systems by treating their interior solutions as zeros of Banach-space mappings, establishing high-probability proximity and loss-value error estimates under stability assumptions.

  • 8 From mean-field PMP to sampled PMP: The sampled PMP is introduced because solving the mean-field PMP requires maximizing an expectation in practice.The sampled system provides necessary conditions for the empirical optimal control problem.
  • 8 From mean-field PMP to sampled PMP: Mean-field and sampled PMP solutions are represented as zeros of mappings F and F_N, with F_N converging almost surely to F by the law of large numbers.The approximation problem is therefore reduced to approximating zeros of Banach-space mappings.
  • 8.1 Excursion: random approximations of zeros of Banach space mappings: Stability of F locally ensures uniqueness of a nearby zero and requires non-singular derivatives; Lipschitz continuity of DF provides a sufficient condition for stability.These properties support transferring solutions from the population system to its sampled approximation.
  • 8.2 Error estimate for sampled PMP: Under stability, Theorem 6 provides a sampled PMP solution θ_N near a mean-field solution θ* with high-probability convergence.The result is obtained by applying the random-zero approximation theory to the PMP mappings.
  • 8.2 Error estimate for sampled PMP: Under local strong concavity, θ_N is a strict local maximum of the sampled Hamiltonian with probability at least 1 − 6 exp[−(Nλ_0^2)/(K_1 + K_2λ_0)].If the finite-sampled Hamiltonian has a unique local maximizer, θ_N is a sampled PMP solution with the same high probability.
  • 8.2 Error estimate for sampled PMP: The resulting loss-function convergence quantifies a difference between population and empirical risk-minimization solutions, providing an a priori estimate relevant to generalization.The concavity assumption is local along paths involving θ*, whereas stronger assumptions may be global.

9 Conclusion

The paper formulates continuous-time deep learning population-risk minimization as mean-field optimal control and analyzes it through HJB and PMP perspectives. It also relates batch normalization to mean-field dynamics while identifying distribution-independent state dynamics as a scope boundary.

  • 9 Conclusion: The framework models continuous-time deep learning population-risk minimization as mean-field optimal control, explicitly representing network composition through time evolution.The state variable in the HJB formulation is the joint distribution of input-target pairs.
  • 9 Conclusion: The dynamic-programming approach derives an infinite-dimensional HJB equation whose viscosity solution characterizes the optimal loss value and an optimal feedback policy.The equation is posed with probability-measure state variables.
  • 9 Conclusion: The maximum-principle approach proves a mean-field PMP giving necessary conditions for optimality and a local characterization of optimal trajectories.This complements the global value-function characterization supplied by dynamic programming.
  • 9 Conclusion: The analysis assumes that state dynamics depend on the state and control but not on the distribution law of the state.More complex structures such as batch normalization therefore motivate extensions of the framework.
  • 9 Conclusion: Batch normalization can be represented using population state distributions and incorporated into the parameter vector, making it a general form of mean-field dynamics.The paper leaves further theoretical implications of this connection for future work.
  • 9 Conclusion: The framework offers a mathematical basis for developing further approaches to practical and theoretical deep-learning problems.This is stated as the paper’s overall supported consequence.
Loading 1807.01083v1…