Tensor Gaussian Processes: Efficient Solvers for Nonlinear PDEs

AI summary

Overview

• Research area: Machine learning for scientific computing, specifically Gaussian-process and kernel-based numerical solvers for partial differential equations (PDEs).
• Technical level: Advanced. The paper combines tensor decomposition (CP, tensor-ring), reproducing kernel Hilbert spaces, Newton/Fréchet linearization, and Sobolev-space convergence analysis.
• Scope in one sentence: The paper introduces TGPS, a tensor Gaussian process that represents a PDE solution as a multilinear combination of one-dimensional GP factor functions, and solves it with closed-form alternating least squares rather than stochastic training.

What This Paper Is About

Machine-learning PDE solvers such as physics-informed neural networks (PINNs) approximate a PDE solution with a model trained by minimizing residual and boundary losses at collocation points, but they rely on stochastic optimization and typically need tens of thousands of iterations. Existing Gaussian process / kernel solvers are mathematically principled but their time and memory costs exceed 𝒪(M³) and 𝒪(M²), where M is the number of collocation points, and the Kronecker-based alternative requires estimating solution values at all grid points — a quantity that grows exponentially with input dimension. TGPS addresses both bottlenecks by factorizing the solution across dimensions so that the number of trainable parameters, storage, and computation grow linearly with the PDE dimension and with M.

Key Contributions

1. Model. The solution is written as a set of one-dimensional factor functions, each modeled as a Gaussian process f_r^i ~ GP(0, κ_i(·,·)) over a RKHS G^i, combined by either CANDECOMP/PARAFAC (CP) decomposition, u(x) = Σ_{r=1}^{R} Π_{i=1}^{d} f_r^i(x_i), or tensor-ring (TR) decomposition, u(x) = Trace(F^1(x_1) ... F^d(x_d)). Each factor function is represented as kernel interpolation (the GP conditional mean) f_r^i(x_i) = κ_i(x_i, γ_i) K_i^{-1} η_r^i with inducing points γ_i and trainable inducing values η_r^i; collocation points are sampled freely rather than placed on a grid.
2. Algorithm. To handle nonlinear PDEs, the authors linearize nonlinear terms either by partial freezing (freezing part of a nonlinear term at values from the previous iteration, e.g. u(u²−1) ≈ u·((u^prev)²−1)) or by Newton's method (first-order Taylor/Fréchet linearization, e.g. u³ replaced by (u^prev)³ + 3(u^prev)²(u − u^prev)). Because the solution and all of its partial derivatives are multilinear in each block of inducing values H_i, an alternating least squares (ALS) scheme updates each H_i in closed form, replacing stochastic optimization entirely.
3. Theory. The paper proves expressivity results for both decomposition formats, a convergence result for the constrained training problem under the CP format, and a proposition showing that the soft-regularized unconstrained objective shares the same minimizer (and therefore the same convergence guarantee) as the constrained problem for appropriately chosen regularization strengths.
4. Experiments. Five benchmark PDE families are used — viscous Burgers' equations, nonlinear elliptic PDEs, Eikonal PDEs, Allen-Cahn equations, and nonlinear Darcy flow — against DAKS, SKS, and PINN. Code is released at https://github.com/BayesianAIGroup/TGPSolve-NonLinear-PDEs.

Main Findings

