SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chuanbo Hua*
DiffeqML, AI4CO
Federico Berto*
Michael Poli
DiffeqML
Stefano Massaroli
Mila and University of Montreal
Jinkyoo Park
OMELET
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.
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; Obergkampf, 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
Figure 1.1: Wind speed forecasts on the Black Sea. GRAPHSPLINENETS can learn fast and accurate surrogate models for complex dynamics on irregular domains.
*Equal contribution authors.
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 Gunnemann (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 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.
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  $\mathcal{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  $\mathcal{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.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  $\mathbf{X} = \{\mathbf{x}_i,i = 1,\dots ,N\}$  and time  $t\in \mathbb{R}$  as  $\mathbf{Y}^t = \{y_i^t,i = 1,\dots ,N\}$  where  $N$  represents the number of sample points and  $\Omega \subset \mathbb{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_{i}^{t} = u(\mathbf{x}_{i},t)$ . The objective of a dynamics surrogate model is to estimate future states  $\{\hat{\mathbf{Y}}^t,t\in \mathbb{R}_+\}$  given the initial states  $\mathbf{Y}^0$ .
3.2 Message Passing Neural Networks
We employ a graph  $\mathbf{G}^t = \{\mathbf{v}_i^t, \mathbf{e}_{ij}^t\}$  to encode the status of physical process at sample points  $\mathbf{X}$  at  $t$ .  $\mathbf{v}_i^t$  and  $\mathbf{e}_{ij}^t$  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:
$$
\hat {\mathbf {Y}} ^ {t + 1} = \underbrace {\mathcal {D}} _ {\text {d e c o d e r}} \left(\underbrace {\mathcal {P} _ {m} (\cdots (\mathcal {P} _ {1} (\underbrace {\mathcal {E} (\mathbf {X} , \mathbf {Y} ^ {t})}))}\right)} _ {\text {p r o c e s s o r}} \tag {1}
$$
where  $\mathcal{E}(\cdot)$  is the encoder,  $\mathcal{P}_i(\cdot)$  is the  $i$ -th message passing layer, and  $\mathcal{D}(\cdot)$  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.
${}^{2}$  Collocation and interpolation are terms that are frequently used interchangeably. While interpolation is defined as obtaining unknown values from known ones, collocation is usually defined as a finite solution space satisfying equations dictated by known (collocation) points. Thus, collocation can be considered as a flexible subset of interpolation methods that satisfies certain conditions, such as  ${\mathcal{C}}^{1}$  class continuity.
Figure 3.1: Visualization of an 1-D OSC example with the order  $r = 3$  and partition number  $N = 3$ .
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<sup>34</sup>. Then, we initialize one  $r$ -order polynomial with  $r + 1$  unknown coefficients in each partition. These polynomials have  $N \times (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)\times 2$  equations from the  $C^1$  continuity restriction. Then we need  $N\times (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\times (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. We illustrate 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).
Figure 3.2: (a) The COLROW algorithm has the lowest computational complexity with  $O(n^{2})$  compared to the Gaussian elimination algorithm with  $O(n^{3})$  or lower bounds of the generic sparse solvers with  $O(n^{2}\mathrm{logn})$ . (b) Running time for 100 iterations of A: OSC+LU decomposition, B: OSC+COLROW solver, and C: OSC+COLROW solver with GPU acceleration.
Note that these partitions do not need to be isometric, thus allowing for flexibility in their choice.
<sup>4</sup>We employ the Gauss-Legendre quadrature for choosing the collocation points in regular boundary problems. We further discuss the choice of collocation points in Appendix B.1.
4.1 GRAPHSPLINENETS
Fig. 4.1 depicts the entire architecture of our model. GRAPHSPLINENETS can be divided into three main components:  $\bigcirc$  message passing neural networks (MPNN),  $\bigcirc$  time-oriented collocation and  $\bigcirc$  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.
Figure 4.1:  $\odot$  Message passing neural networks take as inputs current states and employ an encoder-processor-decoder architecture for obtaining the next states autoregressively.  $\odot$  time-oriented OSC obtains continuous predictions in time while  $\odot$  Space-oriented OSC is employed to obtain continuous predictions in space. The model is trained end-to-end efficiently by leveraging sparsity in the almost block diagonal (ABD) matrices of the collocation method.
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  $\{\mathbf{x}_i\}$  for a 2-D domain. We have the states of these sample points at the initialization time frame  $\{y_i^{t_0}\}$ .
$\odot$  Message passing neural networks The neural networks take the states of sample points at the initial time frame  $\{(\mathbf{x}_i, y_i^{t_0})\}$  as input and employ an encoder-processor-decoder architecture for obtaining the next steps states  $\hat{y}_i^{t_j}, j = 1, \dots, N$  autoregressively.
- Time-oriented OSC For each sample point, taking  $\mathbf{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, \dots, t_K\}$  and values at these time steps  $\{y_i^{t_0}, \hat{y}_i^{t_1}, \dots, \hat{y}_i^{t_K}\}$ , and the output is the polynomials  $\hat{u}_i(t), t \in [t_0, t_K]$  which provides a continuous prediction along time; thus, we can obtain predictions at any time.
$\odot$  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  $\{\mathbf{x}_i\}$ , their values at this time frame  $\{\hat{y}_i^{t_k}\}$ , and the output is polynomials  $\hat{u}^{t_k}(\mathbf{x}), \mathbf{x} \in \Omega$  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  $\hat{u}(\mathbf{x}, t), (\mathbf{x}, t) \in \Omega \times [t_0, t_K]$ .
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  $\phi = \mathrm{OSC}(\mathrm{MPNN}(\mathbf{X},\mathbf{Y}^{t_0};\hat{\theta}))$  which allow for spatiotemporal continuous outputs  $\hat{u} (\mathbf{x}_i,t_k;\phi)$ . Note that solving OSC problems is equivalent to predicting the future space and time continuous states, given the sample inputs.
Inherently, training the parameters,  $\theta$  of the MPNN is cast as a bi-level optimization problem as  $\theta^{*} = \mathrm{argmin}_{\theta}\mathcal{L}(\phi ,\theta)$  where  $\phi = \mathrm{OSC}(\mathrm{MPNN}(\mathbf{X},\mathbf{Y}^{t_0};\theta))$ . The loss function  $\mathcal{L}$  is defined as:
$$
L = \overbrace {\sum_ {i = 0} ^ {N} \sum_ {k = 0} ^ {K} \left\| y _ {i} ^ {t _ {k}} - \hat {y} _ {i} ^ {t _ {k}} \right\| ^ {2}} ^ {L _ {\mathrm {s}} \equiv \text {s a m p l e p o i n t s e r r o r}} + \underbrace {\sum_ {i = 0} ^ {N _ {i}} \sum_ {k = 0} ^ {N _ {t}} \left\| y _ {i} ^ {t _ {k}} - \hat {u} \left(\mathbf {x} _ {i} , t _ {k}\right) \right\| ^ {2}} _ {L _ {\mathrm {i}} \equiv \text {i n t e r p o l a t i o n p o i n t s e r r o r}} \tag {2}
$$
The left term  $L_{s}$  of Eq. (2) is the loss of the MPNN output  $\hat{y}_i^{t_k}$ , 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  $\theta$ . The right term  $L_{i}$  is the prediction loss of the high-resolution interpolation points  $(\mathbf{x}_i, y_i^{t_k})_{i=0,k=0}^{N_i, N_t}$  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  $\hat{u}(\mathbf{x}_i, t_k; \phi)$  is constructed based upon the optimized parameters  $\phi = \mathrm{OSC}(\mathrm{MPNN}(\mathbf{X}, \mathbf{Y}^{t_0}; \theta))$ , the parameters that are obtained by solving OSC problems with MPNN predictions,  $\mathrm{MPNN}(\mathbf{X}, \mathbf{Y}^{t_0}; \theta)$ . Thus, both loss terms can be optimized with respect to
Figure 4.2: Illustration of data points used for model training. Interpolation points include  $\odot$  time-oriented OSC and  $\odot$  space-oriented OSC interpolation points. Here,  $t$  can be any time frame in  $(t_0, t_K)$ .
the MPNN parameters  $\theta$ , 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.
More specifically, after getting the continuous simulation function  $\hat{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
Figure 4.3: Adaptive collocation strategy: mesh points converge towards areas with faster changing regions. In this example, collocation points move towards the moving wave crests.
$$
\tilde {\mathbf {x}} _ {i} = \mathbf {x} _ {i} + \beta \mathbf {v}, \mathbf {v} = \operatorname {a r g m a x} _ {\mathbf {v}} \nabla_ {\mathbf {v}} \frac {\partial \hat {u}}{\partial t} (\mathbf {x} _ {i}, t _ {k}) (3)
$$
where  $\beta$  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.
5.1 Datasets and Training Settings