Tensor-Compressed Back-Propagation-Free Training for (Physics-Informed)
  Neural Networks
Backward propagation (BP) is widely used to compute the gradients in neural
network training. However, it is hard to implement BP on edge devices due to
the lack of hardware and software resources to support automatic
differentiation. This has tremendously increased the design complexity and
time-to-market of on-device training accelerators. This paper presents a
completely BP-free framework that only requires forward propagation to train
realistic neural networks. Our technical contributions are three-fold. Firstly,
we present a tensor-compressed variance reduction approach to greatly improve
the scalability of zeroth-order (ZO) optimization, making it feasible to handle
a network size that is beyond the capability of previous ZO approaches.
Secondly, we present a hybrid gradient evaluation approach to improve the
efficiency of ZO training. Finally, we extend our BP-free training framework to
physics-informed neural networks (PINNs) by proposing a sparse-grid approach to
estimate the derivatives in the loss function without using BP. Our BP-free
training only loses little accuracy on the MNIST dataset compared with standard
first-order training. We also demonstrate successful results in training a PINN
for solving a 20-dim Hamiltonian-Jacobi-Bellman PDE. This memory-efficient and
BP-free approach may serve as a foundation for the near-future on-device
training on many resource-constraint platforms (e.g., FPGA, ASIC,
micro-controllers, and photonic chips).
Yequan Zhao, Xinling Yu, Zhixiong Chen, Ziyue Liu, Sijia Liu, Zheng Zhang
# Tensor-Compressed Back-Propagation-Free Training for

###### Abstract

Backward propagation (BP) is widely used to compute the gradients in neural network training. However, it is hard to implement BP on edge devices due to the lack of hardware and software resources to support automatic differentiation. This has tremendously increased the design complexity and time-to-market of on-device training accelerators. This paper presents a completely BP-free framework that only requires forward propagation to train realistic neural networks. Our technical contributions are three-fold. Firstly, we present a tensor-compressed variance reduction approach to greatly improve the scalability of zeroth-order (ZO) optimization, making it feasible to handle a network size that is beyond the capability of previous ZO approaches. Secondly, we present a hybrid gradient evaluation approach to improve the efficiency of ZO training. Finally, we extend our BP-free training framework to physics-informed neural networks (PINNs) by proposing a sparse-grid approach to estimate the derivatives in the loss function without using BP. Our BP-free training only loses little accuracy on the MNIST dataset compared with standard first-order training. We also demonstrate successful results in training a PINN for solving a 20-dim Hamiltonian-Jacobi-Bellman PDE. This memory-efficient and BP-free approach may serve as a foundation for the near-future on-device training on many resource-constraint platforms (e.g., FPGA, ASIC, micro-controllers, and photonic chips).

## 1 Introduction

In neural network training, it is widely assumed that the gradient information can be computed via backward propagation (BP) [12] to perform SGD-type optimization. This is true on HPC and desktop computers, because they support Pytorch/Tensorflow backend and automatic differentiation (AD) packages [1] to compute exact gradients. However, performing BP is infeasible on many edge devices (e.g., FPGA, ASIC, photonic chips, embedded micro-processors) due to the lack of necessary computing and memory resources to support AD libraries. Manually calculating the gradient of a modern neural network on edge hardware is time-consuming due to the high complexity and lots of debug iterations. This can significantly delay the development and deployment of AI training accelerators. For instance, designing an FPGA inference accelerator may be done within one week via high-level synthesis by an experienced AI hardware expert, but designing an FPGA training accelerator can take one or two years due to the high complexity of implementing BP. Some analog edge computing platforms (e.g., photonic chips) even do not have scalable memory arrays to store the final or intermediate results of a chain rule-based gradient computation.

This paper investigates the end-to-end training of neural networks and physics-informed neural networks (PINNs) without using BP, which can potentially make on-device neural network training accelerator design as easy as designing an inference accelerator. The demand for edge-device training has been growing rapidly in recent years. One motivation is to ensure AI model performance under varying data stream and environment, as well as device parameter shift. Another motivation is the increasing concerns about data privacy, which requires end-to-end or incremental learning on edge devices based on local data [16]. In science and engineering, PINNs [14, 15] have been increasingly used to solve the forward and inverse problems of high-dimensional partial differential equations (PDE). A PINN needs to be trained again once the PDE initial conditions, boundary conditions, or measurement data changes. Applications of such on-device PINN training include (but are not limited to) safety-aware verification and control of autonomous systems [2, 23, 24] and MRI-based electrical property tomography [25].

