[Paper Review] Amortized Normalizing Flows for Transcranial Ultrasound with Uncertainty Quantification

Rafael Orozco, Mathias Louboutin+4 more|arXiv (Cornell University)|Mar 6, 2023

Medical Imaging Techniques and Applications4 citations

TL;DR

This paper proposes an amortized normalizing flow framework for transcranial ultrasound computed tomography that accelerates image reconstruction by 98% compared to full-waveform inversion (FWI) while providing calibrated uncertainty quantification. By using a physics-informed summary statistic to compress high-dimensional ultrasound data, the method enables fast, generalizable posterior inference with only one forward/gradient call per source, achieving high-fidelity reconstructions and uncertainty maps that correlate with reconstruction error.

ABSTRACT

We present a novel approach to transcranial ultrasound computed tomography that utilizes normalizing flows to improve the speed of imaging and provide Bayesian uncertainty quantification. Our method combines physics-informed methods and data-driven methods to accelerate the reconstruction of the final image. We make use of a physics-informed summary statistic to incorporate the known ultrasound physics with the goal of compressing large incoming observations. This compression enables efficient training of the normalizing flow and standardizes the size of the data regardless of imaging configurations. The combinations of these methods results in fast uncertainty-aware image reconstruction that generalizes to a variety of transducer configurations. We evaluate our approach with in silico experiments and demonstrate that it can significantly improve the imaging speed while quantifying uncertainty. We validate the quality of our image reconstructions by comparing against the traditional physics-only method and also verify that our provided uncertainty is calibrated with the error.

Motivation & Objective

• Address the computational infeasibility of full-waveform inversion (FWI) in transcranial ultrasound computed tomography (TUCT), which can take up to 36 hours per reconstruction.
• Overcome the challenge of high-dimensional, complex ultrasound data in deep learning-based imaging by incorporating physical wave physics into data compression.
• Enable efficient, uncertainty-aware image reconstruction that generalizes across diverse transducer configurations without retraining.
• Provide calibrated Bayesian uncertainty quantification that correlates with reconstruction error, improving reliability for clinical decision-making.
• Develop a method that combines physics-informed modeling with data-driven normalizing flows to achieve fast inference while preserving image quality and uncertainty fidelity.

Proposed method

• Employ a physics-informed summary statistic that compresses raw ultrasound data (32 sources, 256 receivers, 2377 time steps) by a factor of 70×, reducing input dimensionality while preserving physical consistency.
• Use a conditional normalizing flow to model the posterior distribution of acoustic impedance given the compressed summary statistic, enabling efficient sampling and uncertainty quantification.
• Train the normalizing flow in an amortized manner, where a single forward and gradient evaluation per source is sufficient for posterior inference, drastically reducing online inference cost.
• Formulate the inverse problem in a Bayesian framework, allowing for full posterior sampling and pointwise variance estimation that reflects uncertainty in image reconstruction.
• Leverage a single pretraining phase to learn the normalizing flow, enabling fast online inference across multiple transducer configurations without retraining.
• Integrate the physics of the scalar acoustic wave equation with variable density into the summary function, ensuring that the data compression respects wave propagation physics.

Figure 1: Proposed transcranial image reconstruction framework with normalizing flows for uncertainty quantification.

Experimental results

Research questions

• RQ1 Can physics-informed summary statistics enable efficient, high-dimensional data compression for transcranial ultrasound imaging without sacrificing physical fidelity?
• RQ2 Can amortized normalizing flows provide fast, calibrated uncertainty quantification in TUCT while maintaining image quality comparable to full-waveform inversion?
• RQ3 How well does the proposed method generalize across different transducer source configurations and acquisition geometries?
• RQ4 Does the uncertainty map produced by the model correlate with actual reconstruction error, as required for clinical reliability?
• RQ5 Can the method reduce reconstruction time from hours to minutes while preserving or improving image fidelity compared to physics-only FWI?

Key findings

• The proposed method reduces online inference time to approximately 48 seconds (44.8 + 3.23 seconds), a 98% improvement over full-waveform inversion (FWI), which takes ~2100 seconds per reconstruction.
• The method achieves a peak signal-to-noise ratio (PSNR) of 38.67, significantly higher than FWI (33.25) and a supervised U-Net (35.63), indicating superior image fidelity.
• The structural similarity index (SSIM) of 0.9646 for the posterior mean exceeds both FWI (0.9450) and U-Net (0.9332), confirming better structural preservation.
• The root mean square error (RMSE) of 0.0119 is substantially lower than FWI (0.0215) and U-Net (0.0168), demonstrating improved reconstruction accuracy.
• The pointwise variance from the posterior samples correlates strongly with reconstruction error, validating that the uncertainty estimates are calibrated and clinically meaningful.
• The method generalizes effectively across different transducer configurations: posterior means and uncertainty maps were computed in ~3 minutes for three different setups, while FWI required ~1.5 hours per configuration.

Figure 2: 2D transcranial ultrasound imaging setup. (a) Ground truth acoustic impedance $\mathbf{x}^{\ast}$ including source/receiver layout; (b) Observed data $\mathbf{y}_{i}$ from a single source; (c) Starting model, $\mathbf{x}_{0}$ , which includes detailed acoustic information on the skull; (d)