Solving Non-linear Kolmogorov Equations in Large Dimensions by Using Deep Learning: A Numerical Comparison of Discretization Schemes

Abstract

Non-linear partial differential Kolmogorov equations are successfully used to describe a wide range of time dependent phenomena, in natural sciences, engineering or even finance. For example, in physical systems, the Allen–Cahn equation describes pattern formation associated to phase transitions. In finance, instead, the Black–Scholes equation describes the evolution of the price of derivative investment instruments. Such modern applications often require to solve these equations in high-dimensional regimes in which classical approaches are ineffective. Recently, an interesting new approach based on deep learning has been introduced byby E, Han and Jentzen [1, 2]. The main idea is to construct a deep network which is trained from the samples of discrete stochastic differential equations underlying Kolmogorov’s equation. The network is able to approximate, numerically at least, the solutions of the Kolmogorov equation with polynomial complexity in whole spatial domains. In this contribution we study variants of the deep networks by using different discretizations schemes of the stochastic differential equation. We compare the performance of the associated networks, on benchmarked examples, and show that, for some discretization schemes, improvements in the accuracy are possible without affecting the observed computational complexity.

Appendices

Appendix A: Derivation of Equation (28)

For completeness we derive Eq. (28) which follows from the Leimkuhler and Matthews discretization of stochastic trajectories given by (21). Note the differences between this equation and (24) implied by the Euler–Maruyama discretization (19). Using (14) and applying Ito’s lemma (i.e., expanding to second order) we get for small enough $\tau$

$$\begin{aligned} g[\vec {Y}^n_{LM}, \tau _{n+1}]-g[\vec {Y}^n_{LM},\tau _n]= & {} \partial _t g[\vec {Y}^n_{LM},\tau _n]\tau + \nabla _{\vec x} g[\vec {Y}^n_{LM},\tau _n]^TA[\vec {Y}^n_{LM},\tau _n]\tau \nonumber \\{} & {} +\frac{1}{2}\nabla g[\vec {Y}^n_{LM},\tau _n]^T B[\vec {Y}^n_{LM},\tau _n] (\Delta \vec {W}^n+\Delta \vec {W}^{n+1})\nonumber \\{} & {} +\frac{1}{8}\,\mathbb {E}\bigl [(\Delta \vec {W}^n+\Delta \vec {W}^{n+1})^T B^T[\vec {Y}^n_{LM}(\tau _{n}),\tau _n] \text {Hess}_{\vec {x}}g\nonumber \\{} & {} [\vec {Y}^n_{LM}(\tau _{n}),\tau _n]B(\Delta \vec {W}^n+\Delta \vec {W}^{n+1})\bigr ]. \end{aligned}$$

(42)

To evaluate the last term we work in components and use $\mathbb {E}[\Delta W_j^n\Delta W_i^n] = \delta _{ij}\tau$ and $\mathbb {E}[\Delta W_j^n\Delta W_i^{n+1}] = 0$ (because increments are independent). This yields

$$\begin{aligned} g[\vec {Y}^n_{LM}, \tau _{n+1}]-g[\vec {Y}^n_{LM},\tau _n]= & {} \partial _t g[\vec {Y}^n_{LM},\tau _n]\tau + \nabla _{\vec x} g[\vec {Y}^n_{LM},\tau _n]^TA[\vec {Y}^n_{LM},\tau _n]\tau \nonumber \\{} & {} +\frac{1}{2}\nabla g[\vec {Y}^n_{LM},\tau _n]^T B[\vec {Y}^n_{LM},\tau _n] (\Delta \vec {W}^n+\Delta \vec {W}^{n+1})\nonumber \\{} & {} + \frac{1}{4}\text {Tr}\{B B^T[\vec {Y}^n_{LM}(\tau _{n}),\tau _n] \text {Hess}_{\vec {x}}g[\vec {Y}^n_{LM}(\tau _{n}),\tau _n]\}\tau . \end{aligned}$$

(43)

Finally, using (1) for the partial derivative with respect to time we obtain (28).

Appendix B: Further Assessment of the LM Scheme for the Heat Equation

