Title: Learning Efficient Surrogate Dynamic Models with Graph Spline Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Graph

Spline Collocation Method

Graph Neural Networks

Simulation

Partial Differential Equations

PDEs

Physics

Scientific Computing

Surrogate Models

Weather Forecasting
Abstract: While complex simulations of physical systems have been widely used in engineering and scientific computing, lowering their often prohibitive computational requirements has only recently been tackled by deep learning approaches. In this paper, we present GraphSplineNets, a novel deep-learning method to speed up the forecasting of physical systems by reducing the grid size and number of iteration steps of deep surrogate models. Our method uses two differentiable orthogonal spline collocation methods to efficiently predict response at any location in time and space. Additionally, we introduce an adaptive collocation strategy in space to prioritize sampling from the most important regions. GraphSplineNets improve the accuracy-speedup tradeoff in forecasting various dynamical systems with increasing complexity, including the heat equation, damped wave propagation, Navier-Stokes equations, and real-world ocean currents in both regular and irregular domains.

Learning Efficient Surrogate Dynamic Models with Graph Spline Networks

Abstract

While complex simulations of physical systems have been widely used in engineering and scientific computing, lowering their often prohibitive computational requirements has only recently been tackled by deep learning approaches. In this paper, we present GRAPHSPLINENETS, a novel deep-learning method to speed up the forecasting of physical systems by reducing the grid size and number of iteration steps of deep surrogate models. Our method uses two differentiable orthogonal spline collocation methods to efficiently predict response at any location in time and space. Additionally, we introduce an adaptive collocation strategy in space to prioritize sampling from the most important regions. GRAPHSPLINENETS improve the accuracy-speedup tradeoff in forecasting various dynamical systems with increasing complexity, including the heat equation, damped wave propagation, Navier-Stokes equations, and real-world ocean currents in both regular and irregular domains.

1 Introduction

For a growing variety of fields, simulations of partial differential equations (PDEs) representing physical processes are an essential tool. PDE-based simulators have been widely employed in a range of practical applications, spanning from astrophysics (Mücke et al., 2000) to biology (Quarteroni and Veneziani, 2003), engineering (Wu and PortéAgel, 2011), finance, (Marriott et al., 2015) or weather forecasting (Bauer et al., 2015).

However, traditional solvers for physics-based simulation often need a significant amount of computational resources (Houska et al., 2012), such as solvers based on first principles and the modified Gauss-Newton methods. Traditional physics-based simulation heavily relies on knowledge of underlying physics and parameters requiring ad-hoc modeling, which is sensitive to design choices. Finally, even the best traditional simulators are often inaccurate due to the difficulty in approximating real dynamics (Kremer and Hancock, 2006; Oberkampf, 2019; Sanchez-Gonzalez et al., 2020).

An attractive alternative to traditional simulators is to use deep learning to train surrogate models directly from observed data. Data-driven surrogate models can generate

KAIST OMELET

predictions based on the knowledge learned automatically from training without the knowledge of the underlying differential equations. Among deep learning methods, graph neural networks (GNNs) have desirable properties such as spatial equivariance and translational invariance, allowing learning representations of dynamical interactions in a generalizable manner (Pfaff et al., 2021; Bronstein et al., 2021) and on unstructured grids. Despite the benefits of these paradigms, by increasing resolution, graph-based models require significantly heavier calculations. Graph models as Poli et al. (2019); Lienen and Günnemann (2022) predict in continuous time or continuous space. However, the substantial computational overhead due to the need for iterative evaluations of a vector field over time requires considerable computation and limits their scalability (Xhonneux et al., 2020).

In this work, we propose GRAPHSPLINENETS, a novel approach that exploits the synergy between GNNs and the orthogonal spline collocation (OSC) method (Bialecki and Fairweather, 2001; Fairweather and Meade, 2020) to both efficiently and accurately forecast continuous responses of a target system. We use the GNNs to efficiently get the future states on coarse spatial and temporal grids and apply OSC with these coarse predictions to predict states at any location in space and time (i.e., continuous). As a result, our approach can generate high-resolution predictions faster than previous approaches, even without explicit prior knowledge of the underlying differential equation. Moreover, the end-to-end training and an adaptive sampling strategy of collocation points help the GRAPHSPLINENETS to improve prediction accuracy in continuous space and time.

We summarize our contributions as follows:

We introduce GRAPHSPLINENETS, a novel learning framework to generate fast and accurate predictions of complex dynamical systems by employing coarse grids to predict time and space continuous responses via the differentiable orthogonal spline collocation.

We propose an adaptive collocation sampling strategy that adjusts the locations of collocation points based on fast-changing regions, enabling dynamic improvements in forecasting accuracy.

We show GRAPHSPLINENETS are competitive against existing methods in improving both accuracy and speed in predicting continuous complex dynamics on both simulated and real data.

2 Related Works

Modeling dynamical systems with deep learning Deep neural networks have recently been successfully employed in accelerating dynamics forecasting, demonstrating their capabilities in predicting complex dynamics often orders of magnitude faster than traditional numerical solvers (Long et al., 2018; Li et al., 2021; Berto et al., 2022; Pathak et al., 2022; Park et al., 2022; Poli et al., 2022; Lin et al., 2023a,b; Boussif et al., 2022, 2023). However, these approaches are typically limited to structured grids and lack the ability to handle irregular grids or varying connectivity. To address this limitation, researchers have turned to graph neural networks (GNNs) as a promising alternative: GNNs enable learning directly on irregular grids and varying connectivity, making them well-suited for predicting system responses with complex geometric structures or interactions (Li et al., 2022a). Additionally, GNNs inherit physical properties derived from geometric deep learning, such as permutation and spatial equivariance (Bronstein et al., 2021), providing further advantages for modeling complex dynamics. Alet et al. (2019), one of the first approaches to model dynamics with GNNs, represent adaptively sampled points in a graph architecture to forecast continuous underlying physical processes without any a priori graph structure. Sanchez-Gonzalez et al. (2020) extend the GNN-based models to particle-based graph surrogate models with dynamically changing connectivity, modeling interactions through message passing. Graph neural networks have also recently been applied to large-scale weather predictions (Keisler, 2022). However, GNNs are limited in scaling in both temporal and spatial resolutions with finer grids as the number of nodes grows.

Accelerating graph-based surrogate models An active research trend focuses on improving the scalability of graph-based surrogate models. Pfaff et al. (2021) extend the particle-based deep models of (Sanchez-Gonzalez et al., 2020) to mesh-based ones. By employing the existing connectivity of mesh edges, Pfaff et al. (2021) demonstrate better scalability than proximity-based online graph creation modeling particle interactions. Li et al. (2022b) extend convolutional operators to irregular grids by learning a mapping from an irregular input domain to a regular grid: under some assumptions on the domain, this allows for faster inference time compared to message passing. Fortunato et al. (2022) build on the mesh-based approach in (Sanchez-Gonzalez et al., 2020) to create multi-

scale GNNs that, by operating on both coarse and fine grids, demonstrate better scalability compared to single-scale GNN models. GRAPHSPLINENETS do not need to learn additional mappings between different domains or need additional parameters in terms of different scales thus reducing computational overheads by operating directly on unstructured coarse meshes. Our method bridges the gaps between coarse and fine grids in time and space via the orthogonal spline collocation (OSC), thus improving GNN-based methods' scalability.

Collocation methods Collocation and interpolation methods 2 are used to estimate unknown data values from known ones (Bourke, 1999). A fast and accurate collocation method with desirable properties in modeling dynamics is the orthogonal spline collocation (OSC) (Bialecki, 1998; Bourke, 1999). Compared to other accurate interpolators with O ( n 3 ) time complexity, the OSC has a complexity of only O ( n 2 log n ) thanks to its sparse structure, allowing for fast inference; moreover, it is extremely accurate, respects C 1 continuity, and has theoretical guarantees on convergence (Bialecki, 1998). Due to such benefits, a few research works have aimed at combining OSC with deep learning approaches. Guo et al. (2019); Brink et al. (2021) combine deep learning and collocation methods by directly learning collocation weights; this provides advantages over traditional methods since the OSC can ensure C 1 continuity compared to mesh-based methods. A recent work Boussif et al. (2022) has studied learnable interpolators and demonstrated that it is possible to learn an implicit neural representation of a collocation function. GRAPHSPLINENETS differ from previous neural collocation methods in two ways: first, we decrease the computational overhead by not needing to learn interpolators or collocation weights; secondly, by interpolating not only in space but also in time, our method enables considerable speedups by employing coarser temporal grids.

3 Background

3.1 Problem Setup and Notation

We first introduce the necessary notation that we will use throughout the paper. Let superscripts denote time, and subscripts denote space indexes as well as other notations. We represent the state of a physical process at space location X = { x i , i = 1 , · · · , N } and time t ∈ R as Y t = { y t i , i = 1 , · · · , N } where N represents the number of sample points and Ω ⊂ R D is a D -dimension physical domain. A physical process in this domain can be described by a solution of PDEs, i.e., y t i = u ( x i , t ) . The objective of a dynamics surrogate model is to estimate future states { ˆ Y t , t ∈ R + } given the initial states Y 0 .

3.2 Message Passing Neural Networks

We employ a graph G t = { v t i , e t ij } to encode the status of physical process at sample points X at t . v t i and e t ij denote the attribute of sample node i and the attribute of the directed edge between node i and j , respectively. Node attributes are encoded from sample points state information while edge attributes are the distance between every two nodes. Message passing neural networks (MPNNs) employ an encoder-processor-decoder structure to predict the states of sample points at the next timestep by giving the previous states:

where E ( · ) is the encoder, P i ( · ) is the i -th message passing layer, and D ( · ) is the decoder.

3.3 Orthogonal Spline Collocation Method

The OSC method consists of three steps in total: (1) partitioning the input domain and selecting collocation points, (2) generating polynomials with parameters to be determined with observations, and (3) solving equations to determine the parameters.

y

0

.

0

.

0

4

5

5

.

0

0

0

0

0

.

.

0

0

0

0

.

.

2

2

0

0

.

.

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

.

0

0

.

5

y

Ground Truth

Polynomial 1 Simulation

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

y

Ground Truth

Polynomial 1 Simulation

Polynomial 2 Simulation

Polynomial 3 Simulation

Collocation Points

Partition Points

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Ground Truth

Polynomial 1 Simulation

Polynomial 2 Simulation

Polynomial 3 Simulation

Collocation Points

Partition Points

0

.

0

0

.

2

0

.

4

0

.

6

x

Ground Truth

Polynomial 1 Simulation

Polynomial 2 Simulation

Polynomial 3 Simulation

Ground Truth

Polynomial 1 Simulation

0

.

0

0

.

2

0

.

4

0

.

6

x

Ground Truth

Polynomial 1 Simulation

Polynomial 2 Simulation

Polynomial 3 Simulation

x

Partitioning and initializing polynomials The OSC aims to find a series of polynomials under order r while satisfying C 1 continuity to approximate the solution of PDEs. To find these polynomials, we split each dimension of the domain into N partitions 34 . Then, we initialize one r -order polynomial with r +1 unknown coefficients in each partition. These polynomials have N × ( r +1) degrees of freedom, i.e., the variables to be specified to determine these polynomials uniquely.

