Tensor-Train Compressed Separable PINNs: A Curvature-Aware Optimization Framework for Parametric PDEs in High Dimensions

2026-09-28

AI Summary

This study addresses the low training efficiency and memory bottlenecks encountered when training physics-informed neural networks (PINNs) for high-dimensional parametric partial differential equations. To overcome these challenges, this work proposes a curvature-aware second-order optimization framework based on the Gauss–Newton metric. By integrating tensor-train/CP decomposition with a separable architecture design, the method enables compressed computation of the residual Jacobian, fundamentally circumventing full-dimensional Jacobian assembly and substantially reducing both computational complexity and GPU memory overhead. Experimental results demonstrate that, compared to conventional first-order baselines, the proposed framework reduces the required number of iterations by several orders of magnitude and significantly shortens computation time while achieving superior solution accuracy.

Abstract

In this work, we develop a second-order optimization framework for physics-informed neural networks (PINNs) applied to high-dimensional parametric partial differential equations (PDEs). The framework is built on the Gauss--Newton pullback metric, which provides an operator-informed notion of curvature in parameter space and connects the method to the broader family of natural gradient schemes. We show that, for coordinate-separable neural architectures and linear differential operators (or linearized operators in the nonlinear case) admitting a finite separable representation, the residual Jacobian inherits a structured separable factorization. This yields an exact compressed formulation of the Gauss--Newton step in a reduced space, without assembling the full residual Jacobian on the exponentially large tensor-product collocation grid. The dimension of the reduced space (the effective compressed dimension) is determined by the local collocation grid sizes, the separable operator structure, and the contraction pattern of the architecture, thereby avoiding dependence on the full tensor-product grid size and replacing dense linear algebra in parameter space by a substantially smaller structured problem. Within our framework, we investigate canonical polyadic and tensor-train parametrizations and derive their full algebraic characterization relevant to the Gauss--Newton method, including the structure of the residual Jacobian, the resulting compressed system, and its effective compressed dimension. Numerical experiments on high-dimensional PDEs, including parametric problems, demonstrate the high efficiency of the proposed compressed Gauss--Newton method, which achieves substantially lower errors than tensor-compressed first-order baselines with orders of magnitude fewer iterations and only a fraction of the computing time.

Problem

• Physics-Informed Neural Networks
• High-dimensional PDEs
• Parametric PDEs
• Second-order Optimization
• Tensor Decomposition

Innovation

• Physics-Informed Neural Networks
• Gauss-Newton Optimization
• Tensor-Train Decomposition
• Separable Architecture
• High-Dimensional PDEs