We intend to perform BP-free training via stochastic zeroth-order (ZO) optimization [26, 27, 28, 29, 30], which uses a few forward evaluations with perturbed model parameters to approximate the true gradient. In the realm of deep learning, ZO optimization was primarily used for crafting black-box adversarial examples to assess neural network robustness [3, 16], and for parameter and memory-efficient model fine-tuning [26, 27]. Notably, ZO optimization has been rarely used in neural network training from scratch, because the variance of the ZO gradient estimation is large when the number of training variables increases.

Leveraging ZO optimization, we present a scalable BP
free approach for training both standard neural networks and PINNs from scratch. Our novel contributions include:

* [leftmargin=*,noitemsep,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,topsep=0pt,parsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,top,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,top,topsep=0pt,top,sep=0pt,top,topsep=0pt,topsep=0pt,top,topsep=0pt,topsep=0pt,top,topsep=0pt,top,topsep=0pt,top,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,top,topsep=0pt,top,0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,top,topsep=0pt,topsep=0pt,topsep=0pt,top,topsep=0pt,top,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,top,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,top,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsepsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsep=0pt,topsepsep=0pt,topsep=0pt,topsepsep=0pt,topsep=0pt,topsepsep=0pt,topsep=0pt,topsep=0pt,topsepsep=0pt,topsepsep=0pt,topsepsep=0pt,topsep=0pt,topsepsep=0pt,topsep=0pt,topsep=0pt,top
### Tensor-Train (TT) Variance Reduction

Our goal is to develop a BP-free training framework based on ZO optimization. As explained before, the ZO gradient estimation error is dominated by \(d\), which is the number of total trainable variables. In realistic neural network training, \(d\) can easily exceed \(10^{5}\) even for very simple image classification problems. This prevents the direct application of ZO optimization in end-to-end neural network training.

To improve the scalability of ZO training, we propose to significantly reduce the dimensionality and gradient MSE error via a _low-rank_ tensor-compressed training. Let \(\mathbf{W}\in\mathbb{R}^{M\times N}\) be a generic weight matrix in a neural network. We factorize its dimension sizes as \(M=\prod_{i=1}^{L}m_{i}\) and \(N=\prod_{j=1}^{L}n_{j}\), fold \(\mathbf{W}\) into a \(2L\)-way tensor \(\mathbf{\mathcal{W}}\in\mathbb{R}^{m_{1}\times m_{2}\times\cdots\times m_{L} \times n_{1}\times n_{2}\times\cdots\times n_{L}}\), and parameterize \(\mathbf{\mathcal{W}}\) with the tensor-train (TT) decomposition [10]:

\[\mathbf{\mathcal{W}}(i_{1},i_{2},\dots,i_{L},j_{1},j_{2},\dots,j_{L})\approx\prod_ {k=1}^{L}\mathbf{G}_{k}(i_{k},j_{k}) \tag{7}\]

Here \(\mathbf{G}_{k}(i_{k},j_{k})\in\mathbb{R}^{r_{k-1}\times r_{k}}\) is the \((i_{k},j_{k})\)-th slice of the TT-core \(\mathbf{\mathcal{G}}_{k}\in\mathbb{R}^{r_{k-1}\times m_{k}\times n_{k}\times r_{k}}\)by fixing its \(2\)nd index as \(i_{k}\) and \(3\)rd index as \(j_{k}\). The vector \((r_{0},r_{1},\dots,r_{L})\) is called TT-ranks with the constraint \(r_{0}=r_{L}=1\). This TT representation reduces the number of unknown variables from \(\prod\limits_{k=1}^{L}m_{k}n_{k}\) to \(\sum\limits_{k=1}^{L}r_{k-1}m_{k}n_{k}r_{k}\). The compression ratio can be controlled by the TT-ranks, which can be learnt automatically via the Bayesian tensor rank determination [11, 12].

In the ZO training process, we change the training variables from \(\mathbf{W}\) to the TT factors \(\{\mathbf{\mathcal{G}}_{k}\}_{k=1}^{L}\). This reduces the problem dimensionality \(d\) by several orders of magnitude, leading to dramatic reduction of the variance in the RGE gradient estimation and of the query complexity in the CGE gradient estimator. In the ZO training, we only need forward evaluations, where the original matrix-vector product is replaced with low-cost tensor-network contraction. This offers both memory and computing cost reduction in the ZO training process. While we assume \(\mathbf{W}\) as a weight matrix, other model parameters like embedding tables and convolution filters can be compressed simulatenously in the same way (with adjustable ranks) in the ZO training process.

### A Hybrid ZO Optimizer

With the above tensor-compressed variance reduction, we can employ either RGE or CGE for ZO gradient estimation and perform BP-free training. In practice, the RGE method converges very slowly in the late stage of training due to the large gradient error. CGE needs fewer training epochs due to more accurate gradient estimation, but it needs many more forward evaluations per gradient estimation since it only perturbs one model parameter in each forward evaluation.

To enhance both the accuracy and efficiency of the whole ZO training process, we employ a hybrid ZO training scheme that involves two stages:

* **ZO-signRGE coarse training:** The RGE method perturbs all model parameters simultaneously with a single random perturbation, requiring significantly fewer forward queries per epoch compared to the CGE method. However, the two-level stochasticity (in SGD and in the gradient estimation, respectively) of ZO via RGE (ZO-RGE) may cause a large gradient variance and lead to divergence in high-dimensional tasks. To address this issue, we adopt the concept from signSGD [13] and its ZO counterpart, ZO-signSGD [11], to de-noise the ZO gradient estimation by preserving only the sign for each update: \[\mathbf{\theta}_{t}\leftarrow\mathbf{\theta}_{t-1}-\alpha\mathrm{sign}\left[\sum_{i=1 }^{N}\frac{1}{N\mu}\left[\mathcal{L}\left(\mathbf{\theta}_{t-1}+\mu\mathbf{\xi}_{i} \right)-\mathcal{L}(\mathbf{\theta}_{t-1})\right]\mathbf{\xi}_{i}\right].\] We refer to this method as ZO-signRGE. ZO-signRGE retains only the sign of each gradient estimation, mitigating the adverse effects of the high variance of RGE. As a result, it exhibits better robustness to gradient noise and demonstrates faster empirical convergence.
* **ZO-CGE fine-tuning:** The CGE method requires \(d+1\) forward evaluations per gradient estimation. Therefore, the ZO training with CGE (ZO-CGE) requires many function queries in the whole training process. To accelerate training, we further adopt the idea of momentum [10]. The momentum method accumulates a velocity vector in directions of persistent reduction in the objective across iterations [21]. The descent direction \(\mathbf{b}^{t}\) is given by an exponential moving average of the past gradients \(\mathbf{b}^{t}\gets mb^{t-1}+\mathbf{g}^{t}\). Here \(m\) is the momentum term, \(\mathbf{g}^{t}\) is the estimated gradient vector at time \(t\), and \(\mathbf{b}^{0}=\mathbf{g}^{0}\). The ZO-signRGE coarse training rapidly explores a roughly converged solution with a small number of loss evaluations. When the coarse training fails to learn (e.g., the training loss exhibits trivial updates for several epochs), the optimizer switches to ZO-CGE to fine-tune the model.

## 4 BP-free Training for PINN

In this section, we extend the proposed TT-compressed hybrid ZO training to PINNs. Training a PINN is more challenging because the loss function in (5) involves first or even high-order derivatives. We intend to also avoid BP computation in the loss evaluation.

### Stein Gradient Estimation

Without loss of generality, for an input \(\mathbf{x}\in\mathbb{R}^{D}\) and an approximated PDE solution \(\mathbf{u}_{\mathbf{\theta}}(\mathbf{x})\in\mathbb{R}^{n}\) parameterized by \(\mathbf{\theta}\), we consider the first-order derivative \(\nabla_{\mathbf{x}}\mathbf{u}_{\mathbf{\theta}}\) and Laplacian \(\Delta\mathbf{u}_{\mathbf{\theta}}\) involved in the loss function of a PINN training. Our implementation leverages the Stein estimator [10]. Specifically, we represent the PDE solution \(\mathbf{u}_{\mathbf{\theta}}(\mathbf{x})\) via a Gaussian smoothed model:

\[\mathbf{u}_{\mathbf{\theta}}(\mathbf{x})=\mathbb{E}_{\mathbf{\delta}\sim\mathcal{N}(\mathbf{0},\sigma^{2}\mathbf{I})}f_{\mathbf{\theta}}(\mathbf{x}+\mathbf{\delta}), \tag{8}\]

where \(f_{\mathbf{\theta}}\) is a neural network with parameters \(\mathbf{\theta}\); \(\mathbf{\delta}\in\mathbb{R}^{D}\) is the random noise sampled from a multivariate Gaussian distribution \(\mathcal{N}(\mathbf{0},\sigma^{2}\mathbf{I})\). With this special formulation, the
first-order derivative and Laplacian of \(\mathbf{u}_{\mathbf{\theta}}(\mathbf{x})\) can be reformulated as the expectation terms:

\[\nabla_{\mathbf{x}}\mathbf{u}_{\mathbf{\theta}} =\mathbb{E}_{\mathbf{\delta}\sim\mathcal{N}(\mathbf{0},\sigma^{2}\mathbf{I})} \left[\frac{\mathbf{\delta}}{2\sigma^{2}}(f_{\mathbf{\theta}}(\mathbf{x}+\mathbf{\delta})-f_{ \mathbf{\theta}}(\mathbf{x}-\mathbf{\delta}))\right], \tag{9}\] \[\Delta\mathbf{u}_{\mathbf{\theta}} =\mathbb{E}_{\mathbf{\delta}\sim\mathcal{N}[\mathbf{0},\sigma^{2}\mathbf{I})} \left[f_{\mathbf{\theta}}(\mathbf{x}+\mathbf{\delta})+f_{\mathbf{\theta}}(\mathbf{x}-\mathbf{ \delta})-2f_{\mathbf{\theta}}(\mathbf{x})\right]\] \[\qquad\qquad\qquad\qquad\times\frac{\|\mathbf{\delta}\|^{2}-\sigma^{ 2}D}{2\sigma^{4}}.\]

In (He et al., 2023), the above expectation is computed by evaluating \(f_{\mathbf{\theta}}(\mathbf{x}+\mathbf{\delta})\) and \(f_{\mathbf{\theta}}(\mathbf{x}-\mathbf{\delta})\) at a set of i.i.d. Monte Carlo samples of \(\mathbf{\delta}\). Compared with finite difference (Chiu et al., 2022; Xiang et al., 2022) which requires repeated calculations of gradients for each dimension, the Stein estimator can directly provide the vectorized gradients for all dimensions, allowing efficient parallel implementations. However, the Monte-Carlo Stein gradient estimator (He et al., 2023) needs a huge number of (e.g., \(>10^{3}\)) function queries even if variance reduction is utilized. Therefore, it is highly desirable to develop a more efficient BP-free method for evaluating the derivative terms in the loss function.

### Sparse-Grid Stein Gradient Estimator

Now we leverage the sparse grid techniques (Garcke et al., 2006; Gerstner and Griebel, 1998) to reduce the number of function queries in the Stein gradient estimator. Sparse grids have been extensively used in the uncertainty quantification (Nobile, Tempone, and Webster, 2008) of stochastic PDEs, but they have not been used in PINN training.

To begin, we define a sequence of univariate quadrature rules \(V=\{V_{l}:l\in\mathbb{N}\}\). Here \(l\) denotes an accuracy level so that any polynomial function of order \(\leq l\) can be exactly integrated with \(V_{l}\). Each rule \(V_{l}\) specifies \(n_{l}\) nodes \(N_{l}=\{\delta_{1},\ldots,\delta_{n_{l}}\}\) and the corresponding weight function \(w_{l}:N_{l}\rightarrow\mathbb{R}\). A univariate quadrature rule \(V_{k}\) for a function \(f\) of a random variable \(\delta\), can be written as:

\[\int_{\mathbb{R}}f(\delta)p(\delta)\,d\delta\approx V_{k}[f]=\sum_{\delta_{j} \in N_{k}}w_{k}(\delta_{j})f(\delta_{j}). \tag{10}\]

Here \(p(\delta)\) is the probability density function (PDF) of \(\delta\).

Next, we consider the multivariate integration of a function \(f\) over a random vector \(\mathbf{\delta}=(\delta^{1},\ldots,\delta^{D})\). We denote the joint PDF of \(\mathbf{\delta}\) as \(p(\mathbf{\delta})=\prod_{m=1}^{D}p(\delta^{m})\) and define the \(D\)-variate quadrature rule with potentially different accuracy levels in each dimension indicated by the multi-index \(\mathbf{l}=(l_{1},l_{2},...,l_{D})\in\mathbb{N}^{D}\). The Smolyak algorithm (Gerstner and Griebel, 1998) can be used to construct sparse grids by combining full tensor-product grids of different accuracy levels and removing redundant points. Specifically, for any non-negative integer \(q\), define \(\mathbb{N}_{q}^{D}=\left\{\mathbf{l}\in\mathbb{N}^{D}:\sum_{m=1}^{D}l_{m}=D+q\right\}\) and \(\mathbb{N}_{q}^{D}=\emptyset\) for \(q<0\). The level-\(k\) Smolyak rule \(A_{D,k}\) for \(D\)-dim integration can be written as (Wasilkowski and Wozniakowski, 1995):

\[A_{D,k}[f]=\sum_{q=k-D}^{k-1}(-1)^{k-1-q}\left(\begin{array}{c}D-1\\ k-1-q\end{array}\right)\times \tag{11}\] \[\sum_{t\in\mathbb{N}_{q}^{D}}\left(V_{l_{1}}\otimes\cdots\otimes V_ {l_{D}}\right)[f].\]

It follows that:

\[A_{D,k}[f] =\sum_{q=k-D}^{k-1}\sum_{\mathbf{l}\in\mathbb{N}_{q}^{D}}\sum_{\mathbf{ \delta}^{1}\in N_{l_{1}}}\cdots\sum_{\mathbf{\delta}^{D}\in N_{l_{D}}}(-1)^{k-1-q}\times\] \[\left(\begin{array}{c}D-1\\ k-1-q\end{array}\right)\prod_{m=1}^{D}w_{l_{m}}(\delta^{m})f(\delta^{1}, \ldots,\delta^{D}),\]

which is a weighted sum of function evaluations \(f(\mathbf{\delta})\) for \(\mathbf{\delta}\in\bigcup_{q=k-D}^{k-1}\bigcup_{\mathbb{N}_{q}^{D}}\left(N_{l_{1}} \times\cdots\times N_{l_{D}}\right)\). The corresponding weight is \((-1)^{k-1-q}\left(\begin{array}{c}D-1\\ k-1-q\end{array}\right)\prod_{m=1}^{D}w_{l_{m}}(\delta^{m})\). For the same \(\mathbf{\delta}\) that appears multiple times for different combinations of values of \(\mathbf{l}\), we only need to evaluate \(f\) once and sum up the respective weights beforehand. The resulting level-\(k\) sparse quadrature rule defines a set of \(n_{L}\) nodes \(S_{L}=\{\mathbf{\delta}_{1},\ldots,\mathbf{\delta}_{n_{L}}\}\) and the corresponding weights \(\{w_{1},\ldots,w_{n_{L}}\}\). The \(D\)-dim integration can then be efficiently computed with the sparse grids as:

\[\int_{\mathbb{R}^{D}}f(\mathbf{\delta})p(\mathbf{\delta})d\mathbf{\delta}\approx A_{D,k}[f]= \sum_{j=1}^{n_{L}}w_{j}f(\mathbf{\delta}_{j}). \tag{12}\]

In practice, since the sparse grids and the weights do not depend on \(f\), they can be pre-computed for the specific quadrature rule, dimension \(D\), and accuracy level \(k\).

Finally, we implement the Stein gradient estimator in (9) via the sparse-grid integration. Noting that \(\mathbf{\delta}\sim\mathcal{N}\left(\mathbf{0},\sigma^{2}\mathbf{I}\right)\), we can use univariate Gaussian quadrature rules as basis to construct a level-\(k\) sparse Gaussian quadrature rule \(A_{D,k}^{\star}\) for \(D\)-variate integration. Then the first-order derivative and Laplacian in (9) is approximated as:

\[\nabla_{\mathbf{x}}\mathbf{u}_{\mathbf{\theta}} \approx\sum_{j=1}^{n_{L}^{\star}}w_{j}^{\star}\left[\frac{\mathbf{ \delta}_{j}^{\star}}{2\sigma^{2}}(f_{\mathbf{\theta}}(\mathbf{x}+\mathbf{\delta}_{j}^{ \star})-f_{\mathbf{\theta}}(\mathbf{x}-\mathbf{\delta}_{j}^{\star}))\right], \tag{13}\] \[\Delta\mathbf{u}_{\mathbf{\theta}} \approx\sum_{j=1}^{n_{L}^{\star}}w_{j}^{\star}\left(\frac{\|\mathbf{ \delta}_{j}^{\star}\|^{2}-\sigma^{2}D}{2\sigma^{4}}\right)\times\] \[\left(f_{\mathbf{\theta}}(\mathbf{x}+\mathbf{\delta}_{j}^{\star})+f_{\mathbf{ \theta}}(\mathbf{x}-\mathbf{\delta}_{j}^{\star})-2f_{\mathbf{\theta}}(\mathbf{x})\right),\]

where the node \(\mathbf{\delta}_{j}^{\star}\) and weight \(w_{j}^{\star}\) are defined by the sparse grid \(A_{D,k}^{\star}\). We remark that \(n_{L}^{\star}\) is usually significantly smaller than the number of Monte Carlo samples required to evaluate (9) when \(D\) is small. For example, we only need \(n_{L}^{\star}=2D^{2}+2D+1\) nodes to approximate a \(D\)-dimensional integral (\(D>1\)) using a level-\(3\) sparse Gaussian quadrature rule \(A_{D,3}^{\star}\).

## 5 Numerical Results

We test our tensor-compressed ZO training method on the MNIST dataset and a PDE benchmark. We compare our method with various baselines including standard first-order training and state-of-the-art BP-free training methods.

### MNIST Image Classification

To evaluate our proposed tensor-train (TT) compressed hybrid ZO training method on realistic image classification
tasks, we train a Multilayer Perceptron (MLP) network with two FC layers (768\(\times\)1024, 1024\(\times\)10) to classify the MNIST dataset. We compare our TT-compressed hybrid ZO training with four baseline methods: 1) first-order (FO) training using BP for accurate gradient evaluation, 2) ZO-RGE, 3) ZO-signRGE, and 4) ZO-CG. We use SGD to update model parameters. We adopt a mini-batch size 64, an initial learning learning rate \(\alpha\) as 1e-3 with exponential decaying rate of 0.9 every 10 epochs. All experiments run for 100 epochs.

