Introducing dynamical constraints into representation learning

The goal of autoencoder-based representation learning algorithms in general is to learn a low-dimensional latent representation $z$ such that the original high-dimensional observations can be reconstructed from the low-dimensional representation. Most of the recent advances in representation learning are built upon the pioneering work of the Variational Autoencoder (VAE).[1, 6] These algorithms consist of at least three important components: an encoder $q\left(z|X\right)$, a decoder $q\left(X|z\right)$, and a prior distribution $r\left(z\right)$ for the latent representation space. The encoder maps the input $X$ to the representation space $z$, while the decoder reconstructs the input data $X$ from the representation. In addition, the encoder is encouraged to satisfy some structure specified by the prior distribution $r_{\omega }\left(z\right)$ on the latent space, where $\omega$ denotes trainable parameters for the prior. The objective function $L_{rep}$ of VAE, where the subscript rep signifies the representation learning process, can be written as follows:

where $\hat{p}\left(X\right)\equiv \frac{1}{N}{\sum }_{n=1}^N{\delta }_{X_n}$ is the empirical data distribution, and $N$ is the number of data points. The stochastic encoder $q_{\varphi }\left(z|X\right)$ and the stochastic decoder $q_{\psi }\left(X|z\right)$ are usually parameterized by neural networks with the model parameters $\varphi$ and $\psi$. This objective function can be divided into two parts: The first term $L_{REC}$ measures the ability of our representation to reconstruct the input, while the second term $L_{REG}$ can be interpreted as the complexity penalty that acts as a regulariser. Such a trade-off between the prediction capacity and model complexity can then be controlled by a hyper-parameter $\beta \in [0,\infty )$. Compared to the original VAE objective where $\beta =1$, a modified version named $\beta$-VAE shows a better performance in terms of disentanglement with an appropriately tuned $\beta$.[5]

Minimizing the objective in Eq. 1 requires a definition of the prior distribution $r_{\omega }\left(z\right)$. However, such a predefined distribution is usually unavailable a priori. Choosing a simplistic prior like the standard Gaussian distribution as done in VAE could lead to over-regularization, resulting in poor latent representations. Therefore, several flexible prior distributions have been proposed, such as a mixture of Gaussians prior,[7] VampPrior,[8] and normalizing flows.[9] However, if the prior distribution is too flexible, it might defeat its purpose and no longer serve as a regularizer for the learned latent space.

To deal with this dilemma, here we propose the dynamics constrained representation learning framework summarized in Fig. 1. Instead of using a predefined prior distribution to regularize the latent space, we enforce the latent representation to follow a specific yet generic class of dynamics. Our inspiration here stems from physics, where even though different systems can and likely will have different probability distribution along the same low-dimensional projection $z$, they will obey the same dynamics equations, such as the Newton’s second law in classical mechanics and Schrödinger’s equation in quantum mechanics. Therefore, we believe this is a more natural and generic constraint for representation learning.

For Markovian processes, their dynamics can be described by their transition density $p\left(\Delta z|z_t\right)$ with $\Delta z\equiv z_{t+\Delta t}-z_t$. For instance, for deterministic dynamical processes, the transition density would equal a delta function. A tempting idea then is to match the encoded transition density $q_{\varphi }\left(\Delta z|z_t\right)$ to some specific prior transition density $r_{\omega }\left(\Delta z|z_t\right)$. This can be done by minimizing the sliced-Wasserstein distance between these two conditional probability distributions $D_{SW}\left(q_{\varphi }\right(\Delta z_t|z_t),r_{\omega }(\Delta z_t|z_t\left)\right)$ if enough samples from each distribution are available (Methods).[10, 11] Samples from the encoded transition density $q_{\varphi }\left(\Delta z|z_t\right)$ can be obtained, for instance, by performing swarms of short trajectories starting from each point $z_t$. In many cases, however, these burst trajectories are usually unavailable. Therefore, it would be desirable if this approach could also be applied to a single long trajectory. Here we approximate the sliced-Wasserstein distance $D_{SW}\left(q_{\varphi }\right(\Delta z_t|z_t),r_{\omega }(\Delta z_t|z_t\left)\right)$ by discretizing the latent space $Z$ into $K$ bins $\{Z_k{\}}_{k=1}^K$. In other words, instead of forcing $\Delta z$ to follow some specific distribution for every point $z$, we encourage $\Delta z$ within each bin $Z_k$ to follow some specific distribution. When the number of bins $K$ is large enough, this should provide a reasonable approximation. The discretization and resampling schemes are discussed in Methods. With this we can propose a dynamics constrained regularizer $L_{REG}$:

