A.2 Applications of the Calculus of Variations

Consider a possibly nonlinear dynamical system, ${\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{\dot{x}}}=\mathbf{f}({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}},t),$ that accrues a running cost at the rate $\ell({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{x}}(t),t)$. Let us define the (augmented) Lagrangian to be the sum of this running cost with Lagrange multipliers enforcing the constraint that the system obey the dynamics:

with $\bm{\lambda}$ the Lagrange multipliers. Notice that this setup differs conceptually from that of the mechanical systems in the previous section, in which the running cost, rather than the constraint, enforces the system dynamics. However, this poses no mathematical obstacle: Integrating this Lagrangian over time gives the total cost, which is a functional of the function $\bm{x}(t)$. Therefore the optimal trajectory is obtainable with the calculus of variations.

We will reformulate this in terms of a Hamiltonian. Applying the definition of the conjugate variables (Eq. A.3) to this Lagrangian, we obtain

That is, the conjugate variables are the Lagrange multipliers. In this context, they are known as the costate^††margin: costate . Next, we form the Hamiltonian according to the recipe in Eq. A.7:

Finally, the optimal trajectories of the states $\bm{x}^{*}(t)$ and costates $\bm{p}^{*}(t)$ must obey the ELEs—in terms of the Hamiltonian, Eq. A.8. Applying it to Eq. A.9, we obtain

The first equation tells us nothing new, but the second tell us how the costate evolves with time. Typically, we are given an initial state (which we can take to be $\bm{0}$ w.l.o.g.), in which case the state dynamics must be integrated forward in time. If we have a terminal cost, then the transversality condition from Eq. A.3 provides the final costate, and therefore the costate dynamics must be integrated backwards in time.^††margin: A word on numerical solution?

[[[In fact, if you insist on using the sign convention for the Hamiltonian from classical mechanics, as you have below, you get the maximum principle: we need to maximize, rather than minimize, the Hamiltonian.]]] We can generalize the foregoing construction along different avenues. One possibility is to let the dynamics—and the running cost—depend on a time-varying control input, $\bm{u}(t)$:

This is an optimal control problem: we seek the controls that steer the system through some desirable trajectory, also perhaps subject to some constraint on the controls (e.g., that they not be too large). The development is essentially identical to that of the method of adjoints just given. The only difference is that we derive an additional PDE, for the controls. It can be treated just like the state and in general could generate a matching PDE. However, we typically do not expect the cost or dynamics to depend directly on the time derivative of the controls, $\dot{\bm{u}}$, so the Euler-Lagrange equations (Eq. A.1) for the controls reduce to

The fact the Lagrangian and Hamiltonian gradients are equal at the optimal control follows from the recipe for constructing Hamiltonians, Eq. A.7.

How do we actually use these equations to solve for optimal controls? The entire optimal system is now governed by three PDEs: Eq. A.10 (modified to include the controls) and Eq. A.11. ….^††margin: FINISH THIS SECTION

Another possible avenue for generalization is to consider optimizing for fixed (time-invariant) parameters $\bm{\theta}$ (rather than time-varying controls). In this framework, the “neural ODEs” of Chen and colleagues [6] can be³³ 3 Those authors took a slightly different approach. derived. The authors consider a certain continuous-time limit of recurrent neural networks (or continuous-layer limit of residual networks), in which forward propagation becomes a nonlinear dynamical system,

What does backpropagation become?

We can use the same Lagrangian as in the method of adjoints above, so the state and costate dynamics are again given by Eq. A.10. The terminal cost $\mathcal{S}$ can be identified with the loss function, applied at the final layer/time, although we can impose costs at other layers/times through the running cost $\ell$. We also have in hand boundary values: the initial state is the input to the network (or some transformation thereof), and the final costate is the loss gradient evaluated at the network outputs (from the transversality condition, Eq. A.2):

Therefore it is possible to compute, or at least numerically simulate, the trajectories of the state and of the costate, by first a forward pass through the state dynamics, followed by a backward pass with the costate (adjoint) dynamics. Comparing Eq. A.10 to Eqs. 7.19 and 7.21, we see that we have derived precisely the continuous-time analogs of the forward and backward passes required for backpropagation through time.

To improve the model, we need the gradient $\mathrm{d}{J}/\mathrm{d}{\bm{\theta}}$ of the (total) loss. Here we make use of an envelope theorem (cf. Eq. A.4). Along optimal trajectories of the states $\bm{x}$ and their costates $\bm{p}$,^††margin: We leave the extension of the envelope theorem to the costate as an exercise for the reader: Exercise LABEL:ex:We_leave_the_extension_of_the_envelope_theorem_to_the_costate_as_an_exercise_for_the_reader the loss function’s total derivative is just its partial derivative:

This equation provides a simple formula for optimization of the parameters. Indeed, it can be numerically integrated, also in the backwards direction, along with $\bm{p}$.

Notice that there is no need to backpropagate gradients through the ODE solver. The authors also suggest re-solving for the states $\bm{x}$ in the backward pass to avoid having to store them during the forward pass. (A forward pass is necessary in any case in order to initialize the final state.)

Consider an exponential family of probability distributions,

It is well known (and easily shown) that the first and second derivatives of the log-partition function $A({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{\eta}})$ are (respectively) the expectation and covariance of the sufficient statistics:

Since covariance matrices are positive definite^††margin: What about degenerate cases with only positive semi-definite? , the first derivative provides an invertible map between the natural parameters and the moment parameters:

and the log-partition function $A$ is itself convex in the natural parameters.

This suggests a certain analogy between the Lagrangian and the log-partition function, with the natural parameters playing the role of velocities and the moment parameters like momenta. Let us take the analogy seriously and carry out a Legendre transformation, exactly on analogy with Eq. A.7

Notice that, as in Eq. A.8,

However, we can also begin from a different point. Let us simply define the “Hamiltonian” to be the relative entropy of the distribution and its base measure:⁴⁴ 4 The base measure need not integrate to one, i.e. be a proper probability density, so the minimum of the Hamiltonian need not be 0.

This yields the same result as the Legendre transformation, and consequently Eqs. A.13 and A.14 still hold.