**(1) Effectiveness of TT Variance Reduction.** The proposed TT variance reduction is the key to ensure the success of ZO training on realistic neural networks. For the 2-layer MLP model, we fold the sizes of the input, hidden, and output layers as size \(7\times 4\times 4\times 7\), \(8\times 4\times 4\times 8\), \(1\times 5\times 2\times 1\), respectively. We preset the TT-ranks as [1,\(r\),\(r\),\(r\),1]. By varying \(r\) we can control the compression ratio. We term this model as TT-MLP. As shown in Table 1, the TT-compressed approach can reduce the model parameters by \(205\times\). Due to the huge parameter and variance reduction, a hybrid ZO training can achieve a high testing accuracy \(96.64\%\), which is very close to the FO training results (\(97.16\%\)). We further compare our method with two other sets of baselines:

* **FO and ZO training on the MLP model.** As shown in Table 1, the original MLP model has a high testing accuracy of \(97.43\%\) when trained with a FO SGD algorithm. Note that this accuracy is very close to that of our TT-compressed hybrid ZO training method. The ZO training methods cannot converge to a descent solution due to the large variance caused by the 814K model parameters.
* **FO and ZO training with pruning.** An alternative method of reducing model parameters and ZO gradient variance is pruning. To train a sparse-pruned MLP (termed SP-MLP), we adopted the Gradient Signal Preserving (GraSP) method [20] at the initialization, which targets preserving the training dynamics after pruning. To ensure a fair comparison, we constrained the total number of parameters to be at the same level of TT-MLP. Table 1 show that pruning can indeed improve the quality of ZO training, but it still underperforms our proposed TT-compressed hybrid ZO training by a \(4.4\%\) accuracy drop.
2021a) assume a random pre-specified sparsity. We also extend the sparse method to ZO training.
* **Forward Gradient (FG):** Standard FG method uses forward-mode AD to evaluate the _forward gradient_ and performs SGD-type iterations (Baydin et al., 2022). Ref. (Ren et al., 2022) further scales up this method by introducing activation perturbation and local loss to reduce the variance of forward gradient evaluation.
* **Feedback Alignment (FA):** Standard FA (Lillicrap et al., 2016) uses random and fixed backward weights. Direct feedback alignment (DFA) propagated the error through fixed random feedback connections directly from the output layer to each hidden layer. (Nokland, 2016).

