Solving for high dimensional committor functions using artificial neural networks

TL;DR

This paper introduces a neural network-based method to approximate high-dimensional committor functions, central to transition path theory, by solving associated Fokker-Planck equations through variational formulations and stochastic optimization.

Contribution

It presents a novel neural network approach for high-dimensional committor functions, enabling feasible numerical solutions where traditional methods struggle.

Findings

Moderate accuracy achieved for high-dimensional problems.

Neural network approach effectively approximates solutions to complex PDEs.

Method leverages variational formulation and stochastic optimization.

Abstract

In this note we propose a method based on artificial neural network to study the transition between states governed by stochastic processes. In particular, we aim for numerical schemes for the committor function, the central object of transition path theory, which satisfies a high-dimensional Fokker-Planck equation. By working with the variational formulation of such partial differential equation and parameterizing the committor function in terms of a neural network, approximations can be obtained via optimizing the neural network weights using stochastic algorithms. The numerical examples show that moderate accuracy can be achieved for high-dimensional problems.

Tables

Table 1. Table 1: Results for the double well potential ( 26 ) between two planes.

$T$	$E_1$	$E_2$	No. of parameters	$\rho$	No. of samples in $\Omega \setminus A\cup B$	$\alpha$	No. of testing samples
0.2	0.0054	0.0063	145	50	2.0e+04	1/20	1e+05
0.05	0.012	0.020	145	0.5	2.0e+04	1/20	1e+05

Table 2. Table 2: Results for the coecentric spheres example. We compare the cases when N s = 0 subscript 𝑁 𝑠 0 N_{s}=0 and N s = 1 subscript 𝑁 𝑠 1 N_{s}=1 .

$N_s$	$E_1$	$E_2$	No. of parameters	$\rho$	No. of samples in $\Omega \setminus A\cup B$	$\alpha$	No. of testing samples
1	0.053	0.015	542	5.3e+02	3e+04	1/30	1e+05
0	0.17	0.078	542	5.3e+02	3e+04	1/30	1e+05

Table 3. Table 3: Results for the rugged Muller potential example.

$(T,\sigma )$	$E_1$	$E_2$	No. of parameters	$\rho$	No. of samples in $\Omega \setminus A\cup B$	$\alpha$	No. of testing samples
(40, 0.05)	0.057	0.035	1011	3.8e+02	7.4e+04	1/74	7.4e+04
(22, 0.05)	0.037	0.036	1011	1.3e+02	1.5e+05	1/150	1.5e+05

Equations

$$d{\mathbf{X}}_t=-\nabla U({\mathbf{X}}_t)dt+\sqrt{2{\beta }^{-1}}d{\mathbf{W}}_t$$

$$q(\mathbf{x})=\mathbb{P}({\tau }_B<{\tau }_A\mid {\mathbf{X}}_0=\mathbf{x}),$$

$$0=-{\beta }^{-1}\Delta q(\mathbf{x})+\nabla U(\mathbf{x})\cdot \nabla q(\mathbf{x})\text{in}\Omega \setminus A\cup B,q(\mathbf{x}){|}_{\partial A}=0,q(\mathbf{x}){|}_{\partial B}=1.$$

$$q(\mathbf{x})\to q_{\theta }(\mathbf{x})$$

$$K({\mathbf{x}}^i,{\mathbf{x}}^j)=\frac{\exp (-|{\mathbf{x}}^i-{\mathbf{x}}^j{|}^2/2{ϵ}^2)}{\sqrt{\exp (-\beta U({\mathbf{x}}^i))\exp (-\beta U({\mathbf{x}}^j))}},D({\mathbf{x}}^i)=\sum _{j}K({\mathbf{x}}^i,{\mathbf{x}}^j),L=D^{-1}K-I$$

$$L(\Omega \setminus A\cup B,\Omega \setminus A\cup B)q_{\Omega \setminus A\cup B}=-L(\Omega \setminus A\cup B,B)q_B.$$

