An Optimisation Framework for the Well-Conditioned Training of Physics-Informed Neural Networks


Abstract

Physics-informed neural networks (PINNs) have emerged as a promising route to solve partial differential equations, yet they have struggled to reach the precision of classical solvers. The obstacle is increasingly understood to be one of optimisation, owing to the severely ill-conditioned loss landscape. We present DSGNAR: Doubly-Sketched Gauss–Newton with Adaptive Ratio, a scalable second-order optimisation framework that confronts this ill-conditioning and, in doing so, obtains unprecedented accuracy and speed. DSGNAR couples a doubly-sketched Gauss–Newton model with a novel strategy that carefully controls both regularisation and step length. Across a suite of problems spanning nonlinear, chaotic, multi-scale, high-dimensional, and Navier–Stokes, the framework greatly improves on the state of the art: able to attain relative \(\ell_2\) errors as low as \(3\times10^{-16}\) in double precision, improve contemporary results by five orders of magnitude on the canonical Burgers’ equation, and as much as eight orders on a high-dimensional Poisson problem, while remaining markedly faster. We further show that, in single precision, solutions at the limit of round-off error can be obtained very quickly: Burgers’ equation to \(\ell_2^{\text{rel}} = 4.75 \times 10^{-7}\) in under ten seconds. The framework is also robust to the choice of architecture, arithmetic precision, and initial hyperparameters.

The code is available at https://www.github.com/wephy/physics-informed-neural-networks

Physics-informed neural networks ,Gauss–Newton ,Conditioning ,Sketching ,SVD

1 Introduction↩︎

Physics-informed neural networks (PINNs) have become one of the most active paradigms in scientific machine learning, offering a promising route to solving partial differential equations (PDEs). In the few years since their modern revival, they have found use across many fields of science and engineering: biomedicine, chemical engineering, dynamical systems, finance, fluid and solid mechanics, and geophysics [1]. There are many motivations for PINNs, among them: (1) their mesh-free approach, which sidesteps mesh generation, irregular geometries, the stability restrictions of finite-difference and finite-element schemes, and the curse of dimensionality [2]; (2) their ease of use, since a problem can be solved purely by stating the governing equations, with little of the bespoke machinery a traditional solver demands; and (3) their natural data-assimilation, fusing sparse or noisy observations with the governing equations and making them well-suited to inverse problems where the physics is only partially known [3].

The idea of casting a system of physical laws as an unconstrained minimisation problem dates to [4] and [5]; later, with the advent of modern machine learning tools and ideas, it was refined and popularised by [6]. PINNs solve PDEs by approximating the solution with a neural network and iteratively evaluating the current violation of the governing equations at sampled points; using automatic differentiation for the required derivatives, steps are then taken to minimise these violations until the network satisfies the problem.

Figure 1: A representative high-accuracy solve with our framework. The viscous Burgers’ equation (9.1) solved in double precision with a SIREN of d_{\boldsymbol{\theta}}=11{,}285 trainable parameters and sketch size s=4{,}000. (a) The recovered solution u_{\boldsymbol{\theta}}(x,t), including the sharp internal layer that forms near x=0. (b) The training loss \mathcal{L}, driven below 10^{-29} in 331 iterations. (c) The relative \ell_2 error against the reference solution, reaching \ell_2^{\text{rel}}=7.97\times10^{-14} in 346.1 seconds on a single NVIDIA H100 GPU – several orders of magnitude beyond the best previously reported accuracy for this benchmark (2).

PINNs are often regarded as a leading candidate for the next generation of scientific computing, yet it is increasingly clear from the existing literature that they have yet to truly rival traditional numerical methods in many areas. The major hurdle has been the inherent ill-conditioning of PINNs, detailed in [7], alongside a growing consensus that the bottleneck lies in the optimiser [8]. PINNs differ considerably from traditional deep learning applications such as computer vision (CV) or natural language processing (NLP): their networks are typically far smaller – tens of thousands of parameters, as seen in 2, against the millions or billions used in CV, NLP, and large language models (LLMs) – yet, unlike these fields, PINNs demand high precision to be competitive with established classical solvers. It is worth noting that, although branded as machine learning, PINNs can exist independent of real data, relying only on their own residual evaluations; a consequence of this is that overfitting is not a genuine concern, and that losses at the limit of machine precision are theoretically possible, and indeed welcome. Standard first-order optimisers (e.g. [9]) therefore often prove insufficient for training PINNs [7], [10].

Current state-of-the-art methods typically employ second-order techniques that approximate the Hessian. One line of research enhances quasi-Newton methods (such as BFGS [11]) with additional scaling factors [12]; an effective hybrid strategy combines an initial phase of Adam with a later switch to a quasi-Newton variant, proving reliable across a broad spectrum of problems [8]. Other methods employ Kronecker-Factored Approximate Curvature (KFAC) [13], [14], building on the seminal work of [15], [16]. A third line, closest to ours, adopts a Gauss–Newton framework [17][19], carefully deriving matrix-free implementations to reduce the cost of direct Jacobian calculations; for the PINN least-squares objective, the energy natural gradient in fact coincides with a Gauss–Newton step [17], [19]. Tools from randomised linear algebra, sketching in particular [20], [21], are also applied to avoid ever forming the full Gauss–Newton matrix [7], [22], [23]. Beyond the optimiser itself, the choice of globalisation scheme – line search [24], trust-region [25], or Levenberg–Marquardt [26], [27] – significantly affects the results [8], as does the multi-scale nature of the PINN loss, whose dynamics, boundary, and initial-condition terms occupy disparate scales that must be balanced [28], [29].

The success of the Gauss–Newton method for nonlinear least squares and other supervised-learning contexts, alongside its relatively low-cost second-order approximation, motivates further exploration in the field of PINNs. We propose DSGNAR1 (/dɪˈzaɪnə/, like “designer”), a novel doubly-sketched Gauss–Newton method that alleviates the ill-conditioning of the PINN objective and permits aggressive sampling of points. It is further combined with a sophisticated regularisation scheme that aims to achieve richer exploration of the loss landscape. We summarise our contributions below.

Figure 2: image.

  • A conditioning-first principle for the step. We argue that a PINN’s accuracy ceiling is set by how well-conditioned the final region of the iterate is, tured by the effective regularisation – the amount of regularisation a Gauss–Newton step needs before it yields a decrease. We turn this into a concrete rule by letting both the regularisation and step length be implicit, fixing each step by comparing the objective against its Gauss–Newton model to hit a target decrease ratio. This target is chosen dynamically, promoting early exploration before shifting to aggressive descent once minimal regularisation is achieved (3.1).

  • A doubly-sketched Gauss–Newton model. We sketch the Jacobian on both sides, using CountSketch [30] on the rows and a subsampled randomised cosine transform [20], [21] on the columns. In contrast to the matrix-free line of work [17], [18], we build the sketched Jacobian explicitly, doing so quickly and memory-efficiently via the proposed methods (3.2.3). The resulting approximation is small, so its SVD yields inexpensive candidate steps.

  • Cheap, dense residual sampling. As the Jacobian sketch is accumulated, notably over low-memory batches, we are able to sample collocation points far more aggressively than is typically feasible (5). This sidesteps the residual-selection problem that is prevalent in the PINN literature [31].

  • Comprehensive, state-of-the-art results. We show that DSGNAR solves a large suite of PDE problems to precisions that rival classical solvers, in few iterations and little compute time. We provide extensive and highly detailed numerics, outperforming state-of-the-art PINNs by orders of magnitude across nonlinear, chaotic, multi-scale, high-dimensional, and Navier–Stokes problems, while remaining robust to architecture, precision, and initial hyperparameters. 1 previews a representative solve, and 2 presents the full set of results.

The rest of the paper is organised as follows. 2.1 formulates the PINN problem as a weighted sum of squares and introduces the running example. 2.2 recalls the Gauss–Newton method, along with the Levenberg–Marquardt and trust-region machinery we build on, situating our approach in the literature. 3 sets out the optimisation philosophy, how we build the doubly-sketched Jacobian, solve the subproblem through a single SVD, and select each step to obtain the target decrease ratio. 4 gives the implementation details, 5 reports results across the suite of problems in 2, and 6 concludes with discussion. The appendices collect the remaining algorithms, architectural details, problem definitions, and the full solution figures, showcasing all available metrics from training the PINNs.

2 Background↩︎

2.1 Problem setup: a physics-informed neural network↩︎

For our purposes, we consider an initial-boundary value problem on a space-time domain \(\Omega \times [0, T]\) as a triplet of operators \((\boldsymbol{\mathcal{P}}, \boldsymbol{\mathcal{I}}, \boldsymbol{\mathcal{B}})\) acting on a solution \(\boldsymbol{u}\) in a suitable function space \(\mathcal{V}\). For a system of \(L\) equations, the problem is expressed as: \[\begin{align} \boldsymbol{\mathcal{P}}\boldsymbol{u} &= \boldsymbol{0}, & (\boldsymbol{x}, t) &\in \Omega \times [0, T], \tag{1} \\ \boldsymbol{\mathcal{I}}\boldsymbol{u} &= \boldsymbol{0}, & (\boldsymbol{x}, t) &\in \Omega \times \{0\}, \tag{2} \\ \boldsymbol{\mathcal{B}}\boldsymbol{u} &= \boldsymbol{0}, & (\boldsymbol{x}, t) &\in \partial\Omega \times [0, T], \tag{3} \end{align}\] with \(\boldsymbol{u}: \Omega \times [0, T] \to \mathbb{R}^L\). By defining a composite operator \(\boldsymbol{\mathcal{F}}\): \[\boldsymbol{\mathcal{F}}\boldsymbol{u} = \begin{pmatrix}\boldsymbol{\mathcal{P}}\boldsymbol{u} \\ \boldsymbol{\mathcal{I}}\boldsymbol{u} \\ \boldsymbol{\mathcal{B}}\boldsymbol{u} \end{pmatrix},\] the solution \(\boldsymbol{u}\) satisfies the operator identity \(\boldsymbol{\mathcal{F}}\boldsymbol{u} = \boldsymbol{0}\), namely, \(\boldsymbol{u} \in \ker(\boldsymbol{\mathcal{P}}) \cap \ker(\boldsymbol{\mathcal{I}}) \cap \ker(\boldsymbol{\mathcal{B}})\). Physics-informed neural networks (PINNs) aim to approximate the solution with a neural network \(\boldsymbol{u}_{\boldsymbol{\theta}}\) by solving \(\boldsymbol{\mathcal{F}} \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}, t) = \boldsymbol{0}\), where \(\boldsymbol{\theta} \in \mathbb{R}^{d_{\boldsymbol{\theta}}}\) collects all of the network’s trainable parameters, flattened into a single vector of dimension \(d_{\boldsymbol{\theta}}\). These are precisely the quantities the optimiser updates, while the inputs \((\boldsymbol{x},t)\) and the architecture itself are held fixed. The explicit parameter set for each architecture used in this work is stated in 8. The training objective is2: \[\mathcal{L}(\boldsymbol{\theta}) := \dfrac{1}{2} \lVert \boldsymbol{\mathcal{F}}\boldsymbol{u}_{\boldsymbol{\theta}} \rVert ^2_{\boldsymbol{\mathcal{X}}, \boldsymbol{w}},\label{eq:simple95objective}\tag{4}\] with a weighted \(\ell^2\) norm given by: \[\lVert \boldsymbol{\mathcal{F}}\boldsymbol{u}_{\boldsymbol{\theta}} \rVert ^2_{\boldsymbol{\mathcal{X}}, \boldsymbol{w}} := \sum_{m=1}^{M} \frac{w_m}{|{\mathcal{X}}_m|} \sum_{(\boldsymbol{x}, t) \in \mathcal{X}_m} \big[ \boldsymbol{\mathcal{F}}_m\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}, t) \big]^2. \label{eq:objective}\tag{5}\] Here, \(M\) is the total number of conditions; each condition \(m\) has an associated weight \(w_m\) and a set of collocation points \(\mathcal{X}_m\) used to compute residuals, giving \(N=\sum_m |\mathcal{X}_m|\) collocation points in total. For a system of \(L\) equations, we may expect \(L\) conditions for each operator, but there may be more; for example, a single PDE that is second-order in time (e.g. the Wave equation) will require two initial conditions, \(\boldsymbol{\mathcal{I}}_1\) and \(\boldsymbol{\mathcal{I}}_2\). Fewer conditions are also possible: for incompressible Navier–Stokes, only pressure gradients appear, so a condition can be dropped (3); conditions can also be dropped when they are analytically enforced into the network architecture, referred to as hard constraints (17). This formulation extends naturally to other problem types, such as time-independent problems, with numerous examples shown later in this paper.

None

Figure 3: No caption.

2.2 Gauss–Newton methods and relevant safeguards↩︎

