SINGER: Stochastic Network Graph Evolving Operator for High Dimensional PDEs
Authors: Mingquan Feng
Yixin Huang
Weixin Liao
Yuhong Liu
Yizhou Liu
Junchi Yan
Venue: ICLR 2025 Poster (2025)
Keywords: PDE
High Dimension
Neural Operator
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: We present a novel framework, StochastIc Network Graph Evolving operatoR (SINGER), for learning the evolution operator of high-dimensional partial differential equations (PDEs). The framework uses a sub-network to approximate the solution at the initial time step and stochastically evolves the sub-network parameters over time by a graph neural network to approximate the solution at later time steps. The framework is designed to inherit the desirable properties of the parametric solution operator, including graph topology, semigroup, and stability, with a theoretical guarantee. Numerical experiments on 8 evolution PDEs of 5,10,15,20-dimensions show that our method outperforms existing baselines in almost all cases (31 out of 32), and that our method generalizes well to unseen initial conditions, equation dimensions, sub-network width, and time steps.

Published as a conference paper at ICLR 2025
SINGER: STOCHASTIC NETWORK GRAPH EVOLVING OPERATOR FOR HIGH DIMENSIONAL PDES
Mingquan Feng, Yixin Huang, Weixin Liao, Yuhong Liu, Yizhou Liu, Junchi Yan"
Sch. of Computer Science & Sch. of Artificial Intelliencee Shanghai Jiao Tong University [fengmi ngquan, ci 1 10,, 1 liacwsi xin, yanjunch i jesjtu. https:/ thub. com tttpa./.oo///eegggggggaaaaarr//a
ABSTRACT
We present a novel framework, Stochastic Network Graph Evolving operatoR (SINGER), for learning the evolution operator of high-dimensional partial dif- ferential equations (PDEs). The framework uses a sub-network to approximate the solution at the initial time step and stochastically evolves the sub-network parameters over time by a graph neural network to approximate the solution at later time steps. The framework is designed inherit the desirable properties of the parametric solution operator, including graph topology, semigroup, and sta- bility, with a theoretical guarantee. Numerical experiments on 8 evolution PDEs of 50000.0iinms show that our method outperforms existing baselines in almost all cases (31 out of 32), and that our method generalizes well to unseen initial conditions, equation dimensions, sub-network width, and time steps.
INTRODUCTION
Partial differential equations (PDEs) are prevalent and have extensive applications in science, engi- necring, economics, and finance (Strauss, 2007; Folland, 2020). Given that many partial diffcrential equations do not have analytical solutions, numerical methods like finite element methods (Hueb- ner et al., 2001) as well as neural network based solvers (Blechschmidt & Ernst, 2021) have been established. Although these methods have been significantly advanced recently, they suffer the well- known issue of curse of dimensionality since the number of data points increases exponentially with respect the dimension d of independent variables (Bellman, 1966). They quickly become compu- tationally intractable for general high dimensional problems, e.g. d 2 5.
Facing the challenge, recent advances in deep-leaming--asee solvers propose incorporating spe- cific PDE problem structure (Han et al., 2017; Wang et al., 2022a; Yu et al., 2018) and designing memory-efficient training strategies (Hu et al., 2024; Shang et al., 2023). These methods succeed in solving variety of PDEs empirically. Nonetheless, they aim to solve specific instances of PDEs. Consequently, they need to retrain from scratch whenever the initial and/or boundary value changes, which is expensive and time-consuming in high-dimensional problems. To address these limitations, (Gaby et al., 2024; Gaby & Ye, 2024) obtain the solution operator of the high-dim PDEs by learning the evolution of the parameters of a sub-network over time. The training is purely based on the PDE residual loss, without requiring any spatial discretization or solutions of the PDE. Once successfully trained, the model can generalize to new initial conditions with single inference pass. In practice, however, the training process often collapses or converges to an incorrect model.
This paper aims to address the challenge of learning tigh-dimensional PDE solution operators. We establish three basic assumptions of an ideal solution operator: graph topology. semigroup, and stability, and design novel framework that inherently satisfies those assumptions with theoreti- cal guarantees. Specifically, we propose a stochastic differential equation driven by graph neu- ral network (GNN) on the sub-network's network graph model the semigroup evolution of the sub-network parameters. The GNN is designed to handle the graph topology of the sub-network, including permutation invariance and modification robustness. The stochastic noise introduced to stabilize the evolution process and prevent divergence. We theoretically prove the semigroup evo-
"Correspondence author who is also affiliated with Shanghai Innovation Institute. This work was partly supported by NSFC (62222607, 623B1009).

---

Published as a conference paper at ICLR 2025
lution and stability of the proposed framework and empirically demonstrate model effectiveness in accuracy, stability, and generalization on several datasets. We highlight the following contributions:
Formally establish three assumptions of an ideal high- nigh-dimensional PDE solution operator: graph topology, semigroup, and stability.
Propose SINGER consists of graph network, continuous evolution and stochastic noise, corresponding to the three assumptions, with theoretical and empirical justification.
Solve eight PDEs on 5-20 dimensions by SINGER accurately and stably and generalize across initial conditions, equation dimensions, input scales, and time horizons.
2 RELATED WORK
Recent developments in neural PDE solvers have shown promise in overcoming the limitations of traditional numerical solvers, especially for high-dim problems. Below we discuss some of the key approaches in this area. The related work on low-dim PDEs are provided in Appendix D.
Neural High-Dim PDEs Solvers. For high-dim PDEs, the curse of dimensionality poses sig- nificant challenge, thus requiring specialized methods. The existing attempts mainly fall in the physics-informed category. (Han et al., 2018; 2017) proposed the DeepBSDE solver, reforming class of hyperbolic PDEs as backward stochastic differential equations (BSDEs) by Feynman-Kac formula. (Wang al., 2022b;a) proposed tensor neural networks with efficient numerical integra- tion and separable structures for solving separable PDEs such as the Schrödinger equation. The deep Ritz method (Yu et al., 2018) considers solving nigh-dimensional variation PDE problems by minimizing the energy functional of the PDE. However, the above techniques are limited to specific types of PDEs and may not generalize well to others. For general PDEs, (Zang et al., 2020) proposed a weak adversarial network that solves PDEs using the weak formulation. (Hu et al., 2024) designed a stochastic dimension gradient descent to reduce the computational cost.
Neural High-Dim PDEs Operators. The data-driven operator methods are less explored in high- dim PDEs, possibly due to the difficulty in data generating and model generalization. The only two existing works that learn the high-dim PDE solution operator proposed by Gaby et al. (2024); Gaby & Ye (2024). The first work reduces the solution operator to the evolution of the parameters of reduced-order model (a neural network) and introduces another network to learn the evolution in discrete time steps by minimizing PDE residual. The second work extends the first work to continuous evolution by adapting Neural ODE (NODE) (Chen al., 2018) framework, improving the sample efficiency. However, the above methods suffer the instability of the evolving operator, crucial for both training and inferencing. Our work follows this line of research and introduces novel framework that inherently maintains stability, generalization, and other desirable properties.
METHODOLOGY
In the SINGER model, we utilize neural network U to approximate the solution u(x, of the PDE. The parameters 0: of network U evolve according to GNN-driven ODE, where the nodes in the graph represent neurons, and the edges represent dependencies between them. The key advantage of this structure is the permutability of neurons within the same layer, which enhances the model's generalization capabilities. To combat the inherent instability in the evolution process, introduce stochastic noise during training. This noise stabilizes the system and prevents divergence during long-term iterations. A theoretical analysis of the stochastic influence on stability is provided.
3.1 MOTIVATION: THREE ASSUMPTIONS
Our model design is based on three key assumptions on the sub-network and the solution operator of PDEs: the graph topology of U, the semigroup operator, and the stability of the solution.
The first assumption is based on the network graph GVV, E) of the sub-network U. The vertex U, denotes the i-th learnable parameter of U, and the vertex feature is the one-hot encoding of its layer index in U. The edge ei, exists if ve, by are the input and output parameters, respectively, of the same neuron in U. We then assume such graph is good representation of the subnetwork U:
Assumption (Graph Topology). For a proper distance metric d(...), there exists a constant C' > such that dy (U,, U2) S (9, G2) for any sub-networks U1, U2 and their network graph G1 G2-
2

---

Published as a conference paper at ICLR 2025
GNN
Asumptions
Graph Topolegy
Modificaton Robasness
d0, V. w)dd (V, w)dB,
Semigroup
u(x, t)
u(x, = T)
Siability
PDE
Residual
r't =
dr, = R(O,, t)dt
Backpropagate
7(t= r(t=T)
Figure 1: Left: SINGER architecture. Right: three underlying assumptions.
This above statement leads to two properties of the sub-network U:
1. Permutation Invariance: changing the relative order of neurons of the same layer (defined as Perm) does not change the graph and sub-network. Formally, if G1 == Permm(2)) then dg(G G2) O, and du (U1, (1) O.
2. Modification Robustness: adding removing nodes or edges in the graph (defined as Modif) does not substantially change the sub-network. Formally, for a small modification budget 91 = (o2)f(G2) and dg (91,92) < e, then du (U1, U2) Cr.
The above assumption holds for common neural network architecture. E.g., the permutation invari- ance is satisfied by fully connected layers, convolutional layers, and attention layers. Additionally, random removal of node, e.g. via dropout, does not drastically alter the network output. Net- works not following this assumption tend to be sensitive to permutation and modification noise. A desirable solution should be robust to noise, which motivates encoding this property into the model.
The second assumption is on the semigroup property of the evolving operator T(t) that evolves the solution function u from time 0 time t. Intuitively, the semigroup property states that the operator T' is consistent over time, i.c. the evolution from time O to is the same as the concatenation of evolutions from to any middle time t and from time & to The formulation is:
Assumption (Semigroup). Suppose T is mapping [0, +00) (X), where X is the space of the PDE solution function u and (XX) denotes the set of bounded operators on X. The mapping T satisfies the semigroup property, i.e.
((t))((t2) = T(t1 + +12, for any 11, 1/ 2 0 T(0) I, where I is the identity operator.
The above assumption holds for initial value problems of PDEs (Farkas et al., 2011). This property arises from the physical interpretation of PDEs, where the time evolution of the solution inde- pendent. Similar to the Markov property, the future state depends only the current state, with memory of the past. Encoding this prior can prevent from learning non-physical solutions.
The third assumption is the stability of the solution function u. defined the boundedness of the evolution of u under a small perturbation initial condition with formulation as follows.
Assumption (Stability). The solution function of the PDE is stable under small perturbation of the initial condition, i.e. >> O, there exists a 5 >> O such that for any initial condition uo and un with Iu w6ll < 8, the solution functions satisfy [[(())uo)) T())(b))) for any ' 0.
3

