[Paper Review] TT-PINN: A Tensor-Compressed Neural PDE Solver for Edge Computing

Ziyue Liu, Xinling Yu+1 more|arXiv (Cornell University)|Jul 4, 2022

Model Reduction and Neural Networks4 citations

TL;DR

This paper proposes TT-PINN, a tensor-compressed physics-informed neural network that leverages Tensor-Train (TT) decomposition to drastically reduce model parameters and memory usage while maintaining high accuracy in solving partial differential equations (PDEs). By replacing dense weight matrices in fully connected layers with low-rank TT-cores, TT-PINN achieves up to 15× parameter reduction and outperforms standard PINNs with significantly fewer parameters, enabling efficient training on resource-constrained edge devices.

ABSTRACT

Physics-informed neural networks (PINNs) have been increasingly employed due to their capability of modeling complex physics systems. To achieve better expressiveness, increasingly larger network sizes are required in many problems. This has caused challenges when we need to train PINNs on edge devices with limited memory, computing and energy resources. To enable training PINNs on edge devices, this paper proposes an end-to-end compressed PINN based on Tensor-Train decomposition. In solving a Helmholtz equation, our proposed model significantly outperforms the original PINNs with few parameters and achieves satisfactory prediction with up to 15$ imes$ overall parameter reduction.

Motivation & Objective

• To address the challenge of training large physics-informed neural networks (PINNs) on edge devices with limited memory, computing, and energy resources.
• To enable end-to-end training of PINNs on edge platforms by reducing model size and memory footprint without sacrificing solution accuracy.
• To explore the feasibility of using Tensor-Train (TT) decomposition as a compression mechanism for PINNs in solving PDEs.
• To demonstrate that a compressed PINN architecture can match or exceed the performance of larger, uncompressed PINNs.

Proposed method

• The TT-PINN framework replaces the dense weight matrices in fully connected layers of a standard PINN with a low-rank Tensor-Train (TT) decomposition, reducing parameter count and memory usage.
• The weight matrix is factorized into a sequence of TT-cores, where each core is a small tensor that stores the compressed representation of the original matrix.
• The model is trained end-to-end using automatic differentiation, with gradients backpropagated directly through the TT-cores to update the parameters.
• The TT-ranks are determined based on desired compression ratios, and the method supports both fixed and adaptive rank determination.
• The architecture uses a Sine activation function and Xavier initialization, and is trained with the Adam optimizer and learning rate decay.
• The physical laws (PDE, boundary, and initial conditions) are embedded directly into the loss function, preserving the physics-informed nature of the network.

Experimental results

Research questions

• RQ1 Can Tensor-Train decomposition effectively compress PINNs while preserving their solution accuracy for PDEs on edge devices?
• RQ2 How does the performance of a TT-compressed PINN compare to a standard PINN of similar or larger size in solving a Helmholtz equation?
• RQ3 To what extent can parameter and memory reduction be achieved in PINNs using TT decomposition without degrading predictive performance?
• RQ4 Can TT-PINN enable training on edge devices with limited computational and memory resources while maintaining high solution fidelity?

Key findings

• TT-PINN with only 3,713 parameters achieved a relative ℓ₂ error of 3.03e-2, matching the performance of a standard PINN with 50,049 parameters.
• A TT-PINN compressing a 256×256 weight matrix by 100× achieved a mean squared error of 2.25e-4 and a relative ℓ₂ error of 3.03e-2, outperforming a standard PINN with 12,737 parameters.
• The TT-PINN with 40× compression of a 256×256 matrix achieved a relative ℓ₂ error of 7.75e-3, significantly better than the standard PINN with 12,737 parameters.
• The TT-PINN achieved up to 15× overall parameter reduction compared to standard PINNs while maintaining or improving solution accuracy.
• Visual comparisons show that TT-PINN with 3,713 parameters produced a more accurate solution than a standard PINN with 12,737 parameters.
• The method demonstrates that low-rank TT structures can preserve the expressive power of larger PINNs, enabling efficient deployment on edge devices.