We first recall the Gauss–Newton model and the trust-region/Levenberg–Marquardt machinery, and then situate the method relative to the literature it both draws from and departs from. Gauss–Newton methods are employed to solve sum-of-squares objectives: \[\mathcal{L}(\boldsymbol{\theta}) = \frac{1}{2}\sum_{j=1}^{N}r_j^{2}(\boldsymbol{\theta}),\] where each \(r_j: \mathbb{R}^{d_{\boldsymbol{\theta}}} \to \mathbb{R}\) is referred to as a residual, whose value is determined by a given choice of the trainable parameters. We assemble these \(N\) residuals into a residual function, \(\boldsymbol{r}: \mathbb{R}^{d_{\boldsymbol{\theta}}} \to \mathbb{R}^N\), with its corresponding Jacobian \(\boldsymbol{J} \in \mathbb{R}^{N \times {d_{\boldsymbol{\theta}}}}\) defined as: \[\boldsymbol{r}(\boldsymbol{\theta}) = \begin{pmatrix} r_1(\boldsymbol{\theta}) \\ r_2(\boldsymbol{\theta}) \\ \vdots \\ r_N(\boldsymbol{\theta}) \end{pmatrix}, \qquad \boldsymbol{J}(\boldsymbol{\theta}) = \begin{bmatrix} {\nabla r_1(\boldsymbol{\theta})}^{\mathsf{T}} \\ {\nabla r_2(\boldsymbol{\theta})}^{\mathsf{T}} \\ \vdots \\ {\nabla r_N(\boldsymbol{\theta})}^{\mathsf{T}} \end{bmatrix}.\] The gradient and Hessian of \(\mathcal{L}\) are therefore written as: \[\nabla \mathcal{L}(\boldsymbol{\theta}) = {\boldsymbol{J}(\boldsymbol{\theta})}^{\mathsf{T}} \boldsymbol{r}(\boldsymbol{\theta}), \qquad \nabla^2 \mathcal{L}(\boldsymbol{\theta}) = {\boldsymbol{J}(\boldsymbol{\theta})}^{\mathsf{T}} \boldsymbol{J}(\boldsymbol{\theta}) + \sum_{j=1}^{N} r_j(\boldsymbol{\theta}) \nabla^2 r_j(\boldsymbol{\theta}). \label{eq:gradandhess}\tag{6}\] The Gauss–Newton approximation replaces the Hessian in 6 with \({\boldsymbol{J}(\boldsymbol{\theta})}^{\mathsf{T}} \boldsymbol{J}(\boldsymbol{\theta})\), discarding the second-order term \(\sum_j r_j \nabla^2 r_j\) and retaining only the first-order term, which is positive semi-definite by construction. For the PINN least-squares objective, this model is especially natural, and the energy natural gradient direction of [19] coincides with the parameter-space Gauss–Newton step when the same collocation points are used [17]. When using a weighted objective, such as the one defined in 5 , we consider the residual function to absorb the respective weights and normalisations for all of the conditions. However, in practice, it is more convenient to build components from pure, unweighted residuals, \(\boldsymbol{\mathcal{F}}_m \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i, t_i)\). The numerous conditions naturally lead to a structured block Jacobian, where each block \(\boldsymbol{J}_m \in \mathbb{R}^{|\mathcal{X}_m| \times d_{\boldsymbol{\theta}}}\) has entries: \[\left( \boldsymbol{J}_m \right)_{i, \cdot} = \sqrt{\frac{w_m}{|\mathcal{X}_m|}} \nabla_{\boldsymbol{\theta}} \left[ \boldsymbol{\mathcal{F}}_m \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i, t_i) \right].\] Similarly, the residual of the \(m\)-th condition, at the \(i\)-th collocation point \((\boldsymbol{x}_i, t_i) \in \mathcal{X}_m\), is given by: \[\left(\boldsymbol{r}_m\right)_{i} = \sqrt{\frac{w_m}{|\mathcal{X}_m|}} \boldsymbol{\mathcal{F}}_m \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i, t_i).\] The Levenberg–Marquardt method [26], [27] updates parameters \(\boldsymbol{\theta}\) using a Gauss–Newton approximation alongside Tikhonov regularisation [32], yielding an update step: \[\boldsymbol{p}_k^{\textsf{LM}}(\lambda) = -\left[{\boldsymbol{J}}^{\mathsf{T}}_k \boldsymbol{J}_k + \lambda \boldsymbol{I}\right]^{-1} {\boldsymbol{J}}^{\mathsf{T}}_k \boldsymbol{r}_k, \label{eq:lm}\tag{7}\] where \(\boldsymbol{r}_k = \boldsymbol{r}(\boldsymbol{\theta}_k)\) and \(\boldsymbol{J}_k = \boldsymbol{J}(\boldsymbol{\theta}_k)\). The regularisation, with an appropriate \(\lambda\) update strategy, ensures global convergence to a first-order stationary point from any starting point under standard assumptions [25]. In contrast, unregularised Gauss–Newton has no such global guarantee; locally it converges quadratically only in the zero-residual limit, degrading to linear convergence for small but non-zero residuals – at a rate that worsens as the residual at the solution grows – and may fail to converge altogether for large-residual problems [11]. Note that the step corresponds to the minimiser of a regularised quadratic model: \[m^\mathsf{LM}_k(\boldsymbol{p}; \lambda) := \frac{1}{2} \lVert \boldsymbol{r}_k + \boldsymbol{J}_k \boldsymbol{p} \rVert^2 + \frac{\lambda}{2} \lVert \boldsymbol{p} \rVert^2. \label{eq:LM95model}\tag{8}\] From the trust-region perspective, and by duality of the trust-region subproblem [25], the step \(\boldsymbol{p}_k\) is equivalently the minimiser of the unregularised quadratic model, \(m^\mathsf{Q}_k(\boldsymbol{p}) := \frac{1}{2} \lVert \boldsymbol{r}_k + \boldsymbol{J}_k \boldsymbol{p} \rVert^2\), subject to a constraint on the step length: \[\boldsymbol{p}_k^{\textsf{TR}}(\Delta) = \mathop{\mathrm{arg\,min}}_{\boldsymbol{p}} m^\mathsf{Q}_k(\boldsymbol{p}) \qquad \text{subject to}\;\lVert \boldsymbol{p} \rVert \le \Delta.\] The regularisation parameter \(\lambda\) here serves as the Lagrange multiplier for the constraint, establishing an implicit mapping \(\lambda(\Delta)\), or \(\Delta(\lambda)\) in Levenberg–Marquardt, such that \(\lambda \cdot \left(\Delta - \lVert \boldsymbol{p}_k \rVert\right) = 0\) with \(\lambda \geq 0\). Often, a scale-invariant metric for the quality of a local model is required. A popular choice is the decrease ratio, hereafter referred to simply as the ratio, considered in various nonlinear optimisation schemes [11], [25], obtained by comparing the actual reduction in the objective to the reduction predicted by the model: \[\frac{\mathcal{L}(\boldsymbol{\theta}_k) - \mathcal{L}(\boldsymbol{\theta}_k + \boldsymbol{p}_k)}{\mathcal{L}(\boldsymbol{\theta}_k) - m(\boldsymbol{p}_k)}. \label{eq:true95rho}\tag{9}\] A ratio \(\varrho_k \approx 1\) indicates that the model \(m\) is a good local proxy for the true objective, and \(\varrho_k \approx 0\) indicates that it predicted a significantly larger decrease than was actually achieved. A low decrease ratio can be interpreted as having ‘trusted’ the model over too large a region; in a trust-region method, we would subsequently reduce the radius to ensure sufficient decrease is obtained [25].

A direct Levenberg–Marquardt or trust-region solve is, however, out of reach for reasonably sized PINNs: a naive implementation of 7 costs \(\mathcal{O}(N d_{\boldsymbol{\theta}}^2 + d_{\boldsymbol{\theta}}^3)\), and even storing the Gauss–Newton matrix \({\boldsymbol{J}}^{\mathsf{T}}\boldsymbol{J}\) is infeasible at the problem sizes involved (3.2.3). Further enhancements of Gauss–Newton have been proposed using momentum and by solving a smaller linear system [18]. Our work differs from the literature in both our sketching methodology, which acts on both rows and columns (laid out in Sketch, 7), and a novel method that removes the need to hand-tune the regularisation \(\lambda\), deriving it instead via the principled approach detailed in 8. These methods require further considerations, for which we provide accommodations, such as the a tailored per iteration condition weight update in 16 (15). Two questions therefore remain, which the rest of the paper answers. How do we obtain a step cheaply, and in reasonable memory? We replace \(\boldsymbol{J}\) with a small, explicitly built sketch, so that a single SVD cheaply yields the Levenberg–Marquardt step for any desired \(\lambda\) – the subject of 3.2.3 And how do we choose among the many possible steps? This is governed by the philosophy we set out next, and realised by the algorithm of [sec:algorithm95overview]. Our trial step is itself a sketched Gauss–Newton step, and so sits within the line of randomised second-order methods initiated by the Newton Sketch of [33].

3 Method: the optimisation framework↩︎

3.1 Overview and key ideas↩︎

Our proposed method, DSGNAR: Doubly-Sketched Gauss–Newton with Adaptive Ratio, is built around a new optimisation philosophy. When using the Gauss–Newton approximation, one traditionally controls the step by choosing either the regularisation parameter \(\lambda\) within Levenberg–Marquardt, or the radius \(\Delta\) within trust-region. We argue that doing either is insufficient, and can in fact be detrimental in practice. The unintuitive point is that a very “successful” step, one that decreases the objective considerably (perhaps due to a good local model, with \(\varrho \approx 1\)), can directly lead to a poor solution, especially if taken too early, before the conditioning of the problem has been addressed. The loss landscape is complex, non-convex, and high-dimensional, and largely unknown to the optimiser; a large decrease in the objective might therefore seem beneficial, yet an optimal (or good) PINN solution may lie far away, and even the best nearby minimum may still be poor. Instead, we propose observing the current effective regularisation: what scale of regularisation is required for the associated step, from 7 , to yield a decrease at all. The regularisation plays a major role: it simultaneously controls the step length, the interplay between pure Gauss–Newton and gradient descent, and, crucially, the extent to which the minimised objective departs from the true PINN problem. When aiming to solve PDE problems to as high an accuracy as possible, ideally to the limits of machine precision, one cannot afford to solve a shifted objective, in which the fidelity of the underlying physics is compromised. After all, in a PINN problem the governing equations are known exactly, and zero residuals (or as close as the network’s expressivity allows) are achievable in principle. To accommodate this perspective, we propose moving away from controlling either \(\lambda\) or \(\Delta\) directly, and instead analysing a scale-independent quantity: the ratio \(\varrho\), defined in 9 . Using this ratio, we can effectively control each step: values in \((0,1)\) can be chosen to determine the step’s aggressiveness, agnostic to step length, regularisation, and complexities of the current landscape.

In summary, rather than chasing decreases in the objective directly, we navigate the loss landscape by: (1) taking large steps; (2) combating ill-conditioning by remaining in better-behaved regions; and (3) searching for a region capable of solving the true problem as closely as possible. Taking steps as large as possible corresponds to keeping the regularisation as small as possible (2.2). We aim to keep the effective regularisation as small as possible throughout the first stage of training; decreases in it indicate that the inherent conditioning of the problem is becoming less severe. This idea – to first seek out a region of parameter space where minimal regularisation is required, and only then descend – manifests as a two-stage optimisation method. The first stage adopts a small target ratio, \(\varrho^\textsf{Stage 1} \le 0.2\), which is conservative and promotes a decrease in the regularisation; the second stage adopts a larger target, \(\varrho^\textsf{Stage 2} \ge 0.5\), which is more aggressive and promotes a decrease in the objective. This leaves two tasks: building a cheap, explicit model for the step, since this philosophy requires repeated access to the SVD for many candidate \(\lambda\) (3.2.3); and finding, from that model, the step that achieves the desired target ratio ([sec:algorithm95overview]).

We combine the solutions to these tasks, together with supporting routines, into a single algorithm, DSGNAR (4). Every iteration, the sketch operators are formed, and subsequently used in Sketch in 7, producing the sketched Jacobian \(\tilde{\boldsymbol{J}}\), sketched residual \(\tilde{\boldsymbol{r}}\), and full residual \(\boldsymbol{r}\) (3.2). The algorithm then takes the SVD of \(\tilde{\boldsymbol{J}}\) and, from it, computes the step whose ratio matches the current target \(\varrho^\star\) () via LambdaSolve. The step is evaluated in the full parameter space and accepted only if it yields a decrease, with the trust-region radius updated accordingly. Training halts once that radius falls below a tolerance \(\Delta_{\min}\). Two quantities are then refreshed before the next iteration: the condition weights, via UpdateWeights (3.4.0.3), which keep the residuals of different conditions on a similar scale; and the target ratio, via UpdateTargetRatio (3.4.0.4), which advances from its first- to second-stage value once the regularisation has bottomed out. We frame DSGNAR as a trust-region method, despite its non-standard step selection, so as to inherit the associated convergence machinery: a step is taken on success (sufficient decrease); otherwise, the radius centring the next iteration’s probes is contracted.

None

Figure 4: No caption.

3.2 The doubly-sketched Gauss–Newton model↩︎

3.2.0.1 Novelty

Prior second-order methods for PINNs are largely matrix-free, applying the Gauss–Newton operator only through matrix–vector products [17]. Where such methods do randomise, the sketch acts on the matrix \(\boldsymbol{J}\boldsymbol{J}^\top+\lambda\boldsymbol{I}\), as in the randomised Nyström method of [18]. We instead form a Jacobian sketch explicitly, never storing it in full (though we can, in fact, build it fully at no extra cost if it fits in memory). To sketch the Jacobian, we apply compression on both sides: the residual dimension and the parameter dimension. The sketch is accumulated over successive row batches (and per problem condition, see 2.1), which removes memory constraints and makes dense sampling possible, further detailed in 3.2.1.

The residual vector \(\boldsymbol{r} \in \mathbb{R}^{N}\) and its Jacobian \(\boldsymbol{J} \in \mathbb{R}^{N \times d_{\boldsymbol{\theta}}}\) are the central objects in DSGNAR: once obtained, they determine the quadratic model in 8 and the Levenberg–Marquardt step in 7 , for a chosen \(\lambda\). For the problem sizes of interest both \(N\) and \(d_{\boldsymbol{\theta}}\) can be large, so even storing the Gauss–Newton matrix \({\boldsymbol{J}}^{\mathsf{T}}\boldsymbol{J}\) is infeasible. We therefore reduce \(\boldsymbol{J}\) to a small square sketch \(\tilde{\boldsymbol{J}} \in \mathbb{R}^{s \times s}\) of dimension \(s\). A square shape is chosen for the fastest possible SVD, with added stability gained from aggregating \(N\) residuals down to \(s \ll N\), which suppresses noise and reflects the over-determination of the system. Solving through the SVD sidesteps any concerns particular to square systems. Empirically, we observed that taller sketches provided no accuracy benefit, only extra computational cost. The construction of this sketch remains cheap, even under exhaustive dense sampling of residuals (5).

3.2.0.2 Choosing the sketch size

The sketch dimension \(s\) is the primary control for this framework, trading the fidelity of the Gauss–Newton model against the cost of the SVD. As a robust default we recommend taking \(s\) between a third and a half of the parameter count, \(s \in [\,\lfloor d_{\boldsymbol{\theta}}/3\rfloor,\;\lfloor d_{\boldsymbol{\theta}}/2 \rfloor \,]\). Intuitively, \(s\) should comfortably exceed the numerical rank of \(\boldsymbol{J}\), so that the dominant singular values and right singular vectors are faithfully preserved. Values much smaller risk discarding curvature directions, while values approaching \(d_{\boldsymbol{\theta}}\) forfeit the speed advantage of sketching for little accuracy gain. These design choices lead to a scalable framework, as we can avoid any memory issues associated with the sizes of \(N\) and \(d_{\boldsymbol{\theta}}\), where all costly mechanisms scale with the sketch dimension \(s\), which can be chosen independently. Further, experiments in this paper highlight that the proposed framework lends itself to few, highly effective iterations; thus, even with a very large sketch dimension, and a consequentially expensive SVD, these methods may still be preferred over other optimisers. In 5 the computation of full and sketched Jacobians with various methods is explored, highlighting that our proposed method allows large problems to be tackled, where other implementations completely fail.

Figure 5: image.

3.2.1 CountSketch↩︎

The step \(\boldsymbol{p}_k\) lives in parameter space, so the quality of a sketch \(\tilde{\boldsymbol{J}}\) is judged solely by how faithfully it preserves the singular values and right singular vectors of \(\boldsymbol{J}\), which also live in \(\mathbb{R}^{d_{\boldsymbol{\theta}}}\) and together determine the step (3.3). The left singular vectors live in the residual space \(\mathbb{R}^{N}\) and carry no geometric meaning for the step, and no quantity is ever lifted back to \(\mathbb{R}^{N}\): the row dimension may therefore be compressed aggressively, by a cheap, non-invertible operation that need not preserve geometry exactly. Further, we wish to have the property: \[\label{eq:batchable} \boldsymbol{C}\boldsymbol{J} = \sum_{\text{batches}} \boldsymbol{C}\boldsymbol{J}_{\text{batch}},\tag{10}\] which is the linearity of the sketch \(\boldsymbol{C}\). A suitable choice is then CountSketch [30], defined by \(K\) independent hash functions \(h_k: \{1,\dots,N\} \to \{1,\dots,s\}\), and sign vectors \(\boldsymbol{\varepsilon}_k \in \{\pm1\}^{N}\), drawn uniformly at random. The CountSketch matrix \(\boldsymbol{C} \in \mathbb{R}^{s \times N}\) is: \[\boldsymbol{C}_{ij} = \dfrac{1}{\sqrt{K}} \displaystyle\sum_{k=1}^{K} (\boldsymbol{\varepsilon}_k)_j\cdot\mathbb{1}[h_k(j) = i].\] This has the effect that each row \(j\) of \(\boldsymbol{J}\) is multiplied by \((\boldsymbol{\varepsilon}_k)_j\), and output to row \(h_k(j)\). To reduce variance, and following theoretical development [34], [35], this process happens \(K\) times for each row of \(\boldsymbol{J}\). This has negligible overhead, as \(\boldsymbol{C}\) is extremely sparse, containing only \(K\) non-zeros per column; as such, \(\boldsymbol{C}\) belongs to the OSNAP family of sparse oblivious subspace embeddings [36], with the single-hash case \(K=1\) recovering the original CountSketch [30], [37].3 In 7, we detail the process by which the sketched Jacobian is constructed. There, the application of \(\boldsymbol{C}\) occurs on a column-sketched sub-batch Jacobian \(\hat{\boldsymbol{J}}_{\text{batch}}\) of size \(b \times s\) (defined formally in 3.2.3), where \(b\) is the batch size (the number of collocation points evaluated). This incurs a negligible cost of \(\mathcal{O}(K \cdot s \cdot b )\) per batch, compared to \(\mathcal{O}(s^2 \cdot b)\) for a dense embedding. Note that \(K=2\) or \(K=4\) is typical.

