From latent dynamics to meaningful representations

TL;DR

This work tackles the challenge of learning meaningful latent representations when prior distributions over latents are unavailable. It introduces DynAE, a representation learner constrained solely by latent dynamics, enforcing overdamped Langevin behavior with a learnable transition density via a sliced-Wasserstein regularizer. The approach yields identifiability of the latent variables up to an isometry and demonstrates successful recovery of physically meaningful factors across diverse systems, including alanine dipeptide and Brownian DNA, while outperforming or matching dynamic baselines on higher-dimensional tasks. By shifting regularization from latent priors to transition dynamics and extending the sliced-Wasserstein auto-encoder to dynamics, the framework provides robust, disentangled representations with broad applicability to stochastic systems. This dynamics-centric perspective offers a practical path to recover ground-truth latent structure in complex datasets without strong prior assumptions about latent distributions.

Abstract

While representation learning has been central to the rise of machine learning and artificial intelligence, a key problem remains in making the learned representations meaningful. For this, the typical approach is to regularize the learned representation through prior probability distributions. However, such priors are usually unavailable or are ad hoc. To deal with this, recent efforts have shifted towards leveraging the insights from physical principles to guide the learning process. In this spirit, we propose a purely dynamics-constrained representation learning framework. Instead of relying on predefined probabilities, we restrict the latent representation to follow overdamped Langevin dynamics with a learnable transition density - a prior driven by statistical mechanics. We show this is a more natural constraint for representation learning in stochastic dynamical systems, with the crucial ability to uniquely identify the ground truth representation. We validate our framework for different systems including a real-world fluorescent DNA movie dataset. We show that our algorithm can uniquely identify orthogonal, isometric and meaningful latent representations.

Figures (8)

• Figure 1: a. Network architecture used for dynamics constrained framework. The encoder $\phi$, the decoder $\psi$, the diffusion $\mathbf{M}_\omega$ and the force field $\mathbf{f}_\omega$ are nonlinear deep neural networks. b. A flowchart illustrating the dynamics constrained framework. To regularize latent dynamics, the latent representation is discretized into bins represented by dashed circles. The samples, depicted as yellow circles, are drawn from each bin based on the encoded transition density $q_\phi(\Delta\mathbf{z}|\mathbf{z}_t)$ and then matched with samples, shown as yellow triangles, drawn from some specific prior transition density $r_\omega(\Delta\mathbf{z}|\mathbf{z}_t)$. The randomly chosen slicing direction is indicated by the red dashed double-headed arrow. See details in the Methods section.
• Figure 2: Recovering the underlying dynamics from the transformed three well model system. The original simulated data (a,b), the transformed data (c,d), and the latent representation learned by our algorithm (e,f) are shown in the figure. The black arrows represent the force field $\mathbf{f}=-\nabla F$ (left) while the ellipses represent the diffusion field $\mathbf{M}$ (right). The ellipses in d are highly stretched due to the extremely anisotropic and inhomogeneous diffusion field caused by the nonlinear mapping function.
• Figure 3: Comparison of the representations learned by different models on the three well model potential. a. $\beta$-VAE. b. SWAE. c. SDE-VAE. d. DynAE. For better comparison, all the representations shown here have been aligned with the ground-truth using the optimum $Q$ from Eq. \ref{['eq:MSE_latent']}. Only DynAE can successfully recover the ground-truth variables up to isometry.
• Figure 4: Learning the underlying dynamics from the thirty-dimensional coordinates of alanine dipeptide in water. a. Structure of alanine dipeptide. The main coordinates describing slow transitions are the torsion angles $\phi$ ($C$-$N$-$C_\alpha$-$C$) and $\psi$ ($N$-$C_\alpha$-$C$-$N$), but the neural network input is only the Cartesian coordinates of the heavy atoms. b. Free energy surface of alanine dipeptide in water at 300K along the dihedral angles $\phi$ and $\psi$. c, d show the latent representation learned by our algorithm. The black arrows represent the force field $\mathbf{f}=-\nabla F$ (c) while the ellipses represent the diffusion field $\mathbf{M}$ (d). e, f illustrate the relationship between our learned latent representation and the ground-truth latent factors $\phi$ and $\psi$.
• Figure 5: Comparison of the results on DNA dataset obtained from $\beta$-VAE (a,e), SWAE (b,f), SDE-VAE (c,g) and DynAE (d,h). a-d. Reconstructions of the DNA molecule position. First row: originals. Second row: reconstructions. Remaining two rows: reconstructions of latent traversals across each latent dimension. It can be seen how $z_0$ and $z_1$ for SDE-VAE and DynAE do a much better job than $\beta$-VAE and SWAE at correlation with the underlying movement along $x$ and $y$ directions. e-h The acquired latent representations of the fluorescent DNA molecule and their correlation with the DNA molecule's $x$ and $y$ movements.