[논문 리뷰] A Novel Training Framework for Physics-informed Neural Networks: Towards Real-time Applications in Ultrafast Ultrasound Blood Flow Imaging

Haotian Guan, Jinping Dong외 1명|arXiv (Cornell University)|2023. 09. 09.

Model Reduction and Neural Networks인용 수 4

한 줄 요약

이 논문은 순차적으로 안정 상태 나비에-스토크스 방정식을 풀고 불확실성 인식 초기화를 통해 병렬 학습을 가능하게 하여 실시간 초고속 초음파 혈류 영상 촬영을 가속화하는 새로운 학습 프레임워크인 SeqPINN와 SP-PINN를 제안한다. 이 방법은 학습 시간을 수개월 정도로 단축시키면서도 직선 혈관에서 RMSE가 각각 1.01 cm/s (SeqPINN)와 1.26 cm/s (SP-PINN)로 낮아져 임상용 혈역학 모델링을 위한 더 빠르고 확장 가능한 추론을 가능하게 한다.

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를 개발하기 위해.
• 스토케스틱 웨이트 평균 가우시안(SWAG)을 이용한 불확실성 추정을 통해 초기화를 개선하고 모델 일반화 능력을 향상시키기 위해.
• 임상 환경에서 환자별 혈류 속도 재구성에 실시간 추론 기능을 달성하기 위해.

제안 방법

• 시간을 순차적인 안정 상태 나비에-스토크스 문제로 분할하고, 전이 학습을 사용하여 한 타임스탬프씩 모델을 학습하는 SeqPINN를 제안한다.
• 모든 타임스탬프에 걸쳐 평균화된 일정한 확률적 경사 하강법(SGD)을 사용하여 초기화하는 병렬 학습 기반의 SP-PINN를 도입한다.
• 일반화 및 수렴 성능 향상을 위해 스토케스틱 웨이트 평균 가우시안(SWAG)을 활용해 불확실성 추정을 수행하고 초기화를 안내한다.
• 유한요소 시뮬레이션과 단일 및 삼분지 혈관 모델의 실내 환상체를 사용하여 프레임워크의 유효성을 검증한다.
• 나비에-스토크스 방정식과 비압축성 조건을 기반으로 한 PDE 제약 손실 함수를 사용하며, 경계 조건은 물리 기반 정규화를 통해 강제한다.
• 기존 CFD의 계산 부담을 피하고 신속한 추론을 가능하게 하기 위해 경량이며 메시 기반의 신경망을 활용한다.

실험 결과

연구 질문

• RQ1 안정 상태 나비에-스토크스 해를 순차적으로 학습함으로써 초고속 초음파 영상 촬영을 위한 PINN 학습 시간을 크게 단축시킬 수 있는가?
• RQ2 모든 타임스탬프에 걸쳐 안정적이고 일반화 가능한 초기화를 통해 PINN의 병렬 학습을 구현할 수 있는가?
• RQ3 SWAG를 통한 불확실성 추정은 PINN 초기화의 일반화 능력과 수렴 성능을 어떻게 향상시키는가?
• RQ4 실시간 혈류 재구성에서 학습 속도, 모델 크기, 정확도 사이의 상호 상충 관계는 어떠한가?
• RQ5 제안된 프레임워크는 고화소율 도플러 초음파 데이터에 대해 실시간 추론을 달성할 수 있는가?

주요 결과

• SeqPINN는 일반 PINN 대비 학습 시간을 다수의 주기로 단축시켰으며, 직선 혈관에서 RMSE는 1.01 cm/s, 삼분지 혈관에서 1.91 cm/s를 달성했다.
• SP-PINN는 직선 혈관에서 RMSE 1.26 cm/s, 삼분지 혈관에서 2.56 cm/s로 SeqPINN와 유사한 정확도를 확보하면서도 병렬 학습을 가능하게 했다.
• SWAG를 통한 불확실성 추정은 안정적인 초기화를 성공적으로 식별하여 모델 일반화 능력 향상과 학습 분산 감소에 기여했다.
• SeqPINN의 순차적 학습 방식은 저장소(N개의 모델)를 투자함으로써 속도를 확보했으며, 개별 모델 크기가 약 461 kB로 작아 GPU 외부 저장소에 적합했다.
• SP-PINN는 GPU에서 더 나은 병렬 계산 활용 덕분에, 제한된 하드웨어 조건에서도 SeqPINN보다 빠른 학습 속도를 기록했다.
• 다른 개선에도 불구하고 GPU 계산 병목 현상으로 인해 실시간 학습은 아직 달성되지 않았으며, 전용 하드웨어 가속이 필요할 것으로 보인다.