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

arXiv: https://arxiv.org/abs/2603.24002 · HTML: https://mlsanity.com/paper/2603.24002/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 41 of its 65 references that are in the map.
- First posted on arXiv 25 Mar 2026.

## Builds on (most cited first, 30 of 41)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Fine-Tuning Language Models with Just Forward Passes](https://mlsanity.com/paper/2305.17333/) (2023)
- [Physics-informed neural networks with hard constraints for inverse design](https://mlsanity.com/paper/2102.04626/) (2021)
- [Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems](https://mlsanity.com/paper/2111.02801/) (2021)
- [Optimal rates for zero-order convex optimization: the power of two function evaluations](https://mlsanity.com/paper/1312.2139/) (2013)
- [A Primer on Zeroth-Order Optimization in Signal Processing and Machine Learning](https://mlsanity.com/paper/2006.06224/) (2020)
- [Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark](https://mlsanity.com/paper/2402.11592/) (2024)
- [Zeroth-Order Stochastic Variance Reduction for Nonconvex Optimization](https://mlsanity.com/paper/1805.10367/) (2018)
- [CAN-PINN: A Fast Physics-Informed Neural Network Based on Coupled-Automatic-Numerical Differentiation Method](https://mlsanity.com/paper/2110.15832/) (2021)
- [Tutorial on amortized optimization](https://mlsanity.com/paper/2202.00665/) (2022)
- [Tackling the Curse of Dimensionality with Physics-Informed Neural Networks](https://mlsanity.com/paper/2307.12306/) (2023)
- [Gradients without Backpropagation](https://mlsanity.com/paper/2202.08587/) (2022)
- [A Theoretical and Empirical Comparison of Gradient Approximations in Derivative-Free Optimization](https://mlsanity.com/paper/1905.01332/) (2019)
- [Sparse MeZO: Less Parameters for Better Performance in Zeroth-Order LLM Fine-Tuning](https://mlsanity.com/paper/2402.15751/) (2024)
- [Variance-reduced Zeroth-Order Methods for Fine-Tuning Language Models](https://mlsanity.com/paper/2404.08080/) (2024)
- [Improved Zeroth-Order Variance Reduced Algorithms and Analysis for Nonconvex Optimization](https://mlsanity.com/paper/1910.12166/) (2019)
- [Zeroth-Order Regularized Optimization (ZORO): Approximately Sparse Gradients and Adaptive Sampling](https://mlsanity.com/paper/2003.13001/) (2020)
- [Randomized Numerical Linear Algebra : A Perspective on the Field With an Eye to Software](https://mlsanity.com/paper/2302.11474/) (2023)
- [Zeroth-Order Fine-Tuning of LLMs in Random Subspaces](https://mlsanity.com/paper/2410.08989/) (2024)
- [ZO-AdaMU Optimizer: Adapting Perturbation by the Momentum and Uncertainty in Zeroth-order Optimization](https://mlsanity.com/paper/2312.15184/) (2023)
- [AdaZeta: Adaptive Zeroth-Order Tensor-Train Adaption for Memory-Efficient Large Language Models Fine-Tuning](https://mlsanity.com/paper/2406.18060/) (2024)
- [Randomized Automatic Differentiation](https://mlsanity.com/paper/2007.10412/) (2020)
- [Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networks](https://mlsanity.com/paper/2312.14499/) (2023)
- [Score-Based Physics-Informed Neural Networks for High-Dimensional Fokker-Planck Equations](https://mlsanity.com/paper/2402.07465/) (2024)
- [Zeroth order optimization with orthogonal random directions](https://mlsanity.com/paper/2107.03941/) (2021)
- [Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators](https://mlsanity.com/paper/2412.00088/) (2024)
- [Unbiased MLMC stochastic gradient-based optimization of Bayesian experimental designs](https://mlsanity.com/paper/2005.08414/) (2020)
- [Bias-Variance Trade-off in Physics-Informed Neural Networks with Randomized Smoothing for High-Dimensional PDEs](https://mlsanity.com/paper/2311.15283/) (2023)

## Abstract (arXiv)

Physics-Informed Neural Networks (PINNs) for high-dimensional and high-order partial differential equations (PDEs) are primarily constrained by the O(dk) spatial derivative complexity and the O(P) memory overhead of backpropagation (BP). While randomized spatial estimators successfully reduce the spatial complexity to 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/varepsilon2), leading to numerical divergence. To address these challenges, we propose the Stochastic Dimension-free Zeroth-order Estimator (SDZE), a unified framework that achieves dimension-independent complexity in both space and memory. Specifically, SDZE leverages Common Random Numbers Synchronization (CRNS) to algebraically cancel the O(1/varepsilon2) variance by locking spatial random seeds across perturbations. Furthermore, an implicit matrix-free subspace projection is introduced to reduce parameter exploration variance from O(P) to O(r) while maintaining an 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.
