arXiv:2512.09202cs.LGcs.AI2025-12被引 1

让物理神经网络在边缘设备上高效求解偏微分方程。

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

  • 用张量压缩与全量化训练降低内存和算力消耗。
  • 在2D泊松、20D HJB等方程上实现5.5到83.5倍加速。
  • 适合需要实时科学计算的边缘部署场景。

物理信息神经网络(PINNs)通过将物理定律嵌入训练目标,成为求解偏微分方程(PDEs)的有前景方法。然而,其在资源受限设备上的部署受限于高阶自动微分、密集张量运算及对全精度算术的依赖所导致的巨大计算与内存开销。为此,我们提出一个可在边缘设备上实现可扩展、节能的PINN训练框架。该框架融合了全量化训练、基于Stein估计器(SE)的残差损失计算以及张量-列车(TT)分解的权重压缩。其三大创新包括:(1)采用平方块混合精度格式(SMX),消除反向传播中的数据冗余;(2)针对Stein估计器设计差异量化方案,缓解下溢问题;(3)提出部分重建方案(PRS)用于TT层,减少量化误差累积。我们进一步设计了可调精度的硬件加速器PINTA,以充分释放框架性能。在二维泊松方程、二十维汉密尔顿-雅各比-贝尔曼(HJB)方程和一百维热方程上的实验表明,该框架在精度上达到或优于全精度无压缩基线,同时实现5.5至83.5倍加速和159.6至2324.1倍能效提升。本工作实现了边缘设备上的实时偏微分方程求解,为大规模节能科学计算铺平道路。

原文摘要 · Abstract (English)

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.

PINNs边缘计算张量压缩量化训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。