Deep Neural Networks with Symplectic Preservation Properties

On $\mathbb{R}^{2n}$, we denote the standard Cartesian coordinates as $q_{1},\cdots,q_{n},p_{1},\cdots,p_{n}$, corresponding to the ”position” and ”momentum” coordinates in Hamiltonian mechanics. The standard symplectic form on $\mathbb{R}^{2n}$ is the differential 2-form

and a transformation $\varphi:\mathbb{R}^{2n}\rightarrow\mathbb{R}^{2n}$ is called a symplectomorphism if $\varphi^{*}\omega=\omega$. This means

where

or equivalently,

where

is the Jacobian matrix of $\varphi$, and

is the matrix of the standard symplectic form $\omega$.

The most essential property of a Hamiltonian system

where

is that its flow map defines a family of symplectomorphisms. This means that if we solve (7) from time $t_{0}$ to time $t_{1}$, then the mapping defined by $(q(t_{0}),p(t_{0}))\to(q(t_{1}),p(t_{1}))$ is an $\mathbb{R}^{2n}\to\mathbb{R}^{2n}$ symplectomorphism. The inverse is also true: If a differential equation system on $\mathbb{R}^{2n}\to\mathbb{R}^{2n}$ satisfies than the flow maps are symlectomorphisms, then there exists a function $H\in C^{2}(\mathbb{R}^{2n+1})$ such that the system can be written as Hamiltonian system (7).

Real NVP (Real-valued Non-Volume Preserving) [3, 1] is a generative model used for density estimation. Real NVP networks use invertible transformations, allowing us to go back and forth between the original and transformed spaces. The structure of real NVP is as follows: The input and output of the network are both $N$-dimensional vectors. An $N$-dimensional vector

received as the input is partitioned in to two parts

A Real NVP transformation keeps one of the parts unchanged and perform an ”entry-wise linear transformation” on the other part, whose coefficients are determined by the unchanged part. Specifically, the input $z$ undergoes the following transformation:

where $s,b:\mathbb{R}^{n}\rightarrow\mathbb{R}^{N-n}$ are two functions which are given as a neural networks in practice, and the symbol ”$\odot$” the Hadamard product (entry-wise product) operator:

The inverse of this mapping (16) is clear:

The transformation (16) is often exhibited as a diagram like .

The apparent limitation of transformation (16) is that it does not change the part $z_{A}$. This can be quickly fixed by appending another real NVP block that keeps the $x_{B}$ part unchanged:

where $\tilde{s},\tilde{b}:\mathbb{R}^{N-n}\rightarrow\mathbb{R}^{n}$ are another two neural network functions, so the composed transformation from $z$ to $y$ given by (16) and (18) do not keep any component unchanged. This can be exhibited as a diagram like .

Of course, we can stack more layers like this to improve the expressivity of the network.

For our goal of building symplectomorphism neural network, the problem of real NVP is directly exhibited in its name: ”NVP” means ”non-volume-preserving”, while a symplectomorphism has to be volume preserving. Indeed, to make real NVP volume preserving (from ”real NVP” to ”real VP”), there is a quick fixation: one only needs to add an extra layer

after the output layer of the network that subtracts the average of the network. Unfortunately, mere volume-preserving property does not guarantee symplecticity. We need further adjustments.

Indeed, we can decompose (16) into two transformations: a ”stretching”

and a ”shearing”

Neither of these two transformations are guaranteed to be symplectic. Nevertheless, we have introduced their symplectic counterparts in the last section: Indeed, we can write (8), (9) and (10) (where (14), (15)) into a more compact form

where $q=(q_{1},\cdots,q_{n})$, $p=(p_{1},\cdots,p_{n})$, $Q=(Q_{1},\cdots,Q_{n})$, $P=(P_{1},\cdots,P_{n})$. And ”$\odot$” is the Hadamard product as before. And now their correspondence with (19) and (20) are clear: (23) is exactly (20) when $\dim x_{A}$ and $\dim x_{B}$ are of the same dimension, and $b$ is the gradient of a function; while (23) is a symmetrized version of (19):

with $s$ being the gradient of a function. We denote the transformations defined by (21), (22), (23) as $\operatorname{qSh}_{F}$, $\operatorname{pSh}_{G}$ and $\operatorname{St_{\Phi}}$, which are short hands for ”q-shearing”, ”p-shearing” and ”stretching”, respectively. These becomes the basic building blocks of the ”symplectic version of real NVP” once we take $F$, $G$ and $\Phi$ in these transformations as trainable neural networks.