where $\Delta z_t\sim q_{\varphi }(\Delta z_t|z_t\in Z_k),$ $\Delta {\tilde{z}}_t\sim r_{\omega }\left(\Delta {\tilde{z}}_t|{\tilde{z}}_t\right)\in Z_k,$ and $r_{\omega }\left(\Delta z_t|z_t\right)$ is a prior transition probability which is related to the dynamics of the latent representation. In this way, even if an accurate prior distribution $r_{\omega }\left(z\right)$ is unavailable, we can still regularize the latent space by forcing its dynamics to follow some specific form of the prior transition probability $r_{\omega }\left(\Delta z_t|z_t\right)$.

Similarly, the reconstruction loss $L_{REC}$ in Eq. 1 can also be reformulated in terms of bins $\left\{Z_k\right\}$ and generalized to learning dynamics. For simplicity, we only consider in this paper the case of mappings where the latent variable $z$ is a deterministic function of the input data $X$ given as $z=\varphi \left(X\right)$, but the algorithm should be easily generalized to stochastic encoders. Let $\psi$ denote the deterministic decoder that maps the latent variable $z$ back to the original data space. In this way, the reconstruction loss $L_{REC}$ can be replaced by the $L_2$-loss, namely:

Combining the reconstruction loss $L_{REC}$ in Eq. 3 and the regularizer $L_{REG}$ in Eq. 13, we can tune the hyper-parameter $\beta$ to force the latent space to follow specific dynamics given the prior transition density $r_{\omega }\left(\Delta z|z_t\right)$. In the next subsection, we will discuss how to choose such a prior transition density.

We focus in this paper on stochastic systems governed by the overdamped Langevin equation or Brownian dynamics as the underlying dynamics, but the framework should be more generally applicable to other dynamical systems. Consider a system where the representation $z$ obeys Brownian dynamics,

where $f\left(z\right)$ is the force field experienced along $z$, $M\left(z\right)$ is the diffusion matrix, and $dw$ is a Brownian motion in $R^d$. Further, the force field $f\left(z\right)$ can also be given by $f\left(z\right)=-\nabla F\left(z\right)$ where we have introduced the free energy $F\left(z\right)=-kT\log r_{\omega }\left(z\right)$. If we assume the diffusion matrix in the learned representation space is diagonal, Eq. 4 can be simplified as:

where ${ϵ}_i\sim N(0,1)$. This assumption of diagonal diffusion matrix allows us to obtain a set of uncorrelated latent variables. For simplicity, in further derivations we absorb the time unit into the diffusion matrix $M$ and set $dt=\Delta t=1$. The prior transition density of the system $r_{\omega }\left(\Delta z_t|z_t\right)$ parameterized by $\omega$ can then be obtained:

where both the force field $f\left(z\right)=f_{\omega }\left(z\right)$ and the diffusion matrix $M=M_{\omega }\left(z\right)$ are nonlinear functions of the representation $z$ parameterized by neural network parameters $\omega$. Given the pairwise samples $\{\Delta z_t^{n_k},z_t^{n_k}{\}}_{n_k=1}^{N_k}$ from each bin $Z_k$, both the force field $f_{\omega }\left(z\right)$ and the diffusion matrix $M_{\omega }\left(z\right)$ can be inferred through likelihood maximization or equivalently by minimizing the following prior loss:

Once the estimated force field $f_{\omega }\left(z\right)$ and the diffusion matrix $M_{\omega }\left(z\right)$ are obtained, we can easily draw samples $\left\{\Delta {\tilde{z}}_t\right\}$ from the prior transition density $r_{\omega }\left(\Delta z_t|z_t\right)$ using Eq. 5, which can then be used to regularize the representation learnt through Eq. 13.

In fact, given this generalizable framework, we can introduce even more constraints into the prior. For instance, in our experiments we observed that enforcing a constant diffusion $M_{ij}^{\ast }\equiv {\delta }_{ij}$ largely stabilized the optimization process of our algorithm and avoided the algorithm learning trivial scaling. Namely, we incorporate an inductive bias into the learning process and regularize the latent representation by directly using the samples $\left\{\Delta {\tilde{z}}_t\right\}$ generated from the simplified Langevin dynamics obtained after using $M_{ij}^{\ast }\equiv {\delta }_{ij}$ in Eq. 5:

By learning a latent space with isotropic and homogenous diffusion as in Eq. 8, we can aim to preserve some underlying geometric properties of the data as we show later through numerical examples.

To sum up, the whole algorithm is aiming to optimize the following two objective functions in Eqs. 9 and 11:

where,

and,

The workflow so constructed can be summarized through Fig. 1 and Extended Alg. 1. Our algorithm iteratively discretizes the latent space into bins, and resamples the training dataset to focus on poorly sampled regions (see Methods). Then the representation learning is implemented via a two-step optimization process: The first step fixes the prior transition density $r_{\omega }\left(\Delta z_t|z_t\right)$, and concentrates on learning a better representation by minimizing the representation loss (Eq. 9); the second step fixes the latent representation and focuses on inferring the transition density $r_{\omega }\left(\Delta z_t|z_t\right)$ through optimizing the prior loss (Eq. 11), which in turn provides the information to guide the representation learning process.