3.2.1.1 A novel consequence

This particular use of CountSketch provides desirable properties for PINN problems that, to our knowledge, have not previously been explored in the literature. With a fixed sketch size \(s\), the expensive mechanisms in the optimiser are fixed, leaving us free to sample as many residuals as we like. Owing to this low cost, we can assemble residuals from very dense sampling into the \(s\) buckets. This causes important, or troublesome, regions of the PDE to naturally dominate their buckets, as these have a higher associated value \(\boldsymbol{\mathcal{F}}_m \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i, t_i)\). The sheer number of residuals aggregated into each bucket also averages out noise from the system, yielding a better-conditioned problem to solve. This effectively removes the usual PINN problem of where best to place (a typically small set of) collocation points.

3.2.2 SRCT↩︎

The residual dimension (3.2.1) is compressed once and never reconstructed. The parameter dimension is different: the step is solved in a sketched space and must be lifted back through the embedding to update \(\boldsymbol{\theta}\), so its compression must instead be a near-isometry, so that this lift is faithful. To reduce the dimension \(d_{\boldsymbol{\theta}}\), we therefore use a random subspace embedding, compressing columns of the Jacobian and thus alleviating the computational cost of Gauss–Newton steps. Historically, Gaussian embeddings were favoured, as they come with strong theoretical guarantees [20]. In our work, we instead use subsampled randomised trigonometric transforms (SRTTs), further detailed in [20]. These have been observed to match the practical performance of Gaussian embeddings, while additionally admitting a Fast Fourier Transform implementation, as pioneered in SRFT [38]; this gives an improved computational complexity over Gaussian embeddings, \(\mathcal{O}(d_{\boldsymbol{\theta}} \log d_{\boldsymbol{\theta}})\) in our case. As a subspace embedding, an SRTT preserves all singular values of a fixed \(d\)-dimensional subspace to a desired relative accuracy, provided enough columns are retained [20].

SRTTs are composed of a random sign flip \(\boldsymbol{D} \in \mathbb{R}^{d_{\boldsymbol{\theta}}\times d_{\boldsymbol{\theta}}}\), a random permutation \(\boldsymbol{\Pi} \in \mathbb{R}^{d_{\boldsymbol{\theta}}\times d_{\boldsymbol{\theta}}}\), a unitary trigonometric transform \(\boldsymbol{F} \in \mathbb{R}^{d_{\boldsymbol{\theta}}\times d_{\boldsymbol{\theta}}}\) – applied analogously to the FFT – and finally, a random restriction \(\boldsymbol{S} \in \mathbb{R}^{d_{\boldsymbol{\theta}}\times s}\), giving the embedding \[\boldsymbol{\Omega}\,\boldsymbol{S} = (\boldsymbol{D}\, \boldsymbol{\Pi}\, \boldsymbol{F})\, \boldsymbol{S}, \qquad \boldsymbol{\Omega} := \boldsymbol{D}\, \boldsymbol{\Pi}\, \boldsymbol{F}.\]

As a product of orthogonal factors, \(\boldsymbol{\Omega} \in \mathbb{R}^{d_{\boldsymbol{\theta}}\times d_{\boldsymbol{\theta}}}\) is itself orthogonal (unitary). We choose the subsampled randomised cosine transform (SRCT) for \(\boldsymbol{F}\), as we deal with real-valued matrices, opting for the discrete cosine type-II variant. Together with the construction of \(\boldsymbol{C}\) in 3.2.1, the initialisation of \(\boldsymbol{\Omega}\) and \(\boldsymbol{S}\) constitutes InitSketch in 4, meaning a single, consistent sketch is used throughout every iteration. This is primarily due to compiler constraints (4.1), though we also observed empirically that the best results were obtained with a consistent sketch size.

Figure 6: Comparison of the full Jacobian (a) with the corresponding doubly-sketched Jacobian (b), using CountSketch and SRCT, for an early iteration solving Burgers’ equation. The full Jacobian has N_\mathcal{P} + N_\mathcal{I} = 2^{15} + 2^{14} residuals, and 1447 trainable parameters. The doubly-sketched Jacobian is a square of size s=700, with clear isotropic structure.

3.2.3 Assembling the doubly-sketched Jacobian↩︎

We now walk through Sketch, shown in 7, which forms the doubly-sketched Jacobian \(\tilde{\boldsymbol{J}}\), the sketched residual \(\tilde{\boldsymbol{r}}\), and the full residual \(\boldsymbol{r}\), without ever forming \(\boldsymbol{J}\) in full. 6 provides an illustrative example of how the sketched Jacobian compares to its full counterpart. The algorithm begins by looping over the \(M\) conditions of the problem – the true residual function \(\boldsymbol{r}\), which maps over all conditions, is functionally complex: if implemented directly, it would require switch cases depending on the location of the collocation point or the equation being considered, each leading to potentially intensive differential operators. Thus, computing the Jacobian from the true residual function would produce a needlessly large computational graph. However, thanks to CountSketch, and particularly its linearity property 10 , we can instead deal with the Jacobians for each condition, \(\{ \boldsymbol{J}_m\}_{m=1}^{M}\), individually; then, we can simply sum the sketches of these Jacobians together. Each of these Jacobians incurs a scalar multiplication of \(\alpha_m = \sqrt{ w_m / |\mathcal{X}_m|}\), applying the normalisation and condition weight of 5 . As this scaling is applied at accumulation, the routine (featuring all automatic differentiation parts) is agnostic to the current weights \(\boldsymbol{w}\), and so requires no change between iterations. Each condition is processed in macro-batches of size \(b\), so that even the block \(\boldsymbol{J}_m\) need never reside in memory in full. Within each macro-batch, residuals and their Jacobians are computed in sub-batches of size \(b'\) by reverse-mode automatic differentiation; as \(b' \ll d_{\boldsymbol{\theta}}\), this is far cheaper than forward-mode (5), and, being independent, the sub-batches are dispatched in parallel. Each sub-batch Jacobian is column-sketched by the SRCT, \(\hat{\boldsymbol{J}}_{\text{batch}} = \boldsymbol{J}_{\text{batch}}\,\boldsymbol{\Omega}\,\boldsymbol{S}\), before the row-compression \(\boldsymbol{C}\) is applied and accumulated. The full residual \(\boldsymbol{r}\) is assembled alongside, uncompressed, since the decrease ratio is evaluated in the full space (3.3).

None

Figure 7: No caption.

3.3 Solving the subproblem using SVD↩︎

A single SVD of the small square \(s \times s\) sketched Jacobian yields inexpensive Levenberg–Marquardt steps, predicted reductions, and the secular map \(\lambda\!\leftrightarrow\!\Delta\). This allows one to obtain the steps associated with a set of candidate (or probe) trust-region radii \(\{\delta_i\}_i\) at little added cost. An implicit, Jacobian-vector-product solver cannot deliver this factorisation, and would instead need to re-solve from scratch for each predetermined \(\lambda\) or \(\Delta\).

Sketching the row dimension of the Jacobian yields the sketched residuals \(\tilde{\boldsymbol{r}}\), while sketching the column dimension means the resulting step \(\tilde{\boldsymbol{p}}\) lives in a sketched space and must be lifted back to the full parameter space for application. Our smaller, inexact model in the sketched space is written as:

\[\tilde{m}^\mathsf{LM}_k(\boldsymbol{\tilde{p}}; \lambda) := \frac{1}{2} \lVert \tilde{\boldsymbol{r}}_k + \tilde{\boldsymbol{J}}_k \tilde{\boldsymbol{p}} \rVert^2 + \frac{\lambda}{2} \lVert \tilde{\boldsymbol{p}} \rVert^2, \label{eq:LM95model95approximated}\tag{11}\] with \(\tilde{\boldsymbol{p}} \in \mathbb{R}^s\), \(\tilde{\boldsymbol{J}}_k \in \mathbb{R}^{s \times s}\) and \(s \ll N\). The analogous trust-region subproblem, with \(\tilde{m}^\mathsf{Q}_k(\tilde{\boldsymbol{p}}) := \frac{1}{2} \lVert \tilde{\boldsymbol{r}}_k + \tilde{\boldsymbol{J}}_k \tilde{\boldsymbol{p}} \rVert^2\), is then: \[\min_{\tilde{\boldsymbol{p}}} \;\tilde{m}^\mathsf{Q}_k(\tilde{\boldsymbol{p}}) \qquad \text{subject to} \;\lVert \tilde{\boldsymbol{p}} \rVert \le \Delta. \label{eq:tr95subproblem95sketched}\tag{12}\] The use of sketching – namely, the avoidance of ever building the true Gauss–Newton model \(m_k^{\textsf{Q}}\) – means we can never evaluate the ratio in 9 purely from the step \(\boldsymbol{p}_k\) in full space. We instead compute it approximately, based on the sketched model \(\tilde{m}_k^{\textsf{Q}}\), as: \[\varrho_k(\tilde{\boldsymbol{p}}_k, \boldsymbol{p}_k; \boldsymbol{\theta}_k) = \frac{\mathcal{L}(\boldsymbol{\theta}_k) - \mathcal{L}(\boldsymbol{\theta}_k + \boldsymbol{p}_k)}{\tilde{m}^\mathsf{Q}_k(0) - \tilde{m}^\mathsf{Q}_k(\tilde{\boldsymbol{p}}_k)}. \label{eq:gain95ratio} \tag{13}\] Here, the denominator is the predicted decrease in the sketched space, while the numerator is the true reduction in the objective; together these well-approximate the true ratio of 9 .

We solve 11 using factorisation-based methods, since 3.4 requires repeatedly re-solving for many candidate \(\lambda\) and \(\Delta\) from the same sketched Jacobian: a single factorisation reuses itself across this whole array of probes, whereas an iterative solver would need a fresh solve for each. This is made affordable by working in the reduced space of dimension \(s\), which also lets us sidestep the cost of Jacobian-vector products and the need for a preconditioner that a Krylov-type method would otherwise require. Given the thin SVD, \(\tilde{\boldsymbol{J}}_k=\boldsymbol{U}_k\boldsymbol{\Sigma}_k{\boldsymbol{V}_k}^{\mathsf{T}}\), the Levenberg–Marquardt step has the closed-form expression: \[\tilde{\boldsymbol{p}}_k^{\textsf{LM}}(\lambda) = -\boldsymbol{V}_k\,\mathrm{diag}\dfrac{(\sigma_k)_i}{(\sigma_k)_i^2+\lambda}\,{\boldsymbol{U}_k}^{\mathsf{T}} \tilde{\boldsymbol{r}}_k, \label{eq:computesolLMproj}\tag{14}\] where \((\sigma_k)_i\) is the \(i\)-th singular value from \(\boldsymbol{\Sigma}_k\). This is the method used to obtain steps in DSGNAR, and is numerically preferable to solving 7 directly, as it never forms \({\boldsymbol{\tilde{J}}}^{\mathsf{T}}\tilde{\boldsymbol{J}}\), which squares the condition number. This makes iterative methods (Krylov, conjugate gradient) on the already ill-conditioned problem even more prohibitive without a well-chosen preconditioner, discussed in [25, Ch. 5]. Crucially, once the SVD is computed, evaluating the step for a choice of \(\lambda\) is extremely cheap – costing only \(\mathcal{O}(s^2)\) in matrix multiplications – and, in particular, steps for an array of \(\lambda\) choices can be computed quickly in parallel. Further, given a trust-region radius \(\Delta\), the corresponding regularisation \(\lambda\) such that \(\lVert \tilde{\boldsymbol{p}}_k^{\textsf{LM}}(\lambda) \rVert = \Delta\) can be obtained efficiently from the SVD via the secular equation: \[\phi(\lambda) := \sqrt{\sum_{i=1}^{s} \frac{(\sigma_k)_i^2\,({\boldsymbol{U}_k}^{\mathsf{T}}\tilde{\boldsymbol{r}}_k)_i^2}{\big((\sigma_k)_i^2 + \lambda\big)^2}} - \Delta = 0, \label{eq:secular}\tag{15}\] iterating Newton’s method as in [25]. It is this ability to re-solve for many \(\lambda\) and \(\Delta\) from a single SVD that the algorithm of [sec:algorithm95overview] exploits, probing many candidate radii at once to find a step with a target ratio. Note that even disregarding the conditioning benefits of the SVD, implicit access to \({\boldsymbol{J}}^{\mathsf{T}}\boldsymbol{J}\) via Jacobian-vector products could not yield such a factorisation.

3.4 Finding a step with a target ratio \(\varrho\)↩︎

None

Figure 8: No caption.

3.4.0.1 Novelty

Rather than devising an update strategy for the regularisation \(\lambda\) or radius \(\Delta\) – which would be blind to the current landscape – we instead select each step according to the target decrease ratio \(\varrho^\star\). This yields a landscape-aware rule for the regularisation, one that does not depend on the history of previous iterations.

Figure 9: The selection of optimisation step. The model \Delta(\varrho) is built from probe triplets (\Delta_i, \lambda_i, \varrho_i) which are geometrically spaced around the trust-region radius \Delta_k. The next iteration’s \Delta_{k+1} is selected by interpolating the model at the target ratio \varrho. The regularisation corresponding to the target ratio \varrho is chosen as \lambda_k which determines the optimisation step, as in 14 (which is subsequently lifted to the full parameter-space). This visual corresponds to the first optimisation stage, whose goal is to decrease the regularisation over a decrease in the objective, ideally achieving \lambda_k \ll \lambda_{k-1}.

As the ratio \(\varrho\) is only known once a step has been computed, it cannot be chosen directly beforehand. Traditionally, therefore, it has only been used post hoc: to decide whether to increase or decrease the trust-region radius for the next iteration. To instead find a step that achieves some target ratio \(\varrho\), probing is required: by interpolating across several candidate steps, one very near to the target can be found.

Since evaluating probes is cheap in DSGNAR (3.3), we can compute a range of them and read off the radius that attains \(\varrho^\star\) by interpolation, illustrated in 9. The process is detailed in LambdaSolve (8), where a set of geometrically spaced points around the current trust-region radius is considered. For each trust-region radius probe, we use the already obtained SVD to cheaply compute the corresponding regularisation \(\lambda\) (3.3). Using \(\lambda\), the projected step \(\tilde{\boldsymbol{p}}\) is then computed from 14 and lifted to the full-space step \(\boldsymbol{p}\). We then use 13 to obtain each probe’s achieved ratio \(\varrho_i\).

The resulting triples \((\Delta_i, \lambda_i, \varrho_i)\) across all probes allow a smooth model to be built (9). We construct this model using a Piecewise Cubic Hermite Interpolating Polynomial (PCHIP) [39]. To handle noise, particularly once close to the solution, we force monotonicity in the fit; while not necessarily injective, this still permits straightforward interpolation.

