SympNet Architecture

This section discusses the symplectic neural network (SympNet) architecture and its implementation in GeometricMachineLearning.

Principle

SympNets [5] are a type of neural network that can model the trajectory of a canonical Hamiltonian system in phase space. Take $(q^T,p^T)^T=(q_1,\ldots,q_d,p_1,\ldots,p_d)^T\in \mathbb{R}^{2d}$ as the coordinates in phase space, where $q=(q_1, \ldots, q_d)^T\in \mathbb{R}^{d}$ is refered to as the position and $p=(p_1, \ldots, p_d)^T\in \mathbb{R}^{d}$ the momentum. Given a point

\[ \begin{pmatrix} q^{(m)} \\ p^{(m)} \end{pmatrix} \in \mathbb{R}^{2d}\]

the SympNet aims to compute the next position

\[ \mathrm{SympNet}\left( \begin{pmatrix} q^{(m)} \\ p^{(m)} \end{pmatrix} \right) = \begin{pmatrix} \tilde{q}^{(m+1)} \\ \tilde{p}^{(m+1)} \end{pmatrix} \in \mathbb{R}^{2d}\]

and thus predicts the trajectory while preserving the symplectic structure of the system. SympNets are enforcing symplecticity strongly, meaning that this property is hard-coded into the network architecture. The layers are reminiscent of traditional neural network feedforward layers, but have a strong restriction imposed on them in order to be symplectic.

SympNets can be viewed as a symplectic integrator or symplectic one-step method^[1] [1, 63]. Their goal is to predict, based on an initial condition $((q^{(0)})^T,(p^{(0)})^T)^T$, a sequence of points in phase space that fit the training data as well as possible:

\[\begin{pmatrix} q^{(0)} \\ p^{(0)} \end{pmatrix}, \begin{pmatrix} \tilde{q}^{(1)} \\ \tilde{p}^{(1)} \end{pmatrix}, \cdots \begin{pmatrix} \tilde{q}^{(n)} \\ \tilde{p}^{(n)} \end{pmatrix}.\]

The tilde in the above equation indicates predicted data. With standard SympNets^[2] the time step between predictions is not a parameter we can choose but is related to the temporal frequency of the training data. This means that if data is recorded in an interval of e.g. 0.1 seconds, then this will be the time step of our integrator.

SympNets preserve symplecticity by exploiting the $(q, p)$ structure of the system. This is visualized below:

In the figure above we see that an update for $q$ is based on data coming from $p$ and an update for $p$ is based on data coming from $q$. $T_i:\mathbb{R}^d\to\mathbb{R}^d$ is an operation that changes $p$ when $i$ is even and changes $q$ when odd. It has the special property that its Jacobian is a symmetric matrix. There are two types of SympNet architectures: $LA$-SympNets and $G$-SympNets.

$LA$-SympNet

The first type of SympNets, $LA$-SympNets, are obtained from composing two types of layers: symplectic linear layers and symplectic activation layers. For a given integer $w$, the linear part of an $LA$-SympNet is

\[\mathcal{L}^{w,\mathrm{up}} \begin{pmatrix} q \\ p \\ \end{pmatrix} = \begin{pmatrix} \mathbb{I} & A^w/\mathbb{O} \\ \mathbb{O}/A^w & \mathbb{I} \\ \end{pmatrix} \cdots \begin{pmatrix} \mathbb{I} & \mathbb{O} \\ A^2 & \mathbb{I} \\ \end{pmatrix} \begin{pmatrix} \mathbb{I} & A^1 \\ \mathbb{O} & \mathbb{I} \\ \end{pmatrix} \begin{pmatrix} q \\ p \\ \end{pmatrix} + b ,\]

or

\[\mathcal{L}^{w,\mathrm{low}} \begin{pmatrix} q \\ p \end{pmatrix} = \begin{pmatrix} \mathbb{I} & \mathbb{O}/A^w \\ A^w/\mathbb{O} & \mathbb{I} \end{pmatrix} \cdots \begin{pmatrix} \mathbb{I} & A^2 \\ \mathbb{O} & \mathbb{I} \end{pmatrix} \begin{pmatrix} \mathbb{I} & \mathbb{O} \\ A^1 & \mathbb{I} \end{pmatrix} \begin{pmatrix} q \\ p \end{pmatrix} + b . \]

