SGM-PINN: Sampling Graphical Models for Faster Training of Physics-Informed Neural Networks

Consider the following general form of a PDE equation under a set of boundary condition (BC) and initial condition (IC) constraints (Hennigh et al., 2021):

where $\mathcal{F}_{i}$ is the general differential operator, $N_{\mathcal{F}}$ ($N_{\mathcal{C}}$) denotes the number of BCs (ICs), x is the set of independent variables defined over a bounded continuous domain $\mathcal{D}\subseteq\mathbb{R}^{d}$, $u(\textbf{x})$ is the solution to the PDE, $\mathcal{C}_{j}$ is the constraint operator—such as the differential, linear, or nonlinear terms representing the boundary and initial conditions—and $\partial\mathcal{D}$ is a subset of the domain boundary. Consider a simple feed-forward fully-connected neural network with nonlinear activation functions, where the approximate solution $\tilde{u}(\textbf{x})$ can be further expressed as:

where $n$ represents the number of neural network (NN) layers, $\phi_{l}(\textbf{x}_{i})=\sigma(\textbf{W}_{l}\textbf{x}_{i}+\textbf{b}_{l})$ represents the $l^{th}$ layer of the NN, $\textbf{W}_{l}\in\mathbb{R}^{d_{l}\times d_{l-1}}$ and $\textbf{b}_{l}\in\mathbb{R}^{d_{l}}$ are the weight and bias of the $l^{th}$ layer, $\phi_{E}$ is an input encoding layer, $\theta=\{\textbf{W}_{1},\textbf{b}_{1},\cdots,\textbf{W}_{n},\textbf{b}_{n}\}$ is the set of the trainable parameters of the network, and $\sigma$ is the nonlinear activation function. Given the approximated solution $\tilde{u}(\textbf{x},\theta)$, we define the following residuals for both PDE equation and boundary (initial) condition constraints:

In order to train the NN model, residuals can be encoded into a constructed loss function such that $\tilde{u}(x,\theta)$ can gradually approach the true solution $u(x)$ through iteratively optimizing the network parameters $\theta$. The loss function can be represented in the following form:

where $w_{\mathcal{F}}^{(i)}$ and $w_{\mathcal{C}}^{(j)}$ are the weight functions to scale the corresponding loss terms. The model may be trained by stochastic gradient descent such that the model parameters are updated by:

A probabilistic graphical model (PGM) expresses probabilistic relationships between variables via links between nodes representing the variables (Bishop, 2023). In traditional PDE solvers, discretization methods operate on the principle that the solution at any given point is influenced by its immediate neighbors, with interactions between adjacent mesh or grid points playing a crucial role in problems such as FEA. In the PINN examples discussed in this paper, a randomized point cloud of $N$ collocation points with $M$ input features $X=\mathbb{R}^{N\times M}$ is generated prior to training as described in (Raissi et al., 2019).

With the assumption that nearby points will be highly correlated in a relatively dense initial point cloud, we can construct a k-nearest-neighbor (kNN) graph using the low-dimensional spatial coordinates (x,y,z) of each point in the dataset to represent a conditional relationship between points inversely proportional to their distance. This graph, created using an existing highly-efficient kNN algorithm (Malkov and Yashunin, 2018), serves as an undirected PGM for the input data $X$. At later stages in training this model can be re-built in parallel while incorporating additional features from the output, such as flow or temperature.

Node clustering is achieved via a low-resistance-diameter (LRD) decomposition method that allows partitioning any given PGM into multiple node (collocation point) clusters with bounded resistance diameter ¹¹1The maximum effective-resistance distances between any two nodes within the same graph cluster will not exceed a given threshold.. Since the resistance distance between two nodes on a PGM encodes the conditional dependence between the two corresponding data samples, LRD will guarantee that only the most similar data samples will be grouped into the same cluster, which thereby allows SGM-PINN to select highly-representative data samples during the training of PINNs.

In (Alev et al., 2018) it is proven possible to decompose a simple graph into multiple node clusters of effective-resistance diameter at most the inverse of the average node degree (up to constant losses) by removing only a constant fraction of edges. This is accomplished without significantly impacting the graph conductance (keeping the global structure of the graph intact).