The Leihmkuhler–Matthews (LM) scheme appears to perform worse than both the Euler–Maruyama and Milstein schemes for the Black-Scholes equation (Sect. 5.1). This is also the case when the diffusion tensor is homogeneous as tested on the Allen–Cahn equation (Sect. 5.2). However we observe slight improvements when the time step $\tau$ of the discretization becomes smaller, and we therefore conclude that the LM scheme is limited by the number of neural networks needed for each time-slice. Taking smaller time steps adds networks which becomes prohibitive in terms of memory requirements.

However, the LM scheme is well suited for sampling from stationary distributions, and it is therefore natural to ask if this quality persists when solving an equation admitting a stationary limiting distribution. The simplest such equation is the heat equation that we discuss here. We find that for the same discretization interval used in the non-linear equations the LM scheme still leads to higher errors than Euler-Maruyama.

Fig. 11

Comparison between Leimkuhler-Matthews (red points) and Euler-Maruyama (blue points) schemes for the simple heat equation in ten dimensions. The comparison is performed by computing the averaged relative approximation error defined in Eq. (44) as a function of the number of initial points (P) for stochastic trajectoties used as inputs to the $\mathcal {DNN}$’s. We observe that the Euler-Maruyama scheme is superior (Color figure online)

We look at the (forward) heat equation $\frac{\partial g(\vec {x}, t)}{\partial t}=\Delta _x g(\vec {x}, t)$ in $d=10$ dimensions, on the interval $t \in [0,T]$, and with initial condition $g(\vec {x}, 0)=||\vec {x}||^2_{\mathbb {R}^d}$. The exact solution is given by $g(\vec {x}, t)=||\vec {x}||^2_{\mathbb {R}^d} + td$.

We only compare the Leimkuhler–Matthews scheme with the Euler-Maruyama one (since the diffusion coefficient is constant). Figure 11 shows the comparison between the two schemes by computing the averaged relative approximation error $\langle \epsilon \rangle$:

$$\begin{aligned} \langle \epsilon \rangle = \frac{1}{N} \sum _{i=1}^N \left| \frac{g(\vec {x}_i, T) - \mathcal {N}\mathcal {N}(\vec {x}_i, T|\vec {\theta })}{g(\vec {x}_i, T)} \right| \end{aligned}$$

(44)

over $N=10^4$ points $\vec {x}_i$, $i=1,\dots ,N$, randomly chosen in $[0,1]^d$, with $\mathcal {N}\mathcal {N}(\vec {x}, T|\vec {\theta })$ the solution returned by the respective $\mathcal {DNN}$. Figure 11 shows that the Euler-Maruyama scheme is able to approximate better the solution of the heat equation than the Leimkuhler–Matthews one.

Appendix C: A Simpler Deep Learning Algorithm for Linear Equations. The Example of Geometric Brownian Motion

The main focus of this paper is on non-linear Kolmogorov equations, however we briefly discuss for completeness a simpler algorithm that applies to linear equations. Physical examples for applications of linear Kolmogorv equations can be found in [42,43,44]. For linear equations the standard Feynman-Kac formula is simply given by

$$\begin{aligned} g(\vec {x}, t) =\mathbb {E}_{\vec x} [\phi (\vec {X}(T))] \end{aligned}$$

(45)

i.e., the first term in (7). In [52] it is beautifully remarked that this expectation minimizing a certain "mean square error" and that this can be taken as the basis of a deep learning algorithm. The main idea is to minimize the following loss:

$$\begin{aligned} \mathcal {L}(\vec {\theta })=\frac{1}{\vert b-a\vert ^d}\int _{[a,b]^d} d\vec {x}\mathbb {E}_{\vec x}[ \vert \phi (\vec {X}(T)) - \mathcal {N}\mathcal {N}(\vec {x}, T\vert \vec {\theta })\vert ^2\bigr ] \end{aligned}$$

(46)

In practice the expectations are replaced by empirical averages over a sample set of discretized trajectories and the minimization over $\vec {\theta }$ carried on by a gradient descent algorithm. One important simplification with respect to the case of non-linear equations is that only one deep network at time T is optimized. We refer to [52] for further details of the algorithm. Our purpose here is to briefly compare the errors incurred by Euler-Maruyama and Milstein discretizations of the trajectories.