Now we have introduced all the basic symplecticomorphism building blocks, and a symplectomorphism neural network (SymplectoNet, or even shorter, SpNN) is a neural network designed as an arbitrary finite composition of $\operatorname{qSh}_{F}$, $\operatorname{pSh}_{G}$ and $\operatorname{St_{\Phi}}$ where $F$, $G$ and $\Phi$ are arbitrary neural networks with $n$-dimensional input and one-dimensional output.

Of course, the expressivity of this network depends on the complexity of the underlying neural networks $F$, $G$ and $H$, and also on the number of the building blocks we stacked. Indeed, the latter can be even more essential: e.g. if we only use less than four symplectic shearing blocks, we cannot even cover all the linear symplectomorphisms no matter how complicated the underlying network $F$ and $G$ are, because the Jacobian of a shearing transformation is of the form

where $B$, $C$ are symmetric $n\times n$ matrices. The degree of freedom of these matrices are $n(n+1)/2$, while $\dim Sp(2n)=n(2n+1)$, which is greater than $3n(n+1)/2$ for $n>1$. This is why I also designed the symplectic stretching layer $\operatorname{Str}_{\Phi}$. A good practice is to include both the p, q-shearing and the symplectic stretching layers in the network for at lease once. A simplest example is a network with structure $\operatorname{pSh}_{G}\circ\operatorname{St}_{\Phi}\circ\operatorname{qSh}_{F}$ (see 3), which is similar to the structure of a real NVP.

One of the most important features of real NVP is that it is explicitly invertible: one can write out (or, in a more techical term, build the computation graph of) the explicit expression of the neural network function’s inverse function [5]. Our SymplectoNet is inspired by real NVP, so a natural question is whether the SymplectoNet structure is explicitly invertible like real NVP. Next, we will show that the answer is yes.

Indeed, since the inverse of a composed function $f_{1}\circ f_{2}\circ\cdots\circ f_{k}$ is $f_{k}^{-1}\circ\cdots\circ f_{2}^{-1}\circ f_{1}^{-1}$, so we only need to prove that the basic building blocks, $\operatorname{pSh}_{G}$, $\operatorname{qSh}_{F}$ and $\operatorname{St}_{\Phi}$ are explicitly invertible. The inverse of $\operatorname{pSh}_{G}$, $\operatorname{qSh}_{F}$ are obvious: (21) is equivalent to

(22) is equivalent to

therefore the inverse of $\operatorname{pSh}_{G}$, $\operatorname{qSh}_{F}$ are  $\operatorname{pSh}_{-G}$, $\operatorname{qSh}_{-F}$, respectively. And finally we look at $\operatorname{St}_{\Phi}$. Notice that from (23), we have

therefore

this shows that the inverse of $\operatorname{St}_{\Phi}$ is exactly $\operatorname{St}_{-\Phi}$. In conclusion, we have

These results give a neat expression of inverting the SymplectoNet. E.g. the inverse of the SymplectoNet

This shows that the inverse of SymplectoNet is explicitly available.

This example is learning a symplectic map

A network with structure

where $F_{1},G_{1},F_{2},G_{2}$ are $(2,20,10,1)$ dense neural networks, and $\Phi$ is $(2,10,1)$ dense neural network. The loss is the ordinary MSE loss. Adamax with learning rate 0.25 is applied here, and decay by factor 0.99 every 100 epoch.

Firstly, some uniformly random points for

is sampled. The training went for 40,000 epochs, and the loss dropped from $0.3$ to about $10^{-5}$, and the plot is shown in Figure 4(a), and the loss decay is shown in Figure 4(b).

Anoter numerical experiments concerning also (30) but the domain changed to

is also conducted. This time, the geometry of the transformation is more complicated. Note that we cannot do $p:[0,2\pi]$ because this will make the mapping (30) non-injective, while the model is invertible. Thus the model will have difficulty learning the data near the two lines $p=0$ and $p=2\pi$. The training went for 40,000 epochs, and the loss dropped from $0.3$ to about $10^{-5}$, and the plot is shown in Figure 4(c), and the loss decay is shown in Figure 4(d). The majority of the error comes from $p=3\pi/2$ boundary. This is because the points her are close to the points with $p=0$.