Collocation points selection and equation generation We firstly have 2 equations from the boundary condition and ( N -1) × 2 equations from the C 1 continuity restriction. Then we need N × ( r -1) more equations to uniquely define the polynomials. So we select r -1 collocation points where the observation is obtained to determine the parameters in each partition. By substituting the state of collocation points to the polynomials, we can get N × ( r -1) equations. Now we transfer the OSC problem to an algebraic problem. With properly selected base functions, this algebraic equation will be full rank and then has a unique solution (Lee and Micchelli, 2013).

Solving the equations The coefficient matrix of the generated algebraic problem is almost block diagonal (ABD) (De Boor and De Boor, 1978). This kind of system allows for efficient computational routines (Amodio et al., 2000), that we introduce in § 3.4. By solving the equation, we obtain the parameters for the polynomials that can be used to predict the value of any point in the domain. Weillustrate an example of 1 -D OSC problem in Fig. 3.1. More details and examples of the OSC are provided in Appendix A. The OSC method can also be extended to 2 -D and higher-dimensional domains (Bialecki and Fairweather, 2001). In our work, we employ the 1 -D and 2 -D OSC approaches in the time and space domains, respectively.

3.4 Efficiently solving ABD matrices

Most interpolation methods need to solve linear equations. Gaussian elimination is one of the most widely used methods to solve a dense linear equation, which, however, has a O ( n 3 ) complexity (Strassen et al., 1969). Even with the best algorithms known to date, the lower bound of time complexity to solve such equations is O ( n 2 log n ) (Golub and Van Loan, 2013).

In the OSC method, the coefficient matrix for the linear equation follows the ABD structure, which we can efficiently solve with a time complexity O ( n 2 ) by the COLROW algorithm (Diaz et al., 1983) as shown in Fig. 3.2 (a). The core idea for this method is that by using the pivotal strategy and elimination multipliers, we can decompose the coefficient matrix into a set of a permutation matrix and upper or lower triangular matrix that can be solved in O ( n 2 ) time each. The latest package providing this algorithm is in the FORTRAN programming language: our re-implementation in PyTorch (Paszke et al., 2019) allows for optimized calculation, GPU support, and enabling the use of automatic differentiation, as shown in Fig. 3.2 (b).

Number of Elements

Calculation Time

O

(

n

3

)

(Gaussian)

O

(

n

2

log

n

)

(Sparse)

O

(

n

2

)

(COLROW)

(a)

A

B

C

0

5

10

15

20

100

Iteration Running Time

(b)

0

.

6

x

0

.

6

x

Polynomial 2 Simulation

8

1

0

.

Polynomial 3 Simulation

8

1

0

.

Polynomial 2 Simulation

Polynomial 3 Simulation

Collocation Points

Collocation Points

Partition Points

Partition Points

.

.

0

0

y

y

0

0

0

0

.

.

.

.

5

5

0

0

y

y

0

.

0

.

0

.

0

8

.

8

1

0

.

1

0

.

Collocation Points

Partition Points

Collocation Points

Partition Points

4 Methodology

4.1 GRAPHSPLINENETS

Fig. 4.1 depicts the entire architecture of our model. GRAPHSPLINENETS can be divided into three main components: message passing neural networks (MPNN), time-oriented collocation and space-oriented collocation. The MPNN takes observations at collocation points as input and infers a sequence of discrete future predictions via autoregressive rollouts. We then use the time-oriented and space-oriented collocation methods on these discrete predictions to obtain functions that can provide both time and space-continuous predictions.

ˆ

y

t

1

2

ˆ

y

t

1

3

ˆ

y

t

1

1

t

0

+∆

t

ˆ

y

t

0

+∆

t

3

ˆ

y

t

0

+∆

t

1

ˆ

y

t

0

+∆

t

2

ˆ

y

t

0

+∆

t

4

t

0

(

x

3

, y

t

0

3

)

(

x

1

, y

t

0

1

)

(

x

2

, y

t

0

2

)

NN

NN

NN

Time-oriented OSC

t

ˆ

y

t

0

1

ˆ

y

t

1

1

· · ·

ˆ

y

t

0

+∆

t

1

. . .

. . .

Linear Equation

with ABD matrix

Simulation Equation

ˆ

u

1

(

t

0

+∆

t

)

. . .

Linear Equation

with ABD matrix

Simulation Equation

Space-oriented OSC

ˆ

u

t

0

+∆

t

(

x

4

)

ˆ

y

t

0

+∆

t

3

ˆ

y

t

0

+∆

t

1

ˆ

y

t

0

+∆

t

2

ˆ

y

t

0

+∆

t

4

t

0

+∆

t

t

K

t

1

ˆ

y

t

K

1

ˆ

y

t

K

2

ˆ

y

t

K

3

ˆ

y

t

K

1

(

x

i

, y

t

i

)

v

e

v

t

i,

0

=

E

v

(

x

i

, y

t

i

)

v

t

j,

0

v

t

i,

0

e

t

ij,

0

v

t

j,

0

v

t

i,

1

e

t

ij,

1

v

t

j,

1

v

t

i,m

e

t

ij,

0

=

E

e

(

r

i,j

)

r

ˆ

y

t

+1

i

=

D

(

v

t

i,m

)

G

t

0

G

t

1

G

t

m

-

1

G

t

m

ˆ

Y

t

+1

. . .

(

X

,

Y

t

)

Encoder

E

P

1

P

m

D

Decoder

Processor

Construct Graph

Encoder for collocation point states

Encoder for pairwise properties

Pass Message

Pairwise properties between nodes

Get Updated States

P

1

. . .

O

O

O

O

O

i,j

E

E

Assign collocation points We follow the same rule to select collocation points in a 2 -D domain as described in § 3.3 to select sample points { x i } for a 2 -D domain. We have the states of these sample points at the initialization time frame { y t 0 i } .

Message passing neural networks The neural networks take the states of sample points at the initial time frame { ( x i , y t 0 i ) } as input and employ an encoder-processor-decoder architecture for obtaining the next steps states ˆ y t j i , j = 1 , · · · , N autoregressively.

Time-oriented OSC For each sample point, taking x i as an example, the neural networks generate a sequence of predictions along time which can be considered as a 1 -D OSC problem. In this time-oriented OSC problem, collocation points are time steps { t 0 , t 1 , · · · , t K } and values at these time steps { y t 0 i , ˆ y t 1 i , · · · , ˆ y t K i } , and the output is the polynomials ˆ u i ( t ) , t ∈ [ t 0 , t K ] which provides a continuous prediction along time; thus, we can obtain predictions at any time.

Space-oriented OSC For each time frame, we take t k as an example, the neural network generates prediction values of sample points which can be considered a 2 -D OSC problem. In this space-oriented OSC problem, collocation points are positions of sample points { x i } , their values at this time frame { ˆ y t k i } , and the output is polynomials ˆ u t k ( x ) , x ∈ Ω which provides a continuous prediction in the domain. Now we get the prediction of every time frame at any position. Note that the space-oriented OSC can use the time-oriented OSC as the input, which means that we can use the space-oriented OSC at any time frame and then get a spatio-temporal continuous predictions ˆ u ( x , t ) , ( x , t ) ∈ Ω × [ t 0 , t K ] .

<latexi sh

1_b

64="w

X

LJj

Vn

5TFC

q

+E

dm0yQ

>A

B

c

NS8

3Ur

/o

9

GP

K

g

D

7

vH

2

u

M

f

W

z

k

R

I

Z

p

O

Y

<latexi sh

1_b

64="TgG

HO

K

m5kDc9YI78zR0

qQ

>A

C

n

V

LS

N

FJ3U

2v

M

E

y

Z

Bfu

o

d

j

X+

r

/

W

p

P

w

<latexi sh

1_b

64="f27S

8

pzy

L

m9

>A

B

3

VD

gN

EO

r

q

MKu

PUY

I

HC

X

T

F

Q

v

/

R

Go

j

+

5

J

k

w

<latexi sh

1_b

64="2ZCF

okON8

A

vgrz5d

3

J

>

X

c

VDLS

BE

G

Y

T

I

M

Q9

w

p

y

U

mH

0

/

P

qR

n

W

7

u

j

f

+

<latexi sh

1_b

64="

u

2o

C3A

dDK

c

mOG

gnE

>

X

V

LS

N

FJ

U

vqB

zW

R

k

Q

9

H

Mp

T

I0

5

w

f

y7

+

r

/

Y

8

j

Z

P

<latexi sh

1_b

64="Q

Aq8MED

zpKcU

w

0L

X9r

Y

>

CB

Z

S

N

FI

n

V

W/

g

u

2

m

k7dHJ

5

o

R

3v

+T

O

f

G

y

j

P

<latexi sh

1_b

64="

qXwSdP

N

EBQ/

5+

ZFr

A

>

C

H

c

VDL

J3U

2v

g

KI

W

T

u

9

Mp

O

z

0Y

j

7

R

n

y

8m

G

k

f

o

<latexi sh

1_b

64="zIwrO/

AGPZBCVQ5

9NJ

WE

>

n

c

DLS

F

3U

2vq

M

R

Y

yT

d

pm

k

j

7

o

8K

f+

H

0

u

g

X

<latexi sh

1_b

64="u0cgjGv

KSr

o

CD

Zd

9Y

>A

X

V

L

N

FJ3

PU

O

y

I

RE

7

M

Bf nT

/

z

+

8

H

5

W

w

pq

2

m

Q

k

<latexi sh

1_b

64="IGQ

0X

S

kmV

EH

rN

c

w

>A

C

DL

FJ3

PU

d

O

y

Z

vYBT

n

f+

j

9z5

2

K7

p

W

g

R

u

o

8

q

M

/

<latexi sh

1_b

64="pJ

QOV

L

q3

w

S8d9A

5MU

>

C

n

c

D

N

F

rP

uBo

H

Z

T

jWEy

R

G7

z

0

+

v

f

m

I

/

Yk

X

<latexi sh

1_b

64="pRCKH

Ag

2

EqZ

S0Y

z8Pk

>

B9

c

VDL

N

F

r

fU

du ov

T

QwW

+

O

M

X

j

/

n

m

3

G

7

J

I

5

y

<latexi sh

1_b

64="wvW

fK

Zd83j/N

G

g

>A

B

n

c

VDLS

M

F

TX7

+R

2

k

u

r

Y

Q

I

E

9

P

O

z

H

m

5

p

o

y

q

U

0

<latexi sh

1_b

64="

W0R+EZf57

n

N

pVyMDQg wA

>

CI3

c

JS

BFO

jX

9

G

L

q

o

/

PH

v

r

2

d

Y

z

K

k

m

8

T

U

u

<latexi sh

1_b

64="+kvd7p

8

ygjQ

UJI

VDO

Eo

Pc

>A

C

X

LS

N

F

3

r

m

R

G

B

2

0M

5

H

f

u

z

w

n

9

/

Z

Y

WT

K

q

<latexi sh

1_b

64="F

7vdRm

+O

3

Xz

c

Q

>A

CIH

VDLS

N

J

2

g0

K

B

u

p

9

k

G

r

Y

P

T

y

j

5

o

8

W

/

<latexi sh

1_b

64="z dJHuSX

0

f

T

D

E

>A

c

V

L

N

F

2pr

7

ZR

C

W

9oY5

3

y