We work with a geometric Brownian motion, i.e., the linear part of the (backward) Black–Scholes Eq. (30),

$$\begin{aligned} \frac{\partial g(\vec {x},t)}{\partial t} + \frac{1}{2}\sum ^d_{i=1}|\sigma _i x_i|^2 \frac{\partial ^2 g}{\partial x^2_i}(\vec {x},t)+\sum ^d_{i=1}\mu x_i\frac{\partial g}{\partial x_i}(\vec {x},t) = 0, \qquad g(\vec {x}, T) = \psi (\vec {x}) \end{aligned}$$

(47)

for which the exact solution $g(\vec {x}, t)$, $0\le t\le T$ is known,

$$\begin{aligned} g(\vec {x}, t)= & {} \mathbb {E}\bigg [\psi \bigg (x_1 \exp \bigg (\sigma _1 W^{T-t}_1 + \bigg (\mu -\frac{|\sigma _1|^2}{2}\bigg )(T-t)\bigg ), \cdots , x_d \nonumber \\{} & {} \exp \bigg (\sigma _d W^{T-t}_1 + \bigg (\mu -\frac{|\sigma _1|^2}{2}\bigg )(T-t)\bigg )\bigg )\bigg ] \end{aligned}$$

(48)

where $\vec {W}^t$ is the standard Brownian motion (conditioned to start at the origin $\vec {W}^0 = 0$). This is a one-dimensional integral that can be computed by the Monte Carlo method.

The performance by computing the average relative error defined as

$$\begin{aligned} \epsilon = \int _{[0,1]^d} \left| \frac{g(\vec {x}, 0) - \mathcal{N}\mathcal{N}(\vec {x}, 0|\vec {\theta })}{g(\vec {x}, 0)} \right| d\vec {x} \end{aligned}$$

(49)

where $\mathcal{N}\mathcal{N}(\vec {x}, 0|\vec {\theta })$, at $t=0$, is the approximate solution obtained by the deep learning algorithm. We take $r=\frac{1}{20}$, $\mu =r-\frac{1}{10}$, $\sigma _i=\frac{1}{10}+\frac{i}{200}$, $d=100$, $t\in [0,T=1]$, by using a time size interval $\tau =T/N$, with $N=40$. We fix the the initial condition to $g(\vec {x},0)=\psi (\vec {x})=\exp {(-rT)}\max \{[\max _{i \in {1,2,\ldots ,d}}x_i]-100,0\}$ (as in [52]).

In Table 2 and Fig. 12 we report the results obtained for computing the relative error $\epsilon$ by using the two discretization schemes schemes in order to minimize the empirical averages corresponding to (46). We also average $\epsilon$ over 15 different experiments to reduce the variance. The analysis shows that exists an improvement in choosing an higher order discretization scheme for training the neural network.

Fig. 12

Comparison between Milstein (green points) and Euler–Maruyama (blue points) schemes. The comparison is performed by computing the relative approximation error defined in Eq. (49) as a function of the number of steps done for training the parameters $\vec {\theta }$. The figure shows that the Milstein scheme better, on average, approximates the solution of the linear Black–Scholes equation as compared to the Euler–Maruyama one. Error bars are SEM (Color figure online)

Table 2 The table shows the average values of relative error $\epsilon$

Appendix D: Computational Complexity for the Milstein Scheme

In Sect. 5.3, we presented using the Euler–Maruyama scheme that the deep learning methodology studied in this paper does not suffer from the curse of dimensionality, in the sense that the number of training trajectories used scales polynomially in the dimension of space. In this appendix, however, we present the same analysis for the Milstein scheme, and we show that the number of training trajectories scales polynomially in the dimension of space.

As far as we know, it is not know in literature an exact solution for a Kolmogorov non-linear equation with non constant transport coefficients. For this reason, we perform an accurate sampling of solutions using the multilevel Picard method for the equation (32) in the interval $[49.995:50.005]^d$ and then we use those solutions as reference values for our analysis.