---

Published as a conference paper at ICLR 2025
Table 1: Comparison on assumption satisfaction of neural operators.
Method
Graph Topology
Semigroup
Stability
NODE (Gaby & Yc, 2024)
No
Yes
No
PINO (Li el al., 2024)
No No
Yes
SINGER (ours)
Yes
Yes
Yes
The stability holds for many basic PDEs, such heat equation, reaction-diffusion equation, and wave equation(Shirinaaddi & Talebi, 2011). The stability is rooted from the well-posedness of PDEs (Hilditch, 2013). A PDE is well-posed if its solution exists, is unique, and continuously depends on the initial and boundary conditions. The stability we define refers to the continuous dependence of the PDE solution on the initial condition. Stability is of paramount importance in both theoretical analysis and numerical solutions of PDEs. For example, in our experiments, models without stability would collapse during training (see Table 2).
As compared in Table 1, the three assumptions non-trivial for neural network operator. Con- sider two baselines NODE (Gaby & Ye, 2024) and PINO(Li et al., 2024)mmodifie). The NODE model flatten the initial sub-network parameters into vector and evolve the it by freely learnable neural ODE. The flattening operation violates the graph topology assumption, and free NODE model is not guaranteed to be stable. The semigroup property is satisfied by the integration operator in neural ode forwarding.
The PINO learns network mapping the input data and time directly the solution function values at sampling points and time t. It is not directly applicable to high-dim problems, so we modify it mapping the initial sub-network parameters Og and time i the solution sub-network parameters 0, The PINO violates the graph topology if the 0 is flattened. And even if the graph is preserved, the semigroup property is not guaranteed, since the model is a black-box function with regard the time t. Whereas the stability is generally guaranteed due to the continuity of network.
The violation of the assumption usually results in poor experiment performance or even collapse of the model, as we will show in the experiment section. Our model, SINGER, is designed to satisfy all three assumptions, as introduced and justified in the following subsections.
3.2 FORMULATION AND ALGORITHM
We follow the problem setting in Gaby et al. (2024) and consider the initial value problem of the evolution PDE defined as follows:
()u(xxt) t) Fiu) u(( 0) uo(
E (e,
xE
where Q is an open bounded set in Rd, F is possibly nonlinear differential operator of functions u : Q x T.] + R. uo Ω > R stands for an initial value, and T, > 0 is terminal time. The solution is defined in the strong sense, thus is assumed continuous: u E C2 (0, Ta))
The SINGER parameterises solution ul., t) by a sub-network Uo,, with network graph & and param- eters O, at time f. The parameter evolution is modeled by an operator-network V with parameters w. The evolution of the parameters is given by the following equation:
d0, V t, (d ( NB:),
(4)
where N is noise matrix, and B, is a Brownian motion. The operator-network is implemented as a GNN to handle the graph topology of sub-network U. The stochastic noise term N dB, is introduced stabilize the evolution process and prevent divergence.
The network V instantiate the Message Passing Neural Network (MPNN) (Gilmer et al., 2017), a general framework of GNN, to handle the graph topology of sub-network U. D-layer MPNN consists of D pairs of message and update functions, denoted as M; and P. respectively:
m 1-1
Σ
Mr (H (i))
h Pt (h. (h,,m(+1) 1)
(5)
(EN(()
4

---

Published as a conference paper at ICLR 2025
Algorithm Training pipeline for SINGER
Input: parameter 1), sub-network structure U, network graph & and operator-network V
Output: optimal weights w
Sample (0 [ k= uniformly from the parameter set
2: Initialize Eq.4 and Eq.6 with @ [0]] 0]
3: Calculate loss ((w) = & Σ [rr((() by integrating Eq.4 and Eq.6.
4: Use adjoint method according Eq.7 to calculate ww)() Minimize the loss w.r.t
Algorithm 2 Testing pipeline for SINGER
Input: iniiill condition ug, sub-network U, network graph and trained operator-network V
Output: PDE solution w(..t) ()
Estimate 0 using gradient descent with Bernoulli noise: Oo arg ming. |no Uol122
2: Initialize Eq.4 with On Oo
3: Integrate Eq.4 to obtain network parameters Oz.
4: Place Og into sub-network U and retrieve PDE solution ul,tt) Us.
where hú, is the hidden state of node u layer i, .N (e) is the of neighbors of node v. The initial and final hidden state f.. A.P is the input and output of MPNN, respectively. In our model, the input h is the node feature, i.e., the one-hot vector of the parameter's s layer index. The output h of the same shape as the input and is decoded by reading out the element at the layer index.
In order optimize the parameters the operator-networt V, we augment the above equation with accumulated residual ry in strong sense:
dr: R(,,, t) dt ll VoU,, do, F(Uo,]
(6)
where F' the right-hand side differential operator of the PDE, and Ue, is the solution of the PDE at time f with parameters 0, The initial residual ro is set to zero.
The training optimization problem defined as finding the optimal parameters w that minimize the final residual rr over distribution of initial parameters 0o: "" = arg min,, Ee, [rr]. To optimize the parameters and avoid high memory costs, we adopt the scheme presented by Liu et al. (2019).
L is used here represent the loss function Es, [ry]. The gradient is: ØL AL drz If we denote Fur ärr dw
ar as ar and ae, ow
ww as Be, we can get the following relationship:
da aR(O,, t)
Be dt av ov
a0,
OV Ov 00, dB)
dB3
dw +
B.) di + N 00
Dw +
In a way similar to the adjoint method in Neural ODE, Eq.7 solved alongside the original SDE system. This approach eliminates the need to store intermediate states, effectively reducing the original memory burden. In the training phase, SINGER seeks to find the optimal operator network using the adjoint method given in E The algorithm is described in detail in Algorithm 1.
Given a trained operator network, we can solve nigh-dimensional PDEs in the following steps. Firstly, given initial value ug. find 0o arg ming,, I Ue, 2l2. Secondly, solve Eq.4 to retrieve 0,, Finally, the PDE solution is generated using 0, and sub-network U. The detailed algorithm is provided in Algorithm 2.
3.3 THEORETICAL ANALYSIS
3.3.1 GRAPH TOPOLOGY AND SEMIGROUP PROPERTY OF SINGER
Firstly we will focus on the Graph Topology assumption. Since all our operations in the SINGER preserve permutation invariance (clement-wise operation and MPNN), the model itself preserves permutation invariance. As for Modification Robustness, the GNN holds continuity to graph struc- ture (Han et al., 2024), which means that when edges are added or removed, GNN's output won't t drastically change. Thus, model observes Assumption 1.
5

---

Published as a conference paper at ICLR 2025
Test
PLL (fae)dr
do, w. NV W.Wl
Oe
0r
uCs, 0)
Gradient Descent with
0,
a Theorem Train-time Stability
Theerem Test-time Stabilitt
SDE Stability
Figure 2: SIGNER's training/testing and the algorithmic process considered in each theorem.
Then we shall discuss Assumption on our model. As the solution operator is a stochastic process, the definition of semigroup, especially for equality, needs an extension to the stochastic version.
Definition (Stochastic Semigroup). Suppose T {Ti(w), t E [0. +oo)} is stochastic process, where w being Brownian path, the stochastic variable (w) (XX) is solution operator, X is the space of the solution function u and L(X) denotes the set of bounded operators on X. Denote almost surely by a.s., denote identically distributed by m. Then T is:
1. Strong Semigroup, io,
Wh, 12 2 E
Trite (w) TT,TT (w) as.. Tu(w)(u) a.s.
2. Weak Semigroup, for vt,, 12 0,
T (w)))()((u) ~ (w)](())()) To(w)(u)~
(9)
The SINGER evolution operator a stochastic semigroup with a theoretical guarantee under mild conditions (all proofs provided in the Appendix B).
Theorem 1. (O.,,,, t, G; w) is square-integrable w.r.t. time and sub-network Uo, is twice contin- uousty differentiable w.r.t. parameter 0₁, then the solution operator T of Ue, satisfies both strong and weak semigroup property.
3.3.2 STABILITY OF SINGER
We will theoretically analyze the stability of the proposed SINGER. The overall algorithm consid- ered in this section depicted in Fig.
To examine the stability of an SDE, we initialize the system at two slightly different value (6o ro) and (06, ri) (0n + o ra + where (eî. eT) the initial perturbation. The stability assessed based on the long-term behavior of the perturbed model, i.e., how (e? ef) evolve in the long run.
d
t; t, V to,, w)] di N t.
t, dB,
(10)
= VA(ef, t, G; w) dt t, dB,
dr. = | t) - R(O, t)] dt RA t) dt
(1I)
Here we assume that the Brownian motions have the same sample path for both of the initializations to do subtraction. Given the above description of the problem, we can now formally define stability, which is adapted from Liu et al. (2019):
Definition 2. (Lyapunov stability SDE). The solution = 0 is almost surely exponentially stable a.s. YEo R".
l'lim sup,- 1 log ||=,]]
Liu et al. (2019) has already proved the stability of Eq. 10-like SDE in the following theorem:

