cs.LGMar 25, 2026

Stochastic Dimension Zeroth-Order Estimator: Stable and Memory-Efficient Training of PINNs

Authors: Zhangyong Liang, Huanhuan Gao

Abstract

Physics-Informed Neural Networks (PINNs) for high-dimensional and high-order partial differential equations (PDEs) are primarily constrained by the O(dk)\mathcal{O}(d^k) spatial derivative complexity and the O(P)\mathcal{O}(P) memory overhead of backpropagation (BP). While randomized spatial estimators successfully reduce the spatial complexity to O(1)\mathcal{O}(1), their reliance on first-order optimization still leads to prohibitive memory consumption at scale. Zeroth-order (ZO) optimization offers a BP-free alternative; however, naively combining randomized spatial operators with ZO perturbations triggers a variance explosion of O(1/ε2)\mathcal{O}(1/\varepsilon^2), leading to numerical divergence. To address these challenges, we propose the \textbf{S}tochastic \textbf{D}imension-free \textbf{Z}eroth-order \textbf{E}stimator (\textbf{SDZE}), a unified framework that achieves dimension-independent complexity in both space and memory. Specifically, SDZE leverages \emph{Common Random Numbers Synchronization (CRNS)} to algebraically cancel the O(1/ε2)\mathcal{O}(1/\varepsilon^2) variance by locking spatial random seeds across perturbations. Furthermore, an \emph{implicit matrix-free subspace projection} is introduced to reduce parameter exploration variance from O(P)\mathcal{O}(P) to O(r)\mathcal{O}(r) while maintaining an O(1)\mathcal{O}(1) optimizer memory footprint. Empirical results demonstrate that SDZE enables the training of 10-million-dimensional PINNs on a single NVIDIA A100 GPU, delivering significant improvements in speed and memory efficiency over state-of-the-art baselines.

Explore similar work

CardsList