The deep network of Sect. 4.2 (with Milstein discretization) is used since the diffusion term is non-homogeneous. For the number of initial random points in $[49.995:50.005]^d$ of stochastic input trajectories we take $P=4,8,16,32, 64, 128, 256, 512, 1024$. Once the network is trained we obtain the approximate solution $\mathcal {N}\mathcal {N}(\vec {x},0|\vec {\theta })$. The relative error is then easily computed from (36), using $M=10^2$ uniformly random points $\vec {x}_m\in [49.995:50.005]^d$, $m=1, \cdots , M$, computed with the multilevel Picard method.

Figure 13 shows the behaviour of $\langle \epsilon \rangle$ as a function of P, for various values of d. For small value of d, the averaged relative error needs few points P to approximate well the exact solution. In contrast, as expected, as d grows, the number of points P needed for approximate well the solution becomes larger.

By fixing a pre-specified relative error of $\langle \epsilon \rangle \approx 0.0011$ we observe that P scales polynomially with d, roughly as $P=O(d^{1.49})$. This value of P is comparable with the one obtained for the Euler-Maruyama scheme in Sect. 5.3. We conclude that the Milstein scheme improve the accuracy of the Neural Network, without modifying the computational complexity, as expected.

Fig. 13

Left: Average relative error $\langle \epsilon \rangle$ as a function of number of initial points of stochastic trajectories $P=4,8,16,32, 64, 128, 256, 512, 1024$, on a log-log scale i.e. for different values of $d=2,4,6,8,10$. Right: Computational complexity, given an average relative error $\langle \epsilon \rangle < 0.0011$, measured by P as a function of $d=2,4,6,8,10$, on a log-log scale the slope is $b\approx 1.49$. The computational complexity obtained is proportional to $A(\epsilon )d^{1.49}$, with $A(\epsilon ) \sim 7.6$. The results are obtained using the algorithm described in Sect. 4.2, for the Eq. (32) with $T=0.01$

Appendix E: HJB Equation with non Constant Diffusion Coefficient

In Sect. 5.4, we presented a comparison between Euler–Maruyama and Leihmkuhler–Matthews scheme over the Hamilton–Jacobi–Bellman equation in (40). In this appendix we present a comparison between the Euler-Maruyama and a Mistein scheme over an Hamilton–Jacobi–Bellman equation with non constant diffusion tensor. The equation that we try to solve is:

$$\begin{aligned} \frac{\partial g}{\partial t}(\vec {x},t) + \textbf{D}(\vec {x})\Delta g(\vec {x},t) = \lambda ||\nabla g(\vec {x},t)||^2, \end{aligned}$$

(50)

where $\lambda$ is a positive constant set to 1, and $D(\vec {x})_{ij}=x_i^2 \delta _{ij}$, with $\delta _{ij}$ is the $\delta$-Kronecker. The terminal condition is $g(\vec {x}, T)=\phi (\vec {x})=\ln ((1+||\vec {x}||^2)/2)$ with $\vec {x} \in \mathbb {R}^d$.

In Fig. 14 we present the comparison between the output of $\mathcal {D}\mathcal {N}\mathcal {N}_{\textrm{EM}}$ (blue points) and $\mathcal {D}\mathcal {N}\mathcal {N}_{\textrm{M}}$ (green points). Each point is the average over 5 samples. Blue points identify the output of $\mathcal {D}\mathcal {N}\mathcal {N}_{\textrm{EM}}$ and reach the approximate value $g_{\textrm{EM}}(\vec {x}=(50, \dots , 50), 0)=10.049\pm 0.007$. The green points represent the value obtained by $\mathcal {D}\mathcal {N}\mathcal {N}_{\textrm{M}}$ $g_{M}(\vec {x}=(50, \dots , 50), 0)=10.008 \pm 0.008$. Again, it is well shown the different results of the two schemes. No numerical or analytical result is present in the literature for Eq. (50), but numerical observations show that the true value of the solution of $g(\vec {x}=(50, \dots , 50), 0)$ could be around $\sim 10.008$.

Fig. 14

The figure shows the value of $g(\vec {x}=(50, \dots , 50), 0)$ returned by the neural networks $\mathcal {DNN}_{EM}$ (blue points), $\mathcal {DNN}_{M}$ (green points) as function of the number of training steps done for learning the parameters. Each point is the mean over five independent runs. The number of equidistant time steps was fixed at 40, i.e. for each neural networks $\mathcal {DNN}_{EM}$ and $\mathcal {DNN}_{M}$ the value of N is fixed to 40. All the parameters were initialized randomly uniformly between $[-1,1]$. The total number of training steps, i.e. $t_s$, was fixed to 6000 and the learning rate to $\eta =0.008$ (Color figure online)