---

Published as a conference paper at ICLR 2025
Theorem 2. For 10-lke SDE, if V 9; (:w) is L-Lipschitz continuous w.r.f. dr. then it has unique solution with the property lim supt-soo i log ||Exll S - L) a.s. for any Eo RR". In particular, if N² > 2L, the solution Et = 0 is a.s. exponentially stable.
However, we can not apply Theorem 2 directly to Eq.11 since there is no diffusion term in it. It is worth noting that Eq.11 is non autoreereessivee i.c. drs is not influenced by Tt- Any perturbation de
to the initial value of ro will result only constant-level disturbance time t. Given certain assumptions on the system and results from Liu et al. (2019), we can prove the following theorem:
Theorem 3. If (0,,,, 9; w) and R(O,, 1) is I-Lipsscitzz continuous w.r.l. Og and N2 > 2I, then the solution (ey e 0 is almost surely exponentially stable.
Finally, we shall examine the SINGER stability during training. Adapted from Assumptions 3, the definition of its train-time stability is formulated as follows:
Definition 3. (SINGER train-time stability) Ve > 06, if @ 1 oi < 6, then [Uor - Uo,1 || <
where (0,,r) rz) is denoted O for simplicity. Given Theorem 3 and some assumptions on neural networks, we can prove the train-time stability of the SINGERe.
Theorem 4. If the O evolution process satisfies all conditions in Theorem 3, then SINGER is stable during training.
We further focus on the SINGER stability during testing. In the test phase, the PDE initial con- dition is provided as input. The gradient descent method is applied to seek initial sub-network parameters 0g = arg min [(uo, Ue.), where L is the loss function. We first define the test stability:
Definition 4. (SINGER test-time stability) Vs= > O, 36, llUe. - -Vogll < 6, then ||oe-- -o.ll <
To ensure the stability of the process of seeking Oo. we use a gradient descent algorithm with Bernoulli noise. Following a similar proof scheme of Theorem and Theorem 4, we can further confirm SINGER test-time stability.
Theorem 5. If both the O, fitting process and the O, evolution process satisfies all conditions in Theorem 3, then SINGER is stable during testing.
4 EXPERIMENTS
We validate the proposed architecture on benchmark high-dim PDEs and compare it with SOTA methods. The experiments focus accuracy, generalization, stability, and assumption fulfillment. Our model is tested on 5- to 10-dimensional PDEs and evaluated for its ability to generalize across different initial conditions, dimensions of the PDE, the width of the sub-network, and time steps.
We compare our method with 4 baselines. The first one is the non-noise version of our method, denoted as SINGER(-N). We also choose NODE(Gaby & Ye, 2024) and PINOLi et al. (2024) as baselines due to their closeness to our model. The original PINO learns the mapping of function values, which is not suitable for high-dimensional PDEs. Therefore, lution of sub-network parameters. After modification, PINO can be regarded as a variant of NODE without semigroup property. The last baseline is the NODE with noise, denoted as NODE(+N), regarded as the non-graph version of SINGER.
modified it to learn the evo-
4.1 IMPLEMENTATION DETAILS
In all experiments, we implement the SINGER with a 3-layer message passing GNN, with the mes- sage and update function being 4-layer residual networks (ResNets). The hidden dimension is 100, and the noise type is Bernoulli (i.e. dropout) with a probability of 0.1. The validity of using Bern- uolli noise is justified in Appx We also implement NODE baseline using a ResNet with 1000 hidden dimensions. The ODE is integrated with the Runge- Kutta (Butcher, 1996) with a time step 1/10 of the time horizon. The PINO baseline is also with 1000 hidden dimensions. It is implemented with the trick from the Consistency Model (Song et al., 2023), i.e., embedding the time index by an encoder network, then adding the encoded feature to other inputs and feeding into ResNet. The training set consists of 100000 randomly generated 0 from the normal distribution, and the test set is 100 randomly generated 9 from another normal distribution with slightly different mean and vari- ance, suggested by Gaby & Ye (2024). The training is done with Adam optimizer with learning
7

---

Published as a conference paper at ICLR 2025
Table Relative error on PDEs with explicit solution, the best the
and second best are marked, and
symbol denotes collapse in training.
Eqn
Heat
HJB
Dim
d=5
d=15
d=20
d=5
d=10
d=15
d=20
SINGER
0.0045
0.0041
0.0046
0.0175
0.0036
0.0027
0.0099
0.0052
SINGER(-V)
0.0137"
0.0294
0.0371
0.0507
0.0037
0.0027
0.0106
0.0047
NODE
0.0194
0.0342*
0.0502
0.0957
0.0042
0.0090
0.0280
0.0602
NODE(+N)
0.0213
0.0220
0.0196
0.0751
0.0040
0.0041
0.0197
0.0398
PI-NO
0.0543
0.0229
0.0296
0.0567
0.1075
0.5387
0.9386
3,5078
Table Relative error on various PDEs, the best is marked, and the denotes collapse training.
Eqn
Pricing
Sine-Gorden
Dim
d=5
d=10
d=15
d=20
d=5
d=10
d=15
d=20
SINGER
0.0510
0.0217
0.0308
0.0519
0.00051
0.00054
0.00073
0.00066
NODE(+N)
0.0541
0.0258
0,0473
0.0689
0.05431
0.00065
0.00074
0.00085
Eqn
Burgers
Reaction Ditfusion
SINGER
0.0020
0.0032
0.0079
0.0088
0.0022
0.0016
0.0018
0,0020
NODE(+N)
0.0274
0.0068
0.1729
0.0021
0.0024
0.0245
Eqn
HJBLQ
Allen-Cahe
SINGER
0.0041
0,0033
0.0035
0.0037
0,00052
0.00071
0.00088
0.00077
N)
0.0070
0.0039
0.0041
0.0059
1.16401
0.13258
0,00125
0,00373
rate of 0.0005 and the batch size is 64. The metric used in experiments is the Le relative error (L2RE) defined as: L2RE(u, u²) Iu=v'la The function norm evaluated by the Monte Carlo method with 1000 samples. Unless otherwise stated, the L2RE the main text and tables denotes I.2RE(uref- Upred), where uref is the reference solution and upred is the predicted solution.
We choose two PDEs with explicit solutions, the Heat equation and HJB equation (Evans, 2022):
0,u Au 1(x, (6((g))(l)))) Nyy xx, 2,,2tyydy,, x E [-1,1] (0,0.1]
(12)
0,u A 2 Vaulz u(x,t) = 2n / - /1
(4 f)) dy. x xE E [0
(13)
where stands for the density of Gaussian, g(v) the initial function, and the HJB is time-reversed for consistency with the experiment. The subnetwork architecture for heat and HJB are two-layer networks, reflecting the structure of the solution:
un (z: tanh (a, sin(=(x 3)) - b.) 0 - (3, E Rd x p R x R
(14)
uo(x) =
w. w.e (2) = (a, b)) E Id x Rdx x R
(15)
where is input dimension and w is sub-network width. See remaining six PDEs in Appendix
4.2 RESULTS AND DISCUSSION
The experiments are organized to answer the following research questions:
RQ1: How does the SINGER accurately and stably solve high-dimensional PDEs?
RQ2: How does the SINGER generalize to unseen longer time indexes and higher dimensions?
RQ3: How does the SINGER satisfy the three assumptions in Section 3.1?
8

---

Published as a conference paper at ICLR 2025
Ref.
Ours
NODD+N
PINO
Ref.
Ours
NOCE-+N
PING
Ours
NODE-N
PINO
Figure 3: Examples on 10-dim Heat(left) and HJB(right). Only the first two spacial dim are shown.
NeN
2-1
L2RE
L2RE
0.1
NeN
Ours
Ours(-N)
0.01
Ours{ NJ
NODE(+N)
NODE
NODE
PINO
0.01
0.001
10
15
20
25
0.1
0.2
0.3
0.4
0.5
0.6
0,7
0.8
Train iter (x20)
Figure 4: Models trained on 10-dimensional heat equation with = 0.1. (Left) Relative error in the early stage (560 iterations) of training. (Right) time extrapolation to 0.8.
RQI: Accuracy and Stability We first evaluate the accuracy and stability of SINGER on PDEs with explicit solutions, i.e. the heat and HJB equations, across dimensions from 5 to 20. The results are shown Table Our model outperforms almost all baselines w.r.t. relative error. Adding noise into the training is beneficial for both our model and NODE, effectively improving accuracy and preventing the collapse of training in the heat equation. The PINO model performs poorly in all cases, with relative error of over 0.05 in most cases. Examples of the solution are shown in Fig. 3, where the proposed model accurately captures the solution of the heat equation and HJB equation.
We select strong baseline NODE(+N) as the comparison baseline for the 6 more PDEs. Those PDEs have no explicit solution in general, and the reference solution is generated by the DeepBSDE (Han et al., 2018). The generation details are in Appendix C. As illustrated in Table 3, the SINGER outperforms NODE(+N) with a clear margin in almost all cases. SINGER is also more stable than NODE(+N), which collapses in training in the Burgers and Reaction-Diffusiio equations.
To further investigate the performance of the models, we visualize the early stage of training in Fig. 4(left). The NODE model collapses in the first few epochs, while the proposed model and NODE(+N) are more stable during training. The proposed model converges faster than NODE(+N) in the longer term, which consistent with the results in Table 2 and Table 3. The PINO is stable but converges slower than the proposed model and NODE(+N). The results above demonstrate that the noise is crucial for training stability, and the graph topology and semigroup property can accel- erate the convergence of the model in the early stage, possibly due to the incorporation of the prior knowledge, saving the model from exploring the wrong direction.
RQ2: Generalization We test the zero-shot generalization for unseen time periods, subnetwork widths, and even higher dimensions. The results are shown in Fig. 4(right) and Table 4. Our model outperforms all baselines in all cases, showing its superior generalization ability. In Fig. 4(right), the models are trained on a short time horizon of t = 0.1 and tested on a long horizon The proposed model maintains low error the horizon extends, whereas other models experience a substantial increase in error or break down into NaN values. Notably, removing the noise from our model will lead collapse in training, highlighting the importance of noise in training.
= 0.8.

---

Published as a conference paper at ICLR 2025
Table 4: Zero-shot generalization of models trained in 10-dim and 80-width to higher PDE dimen- sions and larger subnetwork width.
Dim
10
11
12
13
14
15
Ours
0.0041
0.0159
0.0255
0.0448
0.0672
0.1194
NODE+N+1
0.0220
0.0789
0.1492
0.1259
0.1433
0.1646
Width
80
85
90
95
100
105
Ours
0.0041
0.0445
0.0061
0.0049
0.0051
0.0045
NODE+N+I
0.0220
0.0531
0.1079
0.1306
0.1404
0,1857
Table 5: Ablation study on the violationnressduua)) of three assumptions.
Ours
Ablation
Assumption
Before Train
After Train
Model
Before Train
After Train
Graph Topology
1.6 10-14
10-14
NODE(+N) PINO SINGER(-N)
0.00205
0.00478
Semigroup Stability
x 10-13
4.7 10-10
0.14441
0.01484
0.29058
0,00057
0.29199
0.02538
In Table 4, we add neurons the subnetwork in the test time. Since generally NODE can not han- dle modification of subnetwork, the proposed model is compared with NODE(+N+D), where the +I means ignoring the additional neurons in the subnetwork. The models are trained 10-dim and 80- width and can generalize to higher problem dimensions and wider subnetworks without finetuning. The proposed model outperforms NODE(+N+1) all cases, showcasing its generalization abil- ity. As dimension increases, SINGER error increases slower than NODE(+N+I), When subnetwork width increases, SINGER keeps error constantly while NODE(+N+I) increases error significantly. The generalization arises from the preservation of graph topology since the modification of subnet- work is essentially a modification of graph structure. The generalization of problem dimensions is novel metric in the existing high-dimension PDE solver literature, and this ability might be useful in pre-trained foundation model of PDE solver such as PDEformer (Ye et al., 2024).
RQ3: Assumption Satisfaction Lastly, we numerically check the assumptions in Section 3.1. To quantify the violation of assumptions we define the following metrics:
1) Graph Topology: T.2E(((()()) T()) (1)erm)) where Upern is the randomly neuron-permuted input of u.
2) Semigroup: I2RE(((()) -,) (t))(,,) where &, is randomly selected time index.
3) Stability: (2)) (()) (Unisse)) where Unoise is the input u with added noise.
Table shows the metrics of the proposed SINGER and its ablation models with respect to the assumptions, hefore and after training. For the graph topology, the SINGER inherently satisfies the assumption by its graph network module, while the ablation model does not and even can not learn from training. For the semigroup, the SINGER also inherently satisfies the assumption by its ode framework, while the ablation model does not and learns from training with limited improve- ment. For stability, both the SINGER and ablation models are not satisfied initially, but the SINGER improves significantly after training, while the ablation model improves little. The results demon- strate that the proposed model automatically satisfies the first two assumptions and can learn from training to improve the stability property. Whereas the ablation model merely learns from training and generally can not satisfy the assumptions after training.
5 CONCLUSION
We have introduced a novel architecture for learning high-dimensional PDEs solution operators based graph neural network. Our model is designed to satisfy three key assumptions: graph topology, semigroup property, and stability. We demonstrate the effectiveness of our approach through extensive experiments. Future work and broader impact are discussed in Appendix E.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ferran Alet, Adarsh Keshav Jeewajee, Maria Bauza Villalonga, Alberto Rodriguez, Tomas Lozano- Perez, and Leslie Kaelbling. Graph element networks: adaptive, structured computation and memory. In International Conference Machine Learning, pp. 212-222. PMLR, 2019.
Richard Bellman. Dynamic programming. science, 153(3731):34-37, 1966.
Jan Blechschmidt and Oliver G Emst. Three ways to solve partial differential equations with neural networks review. GAMM-Mineilungen, 44(2):2002 2021.
Andrey Bryutkin, Jiahao Huang, Zhongying Deng, Guang Yang, Carola- Bibbane Schönlieb, and An- gelica Aviles-Rivero. Hamlet: Graph transformer neural operator for partial differential equations. arXiv preprint 1rXiv:2402.03541, 2024.
John Charles Butcher. A history of uunge-kutta methods. Applied numerical mathematics, 20(3):
247-260, 1996.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE Transactions Neural Networks, 6(4):91 1-917, 1995.
Tim De Ryck, Siddhartha Mishra, and Roberto Molinaro. wpinns: Weak physics informed neural networks for approximating entropy solutions of hyperbolic conservation laws. SIAM Journal on Numerical Analysis, 62(2):811-841, 2024.
Lawrence Evans. Partial differential equations, volume 19. American Mathematical Society, 2022.
Bálint Farkas, Petra Csomós, and Alexander Ostermann. Operator semigroups for numerical analy- sis. 2011.
Gerald B Folland. Introduction to partial differential equations. Princeton university press, 2020.
Nathan Gaby and Xiaojing Ye. Approximation of solution operators for high-dimensional pdes. arXiv preprint arXiv:2401.10383 2024.
Nathan Gaby, Xiaojing Ye, and Haomin Zhou. Neural control of parametric solutions for high- dimensional evolution pdes. SIAM Journal on Scientific Computing, 46(2):C155-C11 C1885 2024.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
Andi Han, Dai Shi, Lequan Lin, and Junbin Gao. From continuous dynamics to graph neural net- works: Neural diffusion and beyond. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https: h/tppn/eviewe..c///////// net forum? to:ammid=fPooo itt.rumii-d--osxaaa222 Survey Certi- fication.
Jiequn Han, Arnulf Jentzen, Deep learning leasenngbbased numerical methods for nigh-dimensional parabolic partial differential equations and backward stochastic differential equations. Communi- cations in mathematics and statistics, (44):349-380, 2017.
Jiequn Han, Arnuff Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34):8005-8510, 2018.
David Hilditch. An introduction to well-posedness and free-evolution. International Journal of Modern Physics A, 28(22223): 2013.