$$Lu=f,\mathbf{x}\in \Omega ,Bu=g,\mathbf{x}\in \partial \Omega$$

$$u(\mathbf{x})\approx a(\mathbf{x})+n(\mathbf{x})b(\mathbf{x})$$

$$\underset{u}{\inf }\Vert Lu-f{\Vert }_{L_2}^2,$$

$$\underset{q}{\mathrm{arginf}}\frac{1}{Z(\beta )}{\int }_{\Omega \setminus A\cup B}|\nabla q(\mathbf{x}){|}^2\exp (-\beta U(\mathbf{x}))d\mathbf{x},q(\mathbf{x}){|}_{\partial A}=0,q(\mathbf{x}){|}_{\partial B}=1,\text{boundary condition on}\partial \Omega .$$

$$0$$

$$=$$

$${\int }_{\Omega \setminus A\cup B}\nabla \cdot (\nabla q^{\ast }(\mathbf{x})\eta (\mathbf{x})\exp (-\beta U(\mathbf{x})))d\mathbf{x}={\int }_{\partial \Omega }\nabla q^{\ast }(\mathbf{x})\eta (\mathbf{x})\exp (-\beta U(\mathbf{x}))d\mathbf{s}=0$$

$${\int }_{\partial B_R}\nabla q^{\ast }(\mathbf{x})\exp (-\beta U(\mathbf{x}))d\mathbf{s}\to 0\text{as}R\to \infty$$

$$\underset{\theta \in {\mathbb{R}}^p}{\mathrm{argmin}}\frac{1}{Z(\beta )}{\int }_{\Omega \setminus A\cup B}|\nabla q_{\theta }(\mathbf{x}){|}^2\exp (-\beta U(\mathbf{x}))d\mathbf{x}+\rho {\int }_{\partial A}{q}_{\theta }(\mathbf{x}{)}^{2}d{\mu }_{\partial A}(\mathbf{x})+\rho {\int }_{\partial B}(q_{\theta }(\mathbf{x})-1{)}^{2}d{\mu }_{\partial B}(\mathbf{x}),$$

$$\underset{\theta\in\mathbb{R}^{p}}{\text{argmin}}\;\mathbb{E}_{\nu}\Bigl{(}|\nabla q_{\theta}(\mathbf{x})|^{2}\chi_{\Omega\setminus A\cup B}(\mathbf{x})+\frac{\rho}{\alpha}q_{\theta}(\mathbf{x})^{2}\chi_{\partial A}(\mathbf{x})+\frac{\rho}{\alpha}(q_{\theta}(\mathbf{x})-1)^{2}\chi_{\partial B}(\mathbf{x})\Bigr{)},$$

$$\nu(\mathbf{x})=\frac{1}{1+2\alpha}\Bigl{(}\frac{1}{Z(\beta)}e^{-\beta U(\mathbf{x})}d\mathbf{x}+\alpha\mu_{\partial A}(\mathbf{x})+\alpha\mu_{\partial B}(\mathbf{x})\Bigr{)},$$

$$\Delta q(\mathbf{x})=0\text{in}\Omega \setminus A\cup B,q(\mathbf{x}){|}_{\partial A}=0,q(\mathbf{x}){|}_{\partial B}=1,\nabla q(\mathbf{x})\sim o(1/|\mathbf{x}{|}^{d-1})\text{as}|\mathbf{x}|\to \infty .$$

$$\Delta \tilde{q}(\mathbf{x})=0\text{in}\Omega \setminus A\cup B,\tilde{q}(\mathbf{x}){|}_{\partial A}=-1,\tilde{q}(\mathbf{x}){|}_{\partial B}=1$$

$$G(\mathbf{x},\mathbf{y})=\frac{\Gamma (d/2)}{(2\pi {)}^{d/2}|\mathbf{x}-\mathbf{y}{|}^{d-2}},$$

