[论文解读] 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

一句话总结

该论文提出SeqPINN和SP-PINN，一种用于物理信息神经网络的新颖训练框架，通过按顺序求解稳态Navier-Stokes方程并实现基于不确定性感知初始化的并行训练，显著加速了实时超快超声血流成像。该方法将训练时间减少多个数量级，同时在直血管上实现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在超快多普勒超声实时血流动力学建模中训练速度缓慢的问题。
• 通过利用稳态Navier-Stokes方程，实现在高帧率血流成像中PINN的快速、可扩展训练。
• 开发一种并行训练方案（SP-PINN），在显著减少训练时间的同时保持模型精度。
• 通过使用随机权重平均高斯（Stochastic Weight Averaging Gaussian）进行初始化，提升模型泛化能力。
• 实现在临床环境中患者特异性血流速度重建的实时推理能力。

提出的方法

• 提出SeqPINN，将时间离散化为一系列稳态Navier-Stokes问题，通过迁移学习逐个时间戳训练模型。
• 提出SP-PINN，一种并行训练方案，对所有时间戳应用平均常数随机梯度下降（SGD）进行初始化。
• 采用随机权重平均高斯（SWAG）估计不确定性，并指导初始化，从而提升泛化能力和收敛性。
• 使用有限元仿真和单支及分叉血管模型的体外幻象验证该框架。
• 基于Navier-Stokes方程和不可压缩性施加PDE约束损失函数，通过物理信息正则化实现边界条件约束。
• 采用轻量化、无网格神经网络，避免传统计算流体动力学（CFD）的计算负担，实现实时推理。

实验结果

研究问题

• RQ1 对稳态Navier-Stokes解进行顺序训练是否能显著减少超快超声成像中PINN的训练时间？
• RQ2 能否在时间戳之间实现稳定且泛化能力强的PINN并行训练初始化？
• RQ3 通过SWAG进行不确定性估计如何改善PINN初始化的泛化能力和收敛性？
• RQ4 在实时血流重建中，训练速度、模型大小与精度之间的权衡如何？
• RQ5 所提出的框架能否实现实时推理，以处理高帧率多普勒超声数据？

主要发现

• 与原始PINN相比，SeqPINN将训练时间减少了多个数量级，在直血管上实现1.01 cm/s的RMSE，在分叉血管上实现1.91 cm/s的RMSE。
• SP-PINN在直血管上实现1.26 cm/s的RMSE，在分叉血管上实现2.56 cm/s的RMSE，精度与SeqPINN相当，同时支持并行训练。
• 通过SWAG进行不确定性估计成功识别出稳定的初始化方式，提升了模型泛化能力并降低了训练方差。
• SeqPINN的顺序训练方法以存储代价（N个模型）换取速度优势，单个模型体积小（约461 kB），适合离GPU存储。
• 由于更有效地利用并行计算，SP-PINN在GPU上的训练速度优于SeqPINN，即使在硬件受限条件下亦然。
• 尽管有所改进，但由于GPU计算瓶颈，实时训练仍未实现，提示需要专用硬件加速。