---

Published as a conference paper at ICLR 2025
Masanobu Horie and Naoto Mitsume. Physics-embedeee neural networks: Graph neural pde solvers with mixed boundary conditions. Advances in Neural Information Processing Systems, 35:23218- 23229, 2022.
Zheyuan Hu, Khemraj Shukla, George Em Karniadakis, and Kenji Kawaguchi. Tackling the curse of dimensionality with physics-informed neural networks. Neural Networks, 176: 106369, 2024.
Kenneth H Huebner, Donald L. Dewhirst, Douglas E Smith, and Ted G Byrom. The finite element method for engineers. John Wiley & Sons, 2001.
George Em Karniadakis, Joannis Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6):422-440, 2021.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Arizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, Anima Anandkumar, et al. Fourier neural operator for parametric partial differential equa- tions. In International Conference on Learning Representations, 2020.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. prysicc-informed neural operator for learning partial differential equations. ACM/JMS Journal of Dala Science, 1(3):1-27, 2024.
Xuanqing Liu, Tesi Xiao, Si Si, Qin Cao, Sanjiv Kumar, and Cho-Jui Hsieh. Neural sde: Stabilizing neural ode networks with stochastic noise, 2019. URL stttps://aaet org/ abs/1906. 02355.
Winfried Lötzsch, Simon Ohler, and Johannes S Otterbach. Learning the solution operator of bound- ary value problems using graph neural networks. arXiv preprint arXiv:2206. 2022.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):118-23) 2229,, 2021.
Levi McClenny and Ulisses Beto.-Neto. adaf-adaptive physics- informed neural networks using a soft attention mechanism. arXiv preprint arXiv:2009.0453 2020.
Bernt Øksendal and Bert Øksendal. Stochastic differential equations. Springer, 2003.
Eckhard Platen Peter F. Kloeden. Numerical Solution of Stochastic Differential Equations. Springer
Berlin, Heidelberg, 1992.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Ehsan Saleh, Saba Ghaffari, Timothy Bretl, Luke Olson, and Matthew West. Learning from integral losses in physics informed neural networks. arXiv preprint arXiv:2305./7387. 2023.
Yong Shang. Fei Wang, and Jingho Sun. Randomized neural network with petrov- gareekin meth- ods for solving linear and nonlinear partial differential equations. Communications in Nonlinear Science and Numerical Simulation, 127:107518, 2023.
H. Shirinabadi and H. A. Talebi. Lyapunov stability analysis of special class of pde systems. The 2nd International Conference Control, Instrumentation and Automation, pp. 648-653, 1. doi: 11A.0///CCAAAumm22011....51
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint 0:9.002303.01469, 2023,
Walter A Strauss. Partial differential equations: An introduction. John Wiley & Sons, 2007.
Yifan Wang, Pengzhan Jin, and Hehu Xie. Tensor neural network and its numerical integration. arXiv preprint arXiv:2207.0275.. 2022a.
Yifan Wang, Yangfei Liao, and Hehu Xie. Solving schr\"fol dinger equation using tensor neural network. arXiv preprint arXiv:2209. /2572, 2022b.
12

