Real-Time FJ/MAC PDE Solvers via Tensorized, Back-Propagation-Free Optical PINN Training

We focus on the ONN [7] architecture with singular value decomposition (SVD) to implement matrix-vector multiplication (MVM), i.e., $y=\bm{W}x=\bm{U\Sigma V^{*}}x$. The unitary matrices $\bm{U}$ and $\bm{V^{*}}$ are implemented by MZIs in Clements mesh [11]. The parametrization of $\bm{U}$ and $\bm{V^{*}}$ is given by $\bm{U}(n)=\bm{D}\prod_{i=2}^{n}\prod_{j=1}^{i-1}\bm{R}_{ij}\left(\phi_{ij}\right)$ where $\bm{D}$ is a diagonal matrix, and each 2-dimensional rotator $\bm{R}_{ij}(\phi_{ij})$ can be implemented by a $2\times 2$ MZI containing two phase shifters and two 50/50 splitters. We denoted all programmable phases as $\bm{\Phi}$ and $\bm{W}$ is parametrized as $\bm{W}(\bm{\Phi})$.

To increase the scalability of ONN, a tensorized optical neural network (TONN)[19] is proposed to realize large-scale ONNs with reduced hardware resources (i.e., MZIs) using the tensor-train (TT) decomposition algorithm. Let $\bm{W}\in\mathbb{R}^{M\times N}$ be a generic weight matrix in a neural network. We factorize its dimension sizes as $M=\prod^{L}_{i=1}m_{i}$ and $N=\prod^{L}_{j=1}n_{j}$, fold $\bm{W}$ into a $2L$-way tensor $\mathbfcal{W}\in\mathbb{R}^{m_{1}\times m_{2}\times\dots\times m_{L}\times n_{1}\times n_{2}\times\dots\times n_{L}}$, and parameterize $\mathbfcal{W}$ with the TT decomposition [20]:

$${\mathbfcal{W}}(i_{1},i_{2},\dots,i_{L},j_{1},j_{2},\dots,j_{L})\approx\prod^{L}\nolimits_{k=1}\mathbf{G}_{k}(i_{k},j_{k})$$

