Inferring the Langevin Equation with Uncertainty via Bayesian Neural Networks

TL;DR

The paper addresses the challenge of inferring Langevin dynamics from observed trajectories while quantifying uncertainty. It introduces Langevin Bayesian networks (LBN), which learn distributions over neural network weights to simultaneously infer the drift field $\bm{Φ}$ and diffusion matrix $\bm{\mathsf{D}}$ in Itô Langevin equations, providing ensemble predictions and uncertainty bounds. The framework yields unbiased estimators for both overdamped and underdamped regimes and demonstrates strong performance across nonlinear, inhomogeneous, neuronal, and thermodynamic systems, including Hodgkin–Huxley neurons and Brownian Carnot engines. This approach offers a scalable, binning-free, gradient-access method for data-driven stochastic dynamics with practical implications for nonequilibrium thermodynamics and complex systems modeling.

Abstract

Pervasive across diverse domains, stochastic systems exhibit fluctuations in processes ranging from molecular dynamics to climate phenomena. The Langevin equation has served as a common mathematical model for studying such systems, enabling predictions of their temporal evolution and analyses of thermodynamic quantities, including absorbed heat, work done on the system, and entropy production. However, inferring the Langevin equation from observed trajectories is a challenging problem, and assessing the uncertainty associated with the inferred equation has yet to be accomplished. In this study, we present a comprehensive framework that employs Bayesian neural networks for inferring Langevin equations in both overdamped and underdamped regimes. Our framework first provides the drift force and diffusion matrix separately and then combines them to construct the Langevin equation. By providing a distribution of predictions instead of a single value, our approach allows us to assess prediction uncertainties, which can help prevent potential misunderstandings and erroneous decisions about the system. We demonstrate the effectiveness of our framework in inferring Langevin equations for various scenarios including a neuron model and microscopic engine, highlighting its versatility and potential impact.

Figures (17)

