SlowGuess/NeurIPS_2023_OCR huggingface dataset:
if time is discretized with the forward Euler method. Other discretization schemes can be used as well, which then lead to technically more involved problems (9) that remain conceptually similar though. The sparse update is lifted to  $\Delta \pmb{\theta}^{(k-1)} = S_k \Delta \pmb{\theta}_s^{(k-1)}$  so that the update (8) can be computed.
3.4 Computational procedure of RSNG
We describe the proposed RSNG procedure in algorithmic form in Algorithm 1. We iterate over the time steps  $k = 1, \dots, K$ . At each time step, we first sketch the Jacobian matrix by creating a submatrix from randomly sampled columns. Notice that  $S_{k}$  need not actually be assembled as its action on the Jacobian matrix can be accomplished by indexing. We then solve the least-squares problem given in (9) using our sketched Jacobian to obtain  $\Delta \theta_{s}^{(k - 1)}$ . A direct solve of this system dominates the computational cost of making a time step and scales in  $O(ns^2)$  time. The components of  $\Delta \theta^{(k - 1)}$  corresponding to the indices that have not been selected are filled with zeros. We then update the parameter  $\theta^{(k - 1)}$  to  $\theta^{(k)}$  via  $\Delta \theta^{(k - 1)}$ .
The whole integration process scales as  $O\left(\frac{T}{\delta t} ns^2\right)$  in time.
We demonstrate RSNG on a wide range of evolution equations, where speedups of up to two orders of magnitude are achieved compared to comparable schemes with dense updates. We also compare to global-in-time methods, where we achieve up to two orders of magnitude higher accuracy.
4.1 Setup and equations
Examples We now describe the details of the PDEs that we use to evaluate our method. We choose these particular setups to test our method on a diverse set of challenges including problems with global and local dynamics and solutions with sharp gradients and fine grained details. For visualization of the solutions of these equations see the Appendix A.4.
Reaction-diffusion problem modeled by Allen-Cahn (AC) equation: The Allen-Cahn equation models prototypical reaction diffusion phenomena and is given as,
$$
\partial_ {t} u (t, x) = \epsilon \partial_ {x x} u (t, x) + u (t, x) - u (t, x) ^ {3}.
$$
We choose  $\epsilon = 5\times 10^{-3}$ , with periodic boundary condition  $\mathcal{X} = [0,2\pi)$  and initial condition,
$$
u _ {0} (x) = \frac {1}{3} \tanh (2 \sin (x)) - \exp (- 2 3. 5 (x - \frac {\pi}{2}) ^ {2}) + \exp (- 2 7 (x - 4. 2) ^ {2}) + \exp (- 3 8 (x - 5. 4) ^ {2})).
$$
This initial condition results in challenging dynamics that are global over the spatial domain.
Flows with sharp gradients described by Burgers' equation: The Burgers' equation is given by,
$$
\partial_ {t} u (t, x) = \epsilon \partial_ {x x} u (t, x) - u (t, x) \partial_ {x} u (t, x).
$$
We choose  $\epsilon = 1\times 10^{-3}$ , with periodic boundary condition  $\mathcal{X} = [-1,1)$  and initial condition,
$$
u _ {0} (x) = \left(1 - x ^ {2}\right) \exp \left(- 3 0 (x + 0. 5) ^ {2}\right).
$$
The corresponding solution field has sharp gradients that move in the spatial domain over time, which can be challenging to approximate.
Charged particles in electric field: The Vlasov equation describes the time evolution of collisionless charged particles under the influence of an electric field. The equation models the distribution of such particles in terms of their position and velocity. We consider the case of one position dimension and one velocity dimension, making our domain  $\mathcal{X} \subseteq \mathbb{R}^2$ . The equation is given by,
$$
\partial_ {t} u (t, x, v) = - v \partial_ {x} u (t, x, v) + \partial_ {x} \phi (x) \partial_ {v} u (t, x, v)
$$
where  $x$  is the position,  $v$  is the velocity and  $\phi$  is the electric field. We consider the case with periodic boundary condition  $\mathcal{X} = [0,2\pi)\times [-6,6)$  and initial condition,
$$
u _ {0} (x, v) = \frac {1}{\sqrt {2 \pi}} \exp (\frac {- v ^ {2}}{2})
$$
with a fixed electric field  $\phi (x) = \cos (x)$ . This particular setup evolves into a distribution with fine grained details along a separatrix surrounding the potential well.
Setup We parameterize with a feed-forward multi-layer perceptron. All our networks use linear layers of width 25 followed by non-linear activation functions, except the last layer which has no activation and is of width 1 so that  $\hat{u}(\boldsymbol{x},\boldsymbol{\theta}(t)) \in \mathbb{R}$ . To vary the number of total parameters  $p$ , we vary the depth of networks ranging from 3-7 layers. We use rational activation functions which in our experiments allowed for fitting initial conditions faster and more accurately than a standard choice such as tanh or ReLU [5]. To enforce periodic boundary conditions, we modify the first layer so that it outputs periodic embeddings as in [7]; for details see Appendix A.5. The periodic embedding ensures that the boundary conditions are enforced exactly. For additional details on enforcing other types of boundary conditions (e.g. Dirichlet and Neumann) exactly in neural networks see [12, 7, 48]. We sample points from the domain on an equidistant grid. For time integration we use a RK4 scheme with a fixed time step size. The time step sizes are  $5\mathrm{e}-3$ ,  $1\mathrm{e}-3$ ,  $5\mathrm{e}-3$  and we integrate up to end time 4, 4, and 3 for the Allen-Cahn, Burgers', and Vlasov equations, respectively. All error bars show  $+/-$  two standard errors over three random realizations which results in different sketching matrices at each time step. Relative errors are computed over the full space-time domain, unless the plot is explicitly over time.
All gradients and spatial derivatives are computed with automatic differentiation implemented in JAX [6]. All computations are done in single precision arithmetic which is the default in JAX. All runtime statistics were computed on the same hardware, a Nvidia Tesla V100 w/ 32 GB memory. All additional hyperparameters are described in Appendix A.5.
RSNG achieves higher accuracy than schemes with dense updates at same computational costs. In Figure 3 we plot the relative error over time. The curves corresponding to "dense updates" use a 3 layer network and integration is performed using dense updates. For RSNG, we use a 7 layer network and integrate with sparse updates, setting the number of parameters we update,  $s$ , equal to the total number of parameters in the 3 layer network and thus equal to the number of parameters that are updated by "dense updates." Thus, the comparison is at a fixed computational cost. The error achieved with RSNG is one to two orders of magnitude below the error obtained with dense updates, across all examples that we consider. In Figure 4(a), we see that as we increase the network size,
(a) Allen-Cahn
(b) Burgers'
(c) Vlasov
Figure 3: We plot the relative error over time for RSNG versus dense updates at  $s = 757$ . We see RSNG leads to orders of magnitude lower errors than dense updates for the same costs.
(a) fix sparsity  $s$ , vary #network parameters  $p$
Figure 4: (a) RSNG benefits from the additional expressiveness of larger networks (larger  $p$ ) while only using a fixed number of parameters (fixed  $s$ ) at each time step. (b) As we decrease the number of parameters  $s$  in the sparse update, but keep the total number of parameters  $p$  of the network the same, we achieve lower errors than dense updates. Thus, RSNG outperforms dense updates while incurring lower computational costs. Error bars generated over random sketch matrices,  $S_{t}$ .
(b) fix #network parameters  $p$ , vary sparsity  $s$
the relative error decreases as the sparse updates allow us to exploit the greater expressiveness of larger networks while incurring no additional computational cost in computing (9). But we note that increasing the size of the full network will make computations of  $J(\theta)$  and  $f(\theta)$  more expensive because of higher costs of computing gradients. However, for the network sizes that we consider in this work, this effect is negligible compared to the cost of solving (9).
RSNG achieves speedups of up to two orders of magnitude In Figure 5(a), we compare the runtime of RSNG to the runtime of a scheme with dense updates that uses a direct solver and to the runtime of a scheme with dense updates that uses an iterative solver as proposed in [16]. The time is computed for Burgers' equation and the sparsity  $s$  of RSNG is chosen such that all methods reach a comparable level of error. We find that RSNG is faster than direct solves with dense updates by two orders of magnitude and faster than the iterative solver by one order of magnitude.
The results show that while using an iterative solver as in [16] does speed up the method relative to direct solves with dense updates, it can still be quite slow for networks with many parameters  $p$ . Additionally, convergence of the iterative method given in [16] requires a number of hyperparameters to be chosen correctly, which may require an expensive search or a priori knowledge about the solution. Note that our RSNG method does not preclude the use of an iterative method to speed up the least-squares solves further.
(a) runtime at  $2\mathrm{e} - 4$  relative error with  $s = 125$
(b) error versus runtime
Figure 5: RSNG has lower computational costs than dense updates with direct and iterative least-squares solvers. Plots for numerical experiment with Burgers' equation
To address the overfitting problem, the work [16] refits the network to the current approximate solution from a random initialization periodically during time integration. In Figure 5(b), we show the relative error versus the runtime for the iterative solver with various numbers of refits and for RSNG at different sparsity  $s$ . While refitting the network can reduce the relative error, it incurs a high computational cost. By contrast, for appropriate sparsity  $s$ , RSNG outperforms the method given in [16] in both speed and accuracy.
Varying sparsity  $s$  at fixed number of total parameters  $p$  in network We now study the effect of varying the sparsity  $s$  (i.e., number of parameters updated by sparse updates) for a fixed network of total size  $p$ , in this case a 7 layer network. In Figure 4(b), we see that for a network of fixed size, sparse updates can reduce the relative error by about  $2 - 3 \times$  when compared to dense updates. This is notable as the computational costs decrease quadratically with  $s$ . Thus, the combination of sparsity and randomized updates in RSNG can deliver both improved performance and lower computational cost. We see that at the beginning, when the number  $s$  is too small, the expressiveness suffers and the error becomes large. This is because if  $s$  is less than the rank of the dense Jacobian then the sparsified Jacobian will necessarily have less representational power. However, we stress that RSNG is robust with respect to  $s$  in the sense that for a wide range of  $s$  values the error is lower than for dense updates.
The high error when performing dense updates  $s = p$  in Figure 4(b) for Allen-Cahn and Burgers' equation is due to the overfitting problem described in Section 2.2. As updates become denser, the method is more likely to overfit to regions of the parameter space in which the Jacobian,  $J(\theta)$ , is ill suited for approximating the right-hand side  $f$  at future time steps (see Section 2). We can see this explicitly in Figure 6 where we plot the residual over time for sparse and dense updates on the Allen-Cahn equation. Initially, the dense updates lead to a lower residual. This makes sense as they begin at the same region of parameters space. But as the two methods navigate to different regions of parameters space, we see RSNG begins to incur a lower residual relative to dense updates. This indicates that RSNG ameliorates the problem of overfitting and so leads to a lower residual as shown in Figure 6(b).
Comparison with global-in-time methods We compare our method to global-in-time methods which aim to globally minimize the PDE residual over the entire space-time domain. We compare to the original PINN formulation given in [41]. Additionally we compare to a variant termed Causal PINNs, which impose a weak form of time dependence through the loss function [51]. We select this variant as it claims to have state of the art performance among PINNs on problems such as the Allen-Cahn equation. In Table 1, we see that our sequential-in-time RSNG method achieves a higher accuracy by at least one order of magnitude compared to PINNs. Additionally, in terms of computational costs, RSNG outperforms both PINN variants, as their global-in-time training is expensive and requires many residual evaluations. We note that the training time of PINNs is directly
(a) residual at early times
(b) residual over full time
Figure 6: Plot (a) shows the residual of dense and sparse updates at early time steps. Initially, dense updates must have a lower residual as  $JS_{t}$  spans a subspace of the tangent space given by  $J$ . But in plot (b), we see that after a few time steps, dense updates overfit and the residual grows quicker than with sparse updates.
PDE	Method	L2Relative Error	Time(s)	s
Allen-Cahn	PINN	6.85e-2	841	N/A
Allen-Cahn	Causal PINN	3.84e-4	3060	N/A
Allen-Cahn	RSNG	1.66e-5	776	800
Allen-Cahn	RSNG	5.34e-5	63	150
Burgers'	PINN	2.34e-3	3451	N/A
Burgers'	Causal PINN	5.19e-4	23027	N/A
Burgers'	RSNG	6.07e-5	2378	800
Burgers'	RSNG	2.05e-4	188	125
Table 1: The sequential-in-time training with RSNG achieves about one order of magnitude higher accuracy than global-in-time methods in our examples. Details on training in Appendix A.6.
dependent on the number of optimization iterations and thus they can be trained faster if one is willing to tolerate even higher relative errors.
5 Conclusions, limitations, and future work
In this work, we introduced RSNG that updates randomized sparse subsets of network parameters in sequential-in-time training with the Dirac-Frenkel variational principle to reduce computational costs while maintaining expressiveness. The randomized sparse updates are motivated by a redundancy of the parameters and by the problem of overfitting. The randomized sparse updates have a low barrier of implementation in existing sequential-in-time solvers. The proposed RSNG achieves speedups of up to two orders of magnitude compared to dense updates for computing an approximate PDE solution with the same accuracy.
Current limitations leave several avenues for future research: first, as discussed in 3.2, uniform sampling is only appropriate when the Jacobian matrix is of low coherence. Future work may investigate more sophisticated sampling methods such as leverage score and pivoting elements of rank revealing QR. Second, there are problems for which overfitting with dense updates is less of an issue; e.g., the charged particles example in our work. Note that due to the sparsity of the updates, RSNG still achieves a speedup compared to dense updates for the same accuracy for this example though. However, more work is needed to better understand and mathematically characterize which properties of the problems influence the overfitting issue.
We make a general comment about using neural networks for numerically solving PDEs: The equations discussed in this paper are standard benchmark examples used in the machine-learning literature; however, for these equations, carefully designed classical methods can succeed and often have lower runtimes than methods based on nonlinear parameterizations [20, 36]. While these equations provide an important testing ground to demonstrate methodological improvements, future work will extend these results to domains where classical linear methods struggle, e.g., high-dimensional problems and problems with slowly decaying Kolmogorov  $n$ -widths [15, 7, 40].
We do not expect that this work has negative societal impacts.
Acknowledgements The authors were partially supported by the National Science Foundation under Grant No. 2046521 and the Office of Naval Research under award N00014-22-1-2728. This work was also supported in part through the NYU IT High Performance Computing resources, services, and staff expertise.
[1] William Anderson and Mohammad Farazmand. Evolution of nonlinear reduced-order solutions for PDEs with conserved quantities. SIAM Journal on Scientific Computing, 44(1):A176-A197, 2022.
[2] Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P. Brenner. Learning data-driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(31):15344-15349, 2019.
[3] Antonio Valerio Miceli Barone, Barry Haddow, Ulrich Germann, and Rico Sennrich. Regularization techniques for fine-tuning in neural machine translation. CoRR, abs/1707.09920, 2017.
[4] Jens Berg and Kaj Nyström. A unified deep artificial neural network approach to partial differential equations in complex geometries. Neurocomputing, 317:28-41, nov 2018.
[5] Nicolas Boulle, Yuji Nakatsukasa, and Alex Townsend. Rational neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 14243-14253. Curran Associates, Inc., 2020.
[6] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs. 2018.
[7] J. Bruna, B. Peherstorfer, and E. Vanden-Eijnden. Neural Galerkin scheme with active learning for high-dimensional evolution equations. Journal of Computational Physics, 2023. (accepted; preprint arXiv:2203.01360).
[8] Peter Yichen Chen, Maurizio M. Chiaramonte, Eitan Grinspun, and Kevin Carlberg. Model reduction for the material point method via an implicit neural representation of the deformation map. Journal of Computational Physics, 478:111908, 2023.
[9] Michael Cogswell, Faruk Ahmed, Ross B. Girshick, Larry Zitnick, and Dhruv Batra. Reducing overfitting in deep networks by decorrelating representations. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016.
[10] P. A. M. Dirac. Note on exchange phenomena in the thomas atom. Mathematical Proceedings of the Cambridge Philosophical Society, 26(3):376-385, 1930.
[11] M. W. M. G. Dissanayake and N. Phan-Thien. Neural-network-based approximations for solving partial differential equations. Communications in Numerical Methods in Engineering, 10(3):195-201, 1994.
[12] Suchuan Dong and Naxian Ni. A method for representing periodic functions and enforcing exactly periodic boundary conditions with deep neural networks. Journal of Computational Physics, 435:110242, jun 2021.
[13] Petros Drineas, Michael W. Mahoney, and S. Muthukrishnan. Relative-error CUR matrix decompositions. CoRR, abs/0708.3696, 2007.
[14] Yifan Du and Tamer A. Zaki. Evolutional deep neural network. Physical Review E, 104(4), oct 2021.
[15] Weinan E, Jiequn Han, and Arnulf Jentzen. Algorithms for solving high dimensional PDEs: from nonlinear Monte Carlo to machine learning. Nonlinearity, 35(1):278, dec 2021.
[53] Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deep potential molecular dynamics: A scalable model with the accuracy of quantum mechanics. Phys. Rev. Lett., 120:143001, Apr 2018.
[54] Tianchen Zhao, Chuhao Sun, Asaf Cohen, James Stokes, and Shravan Veerapaneni. Quantum-inspired variational algorithms for partial differential equations: Application to financial derivative pricing. arXiv, 2207.10838, 2022.
A.1 Details on Figure 2
For Figure 2a-b we look at quantities generated by fitting a network to the true solution at a point in time  $t$ . This is done in the same way we fit initial conditions described in A.5, but in this context the target function is taken to be the true solution at time  $t$ . For Figure 2(a), we compute the Jacobian of the network fitted to the true solution at each point in time and then plot its spectrum. For Figure 2(b), we take the network fitted to the true solution and compute the residual from the least-squares problem 4 to give the data points in the "Direct fit" line.
In Figure 7 we provide versions of Figure 2(a) from the main text but for all the equations considered. Vlasov is a more complex problem that exhibits a less sharp decay but is still distinctly rank deficient.
A.2 Details on Speed-Up
Here we provide additional results on the speed-up provided by RSNG for different equations. In Figure 8 we provide versions of Figure 5(a) from the main text but for all the equations considered. In Figure 9 we show how the runtime of Neural Galerkin schemes scales with  $s$ , thus showing the quadratic speed-up provided by RSNG as we reduce  $s$ .
A.3 Applications to High Dimensional Problems
Neural Galerkin schemes have been shown to be a useful approach to high dimensional PDEs [7]. In Figure 10 we demonstrate that RSNG may be applicable in these settings we well. We fit a neural network to a numerical solution of the Fokker-Planck equation in 5 dimensions; see [7, Section 5.4.1] for a description of the setup. The results show that for a sufficiently large network, the Jacobian has a low-rank structure as in the examples in the paper and a steep decay in the singular values so that random sketching strategies will likely be successful.
A.4 Ground Truth
The ground truth for the Allen-Cahn and Burgers' equations were generated using a spectral method with a fourth order integrator implemented in the spin solver as part of the Chebfun package in Matlab. We used a spatial grid of 10 000 points and time step size of  $1\mathrm{e} - 3$
The ground truth for the Vlasov equation was generated with a 4th order finite difference scheme in space and an RK4 time integration scheme. We sample  $10^{6}$  points over the full 2D space domain and a time step size of  $1\mathrm{e} - 3$  was used for time integration.
In Figure 11,12,13 we show plots of the ground truth solutions at the beginning, middle, and end of integration time for the equations we examine. We can see that they display the characteristics described in Section 4.1.
(c) Vlasov
Figure 7: Decay of singular values of  $J(\pmb{\theta}(t))$
Figure 8: Speedups of RSNG over dense updates with direct and iterative solver.
Figure 9: Speedups of RSNG scale quadratic with sparsity  $s$ .
Figure 10: Singular values of the Jacobian of the network fit to a Fokker-Planck solution in five dimension decays quickly too; providing indication that our RSNG approach is applicable in these settings as well.
Figure 11: True solution  $u(t,x)$  for Allen-Cahn
(a)  $t = 0$
(b)  $t = 2$
(c)  $t = 4$
Figure 12: True solution  $u(t,x)$  for Burgers'
Figure 13: True solution (top) vs RSNG solution (bottom) for Vlasov
A.5 Architecture and Hyperparameters
For an input  $x \in \mathbb{R}^d$  a periodic embedding layer with period  $P$  is defined with the optimization parameters  $a, \phi, b \in \mathbb{R}^d$  as,
$$
\operatorname {P e r i o d i c E m b e d} (x) = \sum_ {i = 1} ^ {d} \left[ a \cos \left(x \frac {2 \pi}{P} + \phi\right) + b \right] _ {i}.
$$
This operation is the repeated  $w$  times for different parameters  $a, \phi, b$  where  $w$  denotes the width of the layer resulting in a output vector  $y \in R^w$
To fit the initial condition, we minimize the  $L^2$  distance between the network and the initial condition as well as the  $L^2$  distance between the first derivative of the network and the first derivative of the initial condition with respect to the spatial domain  $\mathcal{X}$ . We evaluate the loss function over 10 000 equispaced points for Allen-Cahn and Burgers' equation and 200 000 points for Vlasov. We fit our initial conditions with two nonlinear solvers. First we run L-BFGS with 1000 iterations, then we use an Adam optimizer with the following hyperparameters,
- iterations: 100000
- learning rate: 1e-3
- scheduler: cosine decay
decay steps:500
The number of iterations and the number of points we sample are chosen to fit the initial condition to high accuracy to avoid polluting the results in our analysis with errors of fitting the initial condition.