(1)

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 $\mathbfcal{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^{L}_{k=1}m_{k}n_{k}$ to $\sum_{k=1}^{L}r_{k-1}m_{k}n_{k}r_{k}$. The detailed architecture of TONN can be found in [19].

Consider the well-posed initial value partial differential equation (PDE) problem described by:

	$\displaystyle\mathcal{N}[\bm{u}(\bm{x},t)]$	$\displaystyle=l(\bm{x},t),$		$\displaystyle\bm{x}\in\Omega,~{}~{}t\in[0,T],$		(2)
	$\displaystyle\mathcal{I}[\bm{u}(\bm{x},0)]$	$\displaystyle=g(\bm{x}),$		$\displaystyle\bm{x}\in\Omega,$

where $\bm{x}$ and $t$ are the spatial and temporal coordinates; $\Omega\subset\mathbb{R}^{D}$ and $T$ denote the spatial domain and time horizon, respectively; $\mathcal{N}$ is a general nonlinear differential operator; $\mathcal{I}$ represents the initial condition; $\bm{u}\in\mathbb{R}^{n}$ is the solution for the PDE described above. In PINNs [3], a neural network $\bm{u}(\bm{x},t;\bm{\theta})$, parameterized by $\bm{\theta}$, is substituted into PDE (2), resulting in a residual defined as:

$$r(\bm{x},t;\bm{\theta}):=\mathcal{N}[\bm{u}(\bm{x},t;\bm{\theta})]-l(\bm{x},t).$$

(3)

The parameters $\bm{\theta}$ can be obtained by minimizing the loss $\mathcal{L}(\bm{\theta})=\mathcal{L}_{r}(\bm{\theta})+\lambda\mathcal{L}_{0}(\bm{\theta}),$ where

$$\mathcal{L}_{r}(\bm{\theta})=\frac{1}{N_{r}}\sum_{i=1}^{N_{r}}\left\|r(\bm{x_{r}}^{i},t_{r}^{i};\bm{\theta})\right\|_{2}^{2}~{}~{}\text{and}~{}~{}\mathcal{L}_{0}(\bm{\theta})=\frac{1}{N_{0}}\sum_{i=1}^{N_{0}}\left\|\mathcal{I}[\bm{u}(\bm{x_{0}}^{i},0;\bm{\theta})]-g(\bm{x_{0}}^{i})\right\|_{2}^{2},$$

(4)

are the residuals of the PDE and the initial (or terminal) condition, respectively. To adapt PINNs to the edge with constraints in memory, computation, and energy, [21] introduced a tensor-compressed PINN framework, where fully-connected layers are decomposed into a series of TT-cores, as in (1).

The block diagram of our optical PINN training accelerator is shown in Fig. 1. This accelerator does not perform any BP. Instead, it repeatedly call an optical inference accelerator TONN[19] to obtain some loss information. The collected loss information is process in a digital control system, and gradient information is estimated via a zeroth-order optimization to update PINN model parameters. Since propagations are not used, intermediate results will not be computed or stored on the photonic chip.

In the following, we give the details of our TONN design and BP-free training method.

We present two designs of optical neural networks based on tensor-train (TT) compression. The TT-based optical neural network design can greatly reduce the number of photonic devices, latency and energy cost. Furthermore, it can reduce the number of on-chip training variables and improve the convergence of the on-chip training framework.

The first design, called TONN-1, is illustrated in Fig. 2. In this design, the tensor multiplications between the input data and all tensor-train cores (the whole tensorized matrix) are realized in a single clock cycle by cascading the photonic tensor cores in the space domain and adding parallelism in the wavelength domain[19].

The second design, called TONN-2 and shown in Fig. 3, uses a single wavelength-parallel photonic tensor core [22] with time multiplexing. Compared with TONN-1, TONN-2 exhibits a smaller footprint at the expense of higher latency and additional memory requirements. In each clock cycle, the photonic tensor core with parallel processing in the wavelength domain is updated to multiply with the input tensor. Then, the intermediate output data is stored in the buffer for the next cycle.

We adopt the idea of [23] to implement a fully BP-free PINN training method to mitigate the memory bottleneck of on-chip photonic computing. In PINN training, the BP process should be avoided in both the loss function evaluation and in the SGD-type optimization step.

All numerical simulations are based on a software implementation built upon PyTorch [27] backend and TorchONN library [28] to simulate the computational model of an optical computing platform. To show the effectiveness and robustness of our design, we compared our method with different training paradigms. Off-chip Training denotes first pre-training on electrical digital platforms, e.g., CPUs and GPUs, then mapping the trained model to photonic devices. The gradients w.r.t. model parameters and the derivatives w.r.t. the input are computed by BP. Our proposed tensor-compressed BP-free training belongs to On-chip Training as it directly tunes photonic devices (i.e., phase-shifters in MZIs) on-chip and trains from scratch. For off-chip training, we implemented hardware-aware training that incorporates various hardware imperfections and its counterpart hardware-unaware training that runs on an ideal computational model. The hardware-aware training is a hardware-restricted learning problem, where we considered phase-shifter $\gamma$ coefficient drift $\bm{\Gamma}\sim\mathcal{N}(\gamma,\sigma^{2}_{\gamma})$ [13, 29] caused by fabrication variations and thermal cross-talk between adjacent devices $\bm{\Omega}$ [13, 29, 30], and phase bias due to manufacturing error $\bm{\Phi}_{b}\sim\mathcal{U}(0,2\pi)$ and the objective became $\bm{\Phi}^{*}={\arg\min}_{\bm{\Phi}}\mathcal{L}(\bm{W}(\bm{\Omega\Gamma\Phi}+\bm{\Phi}_{b}))$. For on-chip training, we incorporate the same hardware imperfections to mimic the actual analog hardware.

Our results are provided in Table:1. We report the validation loss which is the mean square error (MSE) w.r.t. the ground truth. After training, our proposed BP-free tensor-compressed PINN training achieves a validation loss of 5.53E-3, indicating that the model fits the ground truth well. The tensor-compressed ONN outperforms the un-compressed ONN, indicating that our tensor-compressed training is capable of preserving the expressive power of a wide ONN with greatly reduced model parameters (396$\times$ fewer in this case).

Off-chip training achieves a similar validation loss on the pre-trained model. However, after mapping to real photonic devices, the performance greatly degrades due to the hardware imperfection. The hardware-aware training does not help significantly as the imperfection model in software is not identical to real hardware. Our proposed method inherently circumvents this problem as it directly tunes on the fabricated hardware during on-chip training, thus demonstrating better robustness and better performance.

Energy Consumption: The total energy of the accelerators is mainly consumed in the ADC, DAC, and digital control systems. Here, we focus on the photonic energy consumption per forward, which consists of five parts: laser wall-plug power, microring modulator power, MZI mesh power, microring add-drop filter power, and PD receiver power. The conventional ONN has insurmountable optical loss due to the square scaling rule, so the energy cannot be calculated. The TONN-2 consumes slightly less energy per forward due to lower insertion loss even though it requires 64 cycles.

Latency: The latency per inference in on-chip training is calculated by: $t_{\rm{inference}}=n_{\rm{cycle}}*(t_{\rm{DAC}}+t_{\rm{tuning}}+t_{\rm{opt}}+t_{\rm{ADC}})+t_{\rm{DIG}}$, where $t_{\rm{DAC}}$ is the ADC conversion delay ($\sim$24 ns), $t_{\rm{tuning}}$ is the metal-oxide-semiconductor capacitor (MOSCAP) phase shifter tuning delay ($\sim$0.1 ns), $t_{\rm{opt}}$ is the propagation latency of optical signal ($\sim$51.2 ns for ONN, $\sim$1.6 ns for TONN-1, and $\sim$0.4 ns for TONN-2), $t_{\rm{ADC}}$ is the DAC delay($\sim$24 ns), and $t_{\rm{DIG}}$ is the digital computation overhead ($\sim$500 ns) for gradient calculation and phase updates. The TONN-2 uses 64 cycles for one inference, while ONN and TONN-1 only needs one cycle.

Training Efficiency: In our 20D-HJB example, we need 42 inferences for each loss evaluation and 10 loss evaluations for gradient estimation. Suppose a mini-batch size of 100, 4.20E4 inferences are required for one epoch. The energy consumption per epoch is estimated as 2.71E-04 J and the latency per epoch is estimated as 0.23 ms for TONN-1. On average training reaches a good solution after 5000 epochs, which corresponds to 1.36 J and 1.15 s for solving a 20D-HJB equation.

Footprint: Only the footprint of the photonic devices, which occupy the major area of the accelerator, is used for comparison. The photonic footprint includes the areas of hybrid silicon comb laser, microring resonator (MRR) modulator arrays, photonic tensor cores, MRR add-drop filters, photodiodes, and electrical cross-connects. It can be seen that TONN-2 occupies a much smaller footprint than TONN-1 at the expense of much higher computational latency.