This model allows us to find the trust-region radius and regularisation that approximately achieve the target ratio \(\varrho^\star\). If the model does not cover \(\varrho^\star\), rather than extrapolate, we select a new set of geometrically spaced probes in the direction of \(\varrho^\star\) and repeat. Due to how this is implemented in efficient compiled code, it is impractical to scan indefinitely until the step corresponding to \(\Delta^\star\) is found – at least, not within the compiled machine-learning frameworks (e.g. JAX [40]) we use for implementation. Hence, if \(\varrho^\star\) is not found within the probe range, we instead test the step associated with the nearest extreme probe. As in 4, every such step is still evaluated in full space before acceptance. If the extreme probe with radius \(\Delta^\star = \Delta_k / 3\) also fails to yield a decrease, we use \(\Delta^\star / 3\) as the centre for the next iteration’s probes, covering a region below and overlapping the previous one. We centre probes around \(\Delta_k\), rather than around the regularisation parameter \(\lambda\), because the trust-region radius is far more stable across iterations (aided by the isotropic parameter embedding of SRCT), whereas \(\lambda\) is not. This is illustrated in 22: for iterations 0–100, the trust-region radius remains consistently close to \(1\), while over the same period the regularisation changes by more than 25 orders of magnitude. Because of this stability, the full range of values of \(\varrho\) (from less than zero to close to one) is typically obtainable using probes geometrically spaced within \(\bigl[\tfrac{1}{3}\Delta_k,\;3\Delta_k\bigr]\). Depending on the architecture, this range can be widened further – to, say, a factor of \(5\) in each direction – and, depending on GPU budget and the memory cost of a loss evaluation, one may use fewer or more probes.

3.4.0.2 Accepting the step and updating the radius

Each candidate step is evaluated in the full parameter space before it is taken. A step is accepted when it yields a smaller objective and a positive ratio. In keeping with the philosophy of 3.1, a step whose achieved ratio overshoots the target by more than a small margin is also rejected: an over-aggressive step signals that the trust region is too small, so the radius is expanded, making the subsequent step less damped and giving finer control over the next iteration’s probes. On a normal accept, the radius selected by the probe interpolation of 3.4 is carried forward to centre the next iteration’s probes. On a rejection for insufficient decrease, the radius is shrunk. Training stops once the radius falls below the convergence tolerance \(\Delta_{\min}\).

3.4.0.3 Automatic reweighting

A further issue with the objective in 5 is that the condition weights, \((w_1, \dots , w_M)\), which balance the contributions of each condition (e.g. dynamics \(\mathcal{P}\), initial condition \(\mathcal{I}\), boundary conditions \(\mathcal{B}\)), are fixed. In machine learning and PINNs alike, these hyperparameters are notoriously difficult to tune to obtain satisfactory solutions. Each operator \((\boldsymbol{\mathcal{P}}, \boldsymbol{\mathcal{I}}, \boldsymbol{\mathcal{B}})\) has its own scale with complex interactions [7], which has motivated a wide array of proposed dynamic and adaptive weighting schemes [28], [41][43]. By allowing weights to vary at each iteration \(k\), the objective becomes: \[\mathcal{L}_k(\boldsymbol{\theta}) := \dfrac{1}{2} \lVert \boldsymbol{\mathcal{F}}\boldsymbol{u}_{\boldsymbol{\theta}} \rVert ^2_{\boldsymbol{\mathcal{X}}, \boldsymbol{w}_k} = \dfrac{1}{2} \sum_{m=1}^{M} \frac{(\boldsymbol{w}_k)_m}{|\mathcal{X}_m|} \sum_{(\boldsymbol{x}, t) \in \mathcal{X}_m} \big[ \boldsymbol{\mathcal{F}}_m\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}, t) \big]^2.\label{eq:variation95objective}\tag{16}\] We provide our own adaptive weighting method, UpdateWeights (15). This is residual-based, and so similar in spirit to [29], [42]. We do not claim it is intrinsically better than existing work, but the particular use of CountSketch in this paper requires accounting for interactions that do not arise in standard PINN implementations. A second-order method should, in principle, adequately resolve residuals at different scales; however, because CountSketch aggregates residuals from every condition into a shared set of \(s\) buckets (3.2.1), it is essential that all conditions \(m=1,\dots,M\) provide values on a comparable scale. The performance of UpdateWeights, combined with its light computational footprint, makes it well suited to our task.

3.4.0.4 Automatic stage switching

The transition from the first (regularisation-decreasing) stage to the second (objective-descent) stage is made automatically by UpdateTargetRatio (16). It monitors the recent history of the regularisation \(\lambda\): once a log-linear fit over a sliding window shows a sustained upward trend – evidence that \(\lambda\) has passed its minimum and begun to climb – the target ratio is raised from its small first-stage value \(\varrho^{\textsf{Stage 1}}\le0.2\) to its larger second-stage value \(\varrho^{\textsf{Stage 2}}\ge0.5\). At this point, the optimiser shifts from seeking a well-conditioned region to descending towards a nearby minimum. The full procedure, with its window length and trend thresholds, is given in 7.

4 Implementation↩︎

In this section, we discuss the practical application of DSGNAR (4) to minimising the objective of 16 .

4.1 Environment, machine learning framework and compiler↩︎

All experiments in this paper use the machine learning framework JAX [40] for automatic differentiation. JAX is built on the compiler XLA (Accelerated Linear Algebra) [44], which performs just-in-time (JIT) compilation. In practice, this means there is an upfront cost to compiling JAX routines, which optimises operations for their specific sizes. Varying the number of residuals per condition across iterations, or the network architecture across time intervals (4.4), would therefore incur further compilation on each change. This would also complicate any implementation of adaptive sketch sizes within this compiled setting, and would require further consideration. We note that the recorded \(t_\text{wall}\) times for obtaining solutions omit this compile time, particularly as the compiled routines can be cached to disk.

All experiments are run with JAX version 0.9.1, on Ubuntu version 24.04.4 LTS (Noble Numbat). Initial development was carried out on an NVIDIA RTX 2070 SUPER with 8 GB of GDDR6 memory; later development, and all experiments presented here, used an NVIDIA H100 NVL with 96 GB of HBM3 memory, CUDA version 13.0.

4.2 Training configuration↩︎

When applying DSGNAR (4) to a PDE problem, several initial decisions must be made. For the set of collocation points \(\mathcal{X}\), we opt for simple sampling distributions. As detailed in 2.1 and highlighted in 5, unlike most PINN formulations [31], we are not constrained to a small number of residuals. The residual counts typical of other works remain feasible, and are not essential to the results we observe; however, as laid out in 3.2.3, using more residuals brings clear advantages, such as reducing noise and better capturing important regions. We choose \(N^\mathcal{P} = |\mathcal{X}_\mathcal{P}| = 2^{15}\) and \(N^\mathcal{I}=N^\mathcal{B}=2^{14}\) in all reported experiments. Smaller numbers may be preferred for GPUs with less memory, or for higher-dimensional problems (where more data is required per point), but any memory issues can be avoided by controlling the batch size. We always sample \(\mathcal{X}_\mathcal{P}\) uniformly; depending on the dimension of the problem, \(\mathcal{X}_\mathcal{I}\) and \(\mathcal{X}_\mathcal{B}\) may use uniform or linear sampling. Each experiment figure explicitly details the sampling method used, chosen primarily for convenience. The only deviation is the lid-driven cavity problem of 9.7, which we found to particularly benefit from additional density near the lid: \(\mathcal{X}_\mathcal{P}\) is a concatenation of uniform sampling over the full domain and additional uniform sampling over the top ten percent of the domain nearest the lid, with the boundary points sampled likewise.

When comparing our approximate solutions against either exact solutions or results from traditional methods (when no exact solution is available), we evaluate on uniform grids over the relevant space(-time) domain, distinct from the collocation points \(\mathcal{X}\) used during training. For high-dimensional problems, such as the 10D Poisson, this is not feasible, and we instead opt for a very dense uniform sample (not to be confused with the training collocation points). Because PINNs do not suffer from overfitting in the conventional sense (1), keeping evaluation points separate from collocation points is not a genuine concern here; nonetheless, we maintain the separation throughout.

The trainable parameters \(\boldsymbol{\theta}\) are architecture-dependent, and we opt for standard initialisation methods for each \(\boldsymbol{\theta}_0\), detailed in 8. For the initial trust-region radius \(\Delta_0\), since LambdaSolve searches for the step matching \(\varrho^\star\) regardless of the starting point, this choice has little impact on the final result; we nonetheless choose \(\Delta_0 = 1\), as this is often close to \(\Delta^\star\) for the first iteration across most architectures. The optimiser halts once the trust-region radius falls below the convergence tolerance \(\Delta_{\min}\). For single precision we set \(\Delta_\text{min} = 10^{-4}\), and for double precision we set \(\Delta_\text{min} = 10^{-8}\). The sketch size is problem- and architecture-dependent, and chosen roughly between \(\lfloor d_{\boldsymbol{\theta}}/3 \rfloor\) and \(\lfloor d_{\boldsymbol{\theta}}/2 \rfloor\).

One important decision, however, is the initial target ratio \(\varrho_0\). If adopting the two-stage approach we recommend (3.1), this will be \(\varrho^{\textsf{Stage 1}}\), used for all first-stage iterations. Values around \(0.1\) are often ideal: we choose \(\varrho^{\textsf{Stage 1}} = 0.1\) for single-precision training and \(\varrho^{\textsf{Stage 1}} = 0.075\) for double-precision training, where a lower value corresponds to slower training but potentially a higher-quality solution. For the second stage we conservatively set \(\varrho^{\textsf{Stage 2}} = 0.5\); for faster, more aggressive, though potentially more volatile performance, these values can instead be raised to \((\varrho^{\textsf{Stage 1}}, \varrho^{\textsf{Stage 2}}) = (0.2, 0.8)\). We use \(24\) probes to build the model \(\Delta(\varrho)\) (3.4). As explained in 3.1, the target ratio \(\varrho^\star\) strongly shapes the optimisation trajectory: a higher target forces smaller, more conservative steps that nonetheless appear, at least initially, to yield a larger decrease in the objective per iteration; on this highly non-convex landscape, this can substantially change the final minimiser reached – exactly the effect warned against in 3.1. We quantify this dependence with a dedicated ablation in 10.

For choices of initial weights \(\boldsymbol{w}_0\), UpdateWeights (15) alleviates this sensitivity, which is often problematic in PINNs. This is especially highlighted in 23, panel (g) (labelled ‘Weights’), which shows the PDE-condition weight initialised much lower, at \(10^{-4}\), than its desired value of \(10^{-1}\), yet quickly rising and plateauing at the target. As the optimisation problem associated with the boundary and initial conditions is much simpler, we initialise these weights to \(1\), and tend to set the PDE condition, or any other problematic condition, to something small (e.g. \(10^{-4}\)); these problematic conditions are then gently incorporated over iterations. UpdateWeights contains one tunable hyperparameter, \(\alpha\), which controls how quickly the residuals of different conditions are equalised. We set \(\alpha=0.05\) for all experiments, which we found sufficient to equalise residuals adequately given the number of iterations the optimiser typically requires.

Figure 10: Ablation study on the choice of target ratio \varrho, with a comparison against a two-stage approach, when solving Burgers’ equation in double precision. We show that smaller choices of the target ratio \varrho provide smaller values of regularisation \lambda, and consequently higher-quality solutions. This performance comes at the cost of more iterations; however, the two-stage approach shows that if \varrho switches to a high value once the minimum \lambda is achieved, fewer iterations are required for the same performance.

4.2.0.1 Hybrid precision

Single precision requires no special implementation considerations; double precision, however, can (rarely) experience early divergence. We observed that, for poor initialisations of \(\boldsymbol{\theta}\), large residuals combined with the optimiser’s ability to solve to the level of noise can lead to early failure. Empirically, this was entirely resolved by carrying out a few initial iterations of the exact same DSGNAR procedure in single precision. We therefore run the first ten iterations of every double-precision experiment in single precision, to correct any initial problematic elements of \(\boldsymbol{\theta}\) and allow drastically mis-scaled weights to align first.

4.3 Architectures↩︎

The literature contains a great deal of recent work developing architectures tailored to PINNs, or even to specific PDE problems. DSGNAR is architecture-independent, and achieves superior performance over the state of the art even when paired with a simple MLP. However, we observe that some architectures pair better with DSGNAR: the SIREN architecture (8.4), in particular, provides considerably more stable trust-region radii and regularisation. We also observed, especially on higher-dimensional problems, that better results are possible with alternative architectures.

For most of the numerical results, a SIREN with four layers and a mildly varying width – \(40\) to \(60\) on double-precision problems, to accommodate the added complexity of the PDE – is employed, to highlight that DSGNAR is not sensitive, and that the same machinery can be used to tackle most PDE problems with great success. For a select few problems, we show performance with radically different architectures, such as a GaborNet or SPINN, to further display DSGNAR’s flexibility. Many other architectures, such as LSTMs, were deployed in testing, all yielding strong results. Specific architecture definitions can be found in 8. Each network is parameterised by trainable weights \(\boldsymbol{\theta}\), accepts a single coordinate vector \(\boldsymbol{x} \in \mathbb{R}^{n_\text{in}}\), and returns an output \(\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}) \in \mathbb{R}^{n_\text{out}}\), where \(n_\text{in}\) and \(n_\text{out}\) are determined by the PDE problem. Each architecture is optionally paired with an input embedding (8.1), applied before the main network body.

4.4 Time marching↩︎

This paper considers the standard 4 , which imparts no inherent sense of time dependence. Instead, it relies on the PDE problem being well-posed, such that violations of the time-dependent dynamics result in a higher loss than a close approximate solution would achieve. This can be viewed as treating the time input variable identically to the spatial inputs: the entire space-time domain is solved in unison, no different to time-independent problems. However, in some time-dependent settings – such as large time horizons \(T\), or PDEs with stiff dynamics – solving the whole domain in unison can fail. Time marching addresses this by decomposing the domain into sequential time intervals. As discussed in [41], a PINN can instead be trained on many shorter time intervals \([t_n, t_{n+1}]\). At the end of each interval, the network’s predicted solution is saved and used as the initial condition for the next interval, \([t_{n+1}, t_{n+2}]\). This process repeats, marching the solution forward in time, and thereby leverages warm starts of the parameters \(\boldsymbol{\theta}\). In our results, we use a slightly perturbed version of the previous interval’s trainable parameters \(\boldsymbol{\theta}\), both to aid faster convergence and to avoid possible divergence from a poor initialisation.

5 Experiments↩︎

We now evaluate the proposed DSGNAR framework across a diverse set of PDE problems, chosen to test the optimisation philosophy of 3 under markedly different challenges. These include the canonical viscous Burgers’ equation (9.1), a standard benchmark for PINN optimisers featuring a sharp shock; the fourth-order Kuramoto–Sivashinsky equation (9.4), whose chaotic dynamics are challenging even for traditional numerical methods; the ten-dimensional Poisson equation ([app:10dpoisson95eq]), representative of problems beyond the reach of conventional grid-based solvers; and the steady lid-driven cavity Navier–Stokes problem (9.7), a coupled constrained system modelling incompressible fluid flow. Together these problems span many of the principal difficulties encountered in PINNs, namely sharp solution features, long-time dynamics, high dimensionality, and coupled physics systems. Beyond these four representative examples, we report additional benchmarks – the Wave equation (9.3), the Korteweg–de Vries equation (9.2), the multi-scale Poisson problem (9.5), and a five-dimensional Poisson equation ([app:5dpoisson95eq]) – whose complete figures and quantitative results are collected in 10.