Appendix F: Time Scaling of the Deep Learning Algorithm

In this appendix we briefly discuss how, in practice, the time scales with the dimension d, when the deep network is trained to compute an approximate solution. This experiment is performed on a cluster composed of 22 nodes, 200 cores and 250 GB of RAM divided in 22 different machines with different hardware.

We illustrate this for the heat equation and for the exactly solvable non-linear diffusion equation of Sect. 5.3. Since the diffusion tensor is constant we use the algorithm presented in Sect. 4.1 based on the Euler-Maruyama discretization.

We use $P=4096$ random initial points for the set of stochastic trajectories used for training, and keep the number of training steps at 15000. Figure 15 shows how the time needed to train the network and find an approximate a solution scales with the number of dimensions d. Interestingly we see that for both equations the set of points live on the same straight line, on a log-log scale, with a slope $\sim 1.1$.

Fig. 15

The figure shows how the time returned by the algorithm scales with the dimension d. Each point is an average over 3 experiments. Black circles correspond to the non-linear diffusion equation and blue squares to the heat equation. For both equations points live on the same straight line on a log-log scale with slope $\sim 1.1$ (Color figure online)

Appendix G: Derivation of Equation (26)

For completeness we derive Eq. (26) which follows from the Milstein discretization of stochastic trajectories given by (20) [37, 45]. We recall that the $\partial _{x_k}$ operator acts only on B.

For small enough $\tau$, after the discretization of Eq. (3), and using (20), one obtains:

$$\begin{aligned} g[\vec {Y}^{n+1}_{M}, \tau _{n+1}]&= g[\vec {Y}^n_{M},\tau _n]+ (\partial _t g[\vec {Y}^n_{M},\tau _n])\tau \\&\quad + \sum ^d_{i=1}(\partial _{x_i}g[\vec {Y}^n_{M},\tau _n])\Big (A_i[\vec {Y}^n_M,\tau _n]\tau + \sum _{j=1}^d B_{ij}[\vec {Y}^n_M,\tau _n]\Delta W^n_j \\&\quad +\frac{1}{2}\sum _{j,k,l=1}^d B_{kl}[\vec {Y}^n_M,\tau _n] (\partial _{x_{k}}B_{ij}[\vec {Y}^n_M,\tau _n]) (\Delta W^{n}_{j}\Delta W^{n}_{l}-\delta _{jl}\tau )\Big ) \\&\quad +\frac{1}{2}\sum ^d_{i,r=1}(\partial _{x_i}\partial _{x_r}g[\vec {Y}^n_{M},\tau _n])\Big (A_i[\vec {Y}^n_M,\tau _n]\tau + \sum _{j=1}^d B_{ij}[\vec {Y}^n_M,\tau _n]\Delta W^n_j\\&\quad +\frac{1}{2}\sum _{j,k,l=1}^d B_{kl}[\vec {Y}^n_M,\tau _n] (\partial _{x_{k}}B_{ij}[\vec {Y}^n_M,\tau _n]) (\Delta W^{n}_{j}\Delta W^{n}_{l}-\delta _{jl}\tau )\Big ) \\&\quad \Big (A_r[\vec {Y}^n_M,\tau _n]\tau + \sum _{j'=1}^d B_{rj'}[\vec {Y}^n_M,\tau _n]\Delta W^n_{j'} \\&\quad +\frac{1}{2}\sum _{j',k',l'=1}^d B_{k'l'}[\vec {Y}^n_M,\tau _n] (\partial _{x_{k'}}B_{rj'}[\vec {Y}^n_M,\tau _n]) (\Delta W^{n}_{j'}\Delta W^{n}_{l'}-\delta _{j'l'}\tau )\Big ) \end{aligned}$$

Finally, discarding all the terms of order greater than $\tau$, using the assumption that $g[\vec {Y}^n_{M},\tau _n]$ must be a solution of the nonlinear Backward Kolmogorov Eq. (1), we end up with Eq. (26).