$$\Delta G(\mathbf{x},\mathbf{y})=\delta (\mathbf{x}-\mathbf{y})$$

$$\tilde{q}(\mathbf{x})=\sum _{i=0}^{\infty }{c}_{i}G(\mathbf{x},{\mathbf{y}}^{+,i})-\sum _{i=0}^{\infty }{c}_{i}G(\mathbf{x},{\mathbf{y}}^{-,i})$$

$$c_i=c_{i-1}(r/w_i{)}^{d-2},w_i=w_0-r^2/w_{i-1},{\mathbf{y}}^{+,i}=(w_i-w_0/2){\mathbf{e}}_1,{\mathbf{y}}^{-,i}=-(w_i-w_0/2){\mathbf{e}}_1,$$

$$U(\mathbf{x})=(x_1^2-1{)}^2+0.3\sum _{i=2}^d{x}_i^2,$$

$$A=\{x\in {\mathbb{R}}^d|x_1\le -1\},B=\{x\in {\mathbb{R}}^d|x_1\ge 1\}.$$

$$\frac{1}{Z(\beta )}{\int }_{\Omega \setminus A\cup B}|\nabla q_{\theta }(\mathbf{x}){|}^2\exp (-\beta U(\mathbf{x}))d\mathbf{x}.$$

$$q_{\theta }(\mathbf{x})=1/2({\tilde{q}}_{\theta }(\mathbf{x})+1),{\tilde{q}}_{\theta }(\mathbf{x}):=n_{{\theta }_1}(\mathbf{x})G(\mathbf{x},{\mathbf{y}}^{+,0})+n_{{\theta }_2}(\mathbf{x})G(\mathbf{x},{\mathbf{y}}^{-,0})$$

$$q_{\theta }(\mathbf{x}):=\sum _{k=1}^{N_s}{n}_{{\theta }_k}(\mathbf{x})S_k(\mathbf{x},{\mathbf{y}}^k)+n_{{\theta }_0}(\mathbf{x}),$$

$$E_1=\frac{\Vert q_{\theta }-q{\Vert }_{L_2(\mu )}}{\Vert q{\Vert }_{L_2(\mu )}},E_2=\frac{|v_R(q_{\theta })-v_R(q)|}{v_R(q)}$$

Full text

1 Introduction

In this paper, we study the transition between two states described by the overdamped Langevin process

[TABLE]

where $\mathbf{X}_{t}\in\Omega\subset\mathbb{R}^{d}$, $U:\mathbb{R}^{d}\rightarrow\mathbb{R}$, $\beta=1/T$, and $\mathbf{W}_{t}$ is a $d$-dimensional Wiener process, using the transition path theory [7, 8, 15]. The central object in the transition path theory is the committor function. Let $\tau_{D}$ be the first hitting time for region $D\subset\Omega$. For two disjoint regions $A,B\subset\Omega$, the committor function is defined as

[TABLE]

which is the probability of hitting region $B$ before region $A$ with the stochastic process (1) starting at $\mathbf{x}$. The committor function $q(\mathbf{x})$ provides useful statistical description on properties such as density and probability current of reaction trajectory. However, obtaining the committor function $q(\mathbf{x})$ can be a formiddable task, as it involves solving a high-dimensional Fokker-Planck equation

[TABLE]

The high dimensional nature of (3) renders obtaining $q(\mathbf{x})$ via finite-element-type methods intractable. On the other hand, since the transition paths are often localized to a quasi one-dimensional reaction tube, the region of interest is rather small compare to $\Omega$. Under this approximation, the finite temperature string method [6, 18] is developed to simultaneously find the best “tube” and the corresponding committor function. Another approach is based on explicit dimension reduction using e.g., the leading eigenfunctions of the generator $-\beta^{-1}\Delta+\nabla U(\mathbf{x})\cdot\nabla$ by means of diffusion maps [5, 4] to approximate $q(\mathbf{x})$. More recently, a direct point cloud discretization for the Fokker-Planck equation [13] has been also considered.

