Finite Expression Method for Solving High-Dimensional Partial Differential Equations

Mathematical expressions will be introduced to form the function space of FEX.

In the definition of mathematical expressions, we only consider the expression that forms a valid function. In our context, “$\sin(x\times y)+1$” is a mathematical expression, but, for example, “$5>x$” and “$\sin(x\times y)+$” are not a mathematical expression as they do not form a valid function. The operands and operators comprise the structure of an expression. The parentheses play a role in clarifying the operation order.

“$\sin(x\times y)+1$” is a $3$-finite expression since there are three operators in it (“$\times$”, “$\sin$”, and “+”). The series, such as “$1+x^{1}+\frac{x^{2}}{2}+\frac{x^{3}}{6}+\cdots$”, belongs to a mathematical expression, but it is not a finite expression since the amount of the operators is infinite. Formally, with the concept of finite expression, we can define FEX as follows,

We denote $\mathbb{S}_{k}$ as a set of functions that are formed by finite expressions with the number of operators not larger than $k$. $\mathbb{S}_{k}$ forms the function space of FEX. Clearly, $\mathbb{S}_{1}\subset\mathbb{S}_{2}\subset\mathbb{S}_{3}\cdots$.

We denote a functional $\mathcal{L}:\mathbb{S}\to\mathbb{R}$ associated with a given PDE, where $\mathbb{S}$ is a function space and the minimizer of $\mathcal{L}$ is the best solution to solve the PDE in $\mathbb{S}$. In FEX, given the number of operators $k$, the problem of seeking a finite expression solution is formulated as a CO over $\mathbb{S}_{k}$ via

The choice of the functional $\mathcal{L}$ is problem-dependent and one may conceive a better functional for a specific PDE with a specific constraint or domain. Some popular choices include least-square methods [28, 47, 12] and variation formulations [14, 52, 8, 53].

The most important theoretical question in solving high-dimensional problems is whether or not a solver suffers from the curse of dimensionality. It will be shown that the function space of $k$-finite expressions, i.e., $\mathbb{S}_{k}$ in (1), is a powerful function space that avoids the curse of dimensionality in approximating high-dimensional continuous functions, leveraging the recent development of advanced approximation theory of deep neural networks [44, 24]. First of all, it can be proved that $\mathbb{S}_{k}$ is dense in $C([0,1]^{d})$ for an arbitrary $d\in\mathbb{N}$ in the following theorem.

Here “$|\cdot|$”, “sign$(\cdot)$”, and “$\lfloor\cdot\rfloor$” denote the absolute, sign and floor functions [44], respectively. The proof of Thm. 4 can be found in Appx. 6.1. The denseness of $\mathbb{S}_{k}$ means that the function space of $k$-finite expressions can approximate any $d$-dimensional continuous functions to any accuracy, while $k$ is only required to be $\mathcal{O}(d^{4})$ independent of the approximation accuracy. The proof of Thm. 4 takes the advantage of operators “sign$(\cdot)$” and “$\lfloor\cdot\rfloor$”, which might not be frequently used in mathematical expressions. If it is more practical to restrict the operator list to regular operators like “$+$”, “$\times$”, “$\sin(\cdot)$”, exponential functions, and the rectified linear unit (ReLU), then it can be proved that $\mathbb{S}_{k}$ can approximate Hölder functions without the curse of dimensionality in the following theorem.

The proof of Thm. 5 can be found in Appx. 6.1. Although finite expressions have a powerful approximation capacity for high-dimensional functions, it is challenging to theoretically prove that our FEX solver to be introduced in Sec. 3 can identify the desired finite expressions with this power. However, the numerical experiments in Sec. 4 demonstrate that our FEX solver can identify the desired finite expressions to machine accuracy for several classes of high-dimensional PDEs.

A finite expression can be represented as a binary tree $\mathcal{T}$, as depicted in Fig. 3. Each node in the tree is assigned a value from a set of operators, and all these node values collectively form an operator sequence $\bm{e}$, following a predefined order to traverse the tree (e.g., inorder traversal).

Within each node featuring a unary operator, we incorporate two additional parameters, a scaling parameter $\alpha$ and a bias parameter $\beta$, to enhance expressiveness. All these parameters are denoted by $\bm{\theta}$. Therefore, a finite expression can be denoted as $u(\bm{x};\mathcal{T},\mathbf{e},\bm{\theta})$, a function of $\bm{x}$. For a fixed $\mathcal{T}$, the maximal number of operators is upper bounded by a constant denoted $k_{\mathcal{T}}$. In this implementation, $\{u(\bm{x};\mathcal{T},\mathbf{e},\bm{\theta})|\mathbf{e},\bm{\theta}\}$ forms the function space in the CO to solve a PDE. This is a subset of functions expressible with at most $k_{\mathcal{T}}$ finite expressions.

