[論文レビュー] 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 Computing被引用数 4

ひとこと要約

本論文は、テンソル化光学ニューラルネットワーク（TONN）と位相ドメインチューニングを用いて、物理的制約付きニューラルネットワーク（PINN）のための、初めてのオンチップでバックプロパゲーションを不要とする光学的学習フレームワークを提示する。これにより、リアルタイムかつ超低消費電力（1.36 J、1.15 s）で高次元PDEを解くことが可能となり、20次元HJB PDEに対して1,170倍少ないMZIを用いて実証された。

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.

研究の動機と目的

• 自律システムや医療画像診断などのエネルギー制約のある応用分野において、リアルタイムかつ低消費電力なPDE解法を可能にすること。
• 大規模な光デバイスのサイズとスケーラブルな光学メモリの欠如により、光ニューラルネットワークにおけるバックプロパゲーションのハードウェア不適合性を克服すること。
• バックプロパゲーションを不要とする勾配推定とテンソル圧縮を用いて、スケーラブルで頑健かつエネルギー効率の良い光学PINN学習フレームワークを開発すること。
• 統合光回路プラットフォーム上で大規模PINNのオンチップ学習の実現可能性を実証すること。

提案手法

• 誤差フィードバックを回避し、ハードウェアに不適切なバックプロパゲーションを避けるために、追加の推論のみを用いて勾配と微分を推定するバックプロパゲーションフリーな学習手法を提案する。
• 重み行列の圧縮のためのテンソルトレース（TT）分解を採用し、MZI数の削減と収束性・スケーラビリティの向上を実現する。
• MZIを用いた位相ドメインチューニングによるイン・サスティ・最適化とスケーラブルな推論加速を実現するテンソル化光学ニューラルネットワーク（TONN）を採用する。
• 実装済みの光デバイスに対して直接最適化することで、ハードウェアの不確実性に対する耐性を高める、ハードウェアに配慮したチューニングを実施する。
• フォトニック回路の効率的実装を可能にするために、フォトニックコンポーネントのフォトニックフォームファクタと挿入損失を低減する64サイクルのマルチサイクル推論アーキテクチャ（TONN-2で実装）を採用する。
• ハイブリッドシリコンレーザー、ミクロリングモodulator、MZIメッシュ、フォトダイオードを統合し、オンチップ実装に適したコンactなフォームファクタでフォトニックコンポーネントを統合する。

実験結果

リサーチクエスチョン

• RQ1 フォトニックチップ上での大規模PINNのスケーラブルかつ効率的な学習を可能にするバックプロパゲーションフリーな光学的学習フレームワークを設計可能か？
• RQ2 テンソル圧縮は、光学PINN学習におけるMZI数の削減と収束性の向上にどのように寄与するか？
• RQ3 ハードウェアの不確実性下でも、オンチップ推論ベースの勾配推定は、ソフトウェアシミュレーションまたはオフチップ学習を上回る性能を達成できるか、その程度は？
• RQ4 TONNアーキテクチャにおけるマルチサイクル推論の使用が、PDE解法においてエネルギーと遅延のトレードオフに与える影響は？
• RQ5 提案されたフレームワークは、20D HJBのような高次元PDEに対して、リアルタイムかつfJ/MACレベルのエネルギー効率を達成できるか？

主な発見

• 提案された光学PINN学習フレームワークは、20次元HJB PDEに対して、従来のONNと比較してMZI数を1.17×10³倍削減した。
• 20次元HJB PDEを解くために、合計1.36 Jのエネルギー消費と1.15 sの合計遅延を達成し、リアルタイム性能を実現した。
• TONN-2は、64サイクル推論による高い計算遅延にもかかわらず、648 mm²のTONN-1と比較して26 mm²のフォトニックフォームファクタに大幅に縮小した。
• 追加の推論による勾配推定を用いるバックプロパゲーションフリー手法は、ソフトウェアシミュレーションまたはオフチップ学習と比較して、ハードウェアの不確実性に対して優れた耐性を示した。
• 本フレームワークは、最大1024×1024の全結合ネットワークをサポートし、大規模PINNへのスケーラビリティを実証した。
• TONN-2では、1回の推論あたりのエネルギー消費を5.05×10⁻⁹ Jにまで低減し、fJ/MACレベルのフォトニックエネルギー効率を達成した。