We now first demonstrate the utility of our method using data obtained by simulating a 2D three-well model potential. The simulated data is warped with nonlinear mapping functions to mimic distortion from real-world observations before feeding into our algorithm (Supplementary Information or SI). The simulated data and its transformation are shown in Fig. 2. In the original data space $x$-$y$, the diffusion is constant (Fig. 2(b)). But after the transformation, the diffusion varies dramatically with different $x^{\prime }$ values as shown in Fig. 2(d). Feeding this transformed trajectory data into our algorithm, we find our algorithm can successfully recover the original variables up to an invertible linear transformation, wherein the diffusion is once again isotropic and homogeneous. The results also demonstrate the ability of our algorithm to recover the underlying kinetics even in the rarely sampled regions, such as the transition state in the three well model system.

As a slightly more complex example, we further illustrate the power of our algorithm in the well-studied alanine dipeptide molecule with TIP3P water molecules. We obtain a molecular dynamics trajectory for this system (see SI for details). The backbone torsion angles $\varphi$ and $\psi$ are known to be the most important reaction coordinates separating different metastable states of alanine dipeptide in water as shown in Fig. 3(a-b). We completely ignore prior knowledge of these dihedrals, and to make the task more challenging, we directly work with a 30-dimensional space comprising the three-dimensional coordinates of the 10 heavy atoms as input. To avoid issue due to trivial rotation and translation of the whole system, we align all the configurations to the first frame of the molecular dynamics trajectory. Fig. 3(c-d) illustrate that our algorithm is still able to recover the underlying dynamics even with a 2D Euclidean latent space. Even more interestingly, a clear correspondence between our learned latent representation and the torsion angles $\varphi$ and $\psi$ is found, as shown in Fig. 3(e-f).This suggests that DynamicsAE can be used to recover physically meaningful representations such as the dihedral angles in alanine dipeptide.

We then examine our algorithm in a 2D image dataset (dSprites) with known ground-truth generative factors. Only two available transformations (X-Position, Y-Position) of the dSprites dataset are first considered for better visualization purpose. To generate a long sequence of images, we take a random initial image, and sequentially transform it according to a Gaussian random walk along the two available transformations (X-Position, Y-Position). Further details of data generation can be found in SI. As shown in Fig. 4, we directly compare our algorithm with $\beta$-VAE. Fig. 4(a) shows that both algorithms are capable of accurately reconstructing the input images and disentangling the two ground-truth factors x-position and y-position. However, as the ground-truth factors should follow a uniform distribution by construction, it is evident that only our algorithm uniquely identifies the ground truth factors up to a linear transformation in this case. This can be clearly seen by comparing the two free energy surfaces in Fig. 4(b) where the latent representation learned by $\beta$-VAE follows a Gaussian distribution while the latent representation learned by our algorithm follows a uniform distribution. This simple but important result is not surprising. $\beta$-VAE cannot recover the ground-truth due to the false prior assumption of a Gaussian distribution in the latent space. Therefore, instead of constraining the distribution of latent space to a prior distribution, regularizing the latent space via its dynamics exhibits huge advantages.

We push our algorithm further by learning a higher dimensional latent representation. In this case, three transformations (Scale, X-Position, Y-Position) of the dSprites dataset are considered. The results are presented in Fig. 5. We find that our algorithm can still successfully recover the underlying "ground-truth" factors (Scale, X-Position, Y-Position), which $\beta$-VAE fails to identify. While it clearly outperforms $\beta$-VAE, our algorithm however doesn’t entirely decouple the scale from the x and y positions. This is because there are only 6 distinct values along the Scale factor, where the boundary effect can no longer be ignored. This finally results in a curved manifold as shown in the right column of Fig. 5(b).

We finally apply our algorithm to a real world video dataset consisting of the Brownian motion of DNA molecules in solution as described in the work of Ref. [12]. Following Ref. [4], we use the same method to obtain the “ground-truth" latent variables (the $x$ and $y$ coordinates of the center of DNA molecule). The details of the preprocessing protocol are provided in SI. As shown in Fig. 6(a), both algorithms can reconstruct the positions of the molecules with a 2D latent space. However, even with a large $\beta$, $\beta$-VAE fails to disentangle the $x$ and $y$ positions in this case. In comparison, our algorithm can still successfully disentangle the $x$ and $y$ positions. Finally, Fig. 6(b) clearly illustrates that $\beta$-VAE can only learn a complex nonlinear warping of the ground-truth latent variables, while our algorithm can still uniquely recover the ground-truth up to an isometry.