The configuration of the tree of various depths can be designed as in Fig. 3. Each tree node is either a binary operator or a unary operator that takes value from the corresponding binary or unary set. The binary set can be $\mathbb{B}:=\{+,-,\times,\div,\cdots\}$. The unary set can be $\mathbb{U}:=\{\sin,\exp,\log,\text{Id},(\cdot)^{2},\int\cdot\text{d}x_{i},\frac{\partial\cdot}{\partial x_{i}},\cdots\}$, which contains elementary functions (e.g., polynomial and trigonometric function), antiderivative and differentiation operators. Here “Id” denotes the identity map. Notice that if an integration or a derivative is used in the expression, the operator can be applied numerically. Each entry of the operator sequence $\bm{e}$ corresponds one-to-one with traversing the tree nodes. The length of $\bm{e}$ equals the total number of tree nodes. For example, when inorder traversal is used, Fig. 2b depicts a tree with 4 nodes and $\bm{e}=(\text{Id},\times,\exp,\sin).$

The computation flow of the binary tree is conducted recursively. The operator of a node is granted higher precedence than that of its parent node. First, as in Fig. 3, we present the computation flow of the basic trees (a tree has a depth of 1 with only 1 operator). For a basic tree with a unary operator $u_{1}$, when the input is $i_{0}$, then the output $o_{1}=\alpha u_{1}(i_{0})+\beta$, where $\alpha$ and $\beta$ are scaling and bias parameters, respectively. For a basic tree with a binary operator $b_{1}$, when the input is $i_{01}$ and $i_{02}$, the output becomes $o_{1}=b_{1}(i_{01},i_{02})$. With these two basic trees, we are ready to define the computation for arbitrary depth by recursion, as the examples shown in Fig. 3. Specifically, the input of a parent node is the output of the child node(s). When the tree input at the bottom layer is a $d$-dimensional variable $\bm{x}$ (Fig. 2b), the unary operator directly linked to $\bm{x}$ is applied element-wisely to each entry of $\bm{x}$ and the scaling $\alpha$ becomes a $d$-dimensional vector to perform a linear transformation from $\mathbb{R}^{d}$ to $\mathbb{R}^{1}$.

Given a tree $\mathcal{T}$, we aim to seek the PDE solution from the function space $\{u(\bm{x};\mathcal{T},\mathbf{e},\bm{\theta})|\mathbf{e},\bm{\theta}\}\subset\mathbb{S}_{k_{\mathcal{T}}}$. The mathematical expression can be identified via the minimization of the functional $\mathcal{L}$ associated with a PDE, i.e.,

We introduce the framework for implementing FEX, as displayed in Fig. 2a, to seek a minimizer of (9). The basic idea is to find a good operator sequence $\bm{e}$ that may uncover the structure of the true solution, and then optimize the parameter $\bm{\theta}$ to minimize the functional (9). In our framework, the searching loop consists of four parts: 1) Score computation (i.e., rewards in RL). A mix-order optimization algorithm is proposed to efficiently assess the score of the operator sequence $\bm{e}$ to uncover the true structure. A higher score suggests a higher possibility to help to identify the true solution. 2) Operator sequence generation (i.e., taking actions in RL). A controller is proposed to generate operator sequences with high scores (see Fig. 2b). 3) Controller update (i.e., policy optimization in RL). The controller is updated to increase the probability of producing a good operator sequence via the score feedback of the generated ones. While the controller can be modeled in many ways (e.g., heuristic algorithm), we introduce the policy gradient in RL to optimize the controller. 4) Candidate optimization (i.e., a non-greedy strategy). During searching, we maintain a candidate pool to store the operator sequence with a high score. After searching, the parameters $\bm{\theta}$ of high-score operator sequences are optimized to approximate the PDE solution.

This part provides the setting of FEX and NN-based solvers. The accuracy of a numerical solution $\tilde{u}$ compared with the true solution $u$ is measured by a relative $L^{2}$ error, i.e., $\|\tilde{u}-u\|_{L^{2}(\Omega)}/\|u\|_{L^{2}(\Omega)}$. The integral in the $L^{2}$ norm is estimated by the Monte Carlo integral for high-dimensional problems.