In recent years, the artificial neural-network (NN) has shown great success in representing high-dimensional probablity distributions or classifiers in a variety of machine learning tasks [10, 14, 16]. Motivated by those recent success, in this note we use an NN to provide a low-dimensional parameterization of the committor function

[TABLE]

where $\mathbf{\theta}\in\mathbb{R}^{p}$ is the parameter vector of the NN. $q_{\theta}(\mathbf{x})$ is then obtained by minimizing the variational formulation of (3) over $\theta$ as a nonlinear Galerkin method. In this way, we switch from solving a partial differential equation to solving an optimization problem. Optimizing the cost in the variational formulation of (3) using gradient-descent-based method involves computing an integral with respect to the equilibrium measure. For doing such integration, we use a Monte-Carlo method where the samples are generated according to the stochastic process (1). An NN parameterized committor function can potentially be used to guide further sampling in transition regions wherein sample density is low (in similar spirit as [19] where an NN parameterized force field is used to guide sampling in the application of molecular dynamics).

This paper is organized as the following. In Section 2, we design an NN tailored for solving for the committor function. In Section 3, we demonstrate the success of the proposed method in a few examples. In Section 4, we conclude the note. Before moving on, in the next subsection we survey related methods for solving (3).

1.1 Previous approaches

A popular way to solve (3) is to discretize the generator $-\beta^{-1}\Delta+\nabla U(\mathbf{x})\cdot\nabla$ of the overdamped Langevin process (1) on the sampled points using diffusion map [5, 4]. The lower eigenmodes of such discretized operator in principle can provide a low-dimensional (nonlinear) reparameterization of the committor function. If the transition trajectories lie on a low dimensional manifold, it is possible to discretize the generator accurately via sampling. More precisely, with $N$ samples, let $K\in\mathbb{R}^{N\times N}$, the diagonal matrix $D\in\mathbb{R}^{N\times N}$, the discretized generator $L\in\mathbb{R}^{N\times N}$ be defined as

[TABLE]

respectively. Let $q_{A},q_{B},q_{\Omega\setminus A\cup B}$ be vectors corresponding to committor function values on the points belong to regions $A,B,\Omega\setminus A\cup B$, the committor function $q$ satisfies

[TABLE]

In the presence of a spectral gap the eigenmodes of $L(\Omega\setminus A\cup B,\Omega\setminus A\cup B)$ provides reduced coordinates for the reaction tube, thus solving (6) for $q$ can be seen as expanding $q$ using the reduced coordinates. However, discretizing the generator using diffusion map may suffer from low order of convergence. Therefore [13] improves upon diffusion map by explicitly constructing the tangent plane of each point in the sampled point cloud and discretizing the generator in each of the tangent plane.

On the other hand, recent years have seen usage of machine learning techniques in solving high-dimensional partial differential equations. The success of [3] where an NN parameterized spin wavefunction is used as an ansatz for solving the many-body Schrödinger equation motivates us to consider solving (3) using an NN as well. Our work is also similar to the methods in [12, 17, 2, 9] for solving partial differential equations. [12, 2] demonstrate success of NN-based method for solving boundary values problem

[TABLE]

by assuming

[TABLE]

where $b(\mathbf{x})=0$ on $\partial\Omega$, $a(\mathbf{x})$ is a smooth function that satisfies the boundary conditions on $\partial\Omega$, and $n(\mathbf{x})$ is an NN-parameterized function. Then $u$ is found from solving

[TABLE]

