DeepWKB: Learning WKB Expansions of Invariant Distributions for Stochastic Systems

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ε(x) = Q(ε)-1 Zε(x) exp-V(x)/ε, where V is known as the quasi-potential, ε denotes the noise strength, and Q(ε) is the normalization factor. By utilizing both Monte Carlo data and the partial differential equations satisfied by V and Zε, the DeepWKB method computes V and Zε separately. This enables an approximation of the invariant distribution in the singular regime where ε 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.