Implements of FEX. The depth-3 binary tree (Fig. 2b) with 3 unary operators and 1 binary operator is used to generate mathematical expressions. The binary set is $\mathbbm{B}=\{+,-,\times\}$ and the unary set is $\mathbbm{U}=\{0,1,\text{Id},(\cdot)^{2},(\cdot)^{3},(\cdot)^{4},\exp,\sin,\cos\}$. A fully connected NN is used as a controller $\bm{\chi}_{\Phi}$ with constant input (see Appendix 6.3 for more details). The output size of the controller NN is $n_{1}|\mathbbm{B}|+n_{2}|\mathbbm{U}|$, where $n_{1}=1$ and $n_{2}=3$ represent the number of binary and unary operators, respectively, and $|\cdot|$ denotes the cardinality of a set.

There are four main parts in the implementation of FEX as introduced in Sec. 3.2. We will only briefly describe the key numerical choices here. (1) Score computation. The score is updated first by Adam with a learning rate 0.001 for $T_{1}=20$ iterations and then by BFGS with a learning rate 1 for maximum $T_{2}=20$ iterations. (2) Operator sequence generation. The depth-3 binary tree (Fig. 2b) with 3 unary operators and 1 binary operator is used to generate mathematical expressions. The binary set is $\mathbbm{B}=\{+,-,\times\}$ and the unary set is $\mathbbm{U}=\{0,1,\text{Id},(\cdot)^{2},(\cdot)^{3},(\cdot)^{4},\exp,\sin,\cos\}$. A fully connected NN is used as a controller $\bm{\chi}_{\Phi}$ with constant input. The output size of the controller NN is $n_{1}|\mathbbm{B}|+n_{2}|\mathbbm{U}|$, where $n_{1}=1$ and $n_{2}=3$ represent the number of binary and unary operators, respectively, and $|\cdot|$ denotes the cardinality of a set. (3) Controller update. The batch size for the policy gradient update is $N=10$ and the controller is trained for $1000$ iterations using Adam with a fixed learning rate $0.002$. Especially in Section of “Numerical Convergence”, since the deeper trees are used, the controller is updated for 5000 iterations for trees with different depth. We adopt the $\epsilon$-greedy strategy to increase the exploration of new $\bm{e}$. The probability $\epsilon$ of sampling an $e_{i}$ by random is $0.1$. (4) Candidate optimization. The candidate pool capacity is set to be $K=10$. For any $\bm{e}\in\mathbb{P}$, the parameter $\bm{\theta}$ is optimized using Adam with an initial learning rate $0.01$ for $T_{3}=20,000$ iterations. The learning rate decays according to the cosine decay schedule [23].

Implements of NN-based Solvers. Residual networks (ResNets) [22, 14] $u(\bm{x};\bm{\Theta})$ parameterized by $\bm{\Theta}$ are used to approximate a solution and a minimization problem ${\min}_{\bm{\Theta}}~{}\mathcal{L}(u(\cdot;\bm{\Theta}))$ is solved to identify a numerical solution [47, 49, 14].

The ResNet maps from $\mathbb{R}^{d}$ to $\mathbb{R}^{1}$ and consists of seven fully connected layers with three skip connections. Each hidden layer contains 50 neurons. The neural network is optimized using the Adam optimizer with an initial learning rate of 0.001 for 15,000 iterations. The learning rate is decayed following a cosine decay schedule.

Poisson equation. The coefficient $\lambda$ in the functional (3) is 100. The batch size for the interior and boundary is 5,000 and 1,000, respectively. In the NN method, we use the ReLU² activation, i.e., $(\max\{x,0\})^{2}$, in ResNet to approximate the solution.

Linear conservative law. The coefficient $\lambda$ in the functional (20) is 100. In the NN method, we use ReLU $(\max\{x,0\})$ activation. The batch size for the interior and boundary is 5,000 and 1,000, respectively. We use the same batch size as the NN method except that we increase the batch size to 20,000 for the interior and 4,000 for the boundary when the dimension is not smaller than 36.

Schrödinger equation. The coefficient $\lambda$ in the functional (23) is 1. The batch size for estimating the first term and second term of (23) is 2,000 and 10,000, respectively. Besides, ReLU² is used in ResNet.

Implements of GP. GP is an evolutionary algorithm used for automated symbolic regression and the evolution of computer programs to solve complex problems. Through a process of selection, crossover, and mutation, GP evolves and refines these programs over generations to optimize a specific fitness or objective function. We utilize GP to optimize the CO (1) directly using the Python package GPlearn [1].

