Tensor Gaussian Processes: Efficient Solvers for Nonlinear PDEs

Qiwei Yuan, Zhitong Xu, Yinghao Chen, Yiming Xu, Houman Owhadi, Shandian Zhe
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:3358-3366, 2026.

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 \url{https://github.com/BayesianAIGroup/TGPSolve-NonLinear-PDEs}

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-yuan26a, title = { Tensor Gaussian Processes: Efficient Solvers for Nonlinear PDEs }, author = {Yuan, Qiwei and Xu, Zhitong and Chen, Yinghao and Xu, Yiming and Owhadi, Houman and Zhe, Shandian}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {3358--3366}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/yuan26a/yuan26a.pdf}, url = {https://proceedings.mlr.press/v300/yuan26a.html}, 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 \url{https://github.com/BayesianAIGroup/TGPSolve-NonLinear-PDEs} } }
Endnote
%0 Conference Paper %T Tensor Gaussian Processes: Efficient Solvers for Nonlinear PDEs %A Qiwei Yuan %A Zhitong Xu %A Yinghao Chen %A Yiming Xu %A Houman Owhadi %A Shandian Zhe %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-yuan26a %I PMLR %P 3358--3366 %U https://proceedings.mlr.press/v300/yuan26a.html %V 300 %X 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 \url{https://github.com/BayesianAIGroup/TGPSolve-NonLinear-PDEs}
APA
Yuan, Q., Xu, Z., Chen, Y., Xu, Y., Owhadi, H. & Zhe, S.. (2026). Tensor Gaussian Processes: Efficient Solvers for Nonlinear PDEs . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:3358-3366 Available from https://proceedings.mlr.press/v300/yuan26a.html.

Related Material