Figure 11: No caption

5.1 Evaluation metrics and common optimisation behaviour↩︎

A summary of every solve, in both single and double precision where applicable, is provided in 2, reporting the relative \(\ell_2\) error together with the wall-clock training time and comparisons against the current state of the art. To assess solution quality we use the relative error, \[\ell_2^{\text{rel}} := \frac{\lVert \boldsymbol{u}_{\boldsymbol{\theta}} - \boldsymbol{u}^\star \rVert_2}{\lVert \boldsymbol{u}^\star \rVert_2}, \label{eq:rel95l295error}\tag{17}\] evaluated over the space(-time) domain using points disjoint from the collocation set employed during training. Where practical these evaluation points form a dense uniform grid; otherwise, for higher-dimensional problems, they are taken as a dense uniform sample (4). The reference solution \(\boldsymbol{u}^\star\) is either the exact analytic solution or, where unavailable, a high-fidelity numerical solution generated using a traditional solver. Details for each benchmark are given in 9. For systems with multiple outputs, such as the Navier–Stokes equations, we report \(\ell_2^{\text{rel}}\) for each solution component.

Each benchmark is accompanied by a full-page figure, presented either in this section or in 10, with the numerical precision indicated in the upper-right corner. The upper panels visualise the obtained solution, its absolute error relative to the reference solution, and the evolution of \(\ell_2^{\text{rel}}\) throughout optimisation. For time-dependent problems we additionally report the error as a function of PDE time, \(\ell_2^{\text{rel}}(t)\); for steady problems we instead examine the spatial distribution of the error.

The lower panels report optimisation quantities that are common to every benchmark. Specifically, we show: the objective \(\mathcal{L}\) (16 ); the relative error; the adaptive condition weights updated by UpdateWeights (15); the achieved decrease ratio \(\varrho\); the effective regularisation \(\lambda\) associated with each accepted step; and the trust-region radius \(\Delta\).

These optimisation metrics exhibit remarkably consistent behaviour across all problems. During the first stage of optimisation, the effective regularisation \(\lambda\) decreases steadily over many orders of magnitude while the achieved ratio remains close to its conservative target, indicating that the optimiser is primarily seeking a well-conditioned region of parameter space rather than aggressively reducing the objective (3.1). Once \(\lambda\) reaches its minimum, the target ratio is increased automatically, producing the second stage of optimisation in which the loss and solution error decrease rapidly. Throughout this process, the trust-region radius \(\Delta\) remains comparatively stable – typically varying by less than one order of magnitude while \(\lambda\) changes by twenty or more. This marked difference motivates centring the probe search around \(\Delta\), rather than \(\lambda\), as described in 3.4. Meanwhile, the adaptive weights converge from their initial values to scales that balance the residuals of the different PDE conditions. We therefore use these common optimisation signatures as a template throughout the remainder of this section, highlighting only those behaviours that are specific to each PDE.

5.2 Burgers’ equation↩︎

Burgers’ equation (9.1) remains the canonical benchmark for PINN optimisation The double-precision solution is shown in 11. Panel (a) recovers the solution throughout the space-time domain, including the narrow internal layer that forms near \(x=0\) for viscosity \(\nu=0.01\). Panel (b) shows the absolute error against the reference solution, which remains below \(10^{-13}\) over almost the entire domain. Importantly, whereas PINN solutions typically exhibit their largest errors along the developing shock, our solution displays some of its smallest errors there. This suggests that the dense residual sampling made possible by CountSketch allows the optimiser to resolve precisely the region that is usually the most difficult for PINNs.

The relative error in panel (c) sees a small decrease during the first stage before dropping by approximately nine orders of magnitude during the second stage, ultimately reaching \(\ell_2^{\mathrm{rel}} = 7.97\times10^{-14}\) in 346.1 seconds. This behaviour closely follows the optimisation strategy described in 3.1: the first stage primarily seeks a well-conditioned region of parameter space by driving the regularisation down, while the second stage exploits that conditioning to obtain rapid objective reduction.

The optimisation metrics follow the common pattern described previously. The regularisation \(\lambda\) (i) decreases to approximately \(10^{-30}\) before increasing once the optimiser switches stages, whereas the trust-region radius \(\Delta\) (j) remains close to its initial value throughout, only dropping once progress is not possible. Likewise, the loss (e) undergoes its sharpest reduction only after the stage transition, matched closely by the relative error (f), down to the limits of double precision. Approximately 260 iterations are spent reducing the regularisation before only around 70 further iterations are required to converge, illustrating that the majority of the optimisation effort is devoted to locating a highly well-conditioned region rather than directly minimising the objective. The corresponding single-precision solve (19) exhibits the same optimisation behaviour with a smaller network, and is discussed further in 5.8.

5.3 Kuramoto–Sivashinsky equation↩︎

Figure 12: No caption
Figure 13: No caption

The Kuramoto–Sivashinsky equation (9.4) is a challenging PINN benchmark [8], where competing second- and fourth-order terms produce chaotic spatio-temporal dynamics. We solve it in double precision over sixteen time intervals using time marching (4.4), requiring a total of 4300 optimisation iterations.

In 12, panel (a) shows the recovered chaotic solution, with the corresponding absolute error in panel (b). Due to the chaotic nature of the dynamics, small errors are amplified over time, leading to the growth in relative error visible in panel (d). Nevertheless, the solution remains orders of magnitude more accurate than previous PINN results, achieving \(\ell_2^{\mathrm{rel}} = 5.12\times10^{-7}\) in 5226.6 seconds.

The bottom panels show the first and final time intervals, each following the same two-stage optimisation behaviour described previously. As the time march progresses, accumulated approximation error produces progressively worse initialisations for later intervals. This is reflected in the optimisation metrics: the minimum regularisation increases over time, indicating that later intervals require stronger damping, as a slightly incorrect initial condition leads to worse data; meanwhile, the loss (e) can still decrease without a corresponding reduction in relative error (f). The condition weights (g) likewise adapt to a different balance between residual terms as the solution evolves through time.

5.4 High-dimensional Poisson↩︎

The mesh-free nature of PINNs is particularly valuable in high dimensions, where traditional grid-based methods become impractical due to the curse of dimensionality. We study the ten-dimensional Poisson problem ([app:10dpoisson95eq]), using a separable SPINN architecture (8.6) to represent the high-dimensional solution efficiently. 13 shows a two-dimensional slice of the recovered solution in panel (a), with all remaining coordinates fixed. The corresponding absolute error in panel (b) remains uniformly small throughout the domain, with the largest errors occurring near the boundaries as shown in panel (d). The error distribution in panel (e), taken from samples over the entire domain, shows that the absolute error remains below \(10^{-10}\) everywhere. The final solution achieves a relative error of \(\ell_2^{\mathrm{rel}} = 5.96\times10^{-12}\), improving over previous baselines by more than eight orders of magnitude (2). Despite using a SPINN 8.6, substantially different architecture to the SIREN-based experiments, the optimisation behaviour remains consistent with the common pattern. The regularisation decreases over the first stage before increasing during the final descent phase, while the trust-region radius remains comparatively stable. This demonstrates that the conditioning mechanism underlying DSGNAR is not tied to a particular network architecture, but persists even when the parameterisation is substantially changed.

5.5 Lid-driven cavity (Navier–Stokes)↩︎

Fluid dynamics provides a more demanding test of the framework, as the solution is governed by a coupled system with multiple outputs and constraints. We study the steady lid-driven cavity problem (9.7), a classical benchmark for incompressible Navier–Stokes solvers, in which a moving top wall drives the flow inside a square cavity and produces a primary vortex together with a hierarchy of smaller corner eddies. We solve the equations at Reynolds number \(\mathrm{Re}=100\), giving a three-output system consisting of the two velocity components and pressure.

14 shows the recovered velocity field in (a), with the streamlines capturing the primary vortex and secondary corner structures present in the high-fidelity reference solution. The error distributions in (d–f) show that the remaining error is concentrated around the regions of strongest gradients, particularly in the corners. The solution achieves a relative error of \(1.13\times10^{-4}\) for the velocity field, obtained in 402.9 seconds.

The optimisation behaviour follows the same general pattern observed throughout the paper. However, unlike other problems where \(\lambda\) can often be driven down to the limits of the used precision, the minimum regularisation reached here is considerably higher. This highlights the additional difficulty of the coupled Navier–Stokes system: the standard least-squares PINN objective must simultaneously balance momentum conservation, incompressibility, and boundary conditions, making it harder to identify a region where all residual components can be reduced to the same extent. The trust-region radius remains comparatively stable throughout, reinforcing that the limiting factor is not the step-selection strategy, but the underlying optimisation landscape induced by the PINN formulation itself.

Figure 14: No caption

5.6 Further benchmarks↩︎

The remaining problems from 2, with full figures collected in 10, further demonstrate the flexibility of the framework across different PDE problems. We first consider a second Burgers’ equation solve (20) with a different initial condition, where the moving shock is accurately recovered. The Wave equation is linear and second-order in time, and provides a test of a known weakness of the standard 4 : that time is not explicitly imparted into the optimisation process. The Wave equation solutions in particular show an extremely rapid decrease in the regularisation, while the trust-region radius remains consistent. This highlights the advantage of selecting steps through the target ratio \(\varrho^\star\), as this conditioning change would be difficult to follow with traditional strategies that update \(\lambda\) based only on parameter history. For this problem, we obtain solutions at the numerical precision limit in both single- and double-precision. The Korteweg–de Vries equation, a dispersive third-order PDE, is solved over four time intervals in . Each interval follows the same optimisation behaviour observed throughout the paper, with the first stage reducing the regularisation before transitioning into objective descent. As the solution is propagated forward in time, the minimum achievable regularisation increases and the final error grows, reflecting the increasing difficulty of the later time intervals as they are initialised from approximate rather than exact solutions. We investigate the framework on multi-scale problems in , where Fourier features are leveraged to provide the networks ability to resolve high varied frequencies. In both single- and double-precision we obtain errors close to the numerical limit, and in 27 the intricate multi-scale structure is recovered across the domain. Finally, 28 considers the radically different GaborNet architecture (8.5). This architecture produces a more rapidly changing optimisation landscape and therefore a less constant trust-region radius than the previous architectures. As the radius moves outside the range covered by the current probes, additional iterations are required to locate steps close to the target ratio – grey crosses in (i) indicate rejected achieved ratios. Nevertheless, the framework remains effective: only 100 iterations are required to reduce the regularisation to \(10^{-35}\) and obtain a solution with \(\ell_2^\text{rel}=3.03\times10^{-16}\), in part due to its smoothness – the remaining absolute error is dominated by floating-point effects, which become clearly visible in (b).

5.7 Computational cost and scalability↩︎

A key result illustrating the computational trade-off of DSGNAR is the lid-driven cavity problem in 2. Here, we solve the same Navier–Stokes system using two sketch sizes, \(s=1{,}000\) and \(s=4{,}000\), with the network containing \(9{,}219\) parameters. The larger sketch captures the local model more accurately and therefore approaches the highest-quality solution possible with our framework. The smaller sketch represents a more aggressive setting, sacrificing some accuracy for a substantial reduction in computational cost, achieving a relative error of \(6.34\times10^{-4}\) compared to \(1.13\times10^{-4}\), while reducing the solve time from \(402.9\) seconds to \(80.4\) seconds. This highlights the main computational advantage of the sketched formulation: the sketch size provides a direct control over the balance between accuracy and cost. In problems where the full accuracy of the model is unnecessary, smaller sketches can provide large reductions in runtime while still maintaining accuracy competitive with, or exceeding, existing approaches. As the experiments show errors substantially below many reported PINN results, aggressive sketching provides a practical route towards very fast solves.

5.8 Single vs. double precision↩︎

Double precision is often considered necessary for high-quality PINN solutions, and many recent state-of-the-art results are therefore reported exclusively in double precision. This comes with increased memory requirements and computational cost. Our results demonstrate that this trade-off is not always required. Across four problems in 2, single-precision training achieves relative errors close to the numerical limit of the chosen precision, while requiring substantially less time. For example, the single-precision Burgers’ solve (19) reaches \(\ell_2^{\text{rel}} = 4.75\times10^{-7}\) in \(9.8\) seconds, using only slightly more than \(100\) iterations of 4. This improves on the best comparable results in 2 by approximately two orders of magnitude in error while requiring fewer iterations and less computation.

These results suggest that precision should be treated as a practical choice rather than a strict requirement. For applications where \(10^{-6}\)\(10^{-7}\) relative accuracy is sufficient, single precision provides a substantially cheaper operating point, which has previously been impossible to achieve for PINNs. The effectiveness of DSGNAR in single precision is highlighted in 2: on the KdV problem (9.2), we obtain a relative error of \(9.55 \times 10^{-7}\), comfortably outperforming the best reported double-precision result in the literature.

6 Conclusions and perspectives↩︎

In this paper we introduced DSGNAR, Doubly-Sketched Gauss–Newton with Adaptive Ratio, an optimisation framework built around a conditioning-first view of PINN training. Rather than chasing the largest immediate decrease in the objective, the method first seeks a region of parameter space where minimal regularisation suffices, and only then descends aggressively once that region is found. This is realised through two coupled ingredients. Firstly, a doubly-sketched Gauss–Newton model compresses both the residual dimension, via CountSketch, and the parameter dimension, via a subsampled randomised cosine transform. This creates a small square matrix whose SVD can be computed cheaply and reused to evaluate many candidate steps in parallel (3.2). Secondly, rather than a typical strategy to determine the regularisation \(\lambda\) or trust-region radius \(\Delta\), DSGNAR selects each step by probing for a target decrease ratio \(\varrho^\star\). This ratio is itself adaptive, raised automatically once minimal regularisation has been achieved, producing the two-stage optimisation behaviour observed throughout 5.

Across a suite of nonlinear, chaotic, multi-scale, high-dimensional, and Navier–Stokes problems (2), this combination delivers accuracy far beyond previously reported PINN results while remaining fast. We see a five orders of magnitude improvement over the state of the art on the canonical Burgers’ equation, and as much as eight orders of magnitude on a high-dimensional Poisson problem, with relative \(\ell_2\) errors as low as \(3\times10^{-16}\) in double precision. The same machinery is effective in single precision, reaching \(\ell_2^{\text{rel}}=4.75\times10^{-7}\) on Burgers’ equation in under ten seconds (5.8), suggesting that cheaper single-precision PINNs can now provide highly accurate solves. These results hold with a fixed, small set of hyperparameters across markedly different architectures and problem structures, indicating that DSGNAR is robust.

Open questions remain. Theoretically, while CountSketch and the SRCT column embedding are each classical subspace embeddings with well-understood guarantees ([sec:countsketch] [sec:srct]), these guarantees do not automatically transfer to the doubly-sketched method obtained by composing them; a dedicated subspace-embedding analysis of this combined construction would be desirable. A rigorous analysis of step selection targeting a particular ratio \(\varrho^\star\) is left for future work. When training PINNs, linear convergence is rarely, if ever, obtained beyond the initial iterations, with most reported methods instead exhibiting sublinear convergence [7], [8]; our results (1, 5) instead show linear convergence throughout the run, until high-accuracy termination is achieved; theoretical understanding of these improved rates is left for future work. The coupled Navier–Stokes system of 5.5 also illustrates that not every problem admits regularisation as small as other problems, consequently leading to a relatively poorer solution – the bottleneck instead appears to now sit in the underlying PINN objective itself. Since all experiments here use only the basic 4 , it is therefore possible that alternative PINN formulations, such as FBPINNs [45], [46], could push accuracy further still when combined with DSGNAR.