leading to an optimization problem over the NN parameters. The improvement of [2] over [12] is that $a(\mathbf{x})$ and $b(\mathbf{x})$ are also learned as a neural-network separately from $n(\mathbf{x})$, whereas in [12] they are obtained via explicit construction. Such methods remove the need of specifying basis for discretizing $u$ therefore can complement Galerkin-type method. While these methods obtain rather impressive results in low-dimension, their performance in high dimension is unexplored, which is in fact the most interesting regime. In a very recent work [17], a neural-network is used to parameterize the solution to a high-dimensional parabolic equation. Although such setting is similar the one we consider, in our case the boundary conditions might result singularities in $u$, making it more difficult to be approximated using an NN, which will be address in Section 2. Moreover, since [12, 17, 2] work with the strong form of a partial differential equation, the computational cost can be high as the second order derivative of $u(\mathbf{x})$ is needed, whereas our approach is based on the variational formulation of the PDE. Although using an NN in solving the variational formulation of a PDE [9] has been explored before, [9] does not face the type of singularity issue arises in our application.

2 Proposed method

In this section, we present the general strategy of solving (3) using a neural-network. For simplicity, we let $\Omega=\mathbb{R}^{d}$, $U(\mathbf{x})$ be a confining potential that gives rise to an equilibrium measure $\mu(\mathbf{x}):=\exp(-\beta U(\mathbf{x}))/Z(\beta)$ normalized on the region $\Omega\backslash(A\cup B)$, where $Z(\beta):=\int_{\Omega\setminus A\cup B}\exp(-\beta U(\mathbf{x}))d\mathbf{x}$. Instead of working with the strong form (3), we solve the variational problem

[TABLE]

To see the boundary conditions for $q$ on $\partial\Omega$, let $q^{*}(\mathbf{x})$ be the minimizer of (10) and $q(\mathbf{x},\lambda)=q^{*}(\mathbf{x})+\lambda\eta(\mathbf{x})$. Since $q^{*}(\mathbf{x})$ is a stationary point, for any $\eta(\mathbf{x})$

[TABLE]

The third equality follows from $\eta(\mathbf{x})=0$ on $\partial A,\partial B$, and requiring

[TABLE]

via imposing suitable boundary condition on $\partial\Omega$. Here $\int_{\partial\Omega}d\mathbf{s}$ stands for the surface integral. When the domain is unbounded as the considered case, the condition

[TABLE]

where $B_{R}$ denotes a ball with radius $R$, can ensure (16). Notice that we simply need $\nabla q(\mathbf{x})$ to have subexponential growth as $|\mathbf{x}|\rightarrow\infty$ for (17) to hold, as long as $\exp(-\beta U(\mathbf{x}))\leq\exp(-a|\mathbf{x}|)$ when $|\mathbf{x}|>R$ for some $R,a>0$. The last equality in (11) implies that a solution to (10) provides a solution to (3).

As mentioned earlier, to cope with the high-dimensionality of $q(\mathbf{x})$, the proposed method consists of parameterizing $q(\mathbf{x})$ as an NN function $q_{\theta}(\mathbf{x})$. Instead of (10), we solve

[TABLE]

where the boundary conditions are only enforced as soft-constraints (with hardness tuned by the choice of $\rho$). The first integral is then approximated via sampling according to the overdamped Langevin process (1). To approximate the second and third integrals, for our problems there exist rather convenient scheme for drawing samples from $\mu_{\partial A}(\mathbf{x})$, $\mu_{\partial B}(\mathbf{x})$. The choice of the measures on the boundaries is based on the consideration of sampling convenience. In our examples, we mainly work with regions $A$ and $B$ being balls, therefore the samples on $\partial A$ and $\partial B$ are drawn by normalizing and recentering normally distributed samples. Note that we can rewrite (18) as a single expectation

[TABLE]

if we define a mixture measure

[TABLE]

where $\chi_{\Sigma}(\mathbf{x})$ is the characteristic function of region $\Sigma$, $\alpha$ is a parameter that controls the proportion between the sample size in $\partial A\cup\partial B$ with the sample size in $\Omega\backslash(A\cup B)$. This allows us to solve for (3) as an optimization problem (19) over NN weights using stochastic gradient type methods based on stochastic approximation of the expectation. While this is in principle straightforward, challenges arise due to the specific nature of the high dimensional Fokker-Planck equation we aim to solve. We discuss those challenges below and then the proposed neural-network design to overcome them.

