SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jules Berman
Courant Institute for Mathematical Sciences
jmb1174@nyu.edu
Benjamin Peherstorfer
pehersto@cims.nyu.edu
Training neural networks sequentially in time to approximate solution fields of time-dependent partial differential equations can be beneficial for preserving causality and other physics properties; however, the sequential-in-time training is numerically challenging because training errors quickly accumulate and amplify over time. This work introduces Neural Galerkin schemes that update randomized sparse subsets of network parameters at each time step. The randomization avoids overfitting locally in time and so helps prevent the error from accumulating quickly over the sequential-in-time training, which is motivated by dropout that addresses a similar issue of overfitting due to neuron co-adaptation. The sparsity of the update reduces the computational costs of training without losing expressiveness because many of the network parameters are redundant locally at each time step. In numerical experiments with a wide range of evolution equations, the proposed scheme with randomized sparse updates is up to two orders of magnitude more accurate at a fixed computational budget and up to two orders of magnitude faster at a fixed accuracy than schemes with dense updates.
In science and engineering, partial differential equations (PDEs) are frequently employed to model the behavior of systems of interest. For many PDEs that model complicated processes, an analytic solution remains elusive and so computational techniques are required to compute numerical solutions.
Global-in-time training There have been many developments in using nonlinear parameterizations based on neural networks for numerically approximating PDE solutions. These include techniques such as the Deep Galerkin Method [46], physics-informed neural networks (PINNs) [41], and others [4, 21, 53, 15]; as well as early works such as [11, 42]. In most of these methods, a neural network is used to represent the solution of a (time-dependent) PDE over the whole space-time domain. For this reason they are termed global-in-time methods in the following. To approximate the solution, the neural network is trained to minimize the PDE residual on collocation points sampled from the space-time domain, which requires solving a large-scale optimization problem that can be computationally expensive. Additionally, the solutions learned by global-in-time methods can violate causality, which can become an issue for complex problems that rely on preserving physics [27]. We note that neural networks have been used for approximating PDE solutions in various other ways, such as learning specific component functions [23, 31, 43], finding closure models [2, 25, 50], de-noising [44], and for surrogate modeling [33, 32, 18]. However, we are interested in this work in using neural networks for directly approximating PDE solutions.
Sequential-in-time training with the Dirac-Frenkel variational principle In this work, we follow the Dirac-Frenkel variational principle, which has been used for numerical methods in the field of quantum dynamics for a long time [10, 17, 26, 34, 28] and for dynamic-low rank and related solvers [24, 45, 35, 39, 38, 22]. Instead of globally approximating a PDE solution in time, the Dirac-Frenkel variational principle allows a sequential-in-time training that adapts a nonlinear parameterization, such as a neural network, over time. In contrast to classical numerical methods in vector spaces, the approximate solution in Dirac-Frenkel schemes is allowed to depend nonlinearly on its parameters and so to lie on a smooth manifold. The update to the nonlinear parameterization is calculated at each time step according to the orthogonal projection of the dynamics onto the tangent space of the manifold induced by the nonlinear parameterization. The Dirac-Frenkel variational principle has been adapted for the nonlinear approximation of PDEs with neural networks. In particular [14, 1, 7] formulate a sequential-in-time method based on the Dirac-Frenkel variational principle.
The neural network represents the PDE solution at a point in time. The time-dependence then arises by allowing the parameters—the weights and biases of the network—to vary in time. The network parameters are then evolved forward according to the time dynamics which govern the PDE. This is in contrast to global-in-time methods, in which time enters the network as an additional input variable. By construction, an approximate solution obtained with a sequential-in-time method is causal, in that the solution at future times depends only on the solution at the current time.
Although these methods have demonstrated success in solving various PDEs [34, 14, 7, 54, 16, 8, 30], there are open challenges: First, the sequential-in-time training is prone to overfitting which can lead to a quick accumulation of the residual over time. Second, the local training step has to be repeated at each time step, which can be computationally costly, especially with direct solvers that have costs increase quadratically with the number of network parameters. The work [16] proposes to address the two issues by using iterative solvers, instead of direct ones, and by re-training the network occasionally over the sequential-in-time training. We show with numerical experiments below that the re-training of the network can be computationally expensive. Additionally, the performance of iterative solvers depends on the condition of the problem, which can be poor in the context of sequential-in-time training.
Our approach and contributions: Randomized sparse updates for schemes based on the Dirac-Frenkel variational principle We build on the previous work in sequential-in-time methods following a similar set up as [7] based on the Dirac-Frenkel variational principle. Where all previous methods solve local training problems that update every parameter of the network at each time step, we propose a modification such that only randomized sparse subsets of network parameters are updated at each time step:
(a) The randomization avoids overfitting locally in time and so helps preventing the error from accumulating quickly over the sequential-in-time training, which is motivated by dropout that addresses a similar issue of overfitting due to neuron co-adaptation.
(b) The sparsity of the updates reduces the computational costs of training without losing expressiveness because many of the network parameters are redundant locally at each time step.
Our numerical experiments indicate that the proposed scheme is up to two orders of magnitude more accurate at a fixed computational cost and up to two orders of magnitude faster at a fixed accuracy.
We release our code implementation here: https://github.com/julesberman/RSNG
2 Sequential-in-time training for solving PDEs
2.1 Evolution equations, Dirac-Frenkel variational principle, Neural Galerkin schemes
Given a spatial domain  $\mathcal{X} \subseteq \mathbb{R}^d$ , and a time domain  $\mathcal{T} = [0, T) \subseteq \mathbb{R}$ , we consider a solution field  $u: \mathcal{T} \times \mathcal{X} \to \mathbb{R}$  so that  $u(t, \cdot): \mathcal{X} \to \mathbb{R}$  is in a function space  $\mathcal{U}$  at each time  $t$ , with dynamics
$$
\begin{array}{l l} \partial_ {t} u (t, \boldsymbol {x}) = f (\boldsymbol {x}, u) & \text {f o r} (t, \boldsymbol {x}) \in \mathcal {T} \times \mathcal {X} \\ u (0, \boldsymbol {x}) = u _ {0} (\boldsymbol {x}) & \text {f o r} \boldsymbol {x} \in \mathcal {X} \end{array} \tag {1}
$$
where  $u_0 \in \mathcal{U}$  is the initial condition and  $f$  can include partial derivatives of  $u$  to represent PDEs. We focus in this work on Dirichlet and periodic boundary conditions but the following approach
Figure 1: We propose Neural Galerkin schemes that update randomized sparse subsets of network parameters with the Dirac-Frenkel variational principle. Randomization avoids overfitting locally in time, which leads to more accurate approximations than dense updates. Sparsity reduces the computational costs of training without losing expressiveness because many parameters are redundant locally in time.
can be applied with, e.g., Neumann boundary conditions as well [7]. One approach for imposing Dirichlet boundary conditions is by choosing parameterizations that satisfy the boundary conditions by definition [48].
Sequential-in-time training methods approximate  $u$  with a nonlinear parameterization such as a neural network  $\hat{u} : \mathcal{X} \times \Theta \to \mathbb{R}$ , where the parameter vector  $\pmb{\theta}(t) \in \Theta \subseteq \mathbb{R}^p$  depends on time  $t$ ; the parameter  $\pmb{\theta}(t)$  has  $p$  components and enters nonlinear in the second argument of  $\hat{u}$ . The residual of (1) at time  $t$  is
$$
r (\boldsymbol {x}; \boldsymbol {\theta} (t), \dot {\boldsymbol {\theta}} (t)) = \nabla_ {\boldsymbol {\theta}} \hat {u} (\boldsymbol {x}; \boldsymbol {\theta} (t)) \cdot \dot {\boldsymbol {\theta}} (t) - f (\boldsymbol {x}, \hat {u} (\cdot ; \boldsymbol {\theta} (t))), \tag {2}
$$
where we applied the chain rule to  $\partial_t\hat{u} (\cdot ;\pmb {\theta}(t))$  to formally obtain  $\dot{\pmb{\theta}} (t)$ . Methods based on the Dirac-Frenkel variational principle [10, 17, 34] seek  $\dot{\pmb{\theta}} (t)$  such that the residual norm is minimized, which leads to the least-squares problem
$$
\min  _ {\dot {\boldsymbol {\theta}} (t)} \| \nabla_ {\boldsymbol {\theta} (t)} \hat {u} (\cdot ; \boldsymbol {\theta} (t)) \dot {\boldsymbol {\theta}} (t) - f (\cdot ; \hat {u} (\cdot ; \boldsymbol {\theta} (t))) \| _ {L ^ {2} (\mathcal {X})} ^ {2}, \tag {3}
$$
in the  $L^2(\mathcal{X})$  norm  $\| \cdot \|_{L^2(\mathcal{X})}$  over  $\mathcal{X}$ . The least-squares problem (3) gives a  $\dot{\pmb{\theta}}(t)$  such that the residual is orthogonal to the tangent space at  $\hat{u}(\cdot; \pmb{\theta}(t))$  of the manifold  $\mathcal{M}_{\Theta} = \{\hat{u}(\cdot; \pmb{\theta}) | \pmb{\theta} \in \Theta\}$  induced by the parameterization  $\hat{u}$ ; see Figure 1. Schemes that solve (3) over time have also been termed Neural Galerkin schemes [7] because (3) can be derived via Galerkin projection as well.
The tangent space at  $\hat{u} (\cdot ;\pmb {\theta}(t))$  is spanned by the spanning set  $\{\partial_{\theta_i}\hat{u} (\cdot ;\pmb {\theta}(t))\}_{i = 1}^p$ , which are the component functions of the gradient  $\nabla_{\pmb{\theta}(t)}\hat{u} (\cdot ;\pmb {\theta}(t))$ ; it is important to stress that  $\{\partial_{\theta_i}\hat{u} (\cdot ;\pmb {\theta}(t))\}_{i = 1}^p$  is not necessarily a basis of the tangent space because it can contain linearly dependent functions and be non-minimal. The least-squares problem (3) can be realized by assembling a matrix whose columns are the gradient sampled at  $n\gg p$  points  $\pmb {x}_1,\dots ,\pmb {x}_n\in \mathcal{X}$  resulting in  $J(\pmb {\theta}(t)) = [\nabla_{\pmb{\theta}(t)}\hat{u} (\pmb {x}_1;\pmb {\theta}(t)),\dots ,\nabla_{\pmb{\theta}(t)}\hat{u} (\pmb {x}_n;\pmb {\theta}(t))]^T\in \mathbb{R}^{n\times p}$ , which is a batch Jacobian matrix to which we refer to as Jacobian for convenience in the following. Additionally, we form the right-hand side vector  $\pmb {f}(\pmb {\theta}(t)) = [f(\pmb {x}_1;\pmb {\theta}(t)),\dots ,f(\pmb {x}_n;\pmb {\theta}(t))]^T\in \mathbb{R}^n$  and thus the least-squares problem
$$
\min  _ {\dot {\boldsymbol {\theta}} (t)} \| J (\boldsymbol {\theta} (t)) \dot {\boldsymbol {\theta}} (t) - \boldsymbol {f} (\boldsymbol {\theta} (t)) \| _ {2} ^ {2}. \tag {4}
$$
The choice of the points  $x_{1}, \ldots, x_{n}$  is critical so that solutions of (4) are good approximations of solutions of (3); however, the topic of selecting the points  $x_{1}, \ldots, x_{n}$  goes beyond this work here and so we just note that methods for selecting the points exist [1, 7] and that we assume in the following that we select sufficient points with  $n \gg p$  to ensure that solutions of (4) are good approximations of solutions of (3)
Challenge 1: Parameters are redundant locally in time. Typically, parameterizations  $\hat{u}$  based on deep neural networks lead to Jacobian matrices  $J(\pmb{\theta}(t))$  that are low rank in the least-squares problem (4); see, e.g., [37] and Figure 2(a). In our case, a low-rank matrix  $J(\pmb{\theta}(t))$  means that components in  $\dot{\pmb{\theta}}(t)$  are redundant, because we assume that the samples  $\pmb{x}_1, \dots, \pmb{x}_n$  are sufficiently rich. Even if  $J(\pmb{\theta}(t))$  is low rank and thus the components in  $\dot{\pmb{\theta}}(t)$  are redundant, the problem (4) can still be
(a) low-rankness of Jacobian
(b) residual growth
(c) co-adaptation of neurons
Figure 2: (a) Jacobians that have low rank locally in time imply that there are redundant parameters in the neural network, which motivates the proposed sparse updates that lead to speedups without losing expressiveness. (b) The residual grows quickly with sequential-in-time training (and dense updates). This is not due to a limitation with the expressiveness of the network because directly fitting the network to solutions indicates that there exist other parameters that can lead to lower residuals. (c) Sequential-in-time training (with dense updates) results in co-adapted neurons as indicated by the highly correlated columns of the  $J$  matrix. Plots for experiment with Allen-Cahn equation (Sec. 4).
solved with standard linear algebra methods such as the singular value decomposition (SVD) because they compress the matrix  $J(\pmb{\theta}(t))$  and regularize for, e.g., the minimal-norm solution; however, the costs of performing the SVD to solve (4) scales as  $\mathcal{O}(np^2)$ , and thus is quadratic in the number of parameters  $p$ . This means that a redundancy in  $\dot{\pmb{\theta}}(t)$  of a factor two leads to a  $4 \times$  increase in the computational costs. Note that the problem typically is poorly conditioned because  $J(\pmb{\theta}(t))$  is low rank, which makes the direct application of iterative solvers challenging.
Challenge 2: Overfitting leads to high residual over time. The residual from solving (4) can rapidly increase over time which in turn increases the overall error. This indicates that the tangent space along the trajectory  $\pmb{\theta}(t)$  becomes ill suited for approximating the right-hand side vector  $\pmb{f}(\pmb{\theta}(t))$  in (4). We compare the residual of the least-squares problem (4) that is obtained along a trajectory of  $\pmb{\theta}(t)$  from sequential-in-time training with the schemes above to the residual of (4) from a network that is fit to the true solution at each point in time; details in Appendix A.1. As shown in Figure 2(b), a lower residual is achieved by the network that is fit to the true solution.
We aim to understand this phenomenon through the lens of overfitting: the sequential-in-time training can be thought of as successive fine-tuning, in the sense that at each time step we must make a small update to our parameters to match the solution at the next time step. However, fine-tuning is well known to be prone to over-fitting and model degeneration [3]. In the setting considered in this work, overfitting means that the representation  $\hat{u} (\cdot ;\pmb {\theta}(t))$  does not generalize well to the next time step. Not generalizing well means that a local change to  $\pmb{\theta}(t)$  is insufficient to move  $\hat{u} (\cdot ;\pmb {\theta}(t))$  according to the desired update given by  $\dot{\pmb{\theta}} (t)$  to match the right-hand side  $f(\pmb {\theta}(t))$ , which implies that a large residual is incurred when solving (4). A common approach to prevent overfitting is dropout [47], especially when applied to fine-tuning tasks with dropout variants proposed in [3, 29], while other approaches are formulated specifically around sparse updates [49, 52]. Dropout is motivated by the observation that dense updates to parameters in neural networks can cause overfitting by leading neurons to co-adapt. Typically, co-adaptation is characterized by layer-wise outputs with high covariance [9]. In the case of sequential-in-time training with the schemes discussed above, co-adaptation implies the columns of the Jacobian matrix  $J(\pmb {\theta}(t))$  are correlated and thus close to linearly dependent. So as neurons co-adapt, component functions of the gradient become redundant and may be less suited for approximating  $f(\pmb {\theta}(t))$  causing the high residual for the least-squares problem; see Figure 2(b). This could also be characterized by the ill conditioning issue pointed out in [16]. We see empirical evidence of co-adaptation in Figure 2(c), where we plot component functions of the gradient and see that they are strongly correlated at the end time  $T$ .
3 Randomized Sparse Neural Galerkin (RSNG) schemes
We introduce randomized sparse Neural Galerkin (RSNG) schemes that build on the Dirac-Frenkel variational principle to evolve network parameters  $\pmb{\theta}(t)$  sequentially over time  $t$  but update only sparse subsets of the components of  $\pmb{\theta}(t)$  and randomize which components of  $\pmb{\theta}(t)$  are updated. The sparse updates reduce the computational costs of solving the least-squares problem (4) while taking advantage of the low rank structure of  $J(\pmb{\theta})$  which indicates components of the time derivative  $\dot{\pmb{\theta}}(t)$  are redundant and can be ignored for updating  $\pmb{\theta}(t)$  without losing expressiveness. The randomization of which components of  $\pmb{\theta}(t)$  are updated prevents the overfitting described above.
3.1 Randomized sketch of residual
To define the sketch matrix  $S_{t}$ , let  $e_1, \ldots, e_p$  be the  $p$ -dimensional canonical unit vectors so that  $e_i$  has entry one at component  $i$  and zeros at all other components. We then define  $s$  independent and identically distributed random variables  $\xi_1(t), \ldots, \xi_s(t)$  that depend on time  $t$ . The distribution of  $\xi_i(t)$  is  $\pi$ , which is supported over the set of indices  $\{1, \ldots, p\}$ . The random matrix  $S_t$  of size  $p \times s$  is then defined as  $S_t = [e_{\xi_1(t)}, \ldots, e_{\xi_s(t)}]$ . The corresponding sketched residual analogous to (2) is
$$
r _ {s} (\boldsymbol {x}; \boldsymbol {\theta} (t), \dot {\boldsymbol {\theta}} _ {s} (t)) = \nabla_ {\boldsymbol {\theta}} \hat {u} (\boldsymbol {x}; \boldsymbol {\theta} (t)) S _ {t} \dot {\boldsymbol {\theta}} _ {s} (t) - f (\boldsymbol {x}, \hat {u} (\cdot ; \boldsymbol {\theta} (t)), \tag {5}
$$
where now  $\dot{\theta}_s(t) \in \mathbb{R}^s$  is of dimension  $s \ll p$ .
3.2 Projections onto randomized approximations of tangent spaces
Using the sketch matrix  $S_{t}$ , we obtain from the spanning set  $\{\partial_{\theta_i}\hat{u} (\cdot ;\pmb {\theta}(t))\}_{i = 1}^p$  of component functions of  $\nabla_{\pmb{\theta}}\hat{u} (\cdot ;\pmb {\theta}(t))$  a subset  $\{\partial_{\theta_{\xi_i(t)}}\hat{u} (\cdot ;\pmb {\theta}(t))\}_{i = 1}^s$  with  $s$  functions. The set  $\{\partial_{\theta_{\xi_i(t)}}\hat{u} (\cdot ;\pmb {\theta}(t))\}_{i = 1}^s$  spans at least approximately the tangent space at  $\hat{u} (\cdot ;\pmb {\theta}(t))$  of  $\mathcal{M}_{\Theta}$  but has only  $s\ll p$  elements. The motivation is that the full spanning set  $\{\partial_{\theta_i}\hat{u} (\cdot ;\pmb {\theta}(t))\}_{i = 1}^p$  contains many functions that are close to linearly dependent (Jacobian is low rank) and thus sub-sampling the component functions still gives reasonable tangent space approximations that preserves much of the expressiveness; see Figure 1. While the low rankness depends on the complexity of the problem and parametrization, we observe low rankness in all our examples; see Appendix A.1 for further discussion.
We now introduce a least-squares problem based on the sparse spanning set  $\{\partial_{\theta_{\xi_i(t)}}\hat{u} (\cdot ;\pmb {\theta}(t)\}_{i = 1}^s$  that is analogous to the least-squares problem problem based on the full spanning set given in (4). We seek  $\dot{\pmb{\theta}}_s(t)\in \mathbb{R}^s$  with  $s$  components that solves
$$
\min  _ {\hat {\boldsymbol {\theta}} _ {s} (t) \in \mathbb {R} ^ {s}} \| \nabla_ {\boldsymbol {\theta}} \hat {u} (\cdot ; \boldsymbol {\theta} (t)) S _ {t} \dot {\boldsymbol {\theta}} _ {s} (t) - f (\cdot ; \hat {u} (\cdot ; \boldsymbol {\theta} (t))) \| _ {L ^ {2} (\mathcal {X})} ^ {2}. \tag {6}
$$
To obtain  $\dot{\pmb{\theta}}(t)$  to update  $\pmb{\theta}(t)$ , we set  $\dot{\pmb{\theta}}(t) = S_t\dot{\pmb{\theta}}_s(t)$ . Thus, the components of  $\dot{\pmb{\theta}}(t)$  that are selected by  $S_t$  are set to the corresponding value of the component of  $\dot{\pmb{\theta}}_s(t)$  and all other components are set to zero, which means that the corresponding components of  $\pmb{\theta}(t)$  are not updated. We can realize (6) the same way as the full least-squares problem in (4) by using the full Jacobian matrix and  $S_t$  to define the sparse Jacobian matrix as  $J_s(\pmb{\theta}(t)) = J(\pmb{\theta}(t))S_t$  and the right-hand side vector  $\pmb{f}(\pmb{\theta}(t))$  analogous to Section 2 to obtain the discrete least-squares problem
$$
\min  _ {\dot {\boldsymbol {\theta}} _ {s} (t)} \left\| J _ {s} (\boldsymbol {\theta} (t)) \dot {\boldsymbol {\theta}} _ {s} (t) - \boldsymbol {f} (\boldsymbol {\theta} (t)) \right\| _ {2} ^ {2}. \tag {7}
$$
The choice of the distribution  $\pi$  is critical and depends on properties of the Jacobian matrix  $J(\theta(t))$ . Distributions based on leverage scores provide tight bounds with regard to the number of columns one needs to sample in order for the submatrix to be close to an optimal low rank approximation of the full matrix with high probability [13]. But these distributions can be expensive to sample from. Instead, uniform sampling provides a fast alternative.
The number of columns will not grow too quickly if the full matrix is sufficiently incoherent [19]. This means some columns do not carry a disproportionate amount of information relative to other columns. We numerically see that in our case the Jacobian matrix  $J(\pmb{\theta}(t))$  is sufficiently incoherent. Thus we can choose a uniform distribution over the set of indices  $\{1, \dots, p\}$  to get the benefits of low rank approximation in a computationally efficient way.
Algorithm 1 Randomized Neural Galerkin scheme with sparse updates
Fit parameterization  $\hat{u} (\cdot ;\pmb{\theta}^{(0)})$  to initial condition  $u_{0}$  to obtain  $\pmb{\theta}^{(0)}$  for  $k = 1,\dots ,K$  do Draw realization of sketching matrix  $S_{k}$  as described in Section 3.1 Solve for sparse update  $\Delta \pmb{\theta}_{s}^{(k - 1)}$  with least-squares problem (9) Lift sparse update  $\Delta \pmb{\theta}^{(k - 1)} = S_k\Delta \pmb{\theta}_s^{(k - 1)}$  Update  $\pmb{\theta}^{(k)} = \pmb{\theta}^{(k - 1)} + \delta t\Delta \pmb{\theta}^{(k - 1)}$  end for
3.3 Discretizing in time
We discretize the time interval  $\mathcal{T}$  with  $K\in \mathbb{N}$  regularly spaced time steps  $0 = t_0 < t_1 < \dots < t_K = T$  with  $\delta t = t_k - t_{k - 1}$  for  $k = 1,\ldots ,K$ . At time  $t_0$ , we obtain  $\pmb{\theta}^{(0)}\in \mathbb{R}^p$  by fitting the initial condition  $u_{0}$ . We then update
$$
\boldsymbol {\theta} ^ {(k)} = \boldsymbol {\theta} ^ {(k - 1)} + \delta t \Delta \boldsymbol {\theta} ^ {(k - 1)} \tag {8}
$$
for  $k = 1, \ldots, K$  so that  $\pmb{\theta}^{(k)}$  is the time-discrete approximation of  $\pmb{\theta}(t_k)$  and thus  $\hat{u}(\cdot; \pmb{\theta}^{(k)})$  approximates the solution  $u$  at time  $t_k$ . The sparse update  $\Delta \pmb{\theta}_s^{(k-1)}$  approximates  $\dot{\pmb{\theta}}_s(t_{k-1})$  and is obtained by the time-discrete counterpart of (7), which is given by
$$
\min  _ {\Delta \boldsymbol {\theta} _ {s} ^ {(k - 1)}} \| J _ {s} \left(\boldsymbol {\theta} ^ {(k - 1)}\right) S _ {k} \Delta \boldsymbol {\theta} _ {s} ^ {(k - 1)} - \boldsymbol {f} \left(\boldsymbol {\theta} ^ {(k - 1)}\right) \| _ {2} ^ {2}, \tag {9}
$$