O

G

j

g

8

vQ

P

q

mK

w

I

k

U

<latexi sh

1_b

64="f27S

8

pzy

L

m9

>A

B

3

VD

gN

EO

r

q

MKu

PUY

I

HC

X

T

F

Q

v

/

R

Go

j

+

5

J

k

w

<latexi sh

1_b

64="

u

2o

C3A

dDK

c

mOG

gnE

>

X

V

LS

N

FJ

U

vqB

zW

R

k

Q

9

H

Mp

T

I0

5

w

f

y7

+

r

/

Y

8

j

Z

P

<latexi sh

1_b

64="

j7CSQP

RMIzqV+Y

Lf3

w

>A

B

X

c

D

N

FJ

r

U

d

m

u

K

g

H9

G

O2

T

Z

5

/ok

8

W

n

v

E

y

0

p

<latexi sh

1_b

64="2ZCF

okON8

A

vgrz5d

3

J

>

X

c

VDLS

BE

G

Y

T

I

M

Q9

w

p

y

U

mH

0

/

P

qR

n

W

7

u

j

f

+

<latexi sh

1_b

64="u0cgjGv

KSr

o

CD

Zd

9Y

>A

X

V

L

N

FJ3

PU

O

y

I

RE

7

M

Bf nT

/

z

+

8

H

5

W

w

pq

2

m

Q

k

<latexi sh

1_b

64="20IS

53PO

BY

u9z

T

o

8/W

>A

H

c

VDL

gM

F

X7

+R

C

Zk

G

Uy

v

N

j

J

p

nf

m

q

d

w

Q

r

E

K

<latexi sh

1_b

64="+cV39foYp

Z

7

m

qH

E

M

>A

B

DJSgN

F

wT

U

G

P

K

yI2z

k

d

R

Q

v

8

/

Or

C

u

W

n

5

0

j

L

X

<latexi sh

1_b

64="IGQ

0X

S

kmV

EH

rN

c

w

>A

C

DL

FJ3

PU

d

O

y

Z

vYBT

n

f+

j

9z5

2

K7

p

W

g

R

u

o

8

q

M

/

<latexi sh

1_b

64="OKvjdn7

5

UHSuz2

m+0

w/

q8

>A

B

c

VDL

N

F

3

W

r

kJ

I

X

9g

Pp

T

YQ

f

C

R

M

y

Z

G

o

E

<latexi sh

1_b

64="

07Fr fnc8

AQL

O

BygZ

>

+X

VD

S

N

J3

PU

d

m

u

K

q

H9

GM

2

T

5

j

/ok

W

v

w

E

R

IY

p

C

z

<latexi sh

1_b

64="On98/

JB

y0

C

G

P

d

Q

UM

>A

X

c

VDLS

N

F

2pr

f

z

gE

I

uH

ZwT

jW

R

m

k

Y

K

v

7

53

o

+

q

<latexi sh

1_b

64="z dJHuSX

0

f

T

D

E

>A

c

V

L

N

F

2pr

7

ZR

C

W

9oY5

3

y

O

G

j

g

8

vQ

P

q

mK

w

I

k

U

<latexi sh

1_b

64="N7J

ZOULQTVw

XDWq

Ykf

0c

>A

B

H

S8

F

yp

+R

C

5

I

rm

3

K

/

E

g

2

j

M

d

P

o

9

z

n

u

v

G

<latexi sh

1_b

64="20IS

53PO

BY

u9z

T

o

8/W

>A

H

c

VDL

gM

F

X7

+R

C

Zk

G

Uy

v

N

j

J

p

nf

m

q

d

w

Q

r

E

K

<latexi sh

1_b

64="uo

T

RH

k

L

Nm9

X

zD

EYZM

>A

c

V

F

3U

2v8

y

CG

+

W

O

n

w

Jp

Q

P

0

5

qr jf

7

I

K

d

<latexi sh

1_b

64="2

+

5DmPy

JCf3SHvY/

WXu

>A

B

n

c

V

N

8

E

Ur

q

o

9L

0

G

K

g

7

Q

M

z

d

k

R

FI

Z

p

O

j

<latexi sh

1_b

64="pWQZgDr

Tc

I5C

Jf uYL

w

>A

B+H

V

NS8

E

3U

o/GvX

nk

R

0

F

j

z

O

y

2

P

7

d

qK

M

m

9

<latexi sh

1_b

64="

cCBUDHyo2j3

r0

Fu

N

pkQ

>A

+

V

S8

E

/GvX

Z

In

RT

W

Y

z

L

5

wO

g

P

7

d

f

qK

M

J

m

9

<latexi sh

1_b

64="

u

2o

C3A

dDK

c

mOG

gnE

>

X

V

LS

N

FJ

U

vqB

zW

R

k

Q

9

H

Mp

T

I0

5

w

f

y7

+

r

/

Y

8

j

Z

P

<latexi sh

1_b

64="oj g+

7

Tw2Lf k0

U

nrVC

Q

>A

B/H

c

D

S

N

F

3

W

q

K

J

I

X

9

Pp

z

Y

8

u

R

m

M

y

Z

d

O

5

G

v

E

<latexi sh

1_b

64="UoLW

Tu0

n

9Fz2

PA

I

>

c

VD

S

N

r

w

K

k

X

j

g

C

p

7mY

/

Rf

M

G

d

y

5

q

OE

v

8

<latexi sh

1_b

64="Qk u/

N9LSnjZ

U+GV7

0

c

I

>A

B

X

D

FJ3

2v

5

yC

m

oW

Y

rR

w

HO

z

fM8

p

K

d

g

T

E

P

q

<latexi sh

1_b

64="+ynwO8/

ERrK

TA3u

Gj9W

>

B

H

c

VDLS

N

FJ

U

2v

d

o

0M

m

5

C

X

P

Z

q

z

I

f

7

<latexi sh

1_b

64="

ON

8jrFw/c

GUX0

o2Y

R

JH

k

>A

V

S

E

3

v

Z

In

q

T

W

D

g

z

L

5u yC

p

P

7

d

f

K

Q

M

9

m

<latexi sh

1_b

64="pWQZgDr

Tc

I5C

Jf uYL

w

>A

B+H

V

NS8

E

3U

o/GvX

nk

R

0

F

j

z

O

y

2

P

7

d

qK

M

m

9

<latexi sh

1_b

64="Qq2w dvg

X

j

U

pr

CNzI8

>A

B9H

c

VDLS

M

F

T

7W+R

Zk

u

G5

P

o

Y

y

O

0

/

3J

nf

m

E

K

<latexi sh

1_b

64="/R3N

LoD

g

B

O

Z

YQSqJPw

>A

7X

c

V

8

E

Ur

9

0

G

K

T

2k vH

5

+

u

M

z

nf

m

W

d

I

FC

y

p

j

<latexi sh

1_b

64="N7J

ZOULQTVw

XDWq

Ykf

0c

>A

B

H

S8

F

yp

+R

C

5

I

rm

3

K

/

E

g

2

j

M

d

P

o

9

z

n

u

v

G

<latexi sh

1_b

64="vc5XK0f

T

kyn2NRg

Br

I

JE

>A

+H

V

S8

3U

o/G

Z

W

P

Y

L

m7

Q

M

C

j

Fu

z

d

w

q

O

9

D

p

<latexi sh

1_b

64="

OA

DznUIdWXu

g

y

rTL

Mv8

>

B

H

c

V

NS

EJ3

q/

9

0mk

CF

2

oQ

+

G

Z

j

Y

w

P

f

p

7

5

KR

<latexi sh

1_b

64="3N

ASc

GEFX

W

LmTJ

>

+H

V

8

Ur o/

v

Z

Ink

Yw

5

u2

j

dy

0

M

O

9

g

z

7

R

K

q

D

<latexi sh

1_b

64="yf0vNXnz+

MV

J

g

>A

B

3

c

DLS

F

2

W

du

o

q

T

K9

Pp

k

ZG

E/

O

H

r

R

m

C

Q

5

Y

8

j

<latexi sh

1_b

64="N7J

ZOULQTVw

XDWq

Ykf

0c

>A

B

H

S8

F

yp

+R

C

5

I

rm

3

K

/

E

g

2

j

M

d

P

o

9

z

n

u

v

G

<latexi sh

1_b

64="Py

L

z2

SG

B

q

Y

Ow

>A

/H

c

VD

N

FJ3

r

U

7d

En

g

m

+

0

Z

C

uX

j

9

5

8

Q

M

o

v

R

W

K

I

k

T

p

f

<latexi sh

1_b

64="N7J

ZOULQTVw

XDWq

Ykf

0c

>A

B

H

S8

F

yp

+R

C

5

I

rm

3

K

/

E

g

2

j

M

d

d

<latexi sh

1_b

64="N7J

ZOULQTVw

XDWq

Ykf

0c

>A

B

H

S8

F

yp

+R

C

5

I

rm

3

K

/

E

g

2

j

M

P

o

9

z

n

u

v

G

P

o

9

z

n

u

v

G

<latexi sh

1_b

64="w

u

fOHL2IGk

3

VJ

Er0

D9U

>A

C

c

S

N

F

vq

X

g

B

K

d

nT

o

Q

P

Z

p

z

W

j

7

Y

8

+

M

y

R

5

/

m

<latexi sh

1_b

64="Z

pIG

L

H

w

OWz

M

7mV

DQ

>A

B8n

c

F

TX

+R

2

C

k

u

Y

Uo

j

q9

5

r

3P

J

K

/

d

v

N

0

y

E

f

<latexi sh

1_b

64="KBTCJmZ

Uv7OIX3RM

>A

n

c

V

NS8

E

r

/

9L

0

k

f

d

Y

Q

+

w

P

G

H

p

y

g5

o

D

j

<latexi sh

1_b

64="vc5XK0f

T

kyn2NRg

Br

I

JE

>A

+H

V

S8

3U

o/G

Z

W

P

Y

L

m7

Q

M

C

j

Fu

z

d

w

q

O

9

D

p

<latexi sh

1_b

64="fq7RoTQmdFP0Y

2Vu

EX

I

>A

G

c

ZDLSgM

z9

r

Kk

W8

p

y

U

j

O

H

vB

n

/

5

N

w

J

+

<latexi sh

1_b

64="g

0qp5u yDOoYH+Bfd

C

>A

/3

c

V

LS

N

FJ

U

2v

GzW

R

KEk

m

Z

n

I

8

T

7M

w

Q

P

9

r

j

X

<latexi sh

1_b

64="g

S

Yo73My5

RNj

Uc

8+r

E

>A

B

H

V

W

9L

0m

q

CF

f

Q

Zdu

k

J

w

K

P

n

G

z

/

p

2

O

I

T

D

v

X

<latexi sh

1_b

64="z

3+

m

HARQfWL

TduJ

C/P2E

>

B

c

V

NS8

o

9

U0

K

Y

F

X

v

D

g

7

M

p

nZ

y

Oj

G

I

k

w

5

r

q

4.2 Training strategy and loss function

GRAPHSPLINENETS predicts the fine resolution states in a hierarchical manner. MPNN uses the sample points at the initialized state as input to propagate rollouts on these points as shown in Fig. 4.2. Based on MPNN outputs, two OSC problems are solved to find the polynomial parameters ϕ = OSC(MPNN( X , Y t 0 ; θ )) which allow for spatiotemporal continuous outputs ˆ u ( x i , t k ; ϕ ) . Note that solving OSC problems is equivalent to predicting the future space and time continuous states, given the sample inputs.