The results are summarized in Table 3. For ZO optimizers, we also compare the number of loss computations needed for each iteration. Our method can achieve the best accuracy in weight perturbation-based methods. Compared with SOTA FG and FA algorithms, our method can also achieve comparable accuracy in the MNIST dataset. Note that our method contains over 107\(\times\) fewer model parameters. This can greatly save the memory overhead for on-device training scenarios with restricted memory resources.

We remark that both FG and FA are not exactly black-box optimizers, since they rely on the computational graph of a given deep learning software. In practice, storing the computational graphs on an edge device can be expensive.

### PINN for Solving High-Dim HJB PDE

We further use our BP-free method to train a PINN arising from high-dim optimal control of robots and autonomous systems. We consider the following 20-dim HJB PDE:

\[\begin{split}&\partial_{t}u(\mathbf{x},t)+\Delta u(\mathbf{x},t)-0.05 \left\|\nabla_{\mathbf{x}}u(\mathbf{x},t)\right\|_{2}^{2}=-2,\\ & u(\mathbf{x},1)=\left\|\mathbf{x}\right\|_{1},\quad\mathbf{x}\in[0,1]^{20}, \;\;t\in[0,1].\end{split} \tag{14}\]

Here \(\left\|\cdot\right\|_{p}\) denotes an \(\ell_{p}\) norm. The exact solution is \(u(\mathbf{x},t)=\left\|\mathbf{x}\right\|_{1}+1-t\). The baseline neural network is a 3-layer **MLP** (\(21\times n,n\times n,n\times 1\), \(n\) denotes the number of neurons in the hidden layer) with sine activation. We consider four options for computing derivatives in the loss: 1) finite-difference (**FD**) (Lim, Dutta, and Rotaru, 2022), 2) Monte Carlo-based Stein Estimator (**SE**) (He et al., 2023), 3) our sparse-grid (**SG**) method, and 4) automatic differentiation (**AD**) as a golden reference. We approximate the solution \(u_{\mathbf{\theta}}\) by a transformed neural network \(f^{{}^{\prime}}_{\mathbf{\theta}}(\mathbf{x},t)=f_{\mathbf{\theta}}(\mathbf{x},t)+\left\|\mathbf{x }\right\|_{1}+1-t\), where \(f_{\mathbf{\theta}}(\mathbf{x},t)\) is the base neural network or its TT-compressed version. Specifically, \(u_{\mathbf{\theta}}(\mathbf{x},t)=f^{{}^{\prime}}_{\mathbf{\theta}}(\mathbf{x},t)\) for **FD** and **AD**, and \(u_{\mathbf{\theta}}(\mathbf{x},t)=\mathbb{E}_{(\mathbf{\delta}_{\mathbf{x}},\delta_{t})\sim \mathcal{N}(\mathbf{0},\sigma^{2}\mathbf{I})}f^{{}^{\prime}}_{\mathbf{\theta}}(\mathbf{x}+ \mathbf{\delta}_{\mathbf{x}},t+\delta_{t})\) for **SE** and **SG**. Here the transformed network is designed to ensure that our approximated solution either exactly satisfies (**FD**, **AD**) or closely adheres to the terminal condition (**SE**, **SG**), allowing us to focus solely on minimizing the HJB residual during training.

