Neural Canonical Transformation with Symplectic Flows

Canonical transformation plays a fundamental role in simplifying and solving classical Hamiltonian systems. Intriguingly, it has a natural correspondence to normalizing flows with a symplectic constraint. Building on this key insight, we design a neural canonical transformation approach to automatically identify independent slow collective variables in general physical systems and natural datasets. We present an efficient implementation of symplectic neural coordinate transformations and two ways to train the model based either on the Hamiltonian function or phase-space samples. The learned model maps physical variables onto an independent representation where collective modes with different frequencies are separated, which can be useful for various downstream tasks such as compression, prediction, control, and sampling. We demonstrate the ability of this method first by analyzing toy problems and then by applying it to real-world problems, such as identifying and interpolating slow collective modes of the alanine dipeptide molecule and MNIST database images.

Popular Summary

For centuries, physicists and astronomers have tackled the dynamics of complex interacting systems (such as the Sun-Earth-Moon system) with a mathematical tool known as a canonical transformation. Such a transformation changes the coordinates of the Hamiltonian equations (which describe the time evolution of the system) to simplify computation while preserving the form of the equations. However, despite being a deep concept, its wider application has been limited by cumbersome manual inspection and manipulation. By exploring the inherent connection between canonical transformation and a modern machine-learning method known as the normalizing flow, we have constructed a neural canonical transformation that can be trained automatically using the Hamiltonian function or data.

Normalizing flows are adaptive transformations often implemented as deep neural networks, and they find many real-world applications such as speech synthesis, image generation, and so on. In essence, it is an invertible change of variables that deforms a complex probability distribution into a simpler one. The canonical transformations are normalizing flows, albeit with two crucial twists. First, they are flows in the phase space, which contains both coordinates and momenta. Second, these flows satisfy the symplectic condition, a mathematical property that underlies most intriguing features in classical mechanics.

An immediate application of the neural canonical transformation is to simplify complex dynamics toward independent nonlinear modes, thereby allowing one to identify a small number of slow modes that are essential for applications such as molecule dynamics and dynamical control. Meanwhile, our work also stands as an example of imposing physical principles into the design of deep neural networks for better modeling of natural data.

Figure 1

A neural canonical transformation maps between the latent variables $\mathit{z}$ and physical variables $\mathit{x}$ via a symplectic neural network. The transformation preserves the Hamiltonian equation and connects the phase-space trajectories in the physical and the latent spaces. There are two ways to train the neural canonical transformation: (a) variational free energy based on the Hamiltonian [Eq. (5)]. (b) density of phase-space estimation based on data [Eq. (6)].

Figure 2

The Ringworld samples of Sec. 4a projected to the plane of learned latent coordinates and the polar coordinates.

Figure 3

(a) The learned frequencies of harmonic chain present in Sec. 4b in the prior distribution [Eq. (4)] and the analytical normal-mode frequency. (b) The Jacobian ${\nabla }_{q_i}{Q}_k$ of the two slowest collective coordinates with respect to the original coordinates. Solid lines are the analytical solution.

Figure 4

(a) The learned frequencies of alanine dipeptide presented in Sec. 4c from the MD trajectories. The inset shows the molecule with slow torsion angles. (b) The mutual information between the few slowest modes and the torsion angles.

Figure 5

(a) Latent-space interpolation in the plane of the two slowest collective coordinates ($Q_1$, $Q_2$) generates various molecule conformations. (b) Probability profile of the alanine dipeptide as a function of dihedral angles from generated samples of neural canonical transformation. The paths corresponding to the path with the same color in (a).

Figure 6

(a) A path from the molecular configuration at ($-125°$,150°) to the configuration at ($-75°$,0°) obtained by spherical linear interpolation of the two slowest latent variables. The background is the probability profile of the alanine dipeptide dataset on the plane of the dihedral angles. (b) The molecular conformation along the interpolation path.

Figure 7

(a) The learned frequencies of the MNIST dataset. (b) Classification accuracy on the test dataset based on $n_{\mathrm{slow}}$ slowest modes.

Figure 8

The coordinate transformation Eq. (2) maps the MNIST images to a latent representation with independent frequency modes. We concatenate slow latent variables of an image together with the fast modes of another image and then map the combined latent vector back to the image space with the inverse coordinates transformation. The bottom panel shows the same experiment by concatenating latent vectors at random.

Figure 9

Conceptual compression using the learned collective variables. One performs the coordinate transformation Eq. (2) to the input data, and restores the data based on $n_{\mathrm{slow}}$ slowest modes. The remaining fast modes are thrown away and resampled from the prior distribution.