Finally, integrating DSGNAR into an existing scientific machine learning library [47] would ease adoption by the numerical PDE community and allow more direct comparison against classical solvers on problems of practical interest.

Acknowledgements↩︎

Joseph Webb’s work was supported by NAG Ltd. and a EPSRC Case Award, as well as a Worcester College scholarship (Paul Thornton) and an MMSC MSc course bursary (Mathematical Institute, University of Oxford). Sadok Jerad and Coralia Cartis acknowledge the support of the Hong Kong Innovation and Technology Commission (InnoHK Project CIMDA). Coralia Cartis’ work was also supported by the EPSRC grant EP/Y028872/1, Mathematical Foundations of Intelligence: An “Erlangen Programme” for AI. The authors gratefully thank Prof Endre Sülli for useful discussions and suggestions regarding the numerical experiments.

7 Further algorithms↩︎

7.1 Residual-based condition re-weighting↩︎

15 presents our condition weight update strategy for \(\boldsymbol{w}_k\), referenced from 3.4.0.3. Its purpose is to keep the \(M\) conditions of the problem (dynamics \(\mathcal{P}\), initial conditions \(\mathcal{I}\), boundary conditions \(\mathcal{B}\)) on a comparable residual scale at every iteration, which matters specifically because of how DSGNAR builds its sketch: CountSketch (3.2.1) aggregates residuals from all conditions into the same \(s\) buckets.

The algorithm first computes the per-condition loss \(\ell_m = \|(\boldsymbol{r}_k)_m\|^2\) for each of the \(M\) conditions, and their mean \(\bar\ell_k\) across conditions. Each weight is then rescaled multiplicatively by the ratio of the mean loss to its own condition’s loss, raised to a tunable exponent \(\alpha\): a condition with an above-average loss (\(\ell_m > \bar\ell_k\)) has its weight decreased, and one with a below-average loss has its weight increased, pulling every condition’s contribution back towards the mean. The exponent \(\alpha\) controls how aggressively this correction is applied: \(\alpha = 0\) recovers no re-weighting at all, while larger \(\alpha\) more strongly equalises the per-condition losses. We use \(\alpha = 0.05\) throughout. A small floor \(\varepsilon = 10^{-8}\) guards this ratio against division by zero.

Finally, the whole weight vector is renormalised so that its mean is exactly \(1\). Without this step, the multiplicative updates would allow the overall magnitude of \(\boldsymbol{w}_k\) to drift across iterations.

None

Figure 15: No caption.

7.2 When to increase the target ratio and enter the second optimisation stage↩︎

16 presents our method for determining when to trigger the second stage of the optimiser, referenced from 3.4.0.4. Recall from 3.1 that the first stage deliberately drives the regularisation \(\lambda\) down, seeking a well-conditioned region of parameter space, and that the switch to the second (objective-descent) stage should occur once this pursuit has run its course, namely, once \(\lambda\) stops decreasing and begins to climb again.

The algorithm works with a sliding window of the \(W\) most recent regularisation values, first mapped to a log scale, \(\ell_j = \log(\max(\lambda_{k-W+j}, \varepsilon))\), because \(\lambda\) evolves over many orders of magnitude (3.4). A line is fit to this window by least squares, giving a slope \(\hat{s}\), and the Pearson correlation coefficient \(\hat{c}\). The two act as complementary checks: \(\hat{s}\) measures whether \(\log\lambda\) is climbing on average, while \(\hat{c}\) measures how consistently it is doing so. The stage switch is only triggered when both a sustained upward slope (\(\hat{s} > \tau_s\)) and a strong, consistent trend (\(\hat{c} > \tau_c\)) are observed simultaneously; we use \(W = 30\), \(\tau_s = 10^{-4}\), and \(\tau_c = 0.1\) throughout.

None

Figure 16: No caption.

8 Architectures and features↩︎

All networks in this work share a common calling convention. For a problem posed on a spatial domain of dimension \(n_\text{sp}\), a time-dependent network accepts a space–time coordinate \((\boldsymbol{x}, t) \in \mathbb{R}^{n_\text{sp}} \times [0, T]\), while a steady (time-independent) network accepts a spatial coordinate \(\boldsymbol{x} \in \mathbb{R}^{n_\text{sp}}\) alone. Writing \(n_\text{in}\) for the total input dimension, we have \(n_\text{in} = n_\text{sp} + 1\) for time-dependent problems and \(n_\text{in} = n_\text{sp}\) for steady ones. Each network is parameterised by trainable weights \(\boldsymbol{\theta}\) and returns an output \(\boldsymbol{u}_{\boldsymbol{\theta}} \in \mathbb{R}^{n_\text{out}}\), where \(n_\text{out}\) is fixed by the problem.

Time is always treated independently of space, such that we can apply spatial input embeddings \(\boldsymbol{\phi} : \mathbb{R}^{n_\text{sp}} \to \mathbb{R}^{d_\phi}\) , leaving the temporal coordinate untouched. The vector passed to the network is therefore \[\boldsymbol{e} = \begin{cases} \bigl(\boldsymbol{\phi}(\boldsymbol{x}),\, t\bigr) \in \mathbb{R}^{d_\phi + 1}, & \text{time-dependent},\\[4pt] \boldsymbol{\phi}(\boldsymbol{x}) \in \mathbb{R}^{d_\phi}, & \text{steady}, \end{cases} \label{eq:embedded95input}\tag{18}\] whose dimension we denote \(d_e\). Keeping \(t\) separate mirrors its role in the PDE formulation of 2.1, where temporal and spatial operators act differently, and primarily avoids imposing spatial feature maps (such as periodic embeddings) on a direction where they are inappropriate. The choice of embedding \(\boldsymbol{\phi}\) – the identity, a periodic mapping, or Fourier features – is described in 8.1.

8.1 Input embeddings↩︎

The simplest spatial embedding is the identity, \(\boldsymbol{\phi}(\boldsymbol{x}) = \boldsymbol{x}\), in which the raw spatial coordinates are passed straight to the network body. Two richer embeddings are used in this work, both acting per spatial axis after the coordinate is normalised to the unit interval using the per-axis domain bounds, \[\tilde{x}_j = \frac{x_j - (x_{\min})_j}{(x_{\max})_j - (x_{\min})_j} \in [0, 1], \qquad j = 1, \dots, n_\text{sp}. \label{eq:coord95norm}\tag{19}\] In every case the temporal coordinate, when present, is appended to the embedded spatial vector without transformation, in keeping with 18 .

8.1.0.1 Periodic embedding

To impose exact periodicity in space, each normalised coordinate is mapped to a cosine–sine pair, \[\boldsymbol{\phi}(\boldsymbol{x}) = \bigl( \cos(2\pi \tilde{x}_j),\;\sin(2\pi \tilde{x}_j) \bigr)_{j=1}^{n_\text{sp}} \in \mathbb{R}^{2 n_\text{sp}}, \label{eq:periodic95embed}\tag{20}\] so that the network – and all of its spatial derivatives – are periodic across the domain by construction. This is the embedding underlying the hard-constraint treatment of 8.2.

8.1.0.2 Fourier features

More generally, \(M\) Fourier features per spatial axis are formed from a chosen set of integer modes \(\{m_1, \dots, m_M\}\), \[\boldsymbol{\phi}(\boldsymbol{x}) = \Bigl( \cos(2\pi m_k \tilde{x}_j),\;\sin(2\pi m_k \tilde{x}_j) \Bigr)_{\substack{j=1,\dots,n_\text{sp} \\ k=1,\dots,M}} \in \mathbb{R}^{2 M n_\text{sp}}, \label{eq:fourier95features}\tag{21}\] which recovers the periodic embedding 20 when \(M = 1\) and \(m_1 = 1\). Unlike random Fourier features [48], the modes here are fixed and deterministic; the network learns only how to combine them through its body. The choice of modes can vary. One can choose (linear) modes \(m_k = k\), giving \(\{1, 2, \dots, M\}\), which populate the low-frequency bands. For problems exhibiting features across disparate length scales, it may be favourable to use geometric (dyadic) modes \(m_k = 2^{k-1}\), giving \(\{1, 2, 4, \dots, 2^{M-1}\}\), which span a wide range of spatial frequencies with only \(M\) features per axis.

8.2 Hard constraints in PINNs↩︎

A limitation of training PINNs with the loss of 4 is that boundary conditions are only softly enforced, so the learned solution may not adequately satisfy them; choosing (or adaptively setting) the associated weights can be challenging, leading to poor convergence and accuracy [49]. Following [5], one can instead approximate the solution by a network that enforces boundary conditions analytically – a strategy known as hard constraints. Periodic boundaries are naturally handled by the periodic spatial embedding \(\boldsymbol{\phi}\) of 8.1 [50], as illustrated in 17.

None

Figure 17: No caption.

None

Figure 18: No caption.

8.3 MLP↩︎

The standard multi-layer perceptron (MLP) is a fully connected feed-forward network with \(H\) hidden layers of widths \(d_1, \dots, d_H\), using the hyperbolic tangent \(\tanh\) as its activation function. Writing \(\boldsymbol{h}^{(0)} = \boldsymbol{e}\) for the embedded input of 18 (with \(\boldsymbol{\phi}\) the identity when no spatial embedding is used), the forward pass reads: \[\boldsymbol{h}^{(\ell)} = \tanh\!\left( a_\ell\!\left( \boldsymbol{W}^{(\ell)}\boldsymbol{h}^{(\ell-1)} + \boldsymbol{b}^{(\ell)} \right) \right), \qquad \ell = 1, \dots, H,\] where \(\boldsymbol{W}^{(\ell)} \in \mathbb{R}^{d_\ell \times d_{\ell-1}}\), \(\boldsymbol{b}^{(\ell)} \in \mathbb{R}^{d_\ell}\), and \(a_\ell \in \mathbb{R}\) is a trainable per-layer activation scale, initialised to \(1\). The output is a final affine map onto the \(L\) solution components, \(\boldsymbol{u}_{\boldsymbol{\theta}} = \boldsymbol{W}^{(H+1)}\boldsymbol{h}^{(H)} + \boldsymbol{b}^{(H+1)} \in \mathbb{R}^{L}\), with \(\boldsymbol{W}^{(H+1)} \in \mathbb{R}^{L \times d_H}\). Weights are initialised orthogonally with a \(\tanh\) gain of \(5/3\); output weights use LeCun initialisation [51].

8.3.0.1 Trainable parameters

\[\boldsymbol{\theta} = \{\boldsymbol{W}^{(\ell)}, \boldsymbol{b}^{(\ell)}\}_{\ell=1}^{H+1} \cup \{a_\ell\}_{\ell=1}^{H},\] flattened into a single vector.

8.4 SIREN↩︎

The sinusoidal representation network (SIREN), introduced by [52], replaces the standard nonlinearity with a scaled sine activation. Writing \(\boldsymbol{h}^{(0)} = \boldsymbol{e}\) for the embedded input of 18 , each of the \(H\) hidden layers takes the form: \[\boldsymbol{h}^{(\ell)} = \sin\!\left( \omega_\ell \left( \boldsymbol{W}^{(\ell)}\boldsymbol{h}^{(\ell-1)} + \boldsymbol{b}^{(\ell)} \right) \right), \qquad \ell = 1, \dots, H,\] where \(\omega_\ell > 0\) is a per-layer frequency parameter. A principled initialisation scheme is given in [52]: first-layer weights are drawn from \(\mathcal{U}\!\left[-1/d_0,\, 1/d_0\right]\), and hidden-layer weights from \(\mathcal{U}\!\left[-\sqrt{6/d_{\ell-1}}/\omega_\ell,\, \sqrt{6/d_{\ell-1}}/\omega_\ell\right]\), preserving the distribution of activations across layers. The output is a final affine map onto the \(L\) solution components, identical to the MLP above.

In our implementation each \(\omega_\ell\) is trained as \(\exp(\log \omega_\ell)\), enforcing positivity while allowing the network to tune its own frequency.

8.4.0.1 Trainable parameters

\[\boldsymbol{\theta} = \{\boldsymbol{W}^{(\ell)}, \boldsymbol{b}^{(\ell)}\}_{\ell=1}^{H+1} \cup \{\log\omega_\ell\}_{\ell=1}^{H},\] flattened into a single vector (each \(\omega_\ell>0\) is learned via its log). SIRENs are well suited to PDE problems because sinusoidal activations naturally preserve the smoothness and high-frequency structure of solutions and their derivatives, which \(\tanh\) networks can struggle to represent [52].

8.5 GaborNet↩︎

GaborNet is a multiplicative filter network [53] in which the feature layer is formed from learnable Gabor wavelets rather than sinusoids or standard activations. A Gabor filter is a sinusoid modulated by a Gaussian envelope, \[g(\boldsymbol{e};\, \boldsymbol{\mu},\, \boldsymbol{\omega},\, \sigma) = \exp\!\left(-\dfrac{\lVert \boldsymbol{e} - \boldsymbol{\mu} \rVert^2}{2\sigma^2}\right) \cos(\boldsymbol{\omega}^\mathsf{T}\boldsymbol{e} + \varphi),\] acting on the embedded input \(\boldsymbol{e} \in \mathbb{R}^{d_e}\) of 18 , with trainable filter centres \(\boldsymbol{\mu} \in \mathbb{R}^{d_e}\), frequency vectors \(\boldsymbol{\omega} \in \mathbb{R}^{d_e}\), bandwidth \(\sigma > 0\), and phase \(\varphi \in \mathbb{R}\). Given \(F\) filters, each parameterised by \((\boldsymbol{\mu}_f, \boldsymbol{\omega}_f, \sigma_f, \varphi_f)\), the feature vector is \[\boldsymbol{\psi}(\boldsymbol{e}) = \bigl(g(\boldsymbol{e};\, \boldsymbol{\mu}_f, \boldsymbol{\omega}_f, \sigma_f, \varphi_f)\bigr)_{f=1}^F \in \mathbb{R}^F,\] and the network output is \(\boldsymbol{u}_{\boldsymbol{\theta}} = \boldsymbol{W}\boldsymbol{\psi}(\boldsymbol{e}) + \boldsymbol{b} \in \mathbb{R}^{L}\), with \(\boldsymbol{W} \in \mathbb{R}^{L \times F}\).

8.5.0.1 Trainable parameters

\[\boldsymbol{\theta} = \{\boldsymbol{\mu}_f, \boldsymbol{\omega}_f, \sigma_f, \varphi_f\}_{f=1}^{F} \cup \{\boldsymbol{W}, \boldsymbol{b}\},\] flattened into a single vector.

8.6 SPINN↩︎

The separable physics-informed neural network (SPINN), introduced by [54], addresses the exponential scaling of conventional PINNs in high dimensions by factorising the solution approximation into a sum of rank-one products for each input dimension. Each of the \(n_\text{in}\) input coordinates is given its own branch network \(f_{\boldsymbol{\theta}_i} : \mathbb{R}^{d_{\phi,i}} \to \mathbb{R}^R\) of rank \(R\), acting on the per-axis embedding \(\boldsymbol{\phi}_i\) of that coordinate (a Fourier embedding from 8.1 for spatial axes; the identity for the time axis, so that time remains independent of space). The network output is the rank-\(R\) combination \[\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}, t) = \boldsymbol{W}\!\left( \bigodot_{i=1}^{n_\text{in}} f_{\boldsymbol{\theta}_i}\!\bigl(\boldsymbol{\phi}_i\bigr)\right) + \boldsymbol{b} \in \mathbb{R}^{L},\] where \(\bigodot\) denotes the elementwise (Hadamard) product of the branch outputs and \(\boldsymbol{W} \in \mathbb{R}^{L \times R}\). Each branch is a small MLP with \(\tanh\) activations.