In our hybrid ZO optimization, we set \(\mu=0.1\) and \(N=10\) in the ZO-signRGE coarse training stage. With greatly reduced variance by TT, ZO-signRGE can provide sufficiently accurate results, so fine-tuning is not necessary for this example. In the loss evaluation, we set the step size in **FD** to 0.01 and the noise level \(\sigma\) to 0.1 in **SE** and **SG**. We use 1024 samples in **SE** and 925 samples in **SG** using a level-3 sparse Gaussian quadrature rule to approximate the expectations (8) and (9). Note that for most 2-dim or 3-dim PDEs, our SG method only requires \(10\sim 30\) samples. Our results are summarized below:

* **Effectiveness of BP-free loss computation:** We compare our sparse-grid loss computation with those using AD, FD, and Monte Carlo-based Stein estimator, respectively. To evaluate different loss evaluations, we perform FO training and report the results in Table 4. The BP-free loss computation does not hurt the model performance, and our SG method is competitive compared to the original PINN training using AD for derivative computation.
* **Evaluation of the BP-free PINN Training.** Table 5 shows the MSE error of fully BP-free PINN training, using various BP-free loss computation and the proposed hybrid ZO training. By employing a tensor-train (TT) compressed model to reduce the variance, our method achieved a validation loss similar to standard FO training using AD. This observation clearly demonstrates that our method can bypass BP in both loss evaluation and model parameter updates with little accuracy loss.
* **Comparison of dimensionality/variance reduction approaches.** We compare our TT-compressed model with standard MLPs with dense fully connected (FC) layers and sparse-pruned (SP) MLPs in Table 6. For a fair comparison, all networks have similar model parameters. TT-compressed and SP model can reach better results than FC models as TT and SP can preserve the wide hidden layer.