Inherently, training the parameters, θ of the MPNN is cast as a bi-level optimization problem as θ ∗ = argmin θ L ( ϕ, θ ) where ϕ = OSC(MPNN( X , Y t 0 ; θ )) . The loss function L is defined as:

The left term L s of Eq. (2) is the loss of the MPNN output ˆ y t k i , where N is the number of spatial sample points, and K is the number of MPNN rollout steps.

The left term L s solely depends on the MPNN parameters θ . The right term L i is the prediction loss of the high-resolution interpolation points ( x i , y t k i ) N i ,N t i =0 ,k =0 along time and space, where N i is the number of spatial interpolation points at the one-time frame, and N t is the number of interpolation time frames. Note that the space and time continuous prediction ˆ u ( x i , t k ; ϕ ) is constructed based upon the optimized parameters ϕ = OSC(MPNN( X , Y t 0 ; θ )) , the parameters that are obtained by solving OSC problems with MPNN predictions, MPNN( X , Y t 0 ; θ ) . Thus, both loss terms can be optimized with respect to

· · ·

t

K

· · ·

t

t

0

Sample Points

Interpolation Points

the MPNN parameters θ , and the whole model is trained end-to-end with automatic differentiation through the OSC. This optimization scheme is akin to Bialas and Karwan (1984); Tuy et al. (1993) that solve a bi-level optimization by flattening the loss for the upper- and lower-level problems.

4.3 Adaptive collocation sampling

To allow for prioritized sampling of important location regions, we dynamically optimize the positions of collocation points via their space gradients of the partial derivative values over time.

Static

Adaptive

More specifically, after getting the continuous simulation function ˆ u , the partial derivative equation of time provides the changes of the domain with time. Then, we calculate the maximum gradient vector for each collocation point to move them to more dynamic regions, i.e., with the largest change over time, which causes higher errors in the static collocation points condition. We formulate this process by

where β is the coefficient of the maximum gradient vector. We project collocation points back to the partition if the gradient step moves them outside to ensure a sufficient number of collocation points in each partition cell. We illustrate the adaptive collocation points in Fig. 4.3. We use the states at optimized positions adapted from history rollouts as the subsequent rollouts input. By adjusting the collocation points position, our model can focus on more dynamic parts of the space to get a more accurate prediction.

<latexi sh

1_b

64="w

X

LJj

Vn

5TFC

q

+E

dm0yQ

>A

B

c

NS8

3Ur

/o

9

GP

K

g

D

7

vH

2

u

M

f

W

z

k

R

I

Z

p

O

Y

<latexi sh

1_b

64="/R3N

LoD

g

B

O

Z

YQSqJPw

>A

7X

c

V

8

E

Ur

9

0

G

K

T

2k vH

5

+

u

M

z

nf

m

W

d

I

FC

y

p

j

<latexi sh

1_b

64="

OA

DznUIdWXu

g

y

rTL

Mv8

>

B

H

c

V

NS

EJ3

q/

9

0mk

CF

2

oQ

+

G

Z

j

Y

w

P

f

p

7

5

KR

<latexi sh

1_b

64="/R3N

LoD

g

B

O

Z

YQSqJPw

>A

7X

c

V

8

E

Ur

9

0

G

K

T

2k vH

5

+

u

M

z

nf

m

W

d

I

FC

y

p

j

<latexi sh

1_b

64="DPvu7NZO

9cC

dA

2S

Fw

>

B

H

V

8

EJ3Ur q/

L

0

G

+

gmk

Y

Q

n

K

f

z

p

T

X

W

M

y

R

o

5

Ij

5 Experiments

5.1 Datasets and Training Settings

We evaluate GRAPHSPLINENETS on five dynamical systems of increasing challenge. Simulated datasets consist of three PDEs (partial differential equations), including Heat, Damped Wave, and Navier-Stokes Equations. The two empirical datasets include the Ocean and Black Sea datasets; the Black Sea introduces significant challenges in scalability in the number of nodes, complex dynamics, irregular boundary, and non-uniform meshing. All the models employ the same structure of encoderprocessor-decoder for fair comparisons and the same amount of training data in each testing domain. While inputs of baseline models are directly all of the available data points, inputs of our OSC-based models are only an initialized 12 × 12 collocation point unitary mesh at the initial state and fewer intime sample points. For the Black Sea dataset, we construct a mesh via Delaunay Triangulation with 5000 data points for baselines and 1000 for GRAPHSPLINENETS as done in Lienen and Günnemann (2022). More details on datasets are available in Appendix B. We value open reproducibility and make our code publicly available 5 .

5.2 Evaluation metrics and baselines

We evaluated our model by calculating the mean square error (MSE) of rollout prediction steps with respectively ground truth for t ∈ [0 , 5] seconds. We employ relevant baselines in the field of discrete-step graph models for dynamical system predictions. Graph convolution networks (GCNs) (Kipf and Welling, 2016) and GCN with a hybrid multilayer perceptron (MLP) model are employed as baselines in the ablation study. We also compare our approach with one widely used baseline that employs linear interpolation for physics simulations allowing for continuous predictions in space, i.e., GEN (Alet et al., 2019). A similar setup considering the inherent graph structure and employing MPNNs in mesh space is used by (Pfaff et al., 2021) (MeshGraphNet, MGN in our comparisons). We utilize MGN as the MPNN building block for our GRAPHSPLINENETS.

5.3 Quantitative Analysis

We consider a series of ablation models on four datasets to demonstrate the effectiveness of our model components in our approach in multiple aspects. Quantitative results of ablation study models are shown in Table 5.1 [right]. The ablated models are:

MGN: MeshGraphNet model with 3 message passing layers from Pfaff et al. (2021).

MGN+OSC(Post): model with 3 message passing layers and only post-processing with the OSC method, i.e., we firstly train a MGN model, then we use the OSC method to collocate the prediction as a final result without end-to-end training.

MGN+OSC: MGN with OSC-in-the-loop that allows for end-to-end training.

MGN+OSC+Adaptive: MGN+OSC model that additionally employs our adaptive collocation point sampling strategy.

Post processing vs end-to-end learning We show the effectiveness of end-to-end learning architecture by comparing the MGN+OSC(Post) and MGN+OSC models. Table 5.1 shows that MGN+OSC has a more accurate prediction than MGN+OSC(Post) by more than 8% percent across datasets. This can be explained by the fact that, since the OSC is applied end-to-end, the error between MGN prediction steps is backpropagated to the message passing layers. In contrast, the model has no way of considering such errors in the post-processing steps.

Adaptive collocation points We further show the effectiveness of the adaptive collocation strategy by comparing the MGN+OSC and MGN+OSC+Adaptive: Table 5.1 shows that MGN+OSC+Adaptive has a more accurate prediction than MGN+OSC, i.e., more than 5% improvement on long rollouts in the Black Sea dataset. Adaptive collocation points encourage those points to move to the most dynamic regions in the domain, which is not only able to place

10

10

-

10

10

10

-

-

-

0

1

2

3

4

10

-

1

MGN

±

±

±

±

±

±

±

greater attention on hard-to-learn parts in space but also can let the OSC method implicitly develop a better representation of the domain. Empirical quantitative results on the five datasets are shown in Table 5.1. In the heat equation dataset, our approach reduces long-range prediction errors by 64% with only 20% of the running time compared with the best baseline model. In the damped wave dataset, our approach reduces errors by 42% with a 48% reduction in inference time.

0

2

4

6

Running Time

MGN

MGN+OSC

(Post, CPU)

MGN+OSC

(Post, GPU)

MGN+OSC

(GPU)

MGN+OSC

(GPU)+Adaptive

Neural network inference time

OSC inference time

Adaptive collocation time

In the Navier-Stokes dataset, our method reduces 31% long-range prediction errors while requiring 37% less time to infer solutions compared to the strongest baseline.

Inference times We show the effectiveness of the COLROW algorithm in accelerating the OSC speed by comparing the OSC method with one of the most commonly used algorithms for efficiently solving linear systems 6 and the OSC with the COLROW solver in Fig. 3.2. Our differentiable OSC and adaptive collocation only result in a minor overhead over the graph neural network. The overall scheme is still a fraction of the total in terms of inference time as shown in Fig. 5.4, while making GRAPHSPLINENETS more accurate.

5.4 Sensitivity Analysis

MSE 10 -1 10 -4 10 -3 10 -2 10 -1 10 0 MSE MGN 10 -1 10 0 Prediction Running Time [s] 10 -4 10 -3 10 -2 10 -1 10 0 MSE (a) (8 × 8) (16 × 16) (24 × 10 -2 10 -1 (b) MGN MGN + Linear Number of Collocation Points MGN + Cubic MGN  + B Spline 10 -1 10 -4 10 -3 10 -2 10 -1 10 0 MSE MGN 10 -1 10 0 Prediction Running Time [s] 10 -4 10 -3 10 -2 10 -1 10 0 MSE (a) MGN MGN + Linear 10 0 Prediction Running Time [s] (a) (8 × 8) (16 × 16) (24 × 24) 10 -2 10 -1 (b) 20 0 1 2 3 4 × 10 -3 MGN + Linear Number of Collocation Points MGN + Cubic MGN  + B Spline MGN + OSC Interpolation and collocation methods We demonstrate the efficiency of the OSC method by comparing the combination of MGN with different interpolation and collocation methods, including linear interpolation, cubic interpolation, and B-spline collocation methods. These models are implemented in the end-to-end training loop, and we use these methods in both the time and space dimensions. Results are shown in Fig. 5.2 where we measured the mean square error and running time of 3 second rollouts predictions by varying the number of collocation points from (2 × 2) to (16 × 16) . The model MGN+OSC shows the best accuracy while having a shorter running time. Even though the linear interpolation can be slightly faster than the OSC, it shows a considerable error in the prediction and does not satisfy basic assumptions such as Lipschitz continuity in space.

10

-

1

10

0

Prediction Running Time [s]

10

-

4

10

-

3

10

-

2

10

-

1

10

0

MGN

MGN + Linear

10

0

Prediction Running Time [s]

10

-

2

10

-

1

MGN + Linear

MGN + Cubic

24)

20

40

60

Rollout Steps

0

1

2

3

4

×

10

-

3

(c)

MGN + OSC

10

0

Prediction Running Time [s]

(8

×

10

-

2

10

-

1

MGN + Linear

MGN + Cubic

(8

×

8)

(16

×

16)

(24

10

-

2

10

-

1

(b)

Number of Collocation Points

MGN + Cubic

MGN  + B Spline

40

60

Rollout Steps

(c)

MGN + OSC + Adaptive

(b)

(b)

×

10

4

3

2

1

(8

(8

8)

(16

16)

8)

(16

16)

(24

8)

(16

16)

(24

24)

Number of Collocation Points

×

(c)

×

×

Number of Collocation Po

×

×

×

×

Number of Collocation Points

0

24)

24)

20

4

MGN + Cubic

×

MGN + OSC + Adaptive

MGN + OSC

MGN  + B S

MGN  + B Spline

MGN  + B Spline

Rollout

-

3

