[論文レビュー] A Novel Training Framework for Physics-informed Neural Networks: Towards Real-time Applications in Ultrafast Ultrasound Blood Flow Imaging

Haotian Guan, Jinping Dong+1 more|arXiv (Cornell University)|Sep 9, 2023

Model Reduction and Neural Networks被引用数 4

ひとこと要約

本稿では、ステディステートのナビエ-ストークス方程式を逐次的に解くことで、リアルタイムの超高速ウルトラサウンド血流イメージングを高速化し、不確実性を考慮した初期化による並列学習を可能にする、物理学的に制約されたニューラルネットワーク（PINN）のための新しいトレーニングフレームワーク、SeqPINNおよびSP-PINNを提案する。この手法により、トレーニング時間を複数桁短縮しつつ、直血管では1.01 cm/s（SeqPINN）、1.26 cm/s（SP-PINN）という低いRMSEを達成した。臨床的血行動態モデリングにおける高速でスケーラブルな推論が可能になった。

ABSTRACT

Ultrafast ultrasound blood flow imaging is a state-of-the-art technique for depiction of complex blood flow dynamics in vivo through thousands of full-view image data (or, timestamps) acquired per second. Physics-informed Neural Network (PINN) is one of the most preeminent solvers of the Navier-Stokes equations, widely used as the governing equation of blood flow. However, that current approaches rely on full Navier-Stokes equations is impractical for ultrafast ultrasound. We hereby propose a novel PINN training framework for solving the Navier-Stokes equations. It involves discretizing Navier-Stokes equations into steady state and sequentially solving them with test-time adaptation. The novel training framework is coined as SeqPINN. Upon its success, we propose a parallel training scheme for all timestamps based on averaged constant stochastic gradient descent as initialization. Uncertainty estimation through Stochastic Weight Averaging Gaussian is then used as an indicator of generalizability of the initialization. This algorithm, named SP-PINN, further expedites training of PINN while achieving comparable accuracy with SeqPINN. The performance of SeqPINN and SP-PINN was evaluated through finite-element simulations and in vitro phantoms of single-branch and trifurcate blood vessels. Results show that both algorithms were manyfold faster than the original design of PINN, while respectively achieving Root Mean Square Errors of 0.63 cm/s and 0.81 cm/s on the straight vessel and 1.35 cm/s and 1.63 cm/s on the trifurcate vessel when recovering blood flow velocities. The successful implementation of SeqPINN and SP-PINN open the gate for real-time training of PINN for Navier-Stokes equations and subsequently reliable imaging-based blood flow assessment in clinical practice.

研究の動機と目的

• 超高速ドップラー超音波を用いた臨床的血行動態モデリングにおける、従来のPINNの遅いトレーニング速度を解消すること。
• ステディステートのナビエ-ストークス方程式を活用することで、高フレームレートの血流イメージングに適したPINNの高速かつスケーラブルなトレーニングを可能にすること。
• すべてのタイムスティッピングにわたる並列トレーニング方式（SP-PINN）を構築し、精度を維持しながらトレーニング時間を大幅に短縮すること。
• Stochastic Weight Averaging Gaussian（SWAG）を用いた不確実性推定により、初期化を最適化し、モデルの一般化性能を向上させること。
• 臨床現場における患者特異的血流速度再構築のためのリアルタイム推論を実現すること。

提案手法

• タイムスティッピングごとに時間離散化し、ステディステートのナビエ-ストークス問題に分解し、転移学習を用いて1タイムステップずつモデルをトレーニングするSeqPINNを提案。
• 全タイムステップにわたる平均化された定数の確率的勾配降下法（SGD）を用いた初期化を適用することで、すべてのタイムステップで並列にトレーニング可能なSP-PINNを導入。
• Stochastic Weight Averaging Gaussian（SWAG）を用いて不確実性を推定し、初期化を最適化することで、一般化性能と収束性を向上。
• 有限要素法シミュレーションおよび単一・分岐血管モデルのin vitroプローブを用いてフレームワークの妥当性を検証。
• ナビエ-ストークス方程式および非圧縮性を制約とするPDEに基づく損失関数を採用し、境界条件は物理的制約付き正則化により強制。
• 従来のCFDの計算負荷を回避するため、軽量でメッシュフリーなニューラルネットワークを活用し、高速な推論を実現。

実験結果

リサーチクエスチョン

• RQ1 ステディステートのナビエ-ストークス解の逐次的トレーニングは、超高速ウルトラサウンドイメージングにおけるPINNのトレーニング時間を顕著に短縮できるか？
• RQ2 タイムステップ間でPINNの並列トレーニングが、安定的かつ一般化可能な初期化を用いて実現可能か？
• RQ3 SWAGによる不確実性推定は、PINN初期化の一般化性能と収束性をどのように向上させるか？
• RQ4 リアルタイム血流再構築において、トレーニング速度、モデルサイズ、精度の間にはどのようなトレードオフがあるか？
• RQ5 提案されたフレームワークは、高フレームレートのドップラー超音波データに対してリアルタイム推論を達成できるか？

主な発見

• SeqPINNは、バニラPINNと比較してトレーニング時間を複数桁短縮し、直血管では1.01 cm/sのRMSEを達成。分岐血管では1.91 cm/sのRMSEを記録した。
• SP-PINNはSeqPINNと同等の精度を達成し、直血管で1.26 cm/s、分岐血管で2.56 cm/sのRMSEを示したが、並列トレーニングが可能となった。
• SWAGによる不確実性推定は、安定した初期化を的確に特定し、モデルの一般化性能を向上させるとともに、トレーニングの分散を低減した。
• SeqPINNの逐次トレーニングアプローチは、ストレージ（N個のモデル）を犠牲にして速度を獲得しており、個々のモデルは小さく（約461 kB）GPU外に保存可能であった。
• SP-PINNは、GPU上で限られたハードウェア環境でも、並列計算の効率的活用により、SeqPINNを上回るトレーニング速度を達成した。
• 改善は見られたが、GPU計算のボトル neck が残っており、リアルタイムトレーニングは達成できなかった。これより、専用のハードウェア加速の導入が求められる。