\begin{table}
\begin{tabular}{c c c c} \hline \hline Method & Params & \# of loss eval. & Best Acc. \\ \hline \multicolumn{4}{c}{Hybrid ZO Optimizer} \\ \hline HZO & 814K & 814K & 83.58 \\ TT-HZO (proposed) & **4k** & 4K & **96.64** \\ SP-HZO & **4k** & 4K & 92.24 \\ \hline \multicolumn{4}{c}{Sparse ZO Optimizer} \\ \hline FLOPS & 814K & **0.1K** & 83.50 \\ STP & 814K & 24.4K & 90.20* \\ SZO-SCD & 814K & 18.3K & 93.50* \\ \hline \multicolumn{4}{c}{Forward Gradient} \\ \hline FG-W & 272K & / & 90.75 \\ LG-FG-A & 272K & / & 96.76 \\ FG-W & 429K & / & 91.44 \\ LG-FG-A & 429K & / & **97.45** \\ \hline \multicolumn{4}{c}{Feedback Alignment} \\ \hline FA & 784K & / & 97.90 \\ DFA & 1.26M & / & **98.30** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Comparison of BP-free optimizers on the MNIST dataset. * means fine-tuning on a pre-trained model. In the Forward Gradient group, FG refers to forward gradient, LG refers to local objective functions added, -W refers to weight-perturbed, and -A refers to activity-perturbed.
The SP model shows superior efficiency in FO training and achieves similar result with the TT model in ZO training. Note that the SP model only reduces the trainable parameters but not the model parameters, so TT models can achieve better memory and computation saving.