2.1 Challenge in high $T$ regime

Although this method seems straight-forward, the main difficulty in using an NN to approximate the committor function is that in some situations, singularities are present within region $A$ and $B$. Consider the case where $A$ and $B$ are two balls of radius $r$ and $\Omega=\mathbb{R}^{d},d\geq 3$, centered at $(-w_{0}/2,0,\ldots,0),(w_{0}/2,0,\ldots,0)$. When $T\rightarrow\infty$ and $\beta\rightarrow 0$, (3) becomes

[TABLE]

The last boundary condition is there in order to satisfy (17). A solution to (20) can be obtained by first solving the Laplace’s equation with Dirichlet’s boundary conditions

[TABLE]

and letting $q(\mathbf{x})=(1/2)(\tilde{q}(\mathbf{x})+1)$. A classical way to solve (21) analytically is via method of images. Using the Green’s function

[TABLE]

which solves

[TABLE]

where $\Gamma(\cdot)$ is the gamma function, the solution of (21) can be obtained as

[TABLE]

where

[TABLE]

where $c_{0}=(2\pi)^{d/2}/\Gamma(d/2)$. Due to the singularities in $\tilde{q}$, the committor function may be steep near regions $A,B$ which can present difficulties when using an NN approximation. To illustrate, we let $d=3$, $w_{0}=1,r=0.15$ and we plot the solution of (20) along the $x_{1}$-dimension in Fig. 1. As a contrast, we minimize (19) using an NN with 3 hidden-layers and $\tanh$ nonlinearities, where each hidden-layer has 12 nodes. 3e+04 samples sampled uniformly from the box $[-2,2]^{3}$ are used in the optimization problem. We let $\rho=666$ and $\alpha=1/15$ to enforce the boundary condition on $\partial A$ and $\partial B$. As shown in Fig. 1, the NN has difficulty capturing the behavior of the committor function near the singularities.

2.2 Challenge in low $T$ regime

A different type of singularities can exist in the low temperature regime. Consider the potential

[TABLE]

and

[TABLE]

Here, (26) resembles a double potential well, and $\partial A$ and $\partial B$ are located in the potential wells. When the temperature is low, the equilibrium distribution for such $U(\mathbf{x})$ is concentrated in $A$ and $B$ while the midpoint of $A$ and $B$ has a low density. Therefore, as $q(\mathbf{x})$ goes from 0 to 1 when $x_{1}$ goes from -1 to 1, it is preferrable for $|\nabla q(\mathbf{x})|$ to concentrate around the midpoint of $A$ and $B$ in order to have a low cost

[TABLE]

As an example, we plot the committor function when $T=0.05$ in Fig. 2 when $d=10$. While the committor function is steep around $x_{1}=0$, unlike the case of high $T$, an NN with a single hidden layer with $\tanh$ activation function gives a good approximation. Since our goal is simply to show qualitatively that an NN is capable of handling such singularity issue, we defer the implementation details to Section 3.

2.3 Neural-network architecture

In this subsection, we present an NN network architecture that can deal with the aforementioned challenges. When solving problem (20) via its variational formulation, a natural choice of the NN architecture is to mimic (24) and take

[TABLE]

where $n_{\theta_{1}},n_{\theta_{2}}$ are neural-network parameterized functions, $\theta=[\theta_{1}^{T},\theta_{2}^{T}]^{T}$. By this ansatz, we explicitly remove the dominant singularities at $\mathbf{y}^{+,0},\mathbf{y}^{-,0}$. On the other hand, in order to determine the committor function of the transition process in the potential (26) when $\beta\rightarrow\infty$, $(\tanh(wf(x_{1}))+1)/2$ where $w$ is a large scalar and $f$ is certain smooth function may be used to capture the sharp transition around the midpoint between $\partial A$ and $\partial B$. This suggests using $\tanh$ as nonlinearity in an NN.