Several numerical examples for high-dimensional PDEs including linear and nonlinear cases are provided here. In these tests, the true solutions of PDEs have explicit formulas that can be reproduced by the binary tree defined in Sec. 3.1 and $\mathbbm{B},\mathbbm{U}$ defined in Sec. 4.1.

Three main sets of numerical results for the PDE problems above will be presented. First, the errors of the numerical solutions by NN-based solvers and FEX are compared. Second, a convergence test is analyzed when the tree size of FEX increases. Finally, FEX is applied to design special NN parametrization to solve PDEs in NN-based solvers.

Estimated Solution Error. The depth-3 binary tree (Fig. 3) is used in FEX with four nodes (a root node (R), a middle node (M), and two leave nodes (L1 and L2)). Fig. 4 shows the operator distribution obtained by FEX and the error comparison between the NN method and our FEX. The results show that NN solutions have numerical errors between $\mathcal{O}(10^{-4})$ and $\mathcal{O}(10^{-2})$ and the errors grow in the problem dimension $d$, agreeing with the numerical observation in the literature. Meanwhile, FEX can identify the true solution structure for the Poisson equation and the linear conservation law with errors of order $10^{-7}$, reaching the machine accuracy since the single-float precision is used. In the results of the nonlinear Schrödinger equation, FEX identifies the solutions of the form $\exp(\cos(\cdot))$ but achieves errors of order $10^{-5}$. Note that $\int_{\Omega}\hat{u}(\bm{x})d\bm{x}$ in (23) is estimated by the Monte-Carlo integration with millions of points as an accurate and precomputed constant, but $\int_{\Omega}u(\bm{x})d\bm{x}$ can only be estimated with fixed and small batch size, typically less than 10,000, in the optimization iterations. As the dimension grows, the estimation error of $\int_{\Omega}u(\bm{x})d\bm{x}$ increases, and, hence, even the ground true solution has an increasingly large error according to (23). Therefore, the optimization solver may return an approximate solution without machine accuracy. Designing a functional $\mathcal{L}$ free of the Monte-Carlo error (e.g., Eqns. (3) and (20)) for the nonlinear Schödinger equation could ensure machine accuracy. Furthermore, as shown in Figure 4, GP tends to identify solutions of lower quality that fail to achieve a high level of accuracy.

Numerical Convergence. The numerical convergence analysis is performed using the Poisson equation as an example. Binary trees of depths 2, 3, 4, 6 and 8 are used (see Fig. 3). The square operator $(\cdot)^{2}$ is excluded in $\mathbbm{U}$ so that the binary tree defined in Sec. 3.1 can not reproduce the true solution (sum of the square of coordinates) exactly. This setting can mimic the case of a complicated solution while a small binary tree was used. Fig. 5 shows the error distribution of FEX with the growth of dimensions and the change of tree depths. FEX obtains smaller errors with increasing tree size. Notice that, compared with the errors of NN-based solvers reported in Fig. 4, FEX gets a higher accuracy when a larger tree is used.

FEX-NN. FEX provides interpretable insights of the ground truth PDE solution by the operator distribution obtained from the searching loop. It may be beneficial to design NN models with a special structure to increase the accuracy of NN-based solvers. In the results of the Poisson equation, we observe that the square operator has a high probability to appear at the leave nodes, which suggests that the true solution may contain the structure $\bm{x}^{2}:=(x_{1}^{2},\cdots,x_{d}^{2})$ at the input $\bm{x}$. As a result, we define the FEX-NN by $v(\bm{x}^{2};\bm{\Theta})$ for the Poisson equation. Similarly, we use FEX-NNs $\sin(v(\bm{x};\bm{\Theta}))$ for the linear conservation law and $\exp(v(\bm{x}^{2};\bm{\Theta}))$ for the nonlinear Schrödinger equation. Fig. 4 shows the errors of FEX-NN with the growth of dimensions, and it is clear that FEX-NN outperforms the vanilla NN-based method by a significant margin.

Consider identifying the smallest eigenvalue $\gamma$ and the associated eigenfunction $u$ of the eigenvalue problem [14],

The minimization of the Rayleigh quotient $\mathcal{I}(u)=\frac{\int_{\Omega}\|\nabla u\|_{2}^{2}dx+\int_{\Omega}wu^{2}dx}{\int_{\Omega}u^{2}dx}$, s.t., $u|_{\partial\Omega}=0$, gives the smallest eigenvalue and the corresponding eigenfunction. In the NN-based solver [14], the following functional is defined