The superscripts $\mathrm{up}$ and $\mathrm{low}$ indicate whether the $q$ or the $p$ part is first changed^[3]. The learnable parameters are the symmetric matrices $A^i\in\mathcal{S}_\mathrm{sym}(d)$ and the bias $b\in\mathbb{R}^{2d}$. The integer $w$ is the number of linear layers in one block. It can be shown that five of these layers, i.e. $w\geq{}5$, can represent any linear symplectic map [67], so $w$ need not be larger than five. We denote the set of symplectic linear layers by $\mathcal{M}^L$.

The second type of layer needed for $LA$-SympNets are activation layers.

An $LA$-SympNet is a mapping of the form $\Psi=l_{k} \circ a_{k} \circ l_{k-1} \circ \cdots \circ a_1 \circ l_0$ where $(l_i)_{0\leq i\leq k} \subset \mathcal{M}^L$ and $(a_i)_{1\leq i\leq k} \subset \mathcal{M}^A$. We will refer to $k$ as the number of hidden layers of the SympNet^[4] and the number $w$ above as the depth of the linear layer.

We give an example of calling $LA$-SympNet:

using GeometricMachineLearning

k = 1
w = 2
arch = LASympNet(4;
                    nhidden = k,
                    depth = 2,
                    init_upper_linear = true,
                    init_upper_act = true,
                    activation = tanh)

model = Chain(arch).layers

(LinearLayerQ{4, 4}(), LinearLayerP{4, 4}(), GeometricMachineLearning.BiasLayer{4, 4}(), ActivationLayerQ{4, 4, typeof(tanh)}(tanh), ActivationLayerP{4, 4, typeof(tanh)}(tanh), LinearLayerQ{4, 4}(), LinearLayerP{4, 4}(), GeometricMachineLearning.BiasLayer{4, 4}())

The keywords init_upper_linear and init_upper_act indicate whether the first linear (respectively activation) layer is of $q$ type^[5].

$G$-SympNets

$G$-SympNets are an alternative to $LA$-SympNets. They are built with only one kind of layer, the gradient layer. If we denote by $\mathcal{M}^G$ the set of gradient layers, a $G$-SympNet is a function of the form $\Psi=g_k \circ g_{k-1} \circ \cdots \circ g_1$ where $(g_i)_{1\leq i\leq k} \subset \mathcal{M}^G$. The index $k$ here is the number of layers in the SympNet.

using GeometricMachineLearning

k = 2
n = 10
arch = GSympNet(4; upscaling_dimension = n, n_layers = k, init_upper = true, activation = tanh)

model = Chain(arch).layers

(GradientLayerQ{4, 4, typeof(tanh)}(10, tanh), GradientLayerP{4, 4, typeof(tanh)}(10, tanh))

The keyword init_upper for GSympNet is similar as in the case for LASympNet. The keyword upscaling_dimension is explained in the section on the SympNet gradient layer.

Forced $G$-SympNets

A $G$-SympNet is symplectic: its map preserves the symplectic form, as the flow of a Hamiltonian system does. A damped system is not of that form: a friction or an external force enters the $\dot{p}$ equation as a term that no gradient layer can represent. A forced $G$-SympNet adds a ForcingLayer in front of every gradient layer that changes $p$:

\[\begin{pmatrix} q \\ p \end{pmatrix} \mapsto \begin{pmatrix} q \\ p + \nabla V(q) + f(z) \end{pmatrix},\]

where $f$ is a neural network of its own and $z$ the component it reads. The resulting map is in general not symplectic, which is what lets it model a non-conservative system. ForcedSympNet is the architecture, and ForcingLayer the layer it is built from.

using GeometricMachineLearning
using GeometricMachineLearning: ForcedSympNet

arch = ForcedSympNet(4; upscaling_dimension = 10, n_layers = 2, forcing_type = :P)

model = Chain(arch).layers