## 6 Relevant Work

Training on edge devices.On-device inference has been well studied. However, training requires extra computation for BP and extra memory for intermediate results. Edge devices usually have tight memory and power budget, and run without an operating system, making it infeasible to implement modern deep learning frameworks that support automatic differentiation [1]. Most on-device training frameworks implemented BP by hand and only update within a small subspace (e.g., the last classifier layer [1], biases [1], etc.), only capable of incremental learning or transfer learning by fine-tuning a pre-trained model. Some work leveraged pruning or sparse training to reduce the trainable parameter number[11, 14], but cannot achieve real memory savings. Recent work found tiny, trainable subspace exists in deep neural networks [10, 12], but a pre-training is needed to identify such a subspace, which does not save the training overhead. Tensorized training [10, 11] offers orders-of-magnitude memory reduction in the training process, enabling end-to-end neural network training on FPGA. However, it is non-trivial to implement this FO method to train larger models on edge devices due to the complex gradient computation.

Back-propagation-free trainingDue to the complexity of performing BP on edge devices, several BP-free training algorithms have been proposed. These methods have gained more attentions in recent years as BP is also considered "biologically implausible". ZO optimization [1, 12, 13, 14, 15, 16, 17, 18] plays an important role in tackling signal processing and machine learning problems, where actual gradient information is infeasible. For specific use cases, please refer to the survey by [12]. ZO optimizations were also applied in on-device training for optical neural networks (ONN) [14, 15, 16]. A ZO SGD-based method FLOPS [14] was proposed for end-to-end ONN training for vowel recognition, and a stochastic ZO sparse coordinate descent SZO-SCD [14] was further proposed for efficient fine-tuning for ONN. But most ZO optimizations scale poorly when training real-size neural networks from scratch, due to their dimension-dependent gradient errors. Forward-forward algorithm was proposed for biologically plausible learning [10]. Other BP-free training frameworks include the forward gradient method [1], which updates the weights based on the directional gradient computed by forward-mode AD along a random direction. This method is further scaled up by leveraging activity perturbation and local loss for variance reduction[15]. However, existing BP-free methods focus on training on HPC, lacking special concerns for the special challenges of on-device training. Furthermore, BP-free training of PINN remains an empty field.