• Complexity scales linearly in collocation points and dimension. Time complexity per iteration is 𝒪(M Σ_{i=1}^{d} ((N_i R)² + N_i³ R³)) and space complexity is 𝒪(Σ_{i=1}^{d} (N_i R + N_i²)), where M is the number of collocation points, N_i the inducing points in dimension i, and R the number of factor functions per dimension. Both scale linearly with M and d.
• Better accuracy on simpler 2D problems. On Burgers' equation with viscosity ν = 0.02, nonlinear elliptic PDEs, and the Eikonal PDE (using 600, 1200, 2400, 4800 collocation points for Burgers' and 300, 600, 1200, 2400 for the other two, following Chen et al. and Xu et al.), both TGPS-PF and TGPS-NT achieve the highest solution accuracy. The single reported exception is Burgers' with 600 collocation points, where TGPS-NT is slightly worse than PINN but still ranks third.
• Reported numbers for Burgers' (ν = 0.02), relative L² error (table is truncated in the supplied content): DAKS 3.05E-02, 1.38E-02, 1.51E-03, 1.70E-04; PINN 4.67E-03, 1.17E-03, 6.27E-04, 6.50E-04; SKS 2.51E-02, 9.41E-03, 1.36E-03, 5.59E-04; TGPS-PF 3.56E-03, 5.77E-04, 1.35E-04, 8.83E-05; TGPS-NT 8.42E-03 at 600 points (remaining TGPS-NT entries are not present in the supplied text). SKS grids were 25×25, 35×35, 49×49, and 70×70.
• Scales to hard, high-dimensional settings. For Burgers' equation with viscosity 0.001 and for a 6D Allen-Cahn equation, TGPS scales to tens of thousands of collocation points and reaches errors on the order of 10⁻³ to 10⁻⁶.
• Orders-of-magnitude runtime advantage. Across all benchmarks the method runs orders of magnitude faster than PINNs and recent GP solvers, with comparable or better accuracy; the abstract attributes this to the closed-form ALS updates rather than stochastic optimization. The specific runtime figures are not reported in the supplied content.
• Expressivity of the multilinear model. Lemma 4.1: for u* ∈ H^k(Ω) and any 0 < ε < 1, a CP model with rank up to R̄ = (√d/ε)^{(d−1)/k} achieves ‖u* − u‖_{L²(Ω)} ≲ ε; under a weighted Sobolev assumption u* ∈ H_v^k(Ω) with v_j ≲ j^{−(1+δ')/k}, δ' > δ + k, the same bound holds with R̄ ≲ (1/ε)^{1/k}. Lemma 4.2: for u* ∈ H^{k+1}(Ω), a TR model with R̄ = (√d/ε)^{d/k} achieves the same error, with a weighted-Sobolev rank requirement R̄ ≲ exp{2(1/ε)^{1/k} + log(1/ε)}.
• Training converges to a good approximation. Lemma 4.4 shows that, under domain-regularity, PDE-stability, and RKHS-embedding assumptions (Assumption 4.3), the constrained problem has a minimizer u† with ‖u† − u*‖_{H^l(Ω)} ≲ ε as the fill distance h = max(h_Ω, h_{∂Ω}) goes to zero, provided each factor RKHS uses a universal kernel and U = G¹ ⊗ ... ⊗ G^d. Proposition 4.5 shows the soft-regularized objective (8) can be made to share that minimizer with suitable regularization strengths α_{1M}, α_{2M}.
• TR and CP coincide in the tested 2D setting. For the simpler PDEs the input dimension is d = 2, where TR decomposition reduces to CP, so the paper does not distinguish them there.

Methodology in Plain English

The classical obstacle for GP or kernel PDE solvers is that placing a prior over the solution and all of its derivatives produces a covariance matrix larger than M×M with 𝒪(M³) time and 𝒪(M²) memory. The earlier Kronecker-based fix requires the solution values on a full Cartesian grid, whose size explodes with dimension.

TGPS sidesteps both by assuming the solution can be built from one-dimensional pieces. Each input dimension gets a handful of one-dimensional functions, each a Gaussian process represented by its values at a small set of inducing locations, combined multiplicatively across dimensions. Since each dimension is handled separately, the parameter count, storage, and computation grow linearly with dimension and with the number of collocation points rather than exponentially.

The reason this makes solving fast is a structural property: when the inducing values in every dimension except one are held fixed, the solution and each of its partial derivatives are linear in the remaining block, so the squared PDE residual and boundary error become quadratic in that block. That turns each update into an ordinary least-squares problem with a closed-form answer, and cycling through the dimensions gives the ALS procedure. Nonlinear terms break this linearity, so the authors restore it in two ways: freeze one factor of a nonlinear product at its previous-iteration value, or replace the nonlinear term with its first-order Taylor expansion around the previous iterate. Both discrepancies vanish at convergence. Random collocation points are supported directly, which the grid-based Kronecker approach cannot do.

Why This Matters

• Impact on research. The work shows that a GP/kernel solver does not have to choose between mathematical guarantees and scalability. It replaces stochastic training — which the paper notes can require up to one million iterations for prior GP solvers — with exact closed-form updates, and backs the approach with expressivity, convergence, and error analysis rather than empirical claims alone. It also connects tensor-decomposition solvers for PDEs (a line going back to Beylkin and Mohlenkamp) with the GP/kernel solver literature.
• Real-world applications (as motivated by the benchmark PDE families used).
  • Burgers'-type convection-diffusion problems, which appear in fluid flow and shock modeling.
  • Allen-Cahn equations, which describe phase separation and interface motion in materials science; the paper tests a 6D version.
  • Eikonal equations, which underlie shortest-path, wavefront propagation, and geometric optics problems.
  • Nonlinear Darcy flow, describing flow through porous media, relevant to groundwater and reservoir simulation.
  • Nonlinear elliptic PDEs, a broad class covering steady-state physical models.
• Industry relevance. Fast, differentiable, stochastic-training-free PDE solvers with favorable scaling in dimension matter wherever repeated PDE solves are embedded inside larger workflows — design and shape optimization, uncertainty quantification, surrogate modeling, and simulation-based engineering. Removing the reliance on a structured Cartesian grid also makes the method applicable to irregular domains, which is a practical limitation of prior Kronecker-based GP solvers.

Future Directions

• The convergence and error analysis (Lemma 4.4, Proposition 4.5) is established for the CP format; the paper notes the TR extension follows naturally, but a dedicated convergence treatment for TR is not given in the supplied content.
• The regularization strengths α_{1M}, α_{2M} that make the soft-regularized objective coincide with the constrained minimizer are shown to exist; the paper does not report a practical selection rule in the supplied content.
• No quantitative runtime or wall-clock tables appear in the supplied content, only the qualitative claim of orders-of-magnitude speedups — a natural follow-up is a systematic benchmark of solve time and iteration counts against PINN, DAKS, and SKS.
• The rank R (and tensor-ring ranks R_i) as well as the number of inducing points N_i per dimension are accuracy-controlling hyperparameters; the theory gives bounds on required rank but the practical trade-off is not quantified in the supplied text.
• Beyond the five tested PDE families, applicability to time-dependent, higher-dimensional, and irregular-domain problems with complex boundary geometries remains open.

Target Audience

Researchers and practitioners working at the intersection of probabilistic machine learning and numerical PDEs, particularly those already familiar with PINNs, Gaussian process regression, and kernel methods who want to understand how tensor decomposition can remove the scalability and training-efficiency bottlenecks of kernel-based solvers. The theory sections assume comfort with Sobolev spaces, RKHS, Fréchet derivatives, and tensor decomposition notation. The paper is also relevant to computational scientists and engineers looking for a faster, grid-free alternative to neural PDE solvers, and to applied mathematicians interested in run–time/accuracy trade-offs in tensor-structured solvers.

Authors’ abstract

Machine learning solvers for partial differential equations (PDEs) have attracted growing interest. However, most existing approaches, such as neural network solvers, rely on stochastic training, which is inefficient and typically requires a great many training epochs. Gaussian process (GP)/kernel-based solvers, while mathematical principled, suffer from scalability issues when handling large numbers of collocation points often needed for challenging or higher-dimensional PDEs. To overcome these limitations, we propose TGPS, a tensor-GP-based solver that introduces factor functions along each input dimension using one-dimensional GPs and combines them via tensor decomposition to approximate the full solution. This design reduces the task to learning a collection of one-dimensional GPs, substantially lowering computational complexity, and enabling scalability to massive collocation sets. For efficient nonlinear PDE solving, we use a partial freezing strategy and Newton's method to linerize the nonlinear terms. We then develop an alternating least squares (ALS) approach that admits closed-form updates, thereby substantially enhancing the training efficiency. We establish theoretical guarantees on the expressivity of our model, together with convergence proof and error analysis under standard regularity assumptions. Experiments on several benchmark PDEs demonstrate that our method achieves superior accuracy and efficiency compared to existing approaches. The code is released at https://github.com/BayesianAIGroup/TGPSolve-NonLinear-PDEs