---

Published as a conference paper at ICLR 2025
Zhanhong Ye, Xiang Huang, Leheng Chen, Hongsheng Liu, Zidong Wang, and Bin Dong. Pde- former: Towards a foundation model for one-dimensional partial differential equations. arXiv 2024.
preprint ar arXiv:2402.12666
Bing Yu et al. The deep ritz method: deep learning-basse numerical algorithm for solving varia- tional problems. Communications in Mathematics and Statistics, 6(1):1-12, 2018.
Yaohua Zang, Gang Bao, Xiaojing Ye, and Haomin Zhou. Weak adversarial networks for high- dimensional partial differential equations. Journal of Computational Physics, 411:109409, 2020.
A APPROXIMATION OF DROPOUT
In this section, we form the dropout process as by Liu et al. (2019) to analyze its stability perfor- mance using Itô process. The dropout layer be modeled using Bernoulli process:
pn+1 = On + V(On: n, 9; w) & Yn
Yn p
-n + V( n, 9: + V n, (:
I)
(16)
P
p
Yn
On + V(O n,, (; + V(On, n, n, 9:w) o
1))
1 - Pv n, (;w)] o 2n
On + V(O 9:w) + IV
p
where 7n ~ ((1,7) p) and @ represents the Hadamard product. The random variable 2, formulated approximate the standard Gaussian distribution.
Given the Itō process O, satisfying d0, V(O,, t, () dt V 9; w)) dB,, where B, the standard Brownian motion, we can approximate it using the forward Euler scheme 0n+1 = On + V(O,, n. G; w) At + / 11 n, G; u) AW, The constructed 14. is an order-0.5 weak approximation it (Peter E. Kloeden, 1992) since it satisfies the following property:
[E(AW,)) + (B((A)))) )) At s KA
(17)
where AW, is the approximation to the original Gaussian random variable AW,. which is zn here. weak order 0.5 approximation converges weakly to the Itô process a rate proportional to At Hence, distributionally, can focus Itó process instead regarding stability analysis.
B PROOF DETAILS
Proof of Theorem I. First, rewrite the parameter evolution equation (Eq. 4) as shorthand expres- sion:
do. 14 dt + o: dB.,
(18)
where μe, ae are square-integrable. The sub-network is twice continuously differentiable func- tion w.r.t. Or, denoted Y:(w) Uo.- By Itô formula:
aUe.
1
d
UUO
dy,
do,,
+
d0, d0,,d0,,.
00;
2
00,0;
1,j=1
d aUe.
1
d
82 U2.
aUo. 30,
(19)
Miit +
2
dt +
oitt
dB,
00;
2
00?
1=1
1=1
dldt + + dB,
13