Directly computing effective resistances according to Definition 1 is practically infeasible for large PGMs. However, as the proposed LRD method only requires approximate estimation of edge effective resistances, we use a highly-scalable (linear time) algorithm by exploiting a Krylov subspace approach, discussed in detail in (Aghdaei and Feng, 2022).

Since the proposed decomposition method has a nearly-linear complexity, it can be efficiently applied to large-scale PGMs (graphs) with millions of data samples (nodes). For additional performance, we also decompose the dataset into grids and perform S1 and S2 in independent sub-processes while training continues either with uniform sampling, or a previously calculated distribution. Speedup is roughly linear with the number of available threads.

We use a black-box spectral method for assessing the robustness (stability) of ML models and data, relying on graph-based manifold representations (Cheng et al., 2021), referred to here as the Inverse Stability Rating (ISR). ISR introduces the concept of Distance Mapping Distortion (DMD), denoted as $\gamma^{F}(p,q)$ for a node pair $(p,q)$ transformed by a function $Y=F(X)$, $F(X)$ in this case being the NN $\tilde{u}(\textbf{x},\theta)$. It is defined as the ratio of distances between nodes $p$ and $q$ on the output and input graphs, $\gamma^{F}(p,q)\triangleq\frac{d_{Y}(p,q)}{d_{X}(p,q)}$, where $d_{X}(p,q)$ and $d_{Y}(p,q)$ represent the distances between nodes $p$ and $q$ on the input and output graphs. Intuitively, the DMD can be employed to estimate the change in distance on the output graph (manifold) due to a perturbation on the input graph (manifold), with the largest $\gamma_{max}^{F}$ representing the fastest change in output with respect to input.

Each edge score acts as a surrogate for the directional derivative between two nodes on $F(X)$(Cheng et al., 2021), which in our application are the NN losses, such that:

In (Nabian et al., 2021) additional computational efficiency is achieved by partitioning nodes among a much smaller set of random ”seeds”(Nabian et al., 2021), updating the loss value at each seed to calculate $\mathcal{L}(\theta^{(t)})$ and assigning that value to nearby samples in a piece-wise fashion before calculating $P_{x_{i}}^{(t)}$ for each sample. In our method, we also use nearby neighbors to adjust the probability of sampling points we did not directly update the loss for, but we can add additional weight to clusters with high ISR. As it represents quickly changing local losses, it implies that the loss estimate for that cluster may be poor, and potentially important points may be ignored.

Only $r$ points in each cluster of samples are updated with the latest losses to rank that cluster relative to others, allowing a reduced number of passes through the NN to update $P$. Sampling is maintained with the principle in Equations 6,7 by sampling more from clusters with losses higher than other clusters, as described in Algorithm 1. There is a floor of 1 sample per cluster when determining an epoch, so no cluster is entirely ignored in any epoch. This mitigates the potential for failing to refine or ’forgetting’ the solution (Daw et al., 2023) in areas with relatively low scores as training continues. For parametric examples the ISR discussed in S3 is added as an additional term to the cluster score prior to ranking, using the same subset of samples as—and normalized with—the other PDE losses. S3 is also performed on a background thread, so the only additional overhead on GPU resources and overall wall time is the $r\times N$ loss calculations every $\tau_{e}$ iterations.

Algorithm 1 shows the key steps in SGM-PINN. Step 1 kNN construction using the HNSW algorithm (Malkov and Yashunin, 2018) has a complexity of $O(N\log(N))$, where $N$ is the number of nodes. Step 2 LRD is a nearly-linear time algorithm with a runtime proportional to the number of edges the graph $G$, which is determined by $N\times k$, $O(kN)$. Sampling losses for importance scores are an additional overhead cost, with the number of extra forward passes determined by $\tau_{e}\times r\times N$ for relatively frequent loss updates, as well as $\tau_{G}\times N$ for infrequent $G$ updates. ISR has a complexity $O(N\log(N))$. This gives an overall complexity of $O(N\log(N))+O(kN\tau_{e}+N\tau_{G})$.