Goal-oriented learning of stochastic dynamical systems using error bounds on path-space observables

Abstract

The governing equations of stochastic dynamical systems often become cost-prohibitive for numerical simulation at large scales. Surrogate models of the governing equations, learned from data of the high-fidelity system, are routinely used to predict key observables with greater efficiency. However, standard choices of loss function for learning the surrogate model fail to provide error guarantees in path-dependent observables, such as reaction rates of molecular dynamical systems. This paper introduces an error bound for path-space observables and employs it as a novel variational loss for the goal-oriented learning of a stochastic dynamical system. We show the error bound holds for a broad class of observables, including mean first hitting times on unbounded time domains. We derive an analytical gradient of the goal-oriented loss function by leveraging the formula for Frechet derivatives of expected path functionals, which remains tractable for implementation in stochastic gradient descent schemes. We demonstrate that surrogate models of overdamped Langevin systems developed via goal-oriented learning achieve improved accuracy in predicting the statistics of a first hitting time observable and robustness to distributional shift in the data.

Figures (7)

• Figure 1: Asymmetric double well potential (left) and the associated invariant distribution (center) as it varies with $\theta$. The first exit time from the domain $(\infty, 1]$ with initial condition $x=-1$ and the histogram of states visited by a single sample trajectory (right).
• Figure 2: The absolute error in the mean first exit time (black) of the double well system compared to various forward (solid) and reverse (dashed) quantities: the KL divergence in invariant measures (KL and KL-r), the KL divergence in path measures (path KL and path KL-r), and the goal-oriented error bound (GO and GO-r). The reference model is defined by $\theta^*=0.5$.
• Figure 3: Left: Reference potential $V_{\theta^*}$, initial potential $V_{\theta_0}$, and training data $\mathcal{D}$ from the reference process. Right: Comparison of the error in the observable (dashed) to the loss (solid) over 500 learning iterations with four objectives: forward relative entropy rate (RER), reverse relative entropy rate (RER-r), forward goal-oriented loss (GO), and reverse goal-oriented loss (GO-r).
• Figure 4: We compare the following quantities over gradient descent iterations with the forward GO loss (left column) and reverse GO loss (right column): elements of the loss, the second moment of the observable $\tilde{\mathbb{E}}^x_{\theta,[0,T]}[\phi_T^2]$ and the relative entropy rate $\mathcal{H}$ (top row); and the $\ell_2$ norm of each term in the loss gradient, $\mathcal{H} \cdot \nabla_\theta \tilde{\mathbb{E}}^x_{\theta,[0,T]}[\phi_T^2]$ and $(\mathcal{M}_\phi + \tilde{\mathbb{E}}^x_{\theta,[0,T]}[\phi_T^2])\nabla_\theta \mathcal{H}$ (bottom row).
• Figure 5: The mean (left) and variance (right) of the first exit time associated with 500 model iterates from the four loss functions, compared to their reference values (dotted black line).

Theorems & Definitions (17)

• Definition 2.3: Path-space observable
• Definition 2.4: Kullback-Leibler (KL) divergence
• Definition 2.5: Relative entropy rate
• Lemma 3.1: Goal-oriented error bound
• Remark 3.2
• Remark 3.3
• Definition 3.4: Goal oriented loss functions
• Proposition 3.5: Goal-oriented losses bound error in the observable
• Proposition 3.9: Gradient of path-space observables
• Proposition 3.10: Gradient of goal-oriented loss functions