---

1 Published as a conference paper at ICLR 2025
Given a Brownian path in advance, Denote the solution time with initial condition Y, as ((Y))))) The Itó integral of the SDE gives the following stochastic integral equations for any t1, 12 2 0 and any initial Y, almost surely:
(Y)()( = (/3 (
dr
d,dB, (w) + +Yo
J0
([)())( = 1/' ur f phi
{d, o, (w) + +Y
(20)
Tı (Y)((( -T, (T₁, (W))(
tit
n, dr +
dB (w) + TT, (Y))
follows from the additive property Itô integral that the solution operator a strong semigroup:
P1 Tıp Tr, (V) a.s., (Y)) = Y a.s.
(21)
strong equality implies weak equality (Øksendal & Øksendal, 2003), therefore the solution also satisfies weak semigroup property.
Proof of Theorem 3. Firstly, the task could be separated into two subtasks. We need to evaluate the perturbation's influence on O₂ and r, separately.
1 [t]
0 ||
lim sup log t1000
l|
||
S lim sup log
||
lim sup log || +ox
(22)
+0000
t
[
t
Theorem has already proved the stability of Eq.
limsup - log || <0 10000 t
(23)
And hence 3a > 0,7 0, s.t. VI T, | Sa
As is proved in Liu et al. (2019), it is said that RA aligns with the assumptions if R does, which means that (E t) < L( + 1). Rs here is a scalar function, and the norm represents its absolute value. Substituting the above equations into the SDE we will get:
de; dt
RA(E,, t) S L(e +1)
(24)
By integrating this simple ODE, we obtain an upper bound on the perturbation's 's influence on ry:
E S
( Lie at +1) =6 Lt
_ (1 - eat
(25)
Substituting 25. into the objective function, stable and thus prove that the solution to this SDE system almost surely exponentially stable.
conclude that r; is almost surely exponentially
log leg + Lt + - - --a
lim sup log ||=]]| < lim sup = t-000
10000
t
(26)
1
limsup log Lt| < 0 10000
14

---

Published as a conference paper at ICLR 2025
Proof of Theorem 4. Firstly, since given continuous activation function, neural networks gener- ally be seen as continuous functions of network parameters, we can conclude that VE, Bee, if 10 0,ll < Eg, then ||og lo <
Secondly, Theorem 3 proves a stricter version of stability of the evolution of e Hence, given E 3ee if |0; Ooll < dg. then | O.ll se. and hence | 0.1 < Ee- Conclusively, ve, e if 0 @ill < 60 UVor Ue.ll e. SINGER is stable during training.
Proof of Theorem 5. Firstly, we attempt to approximate the Gradient Descent (GD) with Bernoulli noise SDE. GD could be represented 0,+1 On (O,,) where 7) is the learning rate. we drop random elements of the gradient, we can get the following formula:
Vo,m p
0,,+1 = =0,a- - (7V,L(0,,m) 9
(27)
Y0
0n,n - -V
@
I)
If we make this representation continuous in Appendix A, can get the following SDE:
P d0 dr TV - PVC(O,,))
(28)
Given the conditions provided, Theorem 3 indicates that Eq.28 is stable. Following the proof in Theorem 4, Vgg. 36, if ||Ue, Ueg < 8, then |10 - 0oll < 6g. Since we always set ro 0, |006 Ooll < 6g. conclusion, SINGER is stable during testing.
C EQUATIONS AND EXPERIMENT SETTING DETAILS
C.I BACKWARD STOCHASTIC DIFFERENTIAL EQUATIONS (BSDE) METHOD
The BSDE (Han et al., 2018) method reformulates certain partial differential equations (PDEs) as backward stochastic differential equations (BSDEs), which can then be solved using numerical methods. The general form of BSDE is:
= = g(XT) + [. f(s Y,, k.) ds - f." ZT dWe,
where X, represents the stochastic process, Y, and Z, are adapted processes, and g(Xr) is the terminal condition.
In our experiments, for each type of equation, we generate data by adjusting the terminal condition function (xx) with different parameters. The BSDE solver then used solve these modified equations, providing us with a dataset of triplets (t, x, Upned), where t represents time, x is the state variable of the PDE, and Upred is the predicted solution time t.
C.2 HAMILTON-JACOR--BELLLA LINEAR- LINEAR-QUADER QUADRATIC (HJB-LQ) EQUATION
C.2.1 EQUATION OVERVIEW
The Gammitonn-acobi-Bellle
(HJB) equation arises in control theory and describes the value func- tion for optimal control problem. In the classical finear-uuadratic-Gaussi (LQG) control prob- lem, the equation is given by:
du (1, x) = - -o Au(l())) - te [0, 77. x E Rd
a1
with the terminal condition uT,xx) g(x).
15

---

Published as a conference paper at ICLR 2025
In our experiments, the terminal cost function is defined as:
g(x) In 1 k |001
where & is the adjustable coefficient.
C.2.2 PARAMETER SETTINGS
In our experiments, the following parameters were used for solving the HJB-LQ equation and for training the neural network to approximate the solution:
Total time of the equation: total time = 0.4
Number of time intervals for solving the PDE: num time interval 20
Range of the adjustable coefficient in the terminal condition gi K.range [0.5, 1.5)
Parameters for fitting the generated data using the neural network:
Number of iterations for fitting: fit. iter 20,000
Batch size during fitting: fit batch size
Threshold for stopping fitting based on loss: fit_threshold = 1e-3
Parameters for training the neural network using our method:
Dropout rate during training: train trainddropout 0.3
Batch size during training: train batch size 64
Learning rate during training: train Jeinnnennnnng.re rate = le-4
C.3
SLACTION-DIFFUSION EQUATION
C.3.1
EQUATION OVERVIEW
The reaction-diffusion equation models the behavior of a system undergoing both diffusion and reaction processes. The equation given by:
du
0t (t. x) Au(t, + x,, Vu), Ee [0, x E Rd,
where the function f is defined as:
f(t,2, u, Vu) = min L, ( (u-
-ais-r
sin sin
x4
e
where A is a constant representing the wave frequency, x: a constant related the initial conditions, and 7 is the total time.
The terminal condition g for this sraction-diffssie equation is defined as:
((x) = + +in sin
24 ==1
h
where & is the the adjustable coefficient.
C.3.2 PARAMETER SETTINGS
In our experiments, the following parameters were used for solving the reaction-diffusion equation and training the neural network:
Total time of the equation: total.time 0.6
Number of time intervals for solving the PDE: num time interval 30
16

---

Published as a conference paper at ICLR 2025
Range of the adjustable coefficient in the terminal condition g: r.range [0.5, 1.5]
Parameters for fitting the generated data using the neural network:
Number of iterations for fitting: fit_n mn_ttrr== 30,000
Batch size during fitting: fit_ fit_batch_stt _size =
Threshold for stopping fitting based on loss: fit threshold = le-3
Parameters for training the neural network using our method:
Dropout rate during training: train_dropout 0.4
Batch size during training: bainh_atccssiee 64
Learning rate during training: train_eernningrre 2e-4
C.4 BURGERS-TYPE EQUATION
C.4.1
EQUATION OVERVIEW
The Burgers-type equation is fundamental partial differential equation that models various physical phenomena such as fluid dynamics, gas dynamics, and traffic flow. The general form of the equation is given by:
Ou (t, Øt
2 2
d
Ou
2 o² )²)
Ee [0, x E Rd
2d
Or.
The terminal condition g(x) ) for the Burgers-type equation is defined as:
g 1 + exp k ++ 1 Ei. Ti
where l is the the adjustable coefficient.
C.4.2 PARAMETER SETTINGS
In our experiments, the following parameters were used for solving the Burgers-type equation and training the neural network approximate the solution:
Total time of the equation: total_time 0.3
Number of time intervals for solving the PDE: num num_time intervaterval =
Range of the adjustable coefficient in the terminal condition gi range [0.8, 1.3]
Parameters for fitting the generated data using the neural network:
Number of iterations for fitting: fit_n_iter = 30,000
Batch size during fitting: fit batch size = 64
Thresholl for stopping fitting based on loss: fit. fit_threshold = 1e-3
Parameters for training the neural network using our method:
Dropout rate during training: train dropout 0.3
Batch size during training:
64
Learning rate during training: train learning rate le-4
C.5 PRICING DEFAULT RISK EQUATION
C.5.1 EQUATION OVERVIEW
The Pricing Default Risk equation is variant of the Black-Scholes equation that models the pricing of financial derivatives under default risk. The equation is given by:
Du (t, x)
1 2²2 Au(t, + (-(1 - 6) Beecewiss_.lneerrr((e u, t [0, T), x E Rd 2
Øt
17

---

Published as a conference paper at ICLR 2025
where the piecewise linear term is expressed as:
piecewise linear(u) (u) ReL.U(ReL.UU - v%)) slope + 7/ - 7) + 70
where d represents the recovery rate after default, r is the risk-free interest rate, 2/ is a threshold parameter, 7hh and " are constants representing the different regimes of the default intensity, and slope determines the linear transition between regimes.
The terminal condition g(x) for the Pricing Default Risk equation is defined as: g(ø) = min(x) k
where k is the the adjustable coefficient.
C.5.2 PARAMETER SETTINGS
In our experiments, the following parameters were used for solving the Pricing Default Risk equation and training the neural network to approximate the solution:
Total time of the equation: outaltimee 0.2
Number of time intervals for solving the PDE: num num_timemtt intervvl = 20
Range of the adjustable coefficient in the terminal condition g: rangee [0.5, 1.0]
Parameters for fitting the generated data using the neural network:
Number of iterations for fitting: fit n iter 30,000
Batch size during fitting: fit batch size 64
Threshold for stopping fitting based on loss: fit threshold = 1c-3
Parameters for training the neural network using our method:
Dropout rate during training: train, dropout 0.4
Batch size during training: bainhbatch_size 64
Learning rate during training: waan_learming_rate le-4
C.6 SINE-GORDON EQUATION
C.6.1
EQUATION OVERVIEW
The Sine-Gordon equation a second-order nonlinear partial differential equation frequently uti- lized in the field theory, string theory, and solitons. The equation given by:
du (t, x) Au(t, x) + sin(u(t, [0, T) Rdd
Ōt
The terminal condition g(x) for Sine- Mond-Goroonnoo equation defined as:
5
g(x) 10 + 2|
C.6.2 PARAMETER SETTINGS
In experiments, the following parameters used for solving the Sine-Gordon equation and training the neural network to approximate the solution:
Total time of the equation: total.time 0.3
Number of time intervals for solving the PDE: num sum interval 20
Parameters for fitting the generated data using the neural network:
Number of iterations for fitting: fit n iter = 5,000
Batch size during fitting: fit. stt_bacch_siee sizee = 100
Threshold for stopping fitting based on loss: fit threshold 1e-3
Parameters for training the neural network using our method:
Dropout rate during training: trainudropout = 0.1
Batch size during training: wrin_batch_sinn 100
Learning rate during training: train learning rate le-3
18

---

Published as a conference paper at ICLR 2025
C.7 ALLEN-CAHN EQUATION
C.7.1 EQUATION OVERVIEW
The Allen-Cahn equation is reaction-diffusion equation that models various physical phenomena such as phase separation, interfacial dynamics and pattern formation. The equation is given by:
du at (1, Au(1,x) + u(t, [=] LE [0, IR"
The terminal condition g(x)) for Allen-Cahn equation is defined as:
5
((x)
10 2|
C.7.2 PARAMETER SETTINGS
In experiments, the following parameters were used for solving the Allen-Cahn equation and training the neural network to approximate the solution:
Total time of the equation: total time 0.3
Number of time intervals for solving the PDE: num time interval = 20
Parameters for fitting the generated data using the neural network:
Number of iterations for fitting: fit_n o_jter = 5,000
Batch size during fitting: fit batch size 100
Threshold for stopping fitting based on loss: fit_threshold 1e-3
Parameters for training the neural network using method:
Dropout rate during training: train dropout 0.1
Batch size during training: bein_batch_ste 100
Learning rate during training: train. laiinllerninggrran rate = 1e-3
D RELATED WORK ON NEURAL PDE SOLVERS IN LOW-DIMENSIONS
The neural-network-bare approaches for low dimensional PDEs can be mainly divided into two categories: thysics-informee network and data-driven (perator(Kariiadaki et al., 2021). Specif- ically, Poysics-Informed Neural Networks (PINNs) (Raissi et al., 2019) leverage neural networks to approximate the solution of PDE and enforce boundary loss and PDE residual loss at a few selected points. This approach has various extensions such as weak form De Ryck et al. (2024), adaptive sample weight McClenny Braga-Neto (2020) and integral loss (Saleh et al., 2023). These methods may suffer from slow convergence and re-training requirements for new initial conditions. Alternatively, DeepONet (Chen & Chen, 1995; Lu et al., 2021) leverages the universal approxi- operators and directly fits the PDE solution operators in
mation theorem of iafiniteddimnsion data-driven manner, solving a family of PDEs with a single model. However, this approach requires large amount of data for training and may generalize well to new initial conditions. Physics- informed neural operators (PINO) Li et al. (2024) a hybrid approach incorporating data and PDE residuals in the loss function. There exist many neural operators with other architectures, such as Fourier transform layer (Li et al., 2020), and graph neural networks (GNNs) et al., 2019). The existing GNN-based models describe irregular spacial grids and their relative position as graph (Lötzsch et al., 2022; Horie & Mitsume, 2022; Bryutkin et al., 2024).
E FUTURE WORK AND BROAD IMPACT
Future work will explore extending this method even higher dimensions and more complex sys- tems, potentially incorporating additional techniques of data sampling limitation of approach is the generalization among PDE parameters, which could be addressed
and training. A potential
19

---

Published as a conference paper at ICLR 2025
by incorporating parameters into the model inputs and modifying the training loss. Another limita- tion is that current model is confined to continuous solution, which can be extended to discontinuous cases by adopting weak form residual.
The proposed method has the potential to impact wide range of scientific appli- cations, including quantum mechanics, quantum chemistry and financial mathematics. The ability to learn high-dimensional PDEs solution operator, instead of solving a single problem instance, could lead to significant advancements in these fields. In addition, the generalization ability of the proposed model might pave the way for pre-trained foundation models of high-dimensional PDEs.
Broader Impact:
20


Reviews:
[{"review_id": "YQ73vftkuP", "reviewer": "Reviewer_8dkd", "summary": "In this work the authors propose a stochastic graph neural network based framework for solving high dimension partial differential equations (PDE).  Neural PDE solvers that perform well in low-dimensional settings do not generalize to high-dimensional settings and hence require special attention. However, existing methods in the high dimension regime suffer from instability and do not generalize to different types of PDEs. To overcome these drawback, the authors propose the SINGER framework.\n\nThe SINGER model uses a GNN to approximate the solution of the PDE at the initial time step and then stochastically evolves the network parameters over time according to a GNN-driven ODE. This network is then used to approximate the solution at later time steps. The structure enables permutability of neurons within a layer, this enhances the models capability to generalize. To combat the issue of instability in the evolution process, SINGER introduces noise during the training step.  Further more, SINGER is designed to satisfy three key assumptions: graph topology, semigroup property and stability. The authors theoretically and empirically verify that their proposed framework satisfies these assumptions. Finally, SINGER is validated on 8 benchmark PDEs and its performance is compared with state-of-the-art methods.", "questions": "(page 1) The 3 contributions of the paper can be rewritten clearly and explained in more detail. \n\nThe authors can consider briefly explaining their theoretical contributions earlier in the paper.\n\n(Table 3) Why not replace 'ours' with SINGER", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The authors identify gaps in current literature and do a good job in motivating their research.\n\n2) The contributions of their framework are explained clearly.\n\n3) Theoretical analysis of stability, graph topology and semigroup property of SINGER\n\n4) Computational results are positive", "weaknesses": "Section 3: The authors could do a better job in explaining what the 3 assumptions signify and why They are important for solving PDEs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QWVekL5Wga", "reviewer": "Reviewer_mPx4", "summary": "The paper proposes to solve time-dependent PDEs with a neural network model $U$that approximates the PDE solutions, where the neural network parameters are solutions to a system of stochastic differential equations with a drift term modeled by a Graph Neural Network on a graph induced by the architecture of $U$. \n\nWhile I am unfamiliar with the literature, I think the idea is interesting, especially with the mathematical reasons discussed in the paper. Upon careful inspection, it seems that this work extends NODE with stochastic noise and GNN architecture for the control vector $V$, where they used tools from Liu et al. 2019 (which is a paper that was rejected in ICLR 2020). My slight reservation with this work is due to my confusion as to why the paper Liu et al. 2019 paper has not been published and the work NODE also seems to be under review. However, these facts should not be the basis for rejecting this work. Based on my reading, the math part of this work is rather trivial and I am not against the paper for publication after the authors satisfactorily answer my following questions.", "questions": "1. Looking at Eq. (4), I assumed that the PDE that is to be solved is $u_t = F(U)$? Is there any restriction on $F$ (e.g., class of PDEs) for this approach to work well or to fail? \n2. Can the authors clarify or give a detailed explanation on a sub-network of $U$? At this point, I am just reading the paper as if the entire architecture of $U$ forms a graph, and the parameters in the neural networks are denoted by $\\theta_t$. How do you set $V,E$ a priori? If it is a fully connected network, does it mean $E$ is always 1 between any pair of vertices? How do you choose a sub-network of $U$? How critical is the performance under variations of sub-networks (what if you use the entire full network)?\n3. How do you specify $N$ in (3)? I suspect you need to set it to $N^2>2L$ based on Theorem 1, where $L$ is the Lipschitz constant of $V$, which needs to be estimated in the optimization procedure?\n4. I am not familiar with NODE or PINO, so it is hard for me to say that the comparison in numerical experiments is fair. How does this work compared to any other methods to solve such a high-dimensional PDEs that are cited in the references (such as DeepRitz, PINN).\n5. As noted in p.6, the semigroup property depends on the random seed in the solution of SDEs. I think stating that the method satisfies the semigroup property in Table 1 is overselling. Upon inspecting the setup for the stability (Eq (6)) in the manuscript, the analysis is performed under the assumption that the same path of Brownian noise is used. I think one can avoid this assumption by stating a weaker result (convergence in distribution instead of almost surely convergence). If my conjecture is correct, there is more theoretical work to satisfy the semigroup property in a weaker sense instead of stating \"approximately satisfied\". Table 5 also suggests that the semigroup is weakened after training (unless I misunderstood the reported values in this table)?\n6. Minor: Should $u_{\\theta_t}$ below Eq. (4) be $U_{\\theta_t}$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The central idea of using GNN for modeling the control seems to be novel.", "weaknesses": "This writing needs more clarifications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HpAU2v98fp", "reviewer": "Reviewer_kUCH", "summary": "The authors introduce StochastIc Network Graph Evolving operatoR (SINGER), a novel framework for learning the evolution operator of high-dimensional partial differential equations (PDEs). SINGER employs a sub-network to approximate the initial solution and stochastically evolves its parameters over time using a graph neural network to approximate later solutions. Designed to inherit key properties such as graph topology, semigroup properties, and stability, SINGER comes with theoretical guarantees of performance. Numerical experiments on eight PDEs across various dimensions show that SINGER outperforms existing methods in most cases and generalizes effectively to new conditions and parameters.", "questions": "1. Why is it beneficial to have the stochasticity? Usually they hurts the surrogate performances, is there an ablation on the noise level to illustrate this issue?\n2. The test cases are all very smooth PDEs, any reason why this is the case? Is it possible to either prove the ability on non-smooth PDEs or show with empirical results?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The presentation is clear and concise, and the problem is clearly illustrated. Although if I am correct, the network solves the PDEs at fixed collocation points, it would be great to point that out somewhere.\n2. The empirical results are clear.\n3. The theory of the method are sound and extensive.", "weaknesses": "1. While the empirical results cover most important PDE solvers, it does not compare with the SOTA methods, and merely compare to the vanilla ones. What's more, methods such as neural ordinary differential equations (NODEs) are not meant for solving partial differential equations, it seems a bit unfair to compare to them.\n2. For instance, the proposed method resides in the realm of hypernetworks, it would be nice to see a hypernetwork of PINN or DeepONet.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)