DeepWKB: Learning WKB Expansions of Invariant Distributions for Stochastic Systems

Abstract

This paper introduces a novel deep learning method, called DeepWKB, for estimating the invariant distribution of randomly perturbed systems via its Wentzel-Kramers-Brillouin (WKB) approximation $u_\epsilon(x) = Q(\epsilon)^{-1} Z_\epsilon(x) \exp\{-V(x)/\epsilon\}$, where $V$ is known as the quasi-potential, $\epsilon$ denotes the noise strength, and $Q(\epsilon)$ is the normalization factor. By utilizing both Monte Carlo data and the partial differential equations satisfied by $V$ and $Z_\epsilon$, the DeepWKB method computes $V$ and $Z_\epsilon$ separately. This enables an approximation of the invariant distribution in the singular regime where $\epsilon$ is sufficiently small, which remains a significant challenge for most existing methods. Moreover, the DeepWKB method is applicable to higher-dimensional stochastic systems whose deterministic counterparts admit non-trivial attractors. In particular, it provides a scalable and flexible alternative for computing the quasi-potential, which plays a key role in the analysis of rare events, metastability, and the stochastic stability of complex systems.