This separable structure allows residuals to be evaluated on a tensor grid of \(n_\text{in}\) one-dimensional arrays rather than a full \(n_\text{in}\)-dimensional set of collocation points, reducing the number of network forward passes from \(\mathcal{O}(N^{n_\text{in}})\) to \(\mathcal{O}(N)\). SPINN is therefore particularly effective for the high-dimensional Poisson problems considered in this work (2), where conventional point-wise evaluation would be computationally prohibitive.

8.6.0.1 Trainable parameters

\[\boldsymbol{\theta} = \{\boldsymbol{\theta}_i\}_{i=1}^{n_\text{in}} \cup \{\boldsymbol{W}, \boldsymbol{b}\},\] flattened into a single vector, where each \(\boldsymbol{\theta}_i\) collects the weights and biases of the \(i\)-th branch MLP.

9 Considered PDEs↩︎

The problems below correspond to those reported in 2, presented in the same order.

9.1 Burgers’ Equation↩︎

We consider the viscous Burgers’ equation: \[\label{eq:burgers} \frac{\partial u}{\partial t} + u\,\frac{\partial u}{\partial x} - \nu\,\frac{\partial^2 u}{\partial x^2} = 0,\tag{22}\] on a spatially-periodic space-time domain \((x, t) \in [-1, 1] \times [0, 1]\), with viscosity \(\nu = 0.01\). The initial condition is \(u(x, 0) = -\sin(\pi x)\). Periodic boundary conditions are enforced analytically by mapping the network input \((x, t) \mapsto (\cos(\pi x), \sin(\pi x), t)\), which eliminates \(\boldsymbol{\mathcal{B}}\) from the residual system (see 17). The resulting operator has \(M=2\) conditions: \[\boldsymbol{\mathcal{F}}\,u = \left(\begin{array}{l@{\quad}l} u_t + u\,u_x - \nu\,u_{xx} & \forall\,(x,t)\in[-1,1]\times[0,1] \\[4pt] u(x, 0) + \sin(\pi x) & \forall\, x\in[-1,1] \end{array}\right) = \begin{pmatrix} \mathcal{P}\, u \\ \mathcal{I}\, u \end{pmatrix} = \boldsymbol{0}.\] The reference solution is computed by a Fourier pseudo-spectral discretisation in space (\(1024\) modes), advanced in time with an eighth-order Dormand–Prince adaptive Runge–Kutta integrator (DOP853) at absolute and relative tolerances of \(10^{-15}\) [55], [56].

9.2 Korteweg–De Vries Equation↩︎

We consider the KdV equation: \[\label{eq:KdV95eq} u_t + \eta\,u\,u_x + \mu^2\,u_{xxx} = 0,\tag{23}\] on a spatially-periodic space-time domain \((x, t) \in [-1, 1] \times [0, 1]\), with parameters \(\eta = 1\) and \(\mu = 0.022\). The initial condition is \(u(x, 0) = \cos(\pi x)\), and periodic boundary conditions are again enforced analytically (see 17). The resulting operator has \(M=2\) conditions: \[\boldsymbol{\mathcal{F}}\,u = \left(\begin{array}{l@{\quad}l} u_t + \eta\,u\,u_x + \mu^2\,u_{xxx} & \forall\,(x,t)\in[-1,1]\times[0,1] \\[4pt] u(x, 0) - \cos(\pi x) & \forall\, x\in[-1,1] \end{array}\right) = \begin{pmatrix} \mathcal{P}\, u \\ \mathcal{I}\, u \end{pmatrix} = \boldsymbol{0}.\] The reference solution is computed by a Fourier pseudo-spectral discretisation in space (\(512\) modes), advanced with the classical fourth-order Runge–Kutta scheme at a step of \(5\times10^{-6}\), chosen to respect the dispersive stability restriction \(\Delta t \lesssim \Delta x^3/\mu^2\) [55].

9.3 Wave Equation↩︎

We consider the one-dimensional wave equation: \[u_{tt} - c^2\,u_{xx} = 0,\] with \(c = 2\) on a space-time domain \((x, t) \in [0,1]^2\), with zero-Dirichlet boundaries and a standing-wave initial condition: \[u(x,0) = \sin(\pi x) + \tfrac{1}{2}\sin(4\pi x), \qquad u_t(x,0) = 0.\] The resulting operator has \(M=4\) conditions: \[\boldsymbol{\mathcal{F}}\,u = \left(\begin{array}{l@{\quad}l} u_{tt} - c^2\,u_{xx} & \forall\,(x,t)\in[0,1]^2 \\[4pt] u(x, 0) - \sin(\pi x) - \tfrac{1}{2}\sin(4\pi x) & \forall\, x\in[0,1] \\[4pt] u_t(x, 0) & \forall\, x\in[0,1] \\[4pt] u(\{0,1\}, t) & \forall\, t\in[0,1] \end{array}\right) = \begin{pmatrix} \mathcal{P}\, u \\ \mathcal{I}_1\, u \\ \mathcal{I}_2\, u \\ \mathcal{B}\, u \end{pmatrix} = \boldsymbol{0}.\] The exact solution is \[u(x, t) = \sin(\pi x)\cos(2\pi t) + \tfrac{1}{2}\sin(4\pi x)\cos(8\pi t).\]

9.4 Kuramoto–Sivashinsky Equation↩︎

We consider the one-dimensional Kuramoto–Sivashinsky (KS) equation, a canonical model for spatio-temporal chaos: \[\label{eq:ks} u_t + \alpha\,u\,u_x + \beta\,u_{xx} + \gamma\,u_{xxxx} = 0,\tag{24}\] on a spatially-periodic space-time domain \((x, t) \in [0, 2\pi] \times [0, 1]\), with \[\alpha = \frac{100}{16}, \qquad \beta = \frac{100}{16^2}, \qquad \gamma = \frac{100}{16^4}.\] The initial condition is \(u_0(x) = \cos(x)(1 + \sin(x))\). Periodic boundary conditions are enforced analytically (see 17). The resulting operator has \(M=2\) conditions: \[\boldsymbol{\mathcal{F}}\,u = \left(\begin{array}{l@{\quad}l} u_t + \alpha\,u\,u_x + \beta\,u_{xx} + \gamma\,u_{xxxx} & \forall\,(x,t)\in[0,2\pi]\times[0,1] \\[4pt] u(x,0) - \cos(x)(1 + \sin(x)) & \forall\, x\in[0,2\pi] \end{array}\right) = \begin{pmatrix} \mathcal{P}\, u \\ \mathcal{I}\, u \end{pmatrix} = \boldsymbol{0}.\] The reference solution is computed by a Fourier spectral discretisation in space, advanced with the fourth-order exponential time-differencing Runge–Kutta scheme (ETDRK4) of [57], whose coefficients are evaluated by contour integration over \(64\) Cauchy points [58]; the nonlinear term is dealiased by \(3/2\) zero-padding.

9.5 Multi-Scale Poisson↩︎

We consider a two-dimensional multi-scale Poisson problem (see [46]): \[-\Delta u(x_1, x_2) = f(x_1, x_2), \qquad (x_1, x_2) \in (0,1)^2,\] on the domain \(\Omega = [0,1]^2\) with a right-hand side designed to excite multiple spatial frequencies simultaneously. We have homogeneous Dirichlet boundary conditions \(u = 0\) on \(\partial[0,1]^2\), and right-hand side \[f(x_1, x_2) = \frac{2}{n}\sum_{w \in \mathcal{W}} (w\pi)^2 \sin(w\pi x_1)\sin(w\pi x_2),\] where \(\mathcal{W} = \{4, 8, 12, 16\}\) and \(n = |\mathcal{W}| = 4\). The resulting operator has \(M=2\) conditions: \[\boldsymbol{\mathcal{F}}\,u = \left(\begin{array}{l@{\quad}l} -\Delta u(x_1,x_2) - f(x_1,x_2) & \forall\,(x_1,x_2)\in(0,1)^2 \\[4pt] u(x_1,x_2) & \forall\,(x_1,x_2)\in\partial[0,1]^2 \end{array}\right) = \begin{pmatrix} \mathcal{P}\, u \\ \mathcal{B}\, u \end{pmatrix} = \boldsymbol{0}.\] Since the problem is steady, there is no initial-condition operator \(\boldsymbol{\mathcal{I}}\). The exact solution is \[u_\star(x_1, x_2) = \frac{1}{n}\sum_{w \in \mathcal{W}} \sin(w\pi x_1)\sin(w\pi x_2).\]

9.6 High-Dimensional Poisson↩︎

We consider two high-dimensional Poisson problems following [18].

9.6.0.1 5D Poisson.

The problem is \(-\Delta u(\boldsymbol{x}) = f(\boldsymbol{x})\) on \([0,1]^5\), with right-hand side \[f(\boldsymbol{x}) = \pi^2 \sum_{i=1}^{5} \cos(\pi x_i),\] and Dirichlet boundary conditions \(u(\boldsymbol{x}) = \sum_{i=1}^5 \cos(\pi x_i)\) for \(\boldsymbol{x} \in \partial[0,1]^5\). The resulting operator has \(M=2\) conditions: \[\boldsymbol{\mathcal{F}}\,u = \left(\begin{array}{l@{\quad}l} -\Delta u(\boldsymbol{x}) - \pi^2\displaystyle\sum_{i=1}^{5}\cos(\pi x_i) & \forall\,\boldsymbol{x} \in (0,1)^5 \\[6pt] u(\boldsymbol{x}) - \displaystyle\sum_{i=1}^{5}\cos(\pi x_i) & \forall\,\boldsymbol{x} \in \partial[0,1]^5 \end{array}\right) = \begin{pmatrix} \mathcal{P}\, u \\ \mathcal{B}\, u \end{pmatrix} = \boldsymbol{0}.\] The exact solution is \(u_\star(\boldsymbol{x}) = \sum_{i=1}^5 \cos(\pi x_i)\).

9.6.0.2 10D Poisson.

The problem is \(-\Delta u(\boldsymbol{x}) = 0\) on \([0,1]^{10}\) with harmonic boundary conditions \[u(\boldsymbol{x}) = \sum_{i=1}^{5} x_{2i-1}\,x_{2i}, \qquad \boldsymbol{x} \in \partial[0,1]^{10}.\] The resulting operator has \(M=2\) conditions: \[\boldsymbol{\mathcal{F}}\,u = \left(\begin{array}{l@{\quad}l} -\Delta u(\boldsymbol{x}) & \forall\,\boldsymbol{x} \in (0,1)^{10} \\[4pt] u(\boldsymbol{x}) - \displaystyle\sum_{i=1}^{5} x_{2i-1}\,x_{2i} & \forall\,\boldsymbol{x} \in \partial[0,1]^{10} \end{array}\right) = \begin{pmatrix} \mathcal{P}\, u \\ \mathcal{B}\, u \end{pmatrix} = \boldsymbol{0}.\] The exact solution is \(u_\star(\boldsymbol{x}) = \sum_{i=1}^5 x_{2i-1}\,x_{2i}\).

9.7 Lid-Driven Cavity Flow↩︎

We consider the steady incompressible Navier–Stokes equations in the unit square \(\Omega = [0,1]^2\) at Reynolds number \(\mathrm{Re} = 100\): \[\begin{align} u\,u_x + v\,u_y + p_x - \tfrac{1}{\mathrm{Re}}(u_{xx} + u_{yy}) &= 0, \\ u\,v_x + v\,v_y + p_y - \tfrac{1}{\mathrm{Re}}(v_{xx} + v_{yy}) &= 0, \\ u_x + v_y &= 0. \end{align}\] The network outputs three fields: \(x\)-velocity \(u\), \(y\)-velocity \(v\), and pressure \(p\). Boundary conditions are imposed on all four walls: \[u(x,y) = \begin{cases} 4x(1-x) & y = 1\;\text{(moving lid)},\\ 0 & \text{otherwise,} \end{cases} \qquad v(x,y) = 0 \quad \forall\,(x,y)\in\partial\Omega.\] Since only gradients of \(p\) appear in the momentum equations, the pressure field is determined only up to an additive constant; one has the option to pin it via \(p(0,0) = 0\), but we instead allow an arbitrary additive constant so that we have one fewer condition to solve for. The resulting operator has \(M=5\) conditions \[\boldsymbol{\mathcal{F}}\,\boldsymbol{u} = \left(\begin{array}{l@{\quad}l} u\,u_x + v\,u_y + p_x - \frac{1}{\mathrm{Re}}(u_{xx}+u_{yy}) & \forall\,(x,y)\in\Omega \\[3pt] u\,v_x + v\,v_y + p_y - \frac{1}{\mathrm{Re}}(v_{xx}+v_{yy}) & \forall\,(x,y)\in\Omega \\[3pt] u_x + v_y & \forall\,(x,y)\in\Omega \\[3pt] u(x,y) - \mathbb{I}_{\{y=1\}}\cdot 4x(1-x) & \forall\,(x,y)\in\partial\Omega \\[3pt] v(x,y) & \forall\,(x,y)\in\partial\Omega \end{array}\right) = \begin{pmatrix} \mathcal{P}_1\, \boldsymbol{u} \\ \mathcal{P}_2\, \boldsymbol{u} \\ \mathcal{P}_3\, \boldsymbol{u} \\ \mathcal{B}_1\, \boldsymbol{u} \\ \mathcal{B}_2\, \boldsymbol{u} \end{pmatrix} = \boldsymbol{0},\] where \(\mathcal{P}_1, \mathcal{P}_2, \mathcal{P}_3\) denote the \(x\)-momentum, \(y\)-momentum, and continuity equations respectively, and \(\mathcal{B}_1, \mathcal{B}_2\) the boundary conditions on \(u\) and \(v\); as the problem is steady, there is no initial-condition operator \(\boldsymbol{\mathcal{I}}\). The reference solution is computed with the Dedalus spectral framework [59], using a Chebyshev discretisation in both directions (\(100\times100\)) and a fourth-order semi-implicit backward-difference time-stepper (SBDF4), marched to steady state at \(\mathrm{Re}=100\) with a regularised lid velocity \(u(x,1)=4x(1-x)\).

10 Results and metrics for comprehensive suite of PDEs↩︎

Below are a set of figures showing results and metrics for the rest of the PDEs in 2, with some extra problems to showcase DSGNAR’s capabilities.

Figure 19: No caption
Figure 20: No caption
Figure 21: No caption
Figure 22: No caption
Figure 23: No caption
Figure 24: No caption
Figure 25: No caption
Figure 26: No caption
Figure 27: No caption
Figure 28: No caption

References↩︎