(GradientLayerQ{4, 4, typeof(tanh)}(10, tanh), GeometricMachineLearning.ForcingLayerP{4, 4, false, Chain{Tuple{Dense{2, 10, true, AbstractNeuralNetworks.TanhActivation}, Dense{10, 10, true, AbstractNeuralNetworks.TanhActivation}, Dense{10, 10, true, AbstractNeuralNetworks.TanhActivation}, Linear{10, 2, IdentityActivation}}}}(4, 10, 2, Chain{Tuple{Dense{2, 10, true, AbstractNeuralNetworks.TanhActivation}, Dense{10, 10, true, AbstractNeuralNetworks.TanhActivation}, Dense{10, 10, true, AbstractNeuralNetworks.TanhActivation}, Linear{10, 2, IdentityActivation}}}((Dense{2, 10, true, AbstractNeuralNetworks.TanhActivation}(AbstractNeuralNetworks.TanhActivation()), Dense{10, 10, true, AbstractNeuralNetworks.TanhActivation}(AbstractNeuralNetworks.TanhActivation()), Dense{10, 10, true, AbstractNeuralNetworks.TanhActivation}(AbstractNeuralNetworks.TanhActivation()), Linear{10, 2, IdentityActivation}(IdentityActivation())))), GradientLayerP{4, 4, typeof(tanh)}(10, tanh))

The forcing_type keyword says what the forcing reads: :Q the $q$ component, :P the $p$ component, and :QP both. All three add to $p$, because that is where a forcing enters.

Universal Approximation Theorems

In order to state the universal approximation theorem for both symplectic architectures we first need a few definitions:

Let $U$ be an open set of $\mathbb{R}^{2d}$, and let us denote by $\mathcal{SP}^r(U)$ the set of $C^r$ smooth symplectic maps on $U$. We now define a topology on $C^r(K, \mathbb{R}^{2d})$, the set of $C^r$-smooth maps from a compact set $K\subset{}U$ to $\mathbb{R}^{2d}$ through the norm

\[||f||_{C^r(K,\mathbb{R}^{2d})} = \sum_{|\alpha|\leq r} \underset{1\leq i \leq 2d}{\max} \hspace{2mm} \underset{x\in K}{\sup} |D^\alpha f_i(x)|,\]

where the differential operator $D^\alpha$ is defined by

\[D^\alpha f = \frac{\partial^{|\alpha|} f}{\partial x_1^{\alpha_1}...x_n^{\alpha_n}},\]

with $|\alpha| = \alpha_1 +...+ \alpha_{2d}$. We impose the following condition ($r$ finiteness) on the activation function:

We further consider the topology on $C^r(U, \mathbb{R}^d)$ induced by $||\cdot ||_{C^r(\cdot, \mathbb{R}^d)}$ and the associated notion of denseness:

We can now state the universal approximation theorems:

and

There are many $r$-finite activation functions commonly used in neural networks, for example:

• The sigmoid activation function: $\sigma(x) = {1} / (1+e^{-x})$,
• The hyperbolic tangent function: $\tanh(x) = (e^x-e^{-x}) / (e^x+e^{-x})$.

The universal approximation theorems state that we can, in principle, get arbitrarily close to any symplectomorphism defined on $\mathbb{R}^{2d}$. But this does not tell us anything about how to optimize the network. This is can be done with any common neural network optimizer and these neural network optimizers always rely on a corresponding loss function.

Loss function

To train the SympNet, one needs data along a trajectory such that the model is trained to perform an integration. The loss function is defined as^[6]:

\[\mathrm{loss}(z^\mathrm{c}, z^\mathrm{p}) = \frac{|| z^\mathrm{c} - z^\mathrm{p} ||}{|| z^\mathrm{c} ||},\]

where

\[z^\mathrm{c} = \begin{pmatrix} q^\mathrm{c} \\ p^\mathrm{c} \end{pmatrix}\]

is the current state and

\[z^\mathrm{p} = \begin{pmatrix} q^\mathrm{p} \\ p^\mathrm{p} \end{pmatrix}\]

is the predicted state. In the example section we show how to use SympNets in GeometricMachineLearning.jl and how to modify the loss function.

Library Functions

SympNet <: NeuralNetworkIntegrator

LASympNet(d)

using GeometricMachineLearning
dl = DataLoader(rand(2, 20); suppress_info = true)
LASympNet(dl)

# output

LASympNet{typeof(tanh), true, true}(2, 5, 1, tanh)

GSympNet(d)

ForcedSympNet(d)

ForcingLayer <: AbstractExplicitLayer

check_forcing_type(type::Symbol)

ForcingLayerQ(dim, width = dim, nhidden = 1, activation = tanh; parameters, return_parameters)

ForcingLayerQ(4, 8, 1, tanh; parameters = (m = 1.0,), return_parameters = false)

ForcingLayerP(dim, width = dim, nhidden = 1, activation = tanh; parameters, return_parameters)

ForcingLayerQP(dim, width = dim, nhidden = 1, activation = tanh; parameters, return_parameters)

References