• Figure 1: Schematics of our framework to infer the Langevin dynamics via Langevin Bayesian networks (LBN). To build the Langevin equation from observed trajectories ( a), we feed the trajectories to LBN and obtain the inferred drift vector $\hat{\bm{\Phi}}(\bm{z}, t)$ and diffusion matrix $\hat{\bm{\mathsf{D}}}(\bm{z}, t)$ ( b). The inferred Langevin equation enables us to predict system evolution over time and analyze thermodynamic quantities ( c and d). Here, the shaded areas represent the uncertainty of its prediction.
• Figure 2: Overdamped Langevin equation (OLE) inference for systems with a nonlinear force. (a) Example trajectory and the average of inferred drift fields $\hat{\bm{\Phi}}(\bm{x})$ (red arrows) for the $d=2$ case. The black arrows represent the true drift fields ${\bm{\Phi}}(\bm{x})$ and the colormap indicates the uncertainty of predictions. Inset: $\hat{{\Phi}}_\mu(\bm{x})$ vs. ${{\Phi}}_\mu(\bm{x})$ on $\mathcal{D}_{\rm te}$. (b) The generated trajectory (top) and the cumulative entropy production $\Delta S^t$ (bottom) along the trajectory are compared between the true model (solid lines) and the trained LBN (shaded area, behind the solid line) using the same random seed. (c) Total errors of LBN and stochastic force inference (SFI) for drift field ($\mathcal{E}^2_{\bm{\Phi}}$) and diffusion matrix ($\mathcal{E}^2_{\bm{\mathsf{D}}}$, inset) with increasing dimension $d$. (d) Pointwise errors for drift inference $e^2_{\Phi}$ vs. uncertainties of LBN $\hat{\Sigma}_{\Phi}$ for the $d=2$ case along a trajectory (top), and Pearson correlation coefficient between $e^2_{\Phi}$ and $\hat{\Sigma}_{\Phi}$ with $d$ ($r_{\bm{\Phi}}$, bottom). Error bars indicate the standard deviation of estimates from five independent trajectories and estimators.
• Figure 3: OLE inference for systems with an inhomogeneous diffusion matrix. (a) Example trajectory and diagonal entries of the diffusion matrix (red arrows), $[\hat{\mathsf{D}}_{11}(\bm{x}), \hat{\mathsf{D}}_{22}(\bm{x})]^{\rm T}$, for the $d=2$ case. The black arrows represent the true values $[{\mathsf{D}}_{11}(\bm{x}), {\mathsf{D}}_{22}(\bm{x})]^{\rm T}$ and the colormap indicates the uncertainty of predictions. Insets: $\hat{{\Phi}}_{\mu}(\bm{x})$ vs. ${\Phi}_{\mu}(\bm{x})$ (top) and $\hat{\mathsf{D}}_{\mu\nu}(\bm{x})$ vs. ${\mathsf{D}}_{\mu\nu}(\bm{x})$ (bottom) on $\mathcal{D}_{\rm te}$. (b) Exact and inferred forces, $F_\mu(\bm{x})$ (solid line) and $\hat{F}_\mu(\bm{x})$ (shaded area, behind the solid line), along a trajectory (top) and a scatter plot between them (bottom). (c) Total errors of LBN and SFI for the drift field ($\mathcal{E}^2_{\bm{\Phi}}$) and diffusion matrix ($\mathcal{E}^2_{\bm{\mathsf{D}}}$, inset) with increasing dimension $d$. (d, e) Pearson correlation coefficient $r_{\bm{\Phi}}$ between $e^2_{\Phi}$ and $\hat{\Sigma}_{\Phi}$ ($r_{\bm{\mathsf{D}}}$ between $e^2_{\bm{\mathsf{D}}}$ and $\hat{\Sigma}_{\bm{\mathsf{D}}}$) with increasing $d$ (d) and the duration $\tau$ of a training trajectory for $d=2$ and $d=4$ (e). Error bars indicate the standard deviation of estimates from five independent trajectories and estimators.
• Figure 4: OLE inference for the spiking neuron model. (a) Trajectories of four state variables generated from both the true Hodgkin--Huxley (HH) model (black) and the trained LBN (blue) with the same random seed. (b) Trajectories in the phase space of $V$ and $m$. (c) Autocorrelation function of $V(t)$, denoted by $ACF_{\bm{V}}(t_l)$, for both trajectories with respect to time lag $t_l$. (d) Generated trajectories at $I_{\rm ext}=-15~\text{mV}$ (left) and $I_{\rm ext}=30~\text{mV}$ (right) from the LBN trained at $I_{\rm ext}=0~\text{mV}$ and the true HH model. (e) Mean inter-spike interval ($ISI$), (f) mean of the minimum and maximum values of $V$, and (g) entropy production rate ($\langle \dot{S} \rangle$) with varying external current $I_{\rm ext}$. The blue shaded area in (g) indicates the uncertainty of $\langle \dot{S} \rangle$. When there are no spikes in the generated trajectories, the $ISI$ is set to the length of the trajectory ($=50~\text{ms}$). The red dotted line indicates $I_{\rm ext}=0$, which is the value used for training, and the red shaded area indicates the non-spiking region.
• Figure 5: Underdamped Langevin equation (ULE) inference for a stochastic van der Pol oscillator. (a) Example $xv$ trajectories and inferred drift fields (red arrows) for the $d=1$ case with increasing $\tau$ from left to right ($\tau \simeq 17$, $130$, and $1000$). The black arrows represent the exact fields and the colormap indicates the uncertainties of predictions $\rm{Var}[\bm{\Phi}_{\bm{\theta}}(x, v)]$. Insets: $\hat{\bm{\Phi}}(x, v)$ vs. ${\bm{\Phi}}(x, v)$. (b) Total errors of LBN for drift field ($\mathcal{E}^2_{\bm{\Phi}}$) and diffusion matrix ($\mathcal{E}^2_{\bm{\mathsf{D}}}$, inset) prediction with increasing $\tau$, respectively, for $d=1, \dots, 5$. Error bars indicate the standard deviation of estimates from five independent trajectories and estimators. (c) Trajectories of $x(t)$ generated from both true model (black) and trained LBN (blue) with the same random seed. (d) True effective friction coefficient $\gamma(x)$ and that inferred by LBN $\hat{\gamma}(x)$ with varying $x$. Here, the 'Naive' label indicates the inferred values by a straightforward extension of Eq. \ref{['eq:OLE_estimators']}, $\hat{\bm{\Psi}}_{f}$. (e) Example $xv$ trajectory and drift fields for the $d=1$ case with multiplicative noise. The meanings of the symbols and colormap are the same as in (a). (f) Scatter plots between ${{\Phi}}(x, v)$ and $\hat{{\Phi}}(x, v)$ (top) and ${\mathsf{D}}(x, v)$ and $\hat{\mathsf{D}}(x, v)$ (bottom) for the multiplicative noise case.