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

Yequan Zhao, Xian Xian+6 more|arXiv (Cornell University)|Dec 31, 2023

Neural Networks and Reservoir Computing4 citations

TL;DR

This paper presents the first on-chip, back-propagation-free optical training framework for physics-informed neural networks (PINNs) using tensorized optical neural networks (TONN) and phase-domain tuning, enabling real-time, ultra-low-energy (1.36 J, 1.15 s) solution of high-dimensional PDEs—demonstrated on a 20D HJB PDE with 1,170× fewer MZIs and fJ/MAC photonic energy efficiency.

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 extit{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 imes 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.

Motivation & Objective

• To enable real-time, low-energy PDE solving in energy-constrained applications such as autonomous systems and medical imaging.
• To overcome the hardware incompatibility of back-propagation in optical neural networks due to large photonic device size and lack of scalable optical memory.
• To develop a scalable, robust, and energy-efficient optical PINN training framework using BP-free gradient estimation and tensor compression.
• To demonstrate feasibility of on-chip training for large-scale PINNs on integrated photonic platforms.

Proposed method

• Proposes a back-propagation-free training method that estimates gradients and derivatives using only additional inferences, avoiding error feedback and hardware-unfriendly BP.
• Employs tensor-train (TT) decomposition to compress the weight matrices, reducing MZI count and improving convergence and scalability.
• Uses tensorized optical neural networks (TONN) with MZI-based phase-domain tuning for in-situ optimization and scalable inference acceleration.
• Applies hardware-aware tuning by directly optimizing on fabricated photonic devices, enhancing robustness to hardware imperfections.
• Utilizes a multi-cycle inference architecture (64 cycles in TONN-2) to reduce footprint and insertion loss, enabling efficient photonic computation.
• Integrates photonic components including hybrid silicon lasers, microring modulators, MZI meshes, and photodiodes in a compact footprint for on-chip implementation.

Experimental results

Research questions

• RQ1 Can a BP-free optical training framework be designed to enable scalable and efficient training of large-scale PINNs on photonic chips?
• RQ2 How can tensor compression reduce the number of MZIs and improve convergence in optical PINN training?
• RQ3 To what extent can on-chip, inference-based gradient estimation outperform software-simulated or off-chip training under hardware imperfections?
• RQ4 What are the energy and latency trade-offs of using multi-cycle inference in TONN architectures for PDE solving?
• RQ5 Can the proposed framework achieve real-time, fJ/MAC-level energy efficiency for high-dimensional PDEs like 20D HJB?

Key findings

• The proposed optical PINN training framework reduces the number of MZIs by a factor of 1.17×10³ compared to conventional ONN for a 20D HJB PDE.
• The system achieves a total energy consumption of only 1.36 J and a total latency of 1.15 s to solve the 20D HJB PDE, enabling real-time performance.
• TONN-2 achieves a 26 mm² photonic footprint, significantly smaller than TONN-1 (648 mm²), despite higher computational latency due to 64-cycle inference.
• The BP-free method with additional inferences for gradient estimation shows superior robustness to hardware imperfections compared to software-simulated or off-chip training.
• The framework supports fully connected networks up to 1024×1024, demonstrating scalability for large-scale PINNs.
• Energy consumption per inference is reduced to 5.05×10⁻⁹ J in TONN-2, achieving fJ/MAC-level photonic energy efficiency.