[논문 리뷰] Real-Time FJ/MAC PDE Solvers via Tensorized, Back-Propagation-Free Optical PINN Training

Yequan Zhao, Xian Xian외 6명|arXiv (Cornell University)|2023. 12. 31.

Neural Networks and Reservoir Computing인용 수 4

한 줄 요약

이 논문은 텐서화된 광학 신경망(TONN)과 단계 도메인 튜닝을 사용하여 물리학 기반 신경망(PINN)을 위한 최초의 온칩, 오차 역전파 없음 광학 학습 프레임워크를 제안한다. 이는 실시간, 초저에너지(1.36 J, 1.15 s)로 고차원 PDE를 해결할 수 있게 하며, 20D HJB PDE에 대해 MZI 수를 1,170배 감소시키고 fJ/MAC 광학 에너지 효율성을 달성한다.

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 사이클)를 도입한다.
• 하이브리드 실리콘 레이저, 마이크로링 모odulator, MZI 메시, 광다이오드를 포함한 광학 구성 요소를 밀집된 면적에 통합하여 온칩 구현을 구현한다.

실험 결과

연구 질문

• RQ1 광학 칩에서 대규모 PINN의 확장 가능하고 효율적인 학습을 가능하게 하기 위해 오차 역전파 없음 광학 학습 프레임워크를 설계할 수 있는가?
• RQ2 텐서 압축은 MZI 수를 얼마나 줄이고 광학 PINN 학습의 수렴성은 어떻게 향상시킬 수 있는가?
• RQ3 하드웨어 결함 상황에서 온칩 추론 기반 기울기 추정이 소프트웨어 시뮬레이션 또는 오프칩 학습에 비해 얼마나 뛰어난 성능을 보일 수 있는가?
• RQ4 TONN 아키텍처에서 다중 사이클 추론을 사용할 경우 PDE 해결에 있어 에너지 소비와 지연 시간 간의 상호 교환 관계는 어떠한가?
• RQ5 제안된 프레임워크는 20D HJB와 같은 고차원 PDE에 대해 실시간, fJ/MAC 수준의 에너지 효율성을 달성할 수 있는가?

주요 결과

• 제안된 광학 PINN 학습 프레임워크는 기존 ONN 대비 20D HJB PDE에 대해 MZI 수를 1.17×10³배 감소시켰다.
• 시스템은 20D HJB PDE를 해결하기 위해 총 에너지 소비 1.36 J, 총 지연 시간 1.15 s로 실시간 성능을 달성하였다.
• TONN-2는 64 사이클 추론로 인해 더 높은 계산 지연을 겪음에도 불구하고, TONN-1(648 mm²) 대비 26 mm²의 광학 면적을 확보하여 현저히 작아졌다.
• 추가 추론을 통한 오차 역전파 없음 방법은 소프트웨어 시뮬레이션 또는 오프칩 학습 대비 하드웨어 결함에 대해 뛰어난 견고성을 보였다.
• 프레임워크는 1024×1024까지의 완전 연결 신경망을 지원하여 대규모 PINN에 대한 확장성을 입증하였다.
• TONN-2에서는 추론당 에너지 소비가 5.05×10⁻⁹ J로 감소하여 fJ/MAC 수준의 광학 에너지 효율성을 달성하였다.