(a)

(a)

(a)

10

10

-

1

-

2

(b)

1

N

e [s]

×

10

(b)

10

-

1

-

2

×

(8

8)

(16

16)

(a)

(24

(8

(16

16)

8)

(24

mber of Collocation Points

Number of Collocation Points

MGN + Cubic

×

MGN  + B Spline

10

0

Prediction Running Time [s]

MGN + Linear

×

0

MGN  + B Spline

-

-

10

10

4

3

2

10

0

4

(c)

10

-

1

3

-

1

2

10

2

(b)

4

3

Number of collocation points We study the effect of the number of collocation points on the 3 second rollout prediction error by testing the MGN, MGN+OSC, and MGN+OSC+Adaptive models. The MGN is directly trained with the whole domain data. 1 2 1 2 10 -4 10 -3 0 1

10

0

-

1

(b)

(a)

(8

×

8)

(16

×

16)

10

-

2

10

-

1

Number of Collocation Points

MGN + OSC

20

40

Rollout Steps

MGN + OSC

MGN + OSC + Adaptive

MGN

10

-

1

10

0

Prediction Running Time [s]

10

-

4

10

-

3

10

-

2

10

-

1

10

0

MSE

(24 × 24) 20 40 60 Rollout Steps 0 1 2 3 4 × 10 -3 (c) MGN + OSC MGN + OSC + Adaptive 20 40 60 Rollout Steps MGN + OSC + Adaptive 60 10 0 Prediction Running Time [s] (8 × 8) (16 × 16) (24 × 24) 20 40 Rollout Steps MGN + Linear Number of Collocation Points MGN + Cubic MGN  + B Spline MGN + OSC (8 × 8) (16 × 16) (24 × 24) 10 -2 10 -1 (b) 20 40 60 Rollout Steps 0 1 2 3 4 × 10 -3 (c) Number of Collocation Points MGN + Cubic MGN  + B Spline MGN + OSC MGN + OSC + Adaptive We use a different number of collocation points (from (2 × 2) to (28 × 28) ) in the MPNN process in the rest two models and then compare the output of the OSC with the whole domain to train. With the increase in the number of collocation points, Fig. 5.3 shows that the MGN+OSC and MGN+OSC+Adaptive achieve significant improvements in prediction accuracy over the MGN. The MGN+OSC+Adaptive has a stable better performance than the MGN+OSC, and the improvement is larger when there are fewer collocation points. The reason is that with fewer collocation points, the MGN+OSC has insufficient ability to learn the whole domain. With the adaptive collocation point, the MGN+OSC+Adaptive can focus on hard-tolearning regions during training to obtain overall better predictions.

Number of rollout steps We show the effectiveness of the OSC method in improving long-range prediction accuracy by comparing the MGN and MGN+OSC model. Fig. 5.4 shows the MGN+OSC can keep stable in long-range rollouts compare with the MGN.

10 -3 10 -2 10 -1 10 0 (a) 10 -2 10 -1 (b) 10 -3 10 -2 10 -1 10 0 (a) 10 -2 10 -1 (b) 10 -2 10 -1 10 0 MSE (a) 10 -2 10 -1 10 -2 10 -1 10 0 MSE (a) 10 -2 10 -1 (b) The reason is that with the OSC, we can use fewer neural network rollout steps to obtain longer-range predictions, which avoids error accumulation during the multi-step rollouts and implicitly learns for compensating integration residual errors. In addition, end-to-end learning lets the neural networks in MGN+OSC learn the states between rollout steps, making the prediction stable and accurate.

10

10

-

-

3

4

3

4

-

-

10 -1 10 0 Prediction Running Time [s] (8 × 8) (16 × 16) (24 × 24) MGN MGN + Linear Number of Collocation Points MGN + Cubic MGN  + B Spline 10 -1 10 0 Prediction Running Time [s] (8 × 8) (16 × 16) (24 × 24) MGN MGN + Linear Number of Collocation Points MGN + Cubic MGN  + B Spline 10 -1 10 0 Prediction Running Time [s] (8 × 8) (16 × 16) (24 MGN MGN + Linear Number of Collocation Points MGN + Cubic MGN  + B Spline 10 -1 10 0 Prediction Running Time [s] (8 × 8) (16 × 16) (24 × 24) MGN MGN + Linear Number of Collocation Points MGN + Cubic MGN  + B Spline 10 Number of Message Passing Layers We further showcase the sensitivity to the number of message-passing layers while adopting the same baseline processor for a fair comparison. As Fig. 5.5 shows, a higher number of message-passing layers generally helps in decreasing the MSE at the cost of higher runtime, reflecting find-

20

40

60

Rollout Steps

0

1

2

3

4

×

10

-

3

MGN + OSC

MGN + OSC + Adaptive

20

40

60

Rollout Steps

0

1

2

3

4

×

10

-

3

MGN + OSC

MGN + OSC + Adaptive

×

24)

20

40

60

Rollout Steps

0

1

2

3

4

×

10

-

3

(c)

MGN + OSC

MGN + OSC + Adaptive

20

40

60

Rollout Steps

0

1

2

3

4

×

10

-

3

(c)

MGN + OSC

MGN + OSC + Adaptive

10

-

1

10

0

Prediction Running Time [s]

10

-

4

10

-

3

10

-

2

10

-

1

10

0

MSE

(a)

10

-

2

10

-

1

MGN

MGN + Linear

MGN + Cubic

-

1

10

0

Prediction Running Time [s]

10

-

4

10

-

3

10

-

2

10

-

1

10

0

MSE

MGN

MGN + Linear

4

3

2

1

(16

-

10

2

(8

8)

16)

(2

0

(8

8)

(16

16)

(24

Number of Collocation Points

×

×

60

MGN + OSC + Adaptive

10

-

1

(b)

×

Number of Collocation Point

MGN  + B Spline

MGN + Cubic

ings from prior work such as GNS (Alet et al., 2019) and MGN (Pfaff et al., 2021). Notably, GRAPHSPLINENETS exhibits a Pareto-optimal trade-off curve, outperforming MGN in terms of prediction accuracy and computational efficiency at different numbers of message-passing layers. Downsampling mesh nodes can help in robustness and generalization given the larger spatial extent of communication akin to Fortunato et al. (2022); the OSC plays a crucial part in handling the problem's continuity (e.g., in time) while adaptive collocation effectively helps the model to focus on regions with higher errors thus increasing the model performance.

1

2

3

MSELoss

×

10

-

3

5

10

15

20

25

Run Time

Heat Equation

3 Layers, MGN

3 Layers, Our

15 Layers, MGN

15 Layers, Our

2

3

MSELoss

×

10

-

1

1

2

3

4

5

Run Time

Navier-Stokes Equation

3 Layers, MGN

3 Layers, Our

15 Layers, MGN

15 Layers, Our

MGN. MGN obtain an MSE of (4 . 52 ± 0 . 23) × 10 -1 with a runtime of more than 80 seconds, while GRAPHSPLINENETS remarkably outperforms MGN with an MSE of (3 . 98 ± 0 . 27) × 10 -1 and less than 9 seconds to generate a trajectory - an example of which is shown in Fig. 5.6.

Scalability to Higher Dimensions We showcase GRAPHSPLINENETS on a larger 3D dataset, namely the 3D Navier-Stokes equations dataset from PDEBench (Takamoto et al., 2022) (more details in Appendix B.7). The number of nodes in GraphSplineNets is 4096 , which is more than 4 × compared to the Black Sea Dataset. Moreover, the 3D NavierStokes includes a total of 5 variables: density, pressure, and velocities on 3 axes. For a fair comparison, we keep the model structure the same as in other experiments and compare GRAPHSPLINENETS against

×

×

MGN + OSC

24)

MGN  + B Splin

10

10

-

-

(c)

(c)

(a)

24)

×

MSE

MSE

10

10

4

4

×

10

-

3

×

24)

MSE

×

10

-

3

(a)

(c)

(b)

(b)

×

10

-

3

(c)

(b)

×

10

-

3

<latexi sh

1_b

64="Y+

ON

K

57P3Spj

GJB

D0

A

>

n

c

V

L

g

E

y

r

fU

9

o

F

I

H

u

wmQ2Z

R8

z

W

d

v

k

C

/

M

X

q

T

<latexi sh

1_b

64="0

M

nk

R

V

ADT8

E

3

Kj

>

B

c

NS

J

W

X

O

o

U9

I

Y

5

G

pmu2

Z

dy

/g

F

w

7

Hf

L

r

Q

C

+

z

q

<latexi sh

1_b

64="0

T

7Vrky

QunS

vOB

8

>A

c

DL

gN

E

fUY9

oP

K

F

J

I5

H

z

Z

j

W

R

3

d

+

C

w

m

/

G

p

2

X

M

q

<latexi sh

1_b

64="z

9

S

gWNPQkK

8X

Awo r0q

>

B

n

c

VDL

EOy fUY

F

J

I5

H

u

v

Z

j

7

R

3

d

+

C

m

/

G

p

2

M

T

Initial States

Ground Truth

Simulation

Ground Truth

Simulation

t

= 4

t

= 2

Ground Truth

Simulation

t

= 6

Ground Truth

Simulation

t

= 8

GT

Simulation

Error

GT

Simulation

Error

GT

Simulation

Error

GT

Simulation

Error

GT

5.5 Qualitative analysis

We visualize the prediction of our model with baselines on the Black Sea dataset Fig. 5.7. Our model can accurately predict even with such complex dynamics from real-world data that result in turbulence phenomena while considerably cutting down the computational costs as shown in Table 5.1. We also provide series of qualitative visualizations with different datasets and baselines in Appendix C. Our model has a smoother error distribution and more stable long-range prediction. Thanks to the continuous predictions from GRAPHSPLINENETS, we can simulate high resolutions without needing additional expensive model inference routines, while the other two models can only achieve lower-resolution predictions. Baselines visibly accumulate errors for long-range forecasts; our model can lower the error with smoother and more accurate predictions in space and time.

Ground Truth

GCN

GCN+MLP

GEN

MGN

MGN+OSC(Post)

MGN+OSC

MGN+OSC(Adaptive)

0

.

0

0

.

2

0

.

4

MSE

6 Conclusion and Limitations

In this work, we proposed GRAPHSPLINENETS, a novel approach that employs graph neural networks alongside the orthogonal spline collocation method to enable fast and accurate forecasting of continuous physical processes. We used GNNs to obtain predictions based on coarse spatial and temporal grids, and the OSC method to produce predictions at any location in space and time. Our approach demonstrated the ability to generate high-resolution predictions faster than previous methods, even without explicit prior knowledge of the underlying differential equations. Additionally, we proposed an adaptive collocation sampling strategy that improves the prediction accuracy as the system evolves. We demonstrate how GRAPHSPLINENETS are robust in predicting complex physics in both simulated and real-world data. We believe this work represents a step forward in a new direction in the research area at the intersection of deep learning and dynamical systems that aims at finding fast and accurate learned surrogate models.

Alimitation of our approach is that, by utilizing the orthogonal spline collocation (OSC) method for interpolation, we assume some degree of smoothness in the underlying physical processes between collocation points; as such, our approach may struggle when modeling discontinuity points, which would be smoothed-out. Finally, we do note that despite benefits, deep learning models including GRAPHSPLINENETS often lack interpretability. Understanding how the model arrives at its predictions or explaining the underlying physics in a transparent manner can be challenging, hindering the model's adoption in domains where interpretability is crucial.

References

Alet, A. K. Jeewajee, M. B. Villalonga, A. Rodriguez, T. Lozano-Perez, and L. Kaelbling. Graph element networks: adaptive, structured computation and memory. In International Conference on Machine Learning , pages 212-222. PMLR, 2019.

Amodio, J. Cash, G. Roussos, R. Wright, G. Fairweather, I. Gladwell, G. Kraut, and M. Paprzycki. Almost block diagonal linear systems: sequential and parallel solution techniques, and applications. Numerical linear algebra with applications , 7(5):275-317, 2000.

Bauer, A. Thorpe, and G. Brunet. The quiet revolution of numerical weather prediction. Nature , 525(7567):47-55, 2015.

Berto, S. Massaroli, M. Poli, and J. Park. Neural solvers for fast and accurate numerical optimal control. International Conferences on Learning Representations , 2022.

Berto, C. Hua, J. Park, M. Kim, H. Kim, J. Son, H. Kim, J. Kim, and J. Park. Rl4co: an extensive reinforcement learning for combinatorial optimization benchmark. arXiv preprint arXiv:2306.17100 , 2023.

Bialas and M. Karwan. Two-level linear programming. Management Science , 30(8), 1984.

Bialecki. Convergence analysis of orthogonal spline collocation for elliptic boundary value problems. SIAM journal on numerical analysis , 35(2):617-631, 1998.

Bialecki and G. Fairweather. Orthogonal spline collocation methods for partial differential equations. Journal of Computational and Applied Mathematics , 128(1):55-82, 2001. ISSN 0377-0427. doi: https://doi.org/10.1016/S0377-0427(00)00509-4. URL https:// www.sciencedirect.com/science/article/pii/S0377042700005094 . Numerical Analysis 2000. Vol. VII: Partial Differential Equations.

Bourke. Interpolation methods. Miscellaneous: projection, modelling, rendering , 1(10), 1999.

Boussif, D. Assouline, L. Benabbou, and Y. Bengio. Magnet: Mesh agnostic neural pde solver. Advances in Neural Information Processing Systems , 2022.

Boussif, G. Boukachab, D. Assouline, S. Massaroli, T. Yuan, L. Benabbou, and Y. Bengio. What if we enrich day-ahead solar irradiance time series forecasting with spatio-temporal context? Advances in Neural Information Processing Systems , 2023.

R. Brink, D. A. Najera-Flores, and C. Martinez. The neural network collocation method for solving partial differential equations. Neural Computing and Applications , 33(11):5591-5608, 2021.

M. Bronstein, J. Bruna, T. Cohen, and P. Veliˇ ckovi´ c. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478 , 2021.

A. Ciliberti, M. Grégoire, J. Staneva, A. Palazov, G. Coppini, R. Lecci, E. Peneva, M. Matreata, V. Marinova, S. Masina, et al. Monitoring and forecasting the ocean state and biogeochemical processes in the black sea: recent developments in the copernicus marine service. Journal of Marine Science and Engineering , 9(10):1146, 2021.

De Boor and C. De Boor. A practical guide to splines , volume 27. springer-verlag New York, 1978.

Diaz, G. Fairweather, and P. Keast. Fortran packages for solving certain almost block diagonal linear systems by modified alternate row and column elimination. ACM Transactions on Mathematical Software (TOMS) , 9(3):358-375, 1983.

Fairweather and D. Meade. A survey of spline collocation methods for the numerical solution of differential equations. In Mathematics for large scale computing , pages 297-341. CRC Press, 2020.

Falcon et al. Pytorch lightning. GitHub. Note: https://github. com/PyTorchLightning/pytorchlightning , 3(6), 2019.

Fortunato, T. Pfaff, P. Wirnsberger, A. Pritzel, and P. Battaglia. Multiscale meshgraphnets. arXiv preprint arXiv:2210.00612 , 2022.

H. Golub and C. F. Van Loan. Matrix computations . JHU press, 2013.

Guo, X. Zhuang, and T. Rabczuk. A deep collocation method for the bending analysis of kirchhoff plate. Computers, Materials & Continua , 59(2):433-456, 2019. ISSN 1546-2226. doi: 10.32604/ cmc.2019.06660. URL http://dx.doi.org/10.32604/cmc.2019.06660 .

Houska, F. Logist, M. Diehl, and J. Van Impe. A tutorial on numerical methods for state and parameter estimation in nonlinear dynamic systems. Identification for Automotive Systems , pages 67-88, 2012.

Keisler. Forecasting global weather with graph neural networks. arXiv preprint arXiv:2202.07575 , 2022.

P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907 , 2016.

Kirchmeyer, Y. Yin, J. Donà, N. Baskiotis, A. Rakotomamonjy, and P. Gallinari. Generalizing to new physical systems via context-informed dynamics model. arXiv preprint arXiv:2202.01889 , 2022.

Kremer and B. Hancock. Process simulation in the pharmaceutical industry: a review of some basic physical models. Journal of pharmaceutical sciences , 95(3):517-529, 2006.

J. Lee and C. A. Micchelli. On collocation matrices for interpolation and approximation. Journal of Approximation Theory , 174:148-181, 2013.

Li, C. Hua, J. Park, H. Ma, V. Dax, and M. J. Kochenderfer. Evolvehypergraph: Group-aware dynamic relational reasoning for trajectory prediction. arXiv preprint arXiv:2208.05470 , 2022a.

Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar. Fourier neural operator for parametric partial differential equations. International Conference on Learning Representations , 2021.

Li, D. Z. Huang, B. Liu, and A. Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint arXiv:2207.05209 , 2022b.

Lienen and S. Günnemann. Learning the dynamics of physical systems from sparse observations with finite element networks. International Conference on Learning Representation , 2022.

Lin, J. Kim, C. Hua, S. Kang, and M.-H. Park. Comparing artificial and deep neural network models for prediction of coagulant amount and settled water turbidity: Lessons learned from big data in water treatment operations. Journal of Water Process Engineering , 54:103949, 2023a.

Lin, J. Kim, C. Hua, M.-H. Park, and S. Kang. Coagulant dosage determination using deep learning-based graph attention multivariate time series forecasting model. Water Research , 232: 119665, 2023b.

Logg, K.-A. Mardal, and G. Wells. Automated solution of differential equations by the finite element method: The FEniCS book , volume 84. Springer Science & Business Media, 2012.

Long, Y. Lu, X. Ma, and B. Dong. Pde-net: Learning pdes from data. In International Conference on Machine Learning , pages 3208-3216. PMLR, 2018.

Marriott, S. M. Tan, and N. Marriott. Experiential learning-a case study of the use of computerised stock market trading simulation in finance education. Accounting Education , 24(6):480-497, 2015.

Marullo, R. Santoleri, D. Ciani, P. Le Borgne, S. Péré, N. Pinardi, M. Tonani, and G. Nardone. Combining model and geostationary satellite data to reconstruct hourly sst field over the mediterranean sea. Remote sensing of environment , 146:11-23, 2014.

Mücke, R. Engel, J. P. Rachen, R. J. Protheroe, and T. Stanev. Monte carlo simulations of photohadronic processes in astrophysics. Computer Physics Communications , 124(2-3):290-314, 2000.

B. Nardelli, C. Tronconi, A. Pisano, and R. Santoleri. High and ultra-high resolution processing of satellite sea surface temperature data over southern european seas in the framework of myocean project. Remote Sensing of Environment , 129:1-16, 2013.

L. Oberkampf. Simulation accuracy, uncertainty, and predictive capability: A physical sciences perspective. Computer Simulation Validation: Fundamental Concepts, Methodological Frameworks, and Philosophical Perspectives , pages 69-97, 2019.

Park, F. Berto, A. Jamgochian, M. J. Kochenderfer, and J. Park. Meta-sysid: A meta-learning approach for simultaneous identification and prediction. arXiv preprint arXiv:2206.00694 , 2022.

Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

Pathak, S. Subramanian, P. Harrington, S. Raja, A. Chattopadhyay, M. Mardani, T. Kurth, D. Hall, Z. Li, K. Azizzadenesheli, et al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214 , 2022.

Pfaff, M. Fortunato, A. Sanchez-Gonzalez, and P. W. Battaglia. Learning mesh-based simulation with graph networks. International Conference on Learning Representations , 2021.

Poli, S. Massaroli, J. Park, A. Yamashita, H. Asama, and J. Park. Graph neural ordinary differential equations. arXiv preprint arXiv:1911.07532 , 2019.

Poli, S. Massaroli, F. Berto, J. Park, T. Dao, C. Re, and S. Ermon. Transform once: Efficient operator learning in frequency domain. 2022.

Quarteroni and A. Veneziani. Analysis of a geometrical multiscale model based on the coupling of ode and pde for blood flow simulations. Multiscale Modeling & Simulation , 1(2):173-195, 2003.

Sanchez-Gonzalez, J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning , pages 8459-8468. PMLR, 2020.

Strassen et al. Gaussian elimination is not optimal. Numerische mathematik , 13(4):354-356, 1969.

Takamoto, T. Praditia, R. Leiteritz, D. MacKinlay, F. Alesiani, D. Pflüger, and M. Niepert. Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems , 35:1596-1611, 2022.

Tuy, A. Migdalas, and P. Värbrand. A global optimization approach for the linear two-level program. Journal of Global Optimization , 1(23), 1993.

Wang, D. Zheng, Z. Ye, Q. Gan, M. Li, X. Song, J. Zhou, C. Ma, L. Yu, Y. Gai, T. Xiao, T. He, G. Karypis, J. Li, and Z. Zhang. Deep graph library: A graph-centric, highly-performant package for graph neural networks, 2020.

Y.-T. Wu and F. Porté-Agel. Large-eddy simulation of wind-turbine wakes: evaluation of turbine parametrisations. Boundary-layer meteorology , 138(3):345-366, 2011.

L.-P. Xhonneux, M. Qu, and J. Tang. Continuous graph neural networks. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 10432-10441. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/xhonneux20a.html .

Yadan. Hydra - a framework for elegantly configuring complex applications. Github, 2019. URL https://github.com/facebookresearch/hydra .

Yin, V. Le Guen, J. Dona, E. de Bézenac, I. Ayed, N. Thome, and P. Gallinari. Augmenting physical models with deep networks for complex dynamics forecasting. Journal of Statistical Mechanics: Theory and Experiment , 2021(12):124012, 2021.

Learning Efficient Surrogate Dynamic Models with Graph Spline Networks

Supplementary Material

Table of Contents

A Additional OSC Material

We further illustrate the OSC method by providing numerical examples in this section.

A.1 1-D OSC Example

For simplicity and without loss of generality, we consider the function domain as unit domain [0 , 1] and we set N = 3 , r = 2 , which means we will use a three-order three-piece function to simulate the 1-D ODE problem. We firstly choose the partition points as x i , i = 0 , · · · , 3 , x 0 = 0 , x 3 = 1 . The number of partition points is N +1 = 4 . Distance between partition points can be fixed or not fixed. Then, based on the Gauss-Legendre quadrature rule, we choose the collocation points. The number of collocation point within one partition is r -1 = 1 , so we have in total N × ( r -1) = 3 collocation points ξ i , i = 0 , · · · , 3 .

After getting partition points and collocation points, we will construct the simulator. Here we have three partitions; in each partition, we assign a 2 order polynomial

Note that these three polynomials should be C 1 continuous at the connecting points, i.e., partition points within the domain. For example, Eq. (4a) and Eq. (4b) should be continuous at x 1 , then we

can get two equations

For boundary condition we can also get two equations

Let us summarize the equations we got so far. Firstly, our undefined polynomials have N × ( r +1) = 9 parameters. The C 1 continuous condition will create ( N -1) × 2 = 4 equations, and the boundary condition will create 2 equations. Then we have N × ( r -1) collocation points. For each collocation point, we substitute it to polynomials to get an equation. For example, if the ODE is

By substituting collocation point ξ 0 into the equation, we can get

ˆ

u

(

ξ

0

) + ˆ

u

′

(

ξ

0

) =

f

(

ξ

0

)

Now we can know that the number of equations can meet with the degree of freedom of polynomials

In this example, generated equations will be constructed to an algebra problem Aa = f where the weight matrix is an ABD matrix.

(A)

(B)

0

Not

0

By solving this problem, we can obtain the simulation results. Note that we can employ the COLROW sparse solver to solve the system efficiently (see § 3.4).

A.2 2-D OSC Example

For simplicity and without loss of generality, we consider the function domain as unit domain [0 , 1] × [0 , 1] , and set N x = N y = 2 , r = 3 . Partition points and collocation points selection is similar to the 1-D OSC method; we have N 2 × ( r -1) 2 = 16 collocation points in total. For simplicity, we note the partition points at two dimensions to be the same, i.e., x i , i = 0 , 1 , 2 . Unlike the 1-D OSC

method, we choose Hermite bases to describe as the simulator, which keeps C 1 continuous. As a case, the base function at point x 1 would be

Weseparately assign parameters to basis functions, i.e. H 1 ( x ) = a 1 ,i f 1 ( x )+ b 1 ,i g 1 ( x ) for x variable in [ x 0 , x 1 ] × [ y i -1 , y i ] partition. Then the polynomial in a partition is the multiple combinations of base functions of two dimensions. For example, the polynomial in the partition [ x 0 , x 1 ] × [ y 0 , y 1 ] is

Now we consider the freedom degree of these polynomials. From definition, we have 2 n ( r -1)( n + 1) = 24 parameter. Considering boundary conditions, we have 24 -4 × N = 16 parameters. The number is equal with collocation points N 2 × ( r -1) 2 , which means we can get an algebra equation by substituting collocation points. Solving this equation, we can get the simulator parameters.

We can similarly have multiple basis functions and set parameters to the simulation result for the higher dimension OSC method. And then, select partition points and collocation points using the same strategy as the 2-D OSC method. The rest algebra equation generating and solving equations parts will not be different.

A.3 Simple Numerical Example

We set N = 3 , r = 3 to simulate the problem

we can get a simulation solution as follows, which is visualized in Fig. A.2.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

Simulation

Ground Truth

Partition Points

Collocation Points

A.4 Numerical Analysis for interpolation and collocation methods

Wecompared the OSC with linear, bilinear, 0-D cubic, and 2-D cubic interpolation methods on four types of problems: 1-D linear, 1-D non-linear, 2-D linear, and 2-D non-linear problems. In these experiments, we tested different simulator orders of the OSC method. For example, we set the order of the simulator to 4 for 1-D linear problem and 2 for 2-D linear problem. When the order of the simulator matches the polynomial order of the real solution, OSC can directly find the real solution. For non-linear problems, increasing the order of the simulator would be an ideal way to get lower

loss. For example, we set the order of the simulator to 4 for 1-D non-linear problem and 5 for 2-D non-linear problem. Thanks to the efficient calculation of OSC, even though we use higher-order polynomials to simulate, we use less running time to get results.

B Additional Experimental Details

B.1 Models and implementation details

In this section we introduce additional experimental details.

Model Specific details of model components are introduced below:

All MLPs have ReLU : x → max(0 , x ) nonlinearities between layers.

Specific details of applying the OSC method are introduced as follows:

Wenote that the number of collocation points in one partition is the same for each spatial dimension; features also share collocation points (i.e., wind speeds and temperature measurements share the same partition).

Regarding the choice of collocation points, for regular boundary problems, we initialize collocation points using Gauss-Legendre quadrature, following the recommendation by Bialecki (1998) 7 . This choice is particularly apt for smooth and well-behaved functions over the integration interval, making it suitable for achieving accurate results without oscillatory behavior. In irregular boundary scenarios, we employ mesh points as collocation points. Importantly, the space-oriented OSC technique enables us to obtain values at any spatial position, rendering the entire space available for collocation points, which makes the adaptive collocation points possible.

Training A batch size of 32 was used for all experiments, and the models were trained for up to 5000 epochs with early stopping. We used the Adam optimizer (Kingma and Ba, 2014) with an initial learning rate of 0 . 001 and a step scheduler with a 0 . 85 decay rate every 500 epoch. For all datasets, we used the split of 5 : 1 : 1 for training, validating, and testing for a fair comparison. We set the parameter β for adaptive OSC as constant. In particular, we set β = 1 2 × min partitions width. This ensures that collocations will not shift by more than half of the partition range in one adaptation step, lending stability and robustness to the adaptive mesh.

Space and Time Resolutions We additionally add a more detailed overview of the experiments regarding data setup, especially for space and time resolution, such as total rollout time, space, and time resolution for both OSC space and time in Table B.3.

Note that the rollout time and the other times ∆ t refer to the original dataset prediction horizons. For instance, in the Black Sea dataset, the ground truth data is provided every day (1 d), while the prediction horizon consists of 10 steps (i.e., a total of 10 days). For datasets including more than a single scalar value (Ocean Currents and Black Sea have 3 features, while 3D Navier-Stokes have 5), we the OSC is applied feature-wise.

Example Rollout To further elucidate the rollout procedure with GRAPHSPLINENETS, we showcase a simple example. For the 2D Navier-Stokes experiment, the ground truth trajectories consist of 10-time steps (10 seconds), with a step size of ∆ t = 1 second. In our model, the neural network autoregressively performs the rollout 5 times with ∆ t = 2 seconds, and the OSC (Orthogonal Spline Collocation) method is used to interpolate through the skipped steps. In other words, the ground truth consists of

while the neural networks in our model provide the interpolated outputs as

which are then further refined using the Time-Oriented OSC to produce

B.2 Heat Equation

The heat equation describes the diffusive process of heat conveyance and can be defined by

where u denotes the solution to the equation, and ∆ is the Laplacian operator over the domain. In a n -dimensional space, it can be written as:

Dataset generation We employ FEniCS (Logg et al., 2012) to generate a mesh from the domain and solve the heat equation on these points. The graph neural network then uses the mesh for training.

B.3 Wave Equation

The damped wave equation can be defined by

where c is the wave speed and k is the damping coefficient. The state is X = ( w, ∂w ∂t ) .

Data generation We consider a spatial domain Ω represented as a 64 × 64 grid and discretize the Laplacian operator. ∆ is implemented using a 5 × 5 discrete Laplace operator in simulation; null Neumann boundary condition are imposed for generation. We set c = 330 and k = 50 similarly to the original implementation in Yin et al. (2021).

B.4 2D Incompressible Navier-Stokes

The Navier-Stokes equations describe the dynamics of incompressible flows with a 2-dimensional PDE. They can be described in vorticity form as:

where v is the velocity field and w is the vorticity, ν is the viscosity and f is a forcing term. The domain is subject to periodic boundary conditions.

Data generation Wegenerate trajectories with a temporal resolution of ∆ t = 1 and a time horizon of t = 10 . We use similar settings as in Yin et al. (2021) and Kirchmeyer et al. (2022): the space is discretized on a 64 × 64 grid and we set f ( x, y ) = 0 . 1(sin(2 π ( x + y )) + cos(2 π ( x + y ))) , where x, y are coordinates on the discretized domain. We use a viscosity value ν = 10 -3 .

B.5 Ocean Currents

The ocean currents dataset (Nardelli et al., 2013; Marullo et al., 2014) is composed of hourly data of ocean currents with an Earth grid horizontal resolution of 1 / 12 degrees with regular longitude/latitude equirectangular projection 8 . We employ data from 2019-01-01 to 2022-01-01 at a depth of 0 . 49 meters. We consider an area in the Pacific ocean characterized by turbulent currents corresponding to latitude and longitude (-126 ∼ -120, -36 ∼ -30). We use currents eastward and northward sea water velocities as variables, i.e., uo and vo , respectively as well as water temperature to .

B.6 Black Sea

The Black Sea dataset is composed of daily real-world measurements of ocean currents and temperatures (Ciliberti et al., 2021). The resolution of the raw data is 1 / 27 ◦ x 1 / 36 ◦ . We employ data starting from 01/01/2012 until 01/01/2020; we split training, validation and testing with ratios of 5:1:1 as in the other datasets in the paper sequentially on the temporal axis; i.e., so that the model has not seen data from 2018 or 2019 during training. Of the data points in the grids, less than 50% actually cover the Black Sea due to the irregular shape of the sea. To obtain a mesh, we subsample the grid using Delaunay triangulation to 5000 nodes for baselines and 1000 nodes for GRAPHSPLINENETS, which results in a non-uniform mesh with an irregularly-shaped boundary akin to Lienen and Günnemann (2022). We can observe a detail of the ground truth, non-adaptive and adaptive meshing alongside GRAPHSPLINENETS errors in Figs. B.1 to B.3 respectively. We use currents eastward and northward sea water velocities as variables, i.e., uo and vo , respectively, as well as water temperature to at a single depth of 12.54m. We normalize features by the mean and standard deviation of the training samples. The dataset is available to download from the Copernicus Institute 9 .

B.7 3D Compressible Navier-Stokes Equations

These equations describe the dynamics of compressible fluid flows (Takamoto et al., 2022):

where ρ is the mass density, v is the velocity, p is the gas pressure, ϵ = p/ (Γ -1) is the internal energy, Γ = 5 / 3 , σ ′ is the viscous stress tensor, η the shear and ζ the bulk viscosities. We employ a downsampled version of the dataset from PDEBench of Takamoto et al. (2022), originally 128 × 128 × 128 and subsampled to 64 × 64 × 64 up to N t = 10 for the ground truth data, while the

Wind Speed Ground Truth

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

0

.

6

Without Adaptive Collocation Points

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

Mean Squared Error (max: 0.43)

With Adaptive Collocation Points

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

Mean Squared Error (max: 0.39)

collocation points lay on a dynamic grid of 32 × 32 × 32 = 32768 nodes. We use a total of 600 trajectories which we divide into training and testing with a 90% split as done in PDEBench.

B.8 Hardware and Software

Hardware Experiments were carried out on a machine equipped with an INTEL CORE I9 7900X CPU with 20 threads and a NVIDIA RTX A5000 graphic card with 24 GB of VRAM.

Software Software-wise, we used FEniCS (Logg et al., 2012) for Finite Element simulations for the heat equation experiments and PyTorch (Paszke et al., 2019) for simulating the damped wave and Navier-Stokes equations. We employed the Deep Graph Library (DGL) (Wang et al., 2020) for graph neural networks. We employ as a main template for our codebase the Lightning-HydraTemplate 10 which we believe is a solid starting point for reproducible deep learning even outside of supervised learning (Berto et al., 2023); the Lightning-Hydra-Template combines the PyTorch Lightning library (Falcon et al., 2019) for efficiency with modular Hydra configurations (Yadan, 2019).

C Additional Visualizations

Ground Truth

GCN

GCN+MLP

GEN

MGN

MGN+OSC(Post)

MGN+OSC

MGN+OSC(Adaptive)

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

-

2

-

1

0

1

2

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00  0

.

00

0

.

25

0

.

50

0

.

75

1

.

00  0

.

00

0

.

25

0

.

50

0

.

75

1

.

00  0

.

00

0

.

25

0

.

50

0

.

75

1

.

00  0

.

00

0

.

25

0

.

50

0

.

75

1

.

00  0

.

00

0

.

25

0

.

50

0

.

75

1

.

00  0

.

00

0

.

25

0

.

50

0

.

75

1

.

0

0

Figure C.1: Visualization of the Damped Wave dataset.

Ground Truth

GCN

GCN+MLP

GEN

MGN

MGN

+OSC(Post)

MGN

+OSC

MGN

+OSC(Adaptive)

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Figure C.2: Visualization of the Navier Stokes dataset.

Ground Truth

t=1

t=2

t=3

t=4

t=5

t=7

t=8

GCN

GEN

Our

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

O

O

O

O

O

O

O

0

.

0

Ground Truth

t=0

t=1

t=2

t=3

t=4

t=5

t=7

GCN

GEN

Our

Ground Truth

GCN

GEN

Our

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

GCN

t=0

t=1

t=2

t=3

t=4

t=5

t=7

GEN

Our

GCN

GEN

GCN

0.2

0.4

0.6

0.8

1.0

0.0

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper builds upon a recent advances in deep-learning based collocation methods for physical simulation to improve prediction accuracy in continuous space and time. The authors proposed a hybrid model based on graph neural network and differentiable orthogonal spline collocation method. Interpolation in time as well as in space is adopted to enable employing coarser temporal grids to speed up prediction. The paper also introduces a new adaptive collocation sampling strategy and techniques for training such as not needing to learn interpolators or collocation weights to decrease the computational overhead. Experiments in various PDE settings demonstrate that the proposed model, GRAPHSPLINENETS outperforms state-of-the-art baselines in both accuracy and runtime.\n\nSTRENGTHS:\nThe paper introduces some techniques, a novel loss function comprising OSC function, COLROW algorithm to efficiently solve linear equations associated with OSC methods, and adaptive collocation sampling to allow for prioritized sampling of important location regions. Wide variety of experiments are conducted, and the proposed method outperforms the baselines.\n\nWEAKNESSES:\nAlthough the authors claim that GraphSplineNets forecast continuous responses in time and space, comparison of GraphSplineNets are limited to space continuous baselines. Also, as in the main text, only a fewer points are sampled as collocation points in time — It is still difficult to assess the advantage of incorporating time-oriented OSC as well as space-oriented into the architecture. It is also unclear how is dataset for interpolation points error in the equation (2) generated. \n\nThe explanation on Figure 5.1 is too brief for me to understand. Does the size of input of MGN is larger than that of MGN + OSC (12x12)? In appearance it looks strange to me that the running time of MGN w/o OSC higher than MGN w/ OSC.\n\nQUESTIONS:\n* Is OSC inference time in Figure 5.1 associated with both spacial and time oriented collocation? \n* Does adaptive collocation sampling also perform for important time interval?\n\nLIMITATIONS:\nYes, the limitations are discussed in the main text.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduce GraphSplineNets, providing a framework for combining MPNNs like MGNs with interpolation/collocation methods. Their contributions also include two methods to surpass the performance of a naive combination: end-to-end training and collocation points adaptation. Temporal and spatial downsampling allow the MPNN to run faster, while the collocation method makes possible *increased* inference-time resolution despite the NN's downsampled domain.\n\nSTRENGTHS:\nSignificance:\n\n1. This is an important area, and the paper's core idea is interesting and well presented (as a proof of concept, at least). The paper offers a lot of infrastructure (code and ideas) important to future steps at the intersection of GNN surrogates and collocation.\n2. The inference speed benefit of GSNs is intriguing because applying GNNs at scale can be difficult.\n\nClarity:\n\n3. The motivation for this work is communicated well. This work's connection to other approaches that deal with GNN scalability is clearly established. \n4. The explanation of the GNN-based OSC method is very well done and nicely illustrated. \n\nQuality:\n\n5. OSC seems like a potentially promising approach for dealing with the efficiency challenges of GNN simulations.\n6. The end-to-end optimization approach (differentiating through the OSC) is well justified and intuitive. \n\nOriginality:\n\n7. Plugging GNNs into an OSC method, allowing temporal coarsening of MPNN (e.g. MGN) activity to increase efficiency, and optimizations of this joint GNN+OSC method (end-to-end training, adaptive points).\n8. Giving automatic differentiation support to an ABD matrix equation solving routine.\n\nWEAKNESSES:\nMy main concern is the quality of the evaluation of the method. I think this paper would gain significance if it showed that GSN's performance benefits hold in more challenging settings (by demonstrating why the benefits of GSN will generally be present, by evaluating on harder problems and against better tuned baselines, etc.). In other words, addressing some of the following weaknesses might help other researchers better understand the potential relevance of this method.\n\n1. The paper leaves unclear why its methods beat baseline methods' errors (i.e., is it because downsampled mesh nodes have access to information from farther away?), if the errors it achieves are practically relevant (i.e., sufficiently low for the model to be used), and if it beats stronger baselines (i.e., MGN with more message passing layers or MultiScale MGN). \n\n2. Moreover, no large-scale or 3D problems were considered, and this is at odds with the description of the proposed method as one that addresses efficiency challenges of existing approaches.\n\nMinor concerns: \n\n3. Figure 4.1 is nice, but I think at $t=0$, you should have a $y_3$ instead of a $y_0$.\n\nQUESTIONS:\n1. Why use only 3 MP layers in baseline MGNs? This gives the GraphSplineNets method an unfair advantage in that it can communicate over a larger spatial extent (because it is operating on a coarser mesh). Relatedly, the baseline MGN results (e.g., those in table 5.1) are rather poor compared to results shown in the MGN paper, which recommends using 15 MP layers.\n\n2. If you keep on increasing the collocation points, will GSNs eventually provide the same solution as MGNs? Assuming this is the case, Figure 5.3 suggests that more collocation points is better, but it might be nice to see when the relationship flips and the methods converge.\n\n3. How do GSNs perform on larger scale problems? In particular, can they compete with approaches designed to be more scalable like MultiScale MeshGraphNets (Fortunato et al., 2022)?\n\n4. Line 27: It's true that traditional solvers can struggle, but it might make sense to support this point by citing a paper that analyzes traditional solvers rather than a paper that analyzes ML solvers.\n\n5. In the Appendix, there seem to be some minor typos/errors in the equations. E.g., equations 7 and 11 seem to assume 2 partitions when the example has 3 partitions. This did not affect my score, but cleaning it up would be great!\n\nLIMITATIONS:\nThe authors addressed some key limitations. I might add that no large-scale or 3D problems were considered, which hinders understanding of the benefits of the proposed method at the scales where its contribution (efficiency) might matter most.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method based on deep learning that will speed up the prediction of physical systems by reducing the size of the grid and the number of tests to be carried out. The method consists of combining a graph neural network to predict some values with two orthogonal spline collocation, one spatial and one temporal, to predict values continuously.\n\nSTRENGTHS:\nThe paper is very clear and well written. The graphics are a great complement to the text and help to make it easier to understand. \nThe choice of graph neural networks seems judicious in the context of a physical system that admits links and dynamic interactions between variables. As for orthogonal spline collocation, the theoretical guarantees and computational complexity justify their use. The results show improved predictions and increased speed on various physical problems. The authors are careful to point out the limitation of their approach.\n\nWEAKNESSES:\nSee limitations and questions.\n\nQUESTIONS:\n1) Lines 176-177: I think there is a typo, 2-D domain is written twice\n2) Line 203 and equation (2): you write $L$ in the equation and then you write a rounded \"L\", I think you mix the notation\n3) To what extent is the availability of data for physical systems not too restrictive? \n4) Could we imagine replacing GNNs with physical neural networks?\n\nLIMITATIONS:\nThe impact of GNN on the overall method needs to be quantified a little more, the level of performance required by the neural network needs to be defined. This would make it easier to adapt the method, for example by changing the GNN model for another one.\nWe can make the same point with OSC, trying to combine GNN with an alternative to show what the use of OSC brings specifically, in terms of runtime and performance.\nThis would help justify the choice of these two methods over others.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe presented manuscript proposes a novel architecture to predict the evolution of the dynamical system. The proposed method has two key ingredients: orthogonal spline collocation for both time and spatial dimensions and adaptive selection of the collocation points. Extensive experiments demonstrate the performance of the proposed method over alternatives in both prediction quality and runtime.\n\nSTRENGTHS:\nThe presented method is based on the natural encoder-decoder architecture and intermediate processor via message-passing layers. The main strength of the paper is the extensive experimental evaluation and sensitivity analysis of the results to the adjusting hyperparameters. The authors consider 3 datasets generated from different types of PDEs and 2 datasets with empirical data from monitoring Ocean and black sea. Their method shows the best results in all considered tests. Moreover, the proposed approach is faster than the competitors based on graph neural networks.\n\nWEAKNESSES:\nThe main weakness of the paper is its structure and presentation of the proposed method. In particular,\n1) Too much attention is paid to the standard method for solving structured linear systems for computing interpolation coefficients. I suggest compressing such detailed descriptions since they are mostly textbook facts. At least a comparison with a general-purpose LU solver from PyTorch is strange since the result of such a comparison is obvious. Thus, section 3.4. can be eliminated or significantly compressed.\n2) section 3.2 introduce the notation for the graph, but later there are no details on how this graph is constructed from snapshots\n3) in the paragraph in lines 133-139 no information about how to select collocation points only the number of collocation points is estimated. So, it is unclear how to extract collocation points from the available samples\n4) Figure 4.1 is entirely confusing. I do not find where is input, output, or ingredients of eq (1) and how they should be processed together to get the prediction. Please, specify what operations are done and in what order starting from the beginning and revise section 4.1 respectively.\n5) also, please move to the main text architecture details of the encoder, decoder, and processor blocks and training setup from the supplementary materials  \n\nI will increase the score if the proper revision of the text will be made in the rebuttal phase.\n\nQUESTIONS:\n1) how to select the parameter $\\beta$ in eq (3) ? Is it constant or adaptive too?\n2) do you observe any relationship between the number of processor blocks and the resulting error?\n3) why do you not consider the Chebyshev grid for collocation points and the corresponding barycentric Lagrange interpolation approach? Cheb grid should be more accurate than your static approach (Fig. 4.3 left) but less computationally intensive than the adaptive approach.\n4) what item in the loss function $L_s$ or $L_i$ dominates and what are the possible explanations?\n\nLIMITATIONS:\nThe authors provide the limitations of the proposed approach.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-19"}