Therefore, to deal the issue of singularity, we propose the following neural-network architecture to solve for the committor function：

[TABLE]

where $n_{\theta_{k}}(\mathbf{x})$’s are functions parameterized by the NN, $N_{s}$ is the number of singularities, and each $S_{k}(\mathbf{x},\mathbf{y}^{k})$ is a problem dependent function with singularity at $\mathbf{y}^{k}$. The vector $\theta=[\theta_{0}^{T},\ldots,\theta_{N_{s}}^{T}]^{T}$ contains the parameters of the neural networks. Except $n_{\theta_{0}}(\mathbf{x})$, each $n_{\theta_{k}}(\mathbf{x})$ is an NN with 3 hidden layers where each hidden layer has 6 nodes. $n_{\theta_{0}}(\mathbf{x})$ consists of multiple hidden layers each having 12 nodes. A hyperparameter we tune here is the number of hidden layers in $n_{\theta_{0}}(\mathbf{x})$, where the choice of it is made using cross-validation. More precisely, the NN for committor function should give similar cost $\int_{\mathbf{x}\in\Omega\setminus A\cup B}|\nabla q_{\theta}(\mathbf{x})|^{2}\exp(-\beta U(\mathbf{x}))d\mathbf{x}$ in training and testing samples. We use $\tanh$ as the activation function of the hidden nodes. At high temperature, we expect singularities of $1/|\mathbf{x}|^{d-2}$ type to be dominant, whereas at low temperature the function $n_{\theta_{0}}(\mathbf{x})$ with $\tanh$ nonlinearities should be the main contributor to the committor function. The pipeline of solving for $q_{\theta}(\mathbf{x})$ is depicted in Fig. 3. As shown in Fig. 4, when using such architecture to solve for the variational form of (20), we indeed recover the $1/|\mathbf{x}|$ type behavior near $A$ and $B$.

3 Numerical experiments

In this section, we evaluate the proposed method in a few numerical examples. In these examples, the minimization in (19) using such NN architecture is done using the Adam [11] optimizer, a variant of stochastic gradient descent, in the TensorFlow [1] engine. The ratio $2\alpha$ of samples on $\partial A\cup\partial B$ to samples on $\Omega\setminus\partial A\cup\partial B$ are kept between 1/10 to 1/100. Then $\rho$ is tuned in order to have the boundary conditions satisfied with $10^{-3}$ accuracy. In all of the experiments, 2000 boundary samples are used, and we set the batch size to be 3000 in the Adam optimizer. We evaluate the performance using the following metric

[TABLE]

where $v_{R}(q):=k_{B}T/Z(\beta)\int_{\Omega\setminus A\cup B}|\nabla q(x)|^{2}\exp(-\beta U(x))d\mathbf{x}$ is the rate of reaction. Note that $v_{R}$ is the energy one minimizes for $q$ in the variational formulation. To calculate these errors, we generate samples by simulating the stochastic process (1).

In the first numerical experiment, we solve for the committor function in the potential well (26) with regions $A$ and $B$ being (27) when $d=10$. In this case, $q(\mathbf{x})=f(x_{1})$ where

[TABLE]

To solve this problem using an NN, we set $N_{s}=0$ in (30) as there is no singularity in this problem. In $n_{\theta_{0}}(\mathbf{x})$, only one hidden layer is used. In this example, we sample differently from what is presented in (19). When $T$ is small, it is difficult to obtain sufficient samples near the $\mathbf{x}=0$ saddle point of $U(\mathbf{x})$. Therefore, instead of working with (19) directly, we sample $x_{1}$ uniformly from $[-1,1]$, $(x_{2},\ldots,x_{d})$ from a $d-1$-dimensional gaussian distribution, and change the first term of the integrand in (19) from $|\nabla q_{\theta}(\mathbf{x})|^{2}\chi_{\Omega\setminus A\cup B}(\mathbf{x})$ to

