Tensor-Compressed and Fully-Quantized Training of Neural PDE Solvers

TL;DR

This paper introduces a scalable, energy-efficient framework for training physics-informed neural networks on edge devices by combining quantization, tensor decomposition, and specialized hardware acceleration, achieving high accuracy and significant speedups.

Contribution

It presents a novel integrated framework with mixed-precision training, Stein's estimator quantization, and tensor-train decomposition, enabling efficient PDE solving on resource-constrained platforms.

Findings

Achieves comparable or better accuracy than full-precision baselines.

Delivers 5.5x to 83.5x speedups on various PDE problems.

Provides 159.6x to 2324.1x energy savings.

Abstract

Physics-Informed Neural Networks (PINNs) have emerged as a promising paradigm for solving partial differential equations (PDEs) by embedding physical laws into neural network training objectives. However, their deployment on resource-constrained platforms is hindered by substantial computational and memory overhead, primarily stemming from higher-order automatic differentiation, intensive tensor operations, and reliance on full-precision arithmetic. To address these challenges, we present a framework that enables scalable and energy-efficient PINN training on edge devices. This framework integrates fully quantized training, Stein's estimator (SE)-based residual loss computation, and tensor-train (TT) decomposition for weight compression. It contributes three key innovations: (1) a mixed-precision training method that use a square-block MX (SMX) format to eliminate data duplication during backpropagation; (2) a difference-based quantization scheme for the Stein's estimator that mitigates underflow; and (3) a partial-reconstruction scheme (PRS) for TT-Layers that reduces quantization-error accumulation. We further design PINTA, a precision-scalable hardware accelerator, to fully exploit the performance of the framework. Experiments on the 2-D Poisson, 20-D Hamilton-Jacobi-Bellman (HJB), and 100-D Heat equations demonstrate that the proposed framework achieves accuracy comparable to or better than full-precision, uncompressed baselines while delivering 5.5x to 83.5x speedups and 159.6x to 2324.1x energy savings. This work enables real-time PDE solving on edge devices and paves the way for energy-efficient scientific computing at scale.