提出SDZE框架,让高维偏微分方程神经网络训练更稳定且省内存。
Stochastic Dimension Zeroth-Order Estimator: Stable and Memory-Efficient Training of PINNs
- 用公共随机数同步技术消除零阶优化的方差爆炸问题。
- 在单张A100 GPU上训练1000万维PINN,内存仅需常数级。
- 适合大规模高维物理方程求解,尤其对资源受限场景友好。
高维高阶偏微分方程的物理信息神经网络(PINNs)主要受限于空间导数复杂度$/mathcal{O}(d^k)$和反向传播的$/mathcal{O}(P)$内存开销。虽然随机空间估计器将空间复杂度降至$/mathcal{O}(1)$,但依赖一阶优化仍导致大规模时内存过高。零阶(ZO)优化可免于反向传播,但直接结合随机空间算子与ZO扰动会引发$/mathcal{O}(1/\varepsilon^2)$方差爆炸,导致数值发散。为此,本文提出统一框架SDZE,实现空间与内存的维度无关复杂度。具体地,通过公共随机数同步(CRNS)在扰动间锁定随机种子,代数抵消$/mathcal{O}(1/\varepsilon^2)$方差;并引入隐式无矩阵子空间投影,将参数探索方差从$/mathcal{O}(P)$降至$/mathcal{O}(r)$,同时保持$/mathcal{O}(1)$优化器内存。实验证明,SDZE可在单张NVIDIA A100 GPU上训练1000万维PINN,相比现有最优方法显著提升速度与内存效率。
原文摘要 · Abstract (English)
Physics-Informed Neural Networks (PINNs) for high-dimensional and high-order partial differential equations (PDEs) are primarily constrained by the $\mathcal{O}(d^k)$ spatial derivative complexity and the $\mathcal{O}(P)$ memory overhead of backpropagation (BP). While randomized spatial estimators successfully reduce the spatial complexity to $\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 $\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 $\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 $\mathcal{O}(P)$ to $\mathcal{O}(r)$ while maintaining an $\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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。