Identifiable learning of dissipative dynamics

TL;DR

UID D introduces a universal, identifiable neural framework for learning dissipative, non-equilibrium stochastic dynamics from trajectory data. By structuring the dynamics as an Itô SDE with a learnable potential $V$, diffusion matrix $M$, and banded antisymmetric $W$, it yields a unique stationary density $\rho = \mathcal{Z}^{-1} e^{-V}$ and a clean split of drift into time-reversible and time-irreversible parts, enabling direct computation of entropy production rate (EPR). The authors prove a general identifiability theorem and demonstrate the method on a linear benchmark, polymer stretching under elongational flow, and stochastic gradient Langevin dynamics, revealing scaling laws for barrier heights and EPR with strain rate and batch size, respectively. This data-driven, thermodynamically grounded framework provides a powerful tool for diagnosing irreversibility and comparing non-equilibrium dynamics across systems and conditions, with broad implications for physics, chemistry, and machine learning workflows.

Abstract

Complex dissipative systems appear across science and engineering, from polymers and active matter to learning algorithms. These systems operate far from equilibrium, where energy dissipation and time irreversibility govern their behavior but are difficult to quantify from data. Here, we introduce a universal and identifiable neural framework that learns dissipative stochastic dynamics directly from trajectories while ensuring interpretability, expressiveness, and uniqueness. Our method identifies a unique energy landscape, separates reversible from irreversible motion, and allows direct computation of the entropy production, providing a principled measure of irreversibility and deviations from equilibrium. Applications to polymer stretching in elongational flow and to stochastic gradient Langevin dynamics reveal new insights, including super-linear scaling of barrier heights and sub-linear scaling of entropy production rates with the strain rate, and the suppression of irreversibility with increasing batch size. Our methodology thus establishes a general, data-driven framework for discovering and interpreting non-equilibrium dynamics.

Figures (10)

• Figure 1: Overview of the UIDD. (Blue) Discrete trajectory data serves as training input, and the model is built upon neural networks. (Green) Through a designed parameterization and likelihood-based training on this trajectory data, the framework constructs the UIDD dynamical model. (Orange) Mathematical analysis establishes that UIDD possesses both universal approximation capabilities and unique representation properties. (Grey) The components of the model can be interpreted as thermodynamically meaningful elements, including the stationary density, the time-reversible drift, and the time-irreversible drift. (Purple) Leveraging this physical interpretability and the identifiability, UIDD enables the computation of the energy landscape and the EPR.
• Figure 2: Results for linear system. (a) The learned potential and the negative log of the stationary density, each shifted by subtracting their value at input zero, ensuring that both functions attain zero at the origin. (b) First component for the time-irreversible drift of learned and exact system for $\lambda=1$. (c) Comparison of the global EPR calculated via four approaches. The exact value is computed using Eq. (\ref{['eq:ep_linear']}); Est.$_1$ is calculated by applying a state-space discretization to approximate the definition in Eq. (\ref{['eq:e_p']}); Est.$_2$ and learned $\dot{S}_{\text{tot}}$ are obtained via Monte Carlo approximation of Eq. (\ref{['eq:ep']}), using the exact and learned equations, respectively.
• Figure 3: Simulation setup, variations across random seeds and barrier heights of the learned potential energy for polymer stretching dynamics, (a) Linear bead-rod polymer chains are simulated using a Brownian dynamics framework under a planar elongational flow, whose strength is characterized by the strain rate $\dot{\varepsilon}$. (b) Relative variation ($\times 10^3$) in learned potential forces over $4$ random seeds at different strain rate $\dot{\varepsilon}$. (c-h) Projections of the learned potential energy at $\dot{\varepsilon}=4.63$ onto the $Z_1\text{-}Z_2$ (c,f), $Z_1\text{-}Z_3$ (d,g), and $Z_2\text{-}Z_3$ (e,h) planes, comparing UIDD (c-e) and OnsagerNet (f-h) across $3$ seeds. Here the projections are obtained via minimization (e.g., $V(Z_1, Z_2) = \min_{Z_3} V(Z_1, Z_2, Z_3)$), which closely approximates the marginal potential at low temperatures. (i.j) Energy landscape barrier heights of UIDD and OnsagerNet averaged over $4$ independent experiments; error bars indicate one standard deviation.
• Figure 4: EPR for polymer stretching dynamics. (a) The global EPR obtained via UIDD for varying strain rate $\dot{\varepsilon}$, averaged over $4$ independent experiments with error bars showing one standard deviation. (b) Illustration of the experimental validation and coarse-graining of the global EPR based on device-level observations. (c-f) The total EPR $\dot{s}_{\text{tot}}$ on test trajectories projected onto the $Z_1\text{-}Z_2$ and $Z_1\text{-}Z_3$ planes. The background is the corresponding potential energy surface. One example trajectory from the test dataset are overlaid on the $Z_1\text{-}Z_3$ panels.
• Figure 5: The global EPR and the squared MMD for SGLD with varying mini-batch sizes. Results represent the mean of $4$ independent experiments, with error bars denoting one standard deviation. Two case studies are considered: (a,b) least squares regression and (c,d) independent component analysis. In all panels, the squared MMD between the mini-batch ($p^b$) and full-batch ($p^\infty$) invariant distributions of the SGLD sampler is plotted as a reference. Panels (a) and (c) further display the global EPR, while panels (b) and (d) show the squared MMD between consecutive batch sizes ($p^b$ and $p^{b/2}$), illustrating how both measures vary with increasing batch sizes.

Theorems & Definitions (8)

• Theorem 1
• Lemma 1
• Theorem 2
• Lemma 2
• proof
• proof : Proof of Theorem \ref{['the:appro']}
• Proposition 1: Conditional KL Decomposition
• proof