to seek an NN solution. Considering an example of $w=\|\bm{x}\|_{2}^{2}$ and $\Omega=\mathbb{R}^{d}$, the smallest eigenvalue of (24) is $d$ and the associated eigenfunction is $\exp(-\frac{\|\bm{x}\|_{2}^{2}}{0.5})$. The domain $\Omega$ is truncated from $\mathbb{R}^{d}$ to $[-3,3]^{d}$ for simplification as done in [14].

In FEX, the functional (25) is also used to estimate a solution. FEX discovers a high probability to have the “$\exp$” operator at the tree root ($100$% for $d=2$, $4$, $6$, $8$, and $93.3$% for $d=10$ as shown in Fig. 6). Therefore, it is reasonable to assume that the eigenfunction is of the form $\exp(v(\bm{x}))$. Let $u(\bm{x})$ be $\exp(v(\bm{x}))$ and then Eqn. (24) is simplified to

Eqn. 26 does not have a trivial zero solution so we can avoid the integration constraint used in Eqns.  (23) and (25), which leads to Monte-Carlo errors. Using Equation (26) and the Rayleigh quotient $\mathcal{I}$, the values of $v$ and $\gamma$ are alternatively updated until they reach convergence. The detail of this iterative algorithm is presented in Appx. 6.4. Fig. 6 shows the relative absolute error of the estimated eigenvalues with the growth of the dimensions. We can see that directly optimizing (25) with the NN method produces a large error on the eigenvalue estimation, especially when the dimension is high (e.g., the relative error is up to 20% when $d=10$). With the postprocessing algorithm with FEX, we can identify the eigenvalue with an error close to zero.

We have included the pseudo-code for the proposed FEX implementation in Alg. 1 and Alg. 2. Specifically, Alg. 1 outlines the procedure for solving the CO problem using a predefined fixed tree. On the other hand, Alg. 2 is designed to iteratively expand a tree, thereby enhancing its expressiveness in the pursuit of identifying an improved solution based on the approach described in Alg. 1.

We illustrate the use of a fully connected neural network to model the controller, which generates probability mass functions for each node within a tree, in Figure 7. As an example, let’s assume we have a tree with three nodes: $n_{1}=2$ nodes for the binary set with $|\mathbbm{B}|=2$ operators, and $n_{2}=1$ node for the unary set with $|\mathbbm{U}|=3$ operators. Consequently, the output size of the neural network is $n_{1}|\mathbbm{B}|+n_{2}|\mathbbm{U}|=7$. This output is divided into $n_{1}+n_{2}$ parts, each corresponding to the probability mass function for its respective node. Besides, the input to the neural network is simply a constant vector.

First, by solving Eqn. (25) with our FEX, we obtain an estimated eigenfunction $u(\bm{x};\mathcal{T},\bm{\hat{e}},\hat{\bm{\theta}})$ and get the initial estimation of the eigenvalue through the Rayleigh quotient $\gamma_{0}=\mathcal{I}(u(\cdot;\mathcal{T},\bm{\hat{e}},\hat{\bm{\theta}}))$. Then we can utilize Eqn. (26) to iteratively find the eigenpair. We define the function for Eqn. (26) by

Given $\gamma_{i}$, we aim to find $v$ that is expressed by mathematical expression and minimizes $\mathcal{L}_{2}(v,\gamma_{i})$. Assume $v$ is expressed by a binary tree ($v:=v(\cdot;\mathcal{T},\bm{e},\bm{\theta})$), and then we can search the solution using our FEX with the following optimization,

Next, we can compute the current estimated eigenvalue by $\gamma_{i+1}=\mathcal{I}(\exp(v(\cdot;\mathcal{T},\bm{e}_{i}^{*},\bm{\theta}_{i}^{*})))$.

If continuing this loop for $G$ times, we will obtain the eigenpair $\gamma_{G}$ and $\exp(v(\cdot;\mathcal{T},\bm{e}_{G}^{*},\bm{\theta}_{G}^{*}))$.

Implementation. In our iteration method, the number of the iterative loop is $G=10$. $\lambda_{1}=\lambda_{2}=500$ in  (25). The batch size for estimating the first term and third term of (25) is 10,000 while that of the second term (boundary) is 2,000. ReLU² is used in ResNet, following [14].