As can be seen from Eq. 1, the regularizer $L_{REG}$ in VAE is the Kullback–Leibler divergence between the posterior distribution $q_{\varphi }\left(z|X\right)$ and the prior distribution $r_{\omega }\left(z\right)$, averaged over the empirical data. However, as discussed in Ref. [10], this regularizer $L_{REG}$ encourages the encoded distribution for each training sample to become increasingly similar to the prior distribution, fundamentally at odds with the goal of achieving good reconstruction. This also can been shown in the limiting case where $L_{REG}=D_{KL}\left(q_{\theta }\right(z|X),r_{\omega }(z\left)\right)=0$. This equality occurs if and only if $q_{\theta }\left(z|X\right)=r_{\omega }\left(z\right)$, suggesting the encoder should forget all information about $X$. This can lead to difficulties in reconstruction of the input data. Therefore, Ref. [10] proposed Wasserstein Auto-Encoders (WAE) using generative adversarial networks (GAN) to directly match the distribution of the entire encoding space $q_{\theta }\left(z\right)=\int q_{\theta }\left(z|X\right)\hat{p}\left(X\right)dX$ to the prior $r_{\omega }\left(z\right)$. This prevents the latent representations of different data points from collapsing together, thereby promoting better reconstructions while still maintaining simple latent representations. Subsequent research introduced a much simpler solution known as Sliced Wasserstein Auto-Encoders (SWAE).[11] The trick in SWAE is to project or slice the distribution along multiple randomly chosen directions, and minimize the so-called Wasserstein distance along each of those one-dimensional spaces (which can be done even analytically). According to Ref. [11], let $\{z^n\sim q_{\varphi }(z){\}}_{n=1}^N$ and $\{{\tilde{z}}^n\sim r_{\omega }(\tilde{z}){\}}_{n=1}^N$ be random samples from the encoded input data and the prior distribution respectively. Assuming $z,\tilde{z}\in Z=R^d$, let $\{{\theta }_l{\}}_{l=1}^L$ be $L$ randomly sampled directions from a uniform distribution on the $(d-1)$-sphere $S^{d-1}$. Then the 1D projection of the latent representation $z$ onto ${\theta }_l$ is ${\theta }_l\cdot z$. Therefore, the regularizer $L_{REG}$ in Eq. 1 can be measured by the sliced-Wasserstein distance:

where $i\left[n\right]$ and $j\left[n\right]$ are the indices of sorted ${\theta }_l\cdot {\tilde{z}}^{i\left[n\right]}$ and ${\theta }_l\cdot z^{j\left[n\right]}$ with respect to $n$, correspondingly. We now can generalize this regularizer loss $L_{REG}$ to learning dynamics:

where $\Delta z_t\sim q_{\varphi }(\Delta z_t|z_t\in Z_k),$ $\Delta {\tilde{z}}_t\sim r_{\omega }\left(\Delta {\tilde{z}}_t|{\tilde{z}}_t\right)\in Z_k$.

For simplicity, regular space clustering with Euclidean distance metric is used to discretize the latent space to obtain uniformly distributed bins $\{Z_k{\}}_{k=1}^K$.[23] More advanced distance metrics such as the Mahalanobis distance[24] and others[25] are left for future exploration. The clustering is performed in such a way that cluster centers are at least $d_{min}$ from each other according to the given metric. Then samples can be assigned to cluster centers through Voronoi partitioning. Besides, to make the algorithm pay more attention to rarely sampled regions in the latent space, we resample the dataset according to a well-tempered distribution[26] by introducing a hyperparameter $\gamma \ge 1$:

where ${\hat{N}}_k$ is the actual number of samples within the bin $Z_k$ and $N$ is the total number of samples. For $\gamma =1$, Eq. 14 reduces to $N_k={\hat{N}}_k$ and the original dataset is used for training. For $\gamma \to \infty$, we obtain $N_k=N/K$, which means that we draw samples uniformly from each bin. In practice, we find our results are robust to the choice of $\gamma$, and a moderate value $\gamma =2$ is chosen in this paper.

As shown in Fig. 1(a), the full algorithm consists of four neural networks: the encoder $\varphi$, the decoder $\psi$, and the neural networks learning the diffusion field $M_{\omega }$ and the force field $f_{\omega }$. For non-image datasets, we use fully connected neural networks with 2 hidden layers and ReLU activations to parameterize the encoder and decoder. For image datasets, we use a convolutional neural network for the encoder and a deconvolutional neural network for the decoder. The diffusion and force field in all models are parameterized using the same fully connected neural networks with 3 hidden layers and tanh activation functions. The networks were trained using two Adam optimizers with the same learning rate of 0.001 optimizing the two loss functions shown in Eq. 9 and 11 separately. The implementation details are in SI.