[TABLE]

to ensure sufficient sample coverages along $x_{1}$. In this case,

[TABLE]

For a subset of these samples, we let $x_{1}=1,-1$ to get the samples on the boundaries $\partial A,\partial B$. We use a separate batch of samples, serving as validation dataset, to determine $E_{1}$ and $E_{2}$. In Table 1 we report the error and the number of samples used for solving this problem in dimension $d=10$ with temperature $T=0.2,0.05$.

In the second experiment, we solve for the committor function for the transition process between a pair of coecentric spheres, with potential

[TABLE]

and the regions

[TABLE]

In this example, even with moderate $T$, the committor function still display a singular behavior $q\sim 1/|\mathbf{x}|^{d-2},d\geq 3$. Therefore in (30) we let $N_{s}=1$, $S_{1}=1/|\mathbf{x}|^{d-2}$. We use 3 hidden layers for $n_{\theta_{0}}$. The equilibrium density is proportional to $\exp(-\beta|\mathbf{x}|^{2})$, therefore the samples can be drawn from the gaussian distribution. The samples on the two boundaries are obtained via rescaling samples from the normal distribution to have norm $a$ or $b$. The results for $T=2$, $d=6$, $a=1$, $b=0.25$ are summarized in Table 2. We compare the solution with and without including the $1/|\mathbf{x}|^{d-2}$ type singularity. It is worth noting that the explicit inclusion of singularity is rather important for this example even at moderate temperature. In Fig. 5a, we plot $q_{\theta}(\mathbf{x})$ along several randomly chosen radial directions to check whether $q_{\theta}(\mathbf{x})$ is close to a single-variable function when explicitly including a singular function in the NN architecture. In Fig. 5b, we plot the NN committor function when the singularity is not explicitly taken care of.

In the third experiment, we work with the rugged-Muller potential

[TABLE]

considered in [13]. This is a Muller potential perturbed by a rugged potential in the first two dimension, where the roughness is controlled by $\gamma,k$. In the rest of the dimensions, we place a quadratic potential well where its strength is controlled by $\sigma$. The domain $\Omega=[-1.5,1]\times[-0.5,2]$. The parameters in (37) are taken from [13], for completeness we provide them in the following:

[TABLE]

In this example, we let $T=40,22$, regions $A$ and $B$ being two balls with radius 0.1 centered at $(-0.57,1.43)$ and $(-0.56,0.044)$. The points are again sampled using Euler-Maruyama scheme. The ground truth is obtained via applying finite element method on uniform grid to (3), where the code is provided by the authors of [13]. We use an NN with two singularities of type $\log(|\mathbf{x}-\mathbf{y}|)$ where $\mathbf{y}$ is the position of singularity, and $n_{\theta_{0}}$ that has 3 hidden layers. The results are reported in Table 3 and the contours of the committor function are shown in Fig. 6. As shown in the table, although we can achieve few percents accuracy, for the case with lower temperature more samples are needed to determine the committor function (since the equilibrium distribution is less smooth).

4 Conclusion

In this note, we develop method based on neural-network to represent the high-dimensional committor function. The neural-network parameters are found via optimizing the variational form of the Fokker-Planck equation. In order to better approximate the committor function, the NN function has to be designed carefully in order to deal with the singularities in high and low $T$ regime. Through numerical experiments, we show the usefulness of the proposed alternative approach in dealing with high-dimensional partial differential equations. We remark that the quality of the learned committor function depends crucially on sampling. When the temperature is low, due to the sparsity of samples between regions $A$ and $B$ when a naive sampling scheme is used, the NN approximation to the committor function tends to make a transition that is too sharp compare to the ground truth. The usage of enhanced sampling schemes, for example using the currently learned NN to guide further sampling, is certainly an important future direction to investigate.