Lagrangian Neural Network
A Lagrangian neural network learns a scalar Lagrangian $L(q, \dot{q})$ instead of the dynamics themselves. The dynamics are then whatever that Lagrangian implies, through the Euler–Lagrange equations
\[ \frac{d}{dt}\frac{\partial{}L}{\partial\dot{q}} = \frac{\partial{}L}{\partial{}q}.\]
This is the Lagrangian counterpart of the Hamiltonian neural network: there the network is a Hamiltonian and the dynamics come from a Poisson tensor, here the network is a Lagrangian and the dynamics come from a variational principle.
The loss
Expanding the Euler–Lagrange equations along a trajectory gives
\[ \nabla_{\dot{q}}\nabla_{\dot{q}}L\,\ddot{q} + (\nabla_q\nabla_{\dot{q}}L)^T\dot{q} = \nabla_qL,\]
so the acceleration the learned Lagrangian predicts is
\[ \ddot{q} = (\nabla_{\dot{q}}\nabla_{\dot{q}}L)^{-1} \left( \nabla_qL - (\nabla_q\nabla_{\dot{q}}L)^T\dot{q} \right),\]
which is what GeometricMachineLearning.lagrangian_acceleration computes and what GeometricMachineLearning.LNNLoss compares against the data.
The transpose is not cosmetic. $\nabla_q\nabla_{\dot{q}}L$ is indexed $[i, j] = \partial^2L/\partial{}q_i\partial\dot{q}_j$, while the chain rule contracts the first index with $\dot{q}$. It disappears only for a Lagrangian whose position–velocity coupling happens to be symmetric, which is why the test for this asserts a closed form whose coupling deliberately is not.
Why the derivatives are symbolic
The loss needs the second derivative of the network with respect to its input, and then its own derivative with respect to the network parameters. Taking the inner one with Zygote inside the loss does not work: the parameter gradient then fails with MethodError: no method matching getindex(::IdDict{Any, Any}). The gradient and Hessian are therefore built as symbolic expressions with SymbolicNeuralNetworks and compiled, so that only one differentiation is left for Zygote to do. This is the same route hamiltonian_vector_field takes, and for the same reason.
Training on positions alone
GeometricMachineLearning.LNNLoss needs $\ddot{q}$ in the data. When only positions are available, GeometricMachineLearning.VariationalMidpointLoss trains the same architecture through the discrete Euler–Lagrange equations of the midpoint discrete Lagrangian
\[ L_d(q_n, q_{n+1}) = \Delta{}t\, L\!\left( \frac{q_n + q_{n+1}}{2}, \frac{q_{n+1} - q_n}{\Delta{}t} \right),\]
which ask that $D_2L_d(q_n, q_{n+1}) + D_1L_d(q_{n+1}, q_{n+2}) = 0$ along the data. Only the first derivative of the network is needed for this, because the chain rule through the midpoint is written out rather than taken with a nested Zygote.gradient.
Library Functions
GeometricMachineLearning.LNNLoss — Type
LNNLoss <: NetworkLossThe loss for a Lagrangian neural network.
The network output is a scalar Lagrangian $L(q, \dot{q})$. The loss compares the acceleration its Euler–Lagrange equations predict against the acceleration in the data:
\[ \mathtt{loss}(\mathcal{NN}, \mathtt{input}, \mathtt{output}) = ||\ddot{q}(\mathtt{input}) - \mathtt{output}|| \,/\, ||\mathtt{output}||,\]
where $\mathtt{input}$ stacks $q$ on $\dot{q}$ and is therefore $2n$ rows tall, and $\mathtt{output}$ is $\ddot{q}$ and is $n$ rows tall. See lagrangian_acceleration for $\ddot{q}$.
Constructor
This can be called with a LagrangianNeuralNetwork as its only argument:
LNNLoss(arch)Functor
loss(model, ps, input, output)
loss(ps, input, output) # equivalent to the aboveGeometricMachineLearning.lagrangian_acceleration — Function
lagrangian_acceleration(arch::LagrangianNeuralNetwork)Compute an executable expression of the acceleration that the Lagrangian arch predicts.
The Euler–Lagrange equations of a Lagrangian $L(q, \dot{q})$ are
\[ \frac{d}{dt}\frac{\partial{}L}{\partial\dot{q}} = \frac{\partial{}L}{\partial{}q},\]
which, expanded along a trajectory, reads
\[ \nabla_{\dot{q}}\nabla_{\dot{q}}L\,\ddot{q} + (\nabla_q\nabla_{\dot{q}}L)^T\dot{q} = \nabla_qL,\]
so the acceleration is
\[ \ddot{q} = (\nabla_{\dot{q}}\nabla_{\dot{q}}L)^{-1} \left( \nabla_qL - (\nabla_q\nabla_{\dot{q}}L)^T\dot{q} \right).\]
The transpose is not cosmetic. $\nabla_q\nabla_{\dot{q}}L$ is indexed $[i, j] = \partial^2L/\partial{}q_i\partial\dot{q}_j$, while the chain rule contracts the first index with $\dot{q}$. It disappears only for a Lagrangian whose position–velocity coupling happens to be symmetric.
Implementation
The gradient and the Hessian of the network with respect to its input are built as symbolic expressions with SymbolicNeuralNetworks.Jacobian — applied once, and then again to its own result — and compiled with build_nn_function. They are not taken with Zygote inside the loss: differentiating a loss that itself contains a Zygote.gradient with respect to the network parameters fails with MethodError: no method matching getindex(::IdDict{Any, Any}). This is the route hamiltonian_vector_field takes, and for the same reason.
The compiled functions flatten a batch along the columns — $(1, 2n\cdot{}B)$ for the gradient and $(2n, 2n\cdot{}B)$ for the Hessian — so both are reshaped before use. Trailing input dimensions beyond the first are flattened on the way in and restored on the way out, so an input of size $(2n, s, B)$, which is what Batch hands the loss, yields an acceleration of size $(n, s, B)$.
GeometricMachineLearning._euler_lagrange_acceleration — Function
_euler_lagrange_acceleration(H, g, q̇)Solve the Euler–Lagrange equations of one sample for the acceleration.
H is the full $2n\times{}2n$ Hessian of the Lagrangian with respect to its input, g the full $2n$ gradient, and q̇ the $n$ velocity components of that same input.
GeometricMachineLearning.VariationalMidpointLoss — Type
VariationalMidpointLoss <: NetworkLossThe loss that trains a Lagrangian neural network on positions alone, through the discrete Euler–Lagrange equations of the midpoint discrete Lagrangian.
It needs three consecutive positions and the timestep between them — no velocities and no accelerations.
The midpoint discrete Lagrangian is
\[ L_d(q_n, q_{n+1}) = \Delta{}t\, L\!\left( \frac{q_n + q_{n+1}}{2}, \frac{q_{n+1} - q_n}{\Delta{}t} \right),\]
and the discrete Euler–Lagrange equations ask that
\[ D_2L_d(q_n, q_{n+1}) + D_1L_d(q_{n+1}, q_{n+2}) = 0 .\]
The loss is the norm of that residual, relative to the norm of the two derivatives it is built from, so a Lagrangian scaled by a constant gives the same number.
Implementation
The two slot derivatives are written out with the chain rule rather than taken with a nested Zygote.gradient:
\[ D_1L_d = \frac{\Delta{}t}{2}\nabla_qL - \nabla_{\dot{q}}L, \qquad D_2L_d = \frac{\Delta{}t}{2}\nabla_qL + \nabla_{\dot{q}}L,\]
both evaluated at the midpoint. Only the first derivative of the network with respect to its input is then needed, and that comes from a compiled symbolic expression exactly as hamiltonian_vector_field's does. Written the other way — a Zygote.gradient inside the loss — the parameter gradient fails with MethodError: no method matching getindex(::IdDict{Any, Any}).
Constructor
VariationalMidpointLoss(arch, timestep)where arch is a LagrangianNeuralNetwork.
Functor
loss(model, ps, input, output)
loss(ps, input, output) # equivalent to the aboveinput stacks $q_n$ on $q_{n+1}$ and so has $2n$ rows; output is $q_{n+2}$ and has $n$.
GeometricMachineLearning._discrete_lagrangian_derivatives — Function
_discrete_lagrangian_derivatives(loss, ps, qa, qb)Return (D₁Lᵈ, D₂Lᵈ) of the midpoint discrete Lagrangian on the pair (qa, qb), one column per sample.