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

2024-01-04

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.

Machine Learning,Emerging Technologies,Signal Processing

What problem does this paper attempt to address?

The paper attempts to address the following issues: 1. **The problem of numerically solving partial differential equations (PDEs)**: Numerically solving high-dimensional partial differential equations usually requires a large amount of computational resources and time, especially in real-time application scenarios such as autonomous systems and medical imaging. Traditional methods are limited by computational resources and runtime. 2. **Challenges in training Physics-Informed Neural Networks (PINNs)**: Although PINNs can overcome the curse of dimensionality caused by traditional discretization methods, training PINNs in practical applications is still very expensive, especially in real-time scenarios that require rapid repeated training. 3. **Challenges in training PINNs on photonic chips**: - Photonic multiply-accumulate units (such as Mach-Zehnder interferometers) are relatively large, leading to low integration density. - Implementing the backpropagation (BP) process on photonic chips faces hardware-unfriendly issues. - The efficiency of multiple BP processes required to compute higher-order derivatives is low. To address these issues, the paper proposes a new optical training framework that can handle large-scale PINNs and achieve efficient training on an integrated photonic platform. By introducing a BP-free training method and tensor compression techniques, this framework not only reduces hardware requirements and energy consumption but also improves the scalability and robustness of training. Experimental results show that in solving the 20-dimensional Hamilton-Jacobi-Bellman (HJB) equation, this method demonstrates significant performance improvements, being able to solve complex high-dimensional PDEs with extremely low energy consumption (1.36 joules) and in a short time (1.15 seconds).

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

Experimental Demonstration of an Optical Neural PDE Solver via On-Chip PINN Training

Electrically Programmable Phase-Change Photonic Memory for Optical Neural Networks with Nanoseconds in Situ Training Capability

Easily Scalable Photonic Tensor Core Based on Tunable Units with Single Internal Phase Shifters

Single chip photonic deep neural network with accelerated training

Photonic Neural Networks and Its Applications

Physics-informed Neural Network for Nonlinear Dynamics in Fiber Optics

Accelerated Training of Physics-Informed Neural Networks (PINNs) using Meshless Discretizations

Fourier Domain Physics Informed Neural Network

Optical Neural Engine for Solving Scientific Partial Differential Equations

Control-free and efficient integrated photonic neural networks via hardware-aware training and pruning

Tensor-Compressed Back-Propagation-Free Training for (Physics-Informed) Neural Networks

Efficient On-Chip Learning for Optical Neural Networks Through Power-Aware Sparse Zeroth-Order Optimization

Predicting Ultrafast Nonlinear Dynamics in Fiber Optics by Enhanced Physics-Informed Neural Network

Mirage: An RNS-Based Photonic Accelerator for DNN Training

Multimodal deep learning using on-chip diffractive optics with in situ training capability

Fully forward mode training for optical neural networks

Large-Scale Optical Neural Networks based on Photoelectric Multiplication

Applications of Physics-Informed Neural Network for Optical Fiber Communications.

Photonic Neural Networks and Optics-informed Deep Learning Fundamentals

Optical training of large-scale Transformers and deep neural networks with direct feedback alignment