[1]
J. D. Toscano et al., From PINNs to PIKANs: Recent Advances in Physics-Informed Machine Learning,” Machine Learning for Computational Science and Engineering, vol. 1, no. 1, Jun. 2025, doi: 10.1007/s44379-025-00015-1.
[2]
P. Frey and P.-L. George, Mesh Generation. London, England: Hermes Science Publishing, 2000.
[3]
G. E. Karniadakis, I. G. Kevrekidis, L. Lu, P. Perdikaris, S. Wang, and L. Yang, Physics-informed machine learning,” Nature Reviews Physics 2021 3:6, vol. 3, no. 6, pp. 422–440, Jun. 2021, doi: 10.1038/s42254-021-00314-5.
[4]
M. W. M. G. Dissanayake and N. Phan-Thien, Neural-Network-Based Approximations for Solving Partial Differential Equations,” Communications in Numerical Methods in Engineering, vol. 10, no. 3, pp. 195–201, 1994, doi: 10.1002/cnm.1640100303.
[5]
I. E. Lagaris, A. Likas, and D. I. Fotiadis, Artificial Neural Networks for Solving Ordinary and Partial Differential Equations,” IEEE Transactions on Neural Networks, vol. 9, no. 5, pp. 987–1000, 1998, doi: 10.1109/72.712178.
[6]
M. Raissi, P. Perdikaris, and G. E. Karniadakis, Physics-Informed Neural Networks: A Deep Learning Framework for Solving Forward and Inverse Problems Involving Nonlinear Partial Differential Equations,” Journal of Computational Physics, vol. 378, pp. 686–707, 2019, doi: 10.1016/j.jcp.2018.10.045.
[7]
P. Rathore, W. Lei, Z. Frangella, L. Lu, and M. Udell, Challenges in Training PINNs: A Loss Landscape Perspective,” in Proceedings of the 41st international conference on machine learning, 2024, pp. 42127–42159, doi: 10.48550/arXiv.2402.01868.
[8]
E. Kiyani, K. Shukla, J. F. Urbán, J. Darbon, and G. E. Karniadakis, Optimizing the Optimizer for Physics-Informed Neural Networks and Kolmogorov-Arnold Networks,” Computer Methods in Applied Mechanics and Engineering, vol. 446, p. 118308, 2025, doi: 10.1016/j.cma.2025.118308.
[9]
D. P. Kingma and J. Ba, Adam: A Method for Stochastic Optimization,” in International conference on learning representations, 2015, doi: 10.48550/arXiv.1412.6980.
[10]
J. F. Urbán, P. Stefanou, and J. A. Pons, Unveiling the Optimization Process of Physics-Informed Neural Networks: How Accurate and Competitive Can PINNs Be? Journal of Computational Physics, vol. 523, p. 113656, Jun. 2025, doi: 10.1016/j.jcp.2024.113656.
[11]
J. Nocedal and S. J. Wright, Numerical Optimization, 2nd ed. New York, NY, USA: Springer, 2006.
[12]
M. Al-Baali, Numerical Experience with a Class of Self-Scaling Quasi-Newton Algorithms,” Journal of Optimization Theory and Applications, vol. 96, no. 3, pp. 533–553, 1998, doi: 10.1023/A:1022608410710.
[13]
F. Dangel, J. Müller, and M. Zeinhofer, Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks,” in Advances in neural information processing systems, 2024, vol. 37, pp. 34582–34636, doi: 10.48550/arXiv.2405.15603.
[14]
S. Wang, A. K. Bhartari, B. Li, and P. Perdikaris, Gradient Alignment in Physics-Informed Neural Networks: A Second-Order Optimization Perspective,” in Advances in neural information processing systems, 2025, vol. 38, doi: 10.48550/arXiv.2502.00604.
[15]
J. Martens and R. Grosse, Optimizing Neural Networks with Kronecker-Factored Approximate Curvature,” in Proceedings of the 32nd international conference on machine learning, 2015, vol. 37, pp. 2408–2417, doi: 10.48550/arXiv.1503.05671.
[16]
J. Martens, J. Ba, and M. Johnson, Kronecker-Factored Curvature Approximations for Recurrent Neural Networks,” in International conference on learning representations, 2018, [Online]. Available: https://openreview.net/forum?id=HyMTkQZAb.
[17]
A. Jnini, F. Vella, and M. Zeinhofer, Gauss-Newton Natural Gradient Descent for Physics-Informed Computational Fluid Dynamics,” Computers & Fluids, vol. 307, p. 106955, Jun. 2026, doi: 10.1016/j.compfluid.2025.106955.
[18]
A. Guzmán-Cordero, F. Dangel, G. Goldshlager, and M. Zeinhofer, Improving Energy Natural Gradient Descent through Woodbury, Momentum, and Randomization,” in Advances in neural information processing systems, 2025, vol. 38, pp. 113870–113900, doi: 10.48550/arXiv.2505.12149.
[19]
J. Müller and M. Zeinhofer, Achieving High Accuracy with PINNs via Energy Natural Gradient Descent,” in Proceedings of the 40th international conference on machine learning, Jun. 2023, vol. 202, pp. 25471–25485, doi: 10.5555/3618408.3619465.
[20]
P.-G. Martinsson and J. A. Tropp, Randomized Numerical Linear Algebra: Foundations and Algorithms,” Acta Numerica, vol. 29, pp. 403–572, 2020, doi: 10.1017/S0962492920000021.
[21]
N. Halko, P. G. Martinsson, and J. A. Tropp, Finding Structure with Randomness: Probabilistic Algorithms for Constructing Approximate Matrix Decompositions,” SIAM Review, vol. 53, no. 2, pp. 217–288, 2011, doi: 10.1137/090771806.
[22]
M. Best Mckay, A. Kaur, C. Greif, and B. Wetton, Near-Optimal Sketchy Natural Gradients for Physics-Informed Neural Networks,” in Proceedings of the 42nd international conference on machine learning, Jun. 2025, vol. 267, pp. 4005–4019, [Online]. Available: https://openreview.net/forum?id=bKsZomnmqn.
[23]
M. Yang, D. Xu, Z. Wen, M. Chen, and P. Xu, Sketch-Based Empirical Natural Gradient Methods for Deep Learning,” Journal of Scientific Computing, vol. 92, no. 3, 2022, doi: 10.1007/s10915-022-01939-5.
[24]
L. Armijo, Minimization of Functions Having Lipschitz Continuous First Partial Derivatives,” Pacific Journal of Mathematics, vol. 16, no. 1, pp. 1–3, 1966, doi: 10.2140/pjm.1966.16.1.
[25]
A. R. Conn, N. I. M. Gould, and Ph. L. Toint, Trust Region Methods. Society for Industrial; Applied Mathematics, 2000.
[26]
K. Levenberg, A Method for the Solution of Certain Non-Linear Problems in Least Squares,” Quarterly of Applied Mathematics, vol. 2, pp. 164–168, 1944, doi: 10.1090/qam/10666.
[27]
D. W. Marquardt, An Algorithm for Least-Squares Estimation of Nonlinear Parameters,” Journal of the Society for Industrial and Applied Mathematics, vol. 11, no. 2, pp. 431–441, 1963, doi: 10.1137/0111030.
[28]
G. Wu and Z. Wu, A Multi-Objective Optimization Framework for Adaptive Weighting in Physics-Informed Machine Learning,” Proceedings of the AAAI Conference on Artificial Intelligence, vol. 40, no. 32, pp. 26885–26893, 2026, doi: 10.1609/aaai.v40i32.39900.
[29]
S. J. Anagnostopoulos, J. D. Toscano, N. Stergiopulos, and G. E. Karniadakis, Residual-Based Attention in Physics-Informed Neural Networks,” Computer Methods in Applied Mechanics and Engineering, vol. 421, p. 116805, 2024, doi: 10.1016/j.cma.2024.116805.
[30]
M. Charikar, K. Chen, and M. Farach-Colton, Finding Frequent Items in Data Streams,” Theoretical Computer Science, vol. 312, no. 1, pp. 3–15, 2004, doi: 10.1016/S0304-3975(03)00400-6.
[31]
C. Wu, M. Zhu, Q. Tan, Y. Kartha, and L. Lu, A comprehensive study of non-adaptive and residual-based adaptive sampling for physics-informed neural networks,” Computer Methods in Applied Mechanics and Engineering, vol. 403, p. 115671, Jan. 2023, doi: 10.1016/J.CMA.2022.115671.
[32]
A. N. Tikhonov, Regularization of Incorrectly Posed Problems,” Soviet Mathematics Doklady, vol. 4, pp. 1624–1627, 1963.
[33]
M. Pilanci and M. J. Wainwright, Newton Sketch: A Near Linear-Time Optimization Algorithm with Linear-Quadratic Convergence,” SIAM Journal on Optimization, vol. 27, no. 1, pp. 205–245, Feb. 2017, doi: 10.1137/15M1021106.
[34]
M. Meister, T. Sarlos, and D. Woodruff, Tight Dimensionality Reduction for Sketching Low Degree Polynomial Kernels,” in Advances in neural information processing systems, 2019, vol. 32, doi: 10.5555/3454287.3455137.
[35]
P. Kar and H. Karnick, Random Feature Maps for Dot Product Kernels,” in Proceedings of the 15th international conference on artificial intelligence and statistics, 2012, vol. 22, pp. 583–591, doi: 10.48550/arXiv.1201.6530.
[36]
J. Nelson and H. L. Nguyễn, OSNAP: Faster Numerical Linear Algebra Algorithms via Sparser Subspace Embeddings,” IEEE Annual Symposium on Foundations of Computer Science, pp. 117–126, 2012, doi: 10.1109/FOCS.2013.21.
[37]
K. L. Clarkson and D. P. Woodruff, Low-Rank Approximation and Regression in Input Sparsity Time,” Journal of the ACM, vol. 63, no. 6, pp. 54:1–54:45, 2017, doi: 10.1145/3019134.
[38]
N. Ailon and B. Chazelle, The Fast Johnson–Lindenstrauss Transform and Approximate Nearest Neighbors,” SIAM Journal on Computing, vol. 39, no. 1, pp. 302–322, Jun. 2009, doi: 10.1137/060673096.
[39]
F. N. Fritsch and R. E. Carlson, Monotone Piecewise Cubic Interpolation,” SIAM Journal on Numerical Analysis, vol. 17, no. 2, pp. 238–246, 1980, doi: 10.1137/0717021.
[40]
J. Bradbury et al., JAX: Composable Transformations of Python+NumPy Programs.” 2018, [Online]. Available: http://github.com/jax-ml/jax.
[41]
A. Krishnapriyan, A. Gholami, S. Zhe, R. Kirby, and M. Mahoney, Characterizing Possible Failure Modes in Physics-Informed Neural Networks,” in Advances in neural information processing systems, 2021, vol. 34, pp. 26548–26560, doi: 10.48550/arXiv.2109.01050.
[42]
W. Li, C. Zhang, C. Wang, H. Guan, and D. Tao, Revisiting PINNs: Generative Adversarial Physics-Informed Neural Networks and Point-Weighting Method.” arXiv:2205.08754, 2022, doi: 10.48550/arXiv.2205.08754.
[43]
F. Cao, X. Guo, X. Dong, and D. Yuan, wbPINN: Weight Balanced Physics-Informed Neural Networks for Multi-Objective Learning,” Applied Soft Computing, vol. 170, p. 112632, Jun. 2025, doi: 10.1016/j.asoc.2024.112632.
[44]
A. Sabne, XLA : Compiling Machine Learning for Peak Performance.” 2020, [Online]. Available: https://research.google/pubs/xla-compiling-machine-learning-for-peak-performance/.
[45]
B. Moseley, A. Markham, and T. Nissen-Meyer, Finite basis physics-informed neural networks (FBPINNs): a scalable domain decomposition approach for solving differential equations,” Advances in Computational Mathematics, vol. 49, no. 4, p. 62, 2023, doi: 10.1007/s10444-023-10065-9.
[46]
S. Anderson, V. Dolean, ·. B. Moseley, and J. Pestana, ELM-FBPINNs: an efficient multilevel random feature method,” Machine Learning for Computational Science and Engineering 2026 2:1, vol. 2, no. 1, p. 23, Jun. 2026, doi: 10.1007/S44379-026-00071-1.
[47]
S. Zampini, U. Zerbinati, G. Turkyyiah, and D. Keyes, PETScML: Second-Order Solvers for Training Regression Problems in Scientific Machine Learning,” Platform for Advanced Scientific Computing Conference (PASC ’24), Zurich, Switzerland, vol. 1, pp. 1–12, Jun. 2024, doi: 10.1145/3659914.3659931.
[48]
M. Tancik et al., Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains,” Neural Information Processing Systems, 2020, doi: 10.48550/arXiv.2006.10739.
[49]
S. Wang, X. Yu, and P. Perdikaris, When and why PINNs fail to train: A neural tangent kernel perspective,” Journal of Computational Physics, vol. 449, Jun. 2022, doi: 10.1016/j.jcp.2021.110768.
[50]
S. Dong and N. Ni, A method for representing periodic functions and enforcing exactly periodic boundary conditions with deep neural networks,” Journal of Computational Physics, vol. 435, Jun. 2021, doi: 10.1016/j.jcp.2021.110242.
[51]
Y. LeCun, L. Bottou, G. B. Orr, and K.-R. Müller, Efficient BackProp,” pp. 9–50, 1998, doi: 10.1007/3-540-49430-8_2.
[52]
V. Sitzmann, J. N. P. Martel, A. W. Bergman, D. B. Lindell, and G. Wetzstein, Implicit Neural Representations with Periodic Activation Functions,” in Advances in neural information processing systems, 2020, vol. 33, pp. 7462–7473, doi: 10.48550/arXiv.2006.09661.
[53]
R. Fathony, A. K. Sahu, D. Willmott, and J. Kolter, Multiplicative Filter Networks,” International Conference on Learning Representations, 2021, [Online]. Available: https://openreview.net/forum?id=OmtmcPkkhT.
[54]
J. Cho, S. Nam, H. Yang, S. Yun, Y. Hong, and E. Park, Separable Physics-Informed Neural Networks,” Neural Information Processing Systems, 2023, doi: 10.48550/arXiv.2306.15969.
[55]
L. N. Trefethen, Spectral Methods in MATLAB,” Spectral Methods in MATLAB, Jun. 2000, doi: 10.1137/1.9780898719598.
[56]
E. Hairer, S. P. Nørsett, and G. Wanner, Solving ordinary differential equations i: Nonstiff problems, 2nd ed., vol. 8. Springer, 1993, p. 528.
[57]
S. M. Cox and P. C. Matthews, “Exponential time differencing for stiff systems,” Journal of Computational Physics, vol. 176, no. 2, pp. 430–455, 2002, doi: 10.1006/jcph.2002.6995.
[58]
A. K. Kassam and L. N. Trefethen, Fourth-Order Time-Stepping for Stiff PDEs,” SIAM Journal on Scientific Computing, vol. 26, no. 4, pp. 1214–1233, Jun. 2006, doi: 10.1137/S1064827502410633.
[59]
K. J. Burns, G. M. Vasil, J. Oishi, D. Lecoanet, and B. P. Brown, “Dedalus: A flexible framework for numerical simulations with spectral methods,” Physical Review Research, vol. 2, no. 2, p. 023068, Apr. 2020, doi: 10.1103/PhysRevResearch.2.023068.

  1. Code available at https://www.github.com/wephy/physics-informed-neural-networks↩︎

  2. The factor \(\tfrac12\) is the conventional least-squares scaling.↩︎

  3. A single hash is an isometry only in expectation, \(\mathbb{E}[{\boldsymbol{C}}^{\mathsf{T}}\boldsymbol{C}] = \boldsymbol{I}_N\), with substantial per-draw variance from hash collisions; such embeddings concentrate around this expectation, preserving the geometry of a fixed \(d\)-dimensional subspace to relative accuracy \(\varepsilon\) with high probability once \(K>1\) hashes are used and the sketch has \(\mathcal{O}(d/\varepsilon^2)\) rows. While these results do not directly cover our square, doubly-sketched construction (6), they motivate both our use of multiple hashes for variance reduction and the expectation that a moderate \(s\) suffices.↩︎