## 7 Conclusion

Due to the high complexity and memory overhead of back propagations (BP) on edge devices, this paper has proposed a completely BP-free zeroth-order (ZO) approach to train real-size neural networks from scratch. Our method utilizes tensor compression to reduce the ZO gradient variance and a hybrid approach to improve the efficiency of ZO training. This has enabled highly accurate ZO training on the MNIST dataset, showing superior performance than sparse ZO approach. We have further extended this memory-efficient approach to train physics-informed neural networks (PINN) by developing a sparse-grid Stein gradient estimator for the loss evaluation. Our approach has been successfully solved a 20-dim HJB PDE, achieving similar accuracy with standard training using first-order optimization and automatic differentiation. Due to the huge memory reduction and the BP-free nature, our method can be easily implemented on various resource-constraint edge devices in the near future.

The ZO-CGE fine-tuning stage of our method requires many more function queries than the ZO-signRGE coarse-training stage. This prevents its applications in training large model (transformers) from scratch. We will address this issue in the future.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline Network & Neurons & Params & FO-AD & ZO-SG \\ \hline FC-3 & 768 & 608257 & 3.58E-04 & 1.62E-01 \\ \hline FC-3 & 78 & 7800 & 1.52E-02 & 4.35E-02 \\ FC-2 & 320 & 7700 & 6.55E-04 & 1.00E-02 \\ TT-3 & 768 & 7745 & 4.25E-05 & **4.79E-05** \\ SP-3 & 768 & 7603 & **1.54E-09** & 6.29E-05 \\ \hline \hline \end{tabular}
\end{table}
Table 6: Validation loss of various models for the HJB PDE

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline Network & Params & FO+AD & FO+FD & FO+SE & FO+SG \\  & & & & & (proposed) \\ \hline MLP & 608257 & 3.58E-04 & 5.97E-05 & 8.98E-04 & 4.27E-04 \\ \hline TT-MLP & 7745 & 4.28E-05 & **4.90E-05** & 2.28E-04 & 7.14E-05 \\  & 3209 & 1.63E-05 & 4.56E-05 & 3.82E-04 & **4.49E-05** \\ \hline \hline \end{tabular}
\end{table}
Table 4: FO training with different loss evaluations.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline Network & Params & FO+AD & ZO+FD & ZO+SE & ZO+SG \\  & & & & & (proposed) \\ \hline MLP & 608257 & 3.58E-04 & 9.23E-02 & diverge & 1.62E-01 \\ \hline TT-MLP & 7745 & 4.28E-05 & 3.79E-04 & 4.35E-04 & **4.79E-05** \\  & 3209 & 1.63E-05 & 3.82E-04 & 2.75E-04 & **1.44E-05** \\ \hline \hline \end{tabular}
\end{table}
Table 5: ZO training results for 20-dim HJB PDE.