Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yequan Zhao, Xian Xiao, Xinling Yu, Ziyue Liu, Zhixiong Chen, Geza Kurczveil, Raymond G. Beausoleil, Zheng Zhang

TL;DR

The paper tackles the challenge of real-time, energy-efficient PDE solving by proposing an on-chip training framework for physics-informed neural networks (PINNs) implemented with tensorized optical neural networks (TONNs). It introduces a BP-free training paradigm that relies on repeated optical inferences and zeroth-order gradient estimation, enabled by tensor compression to dramatically reduce photonic hardware. Key contributions include two TT-based inference designs (TONN-1 and TONN-2), a loss/gradient estimation scheme that avoids backpropagation, and a demonstration on a 20-dimensional Hamiltonian-Jacobi-Bellman PDE achieving only 1.36 J of energy and 1.15 s of latency, with over three orders of magnitude fewer Mach-Zehnder interferometers than a dense optical network. The work advances fJ/MAC computing for high-dimensional PDEs, showing robustness to hardware imperfections and enabling scalable, real-time edge PDE solvers.

Abstract

Solving partial differential equations (PDEs) numerically often requires huge computing time, energy cost, and hardware resources in practical applications. This has limited their applications in many scenarios (e.g., autonomous systems, supersonic flows) that have a limited energy budget and require near real-time response. Leveraging optical computing, this paper develops an on-chip training framework for physics-informed neural networks (PINNs), aiming to solve high-dimensional PDEs with fJ/MAC photonic power consumption and ultra-low latency. Despite the ultra-high speed of optical neural networks, training a PINN on an optical chip is hard due to (1) the large size of photonic devices, and (2) the lack of scalable optical memory devices to store the intermediate results of back-propagation (BP). To enable realistic optical PINN training, this paper presents a scalable method to avoid the BP process. We also employ a tensor-compressed approach to improve the convergence and scalability of our optical PINN training. This training framework is designed with tensorized optical neural networks (TONN) for scalable inference acceleration and MZI phase-domain tuning for \textit{in-situ} optimization. Our simulation results of a 20-dim HJB PDE show that our photonic accelerator can reduce the number of MZIs by a factor of $1.17\times 10^3$, with only $1.36$ J and $1.15$ s to solve this equation. This is the first real-size optical PINN training framework that can be applied to solve high-dimensional PDEs.

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

TL;DR

The paper tackles the challenge of real-time, energy-efficient PDE solving by proposing an on-chip training framework for physics-informed neural networks (PINNs) implemented with tensorized optical neural networks (TONNs). It introduces a BP-free training paradigm that relies on repeated optical inferences and zeroth-order gradient estimation, enabled by tensor compression to dramatically reduce photonic hardware. Key contributions include two TT-based inference designs (TONN-1 and TONN-2), a loss/gradient estimation scheme that avoids backpropagation, and a demonstration on a 20-dimensional Hamiltonian-Jacobi-Bellman PDE achieving only 1.36 J of energy and 1.15 s of latency, with over three orders of magnitude fewer Mach-Zehnder interferometers than a dense optical network. The work advances fJ/MAC computing for high-dimensional PDEs, showing robustness to hardware imperfections and enabling scalable, real-time edge PDE solvers.

Abstract

Solving partial differential equations (PDEs) numerically often requires huge computing time, energy cost, and hardware resources in practical applications. This has limited their applications in many scenarios (e.g., autonomous systems, supersonic flows) that have a limited energy budget and require near real-time response. Leveraging optical computing, this paper develops an on-chip training framework for physics-informed neural networks (PINNs), aiming to solve high-dimensional PDEs with fJ/MAC photonic power consumption and ultra-low latency. Despite the ultra-high speed of optical neural networks, training a PINN on an optical chip is hard due to (1) the large size of photonic devices, and (2) the lack of scalable optical memory devices to store the intermediate results of back-propagation (BP). To enable realistic optical PINN training, this paper presents a scalable method to avoid the BP process. We also employ a tensor-compressed approach to improve the convergence and scalability of our optical PINN training. This training framework is designed with tensorized optical neural networks (TONN) for scalable inference acceleration and MZI phase-domain tuning for \textit{in-situ} optimization. Our simulation results of a 20-dim HJB PDE show that our photonic accelerator can reduce the number of MZIs by a factor of , with only J and s to solve this equation. This is the first real-size optical PINN training framework that can be applied to solve high-dimensional PDEs.
Paper Structure (14 sections, 7 equations, 3 figures, 2 tables)

This paper contains 14 sections, 7 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Optical Neural Networks (ONN) and Tensorized Optical Neural Networks (TONN)
  4. Physics-Informed Neural Networks (PINNs) and Tensor-compressed PINNs
  5. BP-Free Optical Training Accelerator for Tensor-Compressed PINN
  6. Overall Architecture
  7. Tensor-compressed Optical Inference Accelerator Design
  8. Tensor-Compressed BP-Free PINNs Training
  9. BP-free Loss Evaluation.
  10. BP-free Gradient Estimation in SGD-type Optimizers.
  11. Experiments Results
  12. Numerical Simulation Results
  13. System Performance
  14. Conclusion

Figures (3)

  • Figure 1: The overall architecture of the BP-free optical training accelerator.
  • Figure 2: TONN-1: The designed tensor-compressed optical inference accelerator based on the TONN architecture with wavelength and space multiplexing.
  • Figure 3: TONN-2: The designed inference accelerator using a single wavelength-parallel photonic tensor core with time multiplexing.