Neural Born Series Operator for Biomedical Ultrasound Computed Tomography

FWI addresses a PDE-constrained nonlinear inverse problem in imaging, offering a breakthrough solutions in reconstructing properties of complex media. Initially prominent in geophysical exploration for oil and gas reservoir imaging[36, 22, 11], FWI integrates wave physics and multiple scattering effects for data interpretation, surpassing traditional ray approximation-based methods [43] in imaging quality. FWI’s utility has recently expanded to the biomedical imaging sector, with notable applications in optical diffraction tomography (ODT) for imaging thick biological samples [16], and USCT for high-resolution imaging of breasts [28, 41], limbs [3, 8], and brains [9, 18]. Despite these advancements, FWI’s translation to clinical settings remains limited due to its significant computational demands, particularly in the scenarios requiring large field of view. Presently, single reconstruction tasks may take hours or even days on high-performance computing clusters [9, 16]. Accelerating FWI is essential for its broader adoption and future progress, especially in biomedical imaging.

Neural operators defines a novel deep-learning paradigm that approximates mappings between infinite-dimensional function spaces, in contrast to conventional neural networks restricted to finite-dimensional Euclidean spaces. This capability positions them as a promising solution for complex, large-scale PDEs, historically challenging in practical applications. A notable implementation in this domain is the Deep Operator Network (DeepONet) [24, 25, 4, 29, 7, 21], featuring a distinctive “branch and trunk” network structure. Alternatively, the Fourier Neural Operator (FNO) [20, 19, 32, 39, 10] creatively parameterizes the integral kernel in Fourier space, enhancing a neural network’s expressiveness and efficiency for operator learning. Both techniques have demonstrated reasonable accuracy, high inference efficiency, and robustness against noise and generalization errors [26]. Spurred by these initial successes, various specialized forms have also been proposed for solving PDE of specific forms. For example, Learned Born Series (LBO)[35] and Born FNO (BFNO)[48] have been proposed for handling Helmholtz equations by modifying FNO based on Born series. Despite their effectiveness in managing complex PDEs, neural operators typically lag behind traditional methods in computational accuracy, which warrants further investigation.

Thanks to their impressive approximation capabilities and negligible computational speed, neural operators are increasingly employed to solve PDE-constrained inverse problems in diverse fields such as turbulent flows [20], solid mechanics [38], weather forecasting [32], material designs [27] and nano-photonics [37]. There are also growing interests in neural operator-based FWI, with many studies focusing on directly learning the mapping from observed data to parametric images of physical properties, as seen in geophysics [46, 50] and ultrasound imaging [23]. However, this data-to-image mapping is often non-unique due to limited observations, leading to generalization challenges. An alternative is to use neural operators as surrogate forward models to accelerate PDE solving, mapping physical properties to data, followed by inversion using a gradient-based optimizer. This approach enhances generalization and robustness, demonstrated in seismic imaging [44], electromagnetic imaging [48], and time-of-flight tomography in ultrasound [6]. Yet, most of these methods focus on limited computational domains (under 30 wavenumbers) to ease learning highly-oscillatory features. The potential of neural operator-based FWI, particularly in biomedical imaging applications like USCT with larger computational domains (around 100 wavenumbers), remains unexplored.

USCT acquires data by first activating point sources with known locations and subsequently measuring the corresponding scattered acoustic signals that have been modulated by internal biological tissues. In steady-state conditions, the propagation of ultrasound waves can be accurately described by a heterogeneous Helmholtz equation, assuming negligible shear motion and nonlinear effects,

where $r=|x|$, $\omega$ is the angular frequency of ultrasound waves, $c(x)$ is the spatial distribution of sound speed, $\rho(x)$ is the source term and $u(x)$ is the complex acoustic field. In the context of USCT, the heterogeneity in sound speed $c(x)$ is confined to a specific domain of interest (DOI), while the sound speed remains constant $c_{0}$ outside of this domain. To account for this boundary condition, the Helmholtz equation must satisfy the Sommerfeld radiation condition:

A full USCT observation employs a sequential activation scheme, where point sources, represented by $\rho_{k}(x)$,are activated one at a time using $N$ transducers arranged in an annular configuration. The corresponding scattered acoustic fields are measured at the same $M$ transducer locations, resulting in a measurement dataset denoted as $y=\{y_{k};y_{k}=u_{k}(x_{f})+n_{k}\}$, where $k=1,\ldots,M$ represents the point source index, $x_{f}\in\mathbb{R}^{M}$ defines the locations of all transducers, and $n_{k}$ is the measurement noise.

The central goal of USCT is to reconstruct the spatial distribution of sound speed within biological tissues, denoted by $c(x)$, using the measurements obtained from the transducers, represented by $y$. This inverse problem can be formulated as a PDE-constrained optimization problem:

Since the Helmholtz equation governs steady-state wave propagation, namely the wavefield in the frequency domain, this inverse problem is commonly referred to as frequency-domain full waveform inversion (FD-FWI). Akin to many other inverse problems, FD-FWI is typically solved using gradient-based optimizers. A prevalent approach for computing the gradient, $\partial L_{k}/\partial c$, in FWI is the adjoint method,

where $i$ denotes the index of USCT transducers and $\delta(\cdot)$ defines a normalized point source at a specific transducer location. The gradient is proportional to the product of two wavefields, where $u_{k}(x)$ is the forward simulation result for source $\rho_{k}$ and $\lambda_{k}(x)$ arises from the backward simulation where the source term is defined based on the discrepancies between forward predictions and measured data. (See Supplementary Material for theoretical details.) Based Eq. 4, a gradient update requires solving the Helmholtz equation twice for each source, in total $2M$ times for a full observation. The computational efficiency of Helmholtz equation solvers is the primary time bottleneck for FWI.

Current studies on neural operators for Helmholtz equations often employ oversimplified setups. These studies typically create heterogeneous sound speeds for operator learning based on public datasets, such as MNIST [42], or Gaussian/Uniform random fields[48, 35]. Additionally, their solution domains are confined to fewer than 20-30 wavelengths, much smaller than the real-world field of view (FOV) in biomedical USCT applications. To develop a more effective neural operator for USCT FWI, we have created a USCT dataset featuring anatomically realistic phantoms and corresponding wavefields simulated using parameters from an actual USCT system. Our dataset comprises 9590 human organ phantoms, including 8000 breasts and 1590 brains. For each phantom, we simulate wavefields from 64 different sources, resulting in a total of 613,760 input-output data pairs. Our dataset will be publicly released upon request.

In this section, we introduce the NBSO, a proficient tool designed to accurately and efficiently approximate the implicit mapping $\mathcal{S}:(c,\rho)\mapsto u$, as defined by the Helmholtz equation in Eq. 1. Drawing inspiration from the CBS solver for the Helmholtz equation, the NBSO utilizes an iterative scheme in developing its neural network architecture. Moreover, the inputs and outputs within each iterative block of the NBSO are carefully designed based on physical concepts, thereby enhancing the operator’s ability to approximate and generalize effectively.

Having established the NBSO as a precise and efficient surrogate model for the Helmholtz equation, we are well-prepared to compute the FWI gradient via the adjoint method, as specified in Eq. 4. This computation involves running the NBSO twice for each transducer source, which includes conducting a forward simulation followed by residual back-propagation. Armed with the calculated gradient, we then employ the Limited-memory BFGS algorithm is to reconstruct the sound speed of target biological tissues.

In this section, we conduct extensive testing of NBSO’s prediction accuracy on our realistic USCT dataset, as described in Sec. 4.1. We implement two distinct settings for our analysis: 1) we train the neural networks using data with a fixed source located on the left-most transducer, and 2) we train neural networks utilizing data from 64 distinct sources. Essentially, the first setting assesses NBSO’s proficiency in mapping sound speed to wavefield, whereas the second evaluates its generalization capability in solving the Helmholtz Equation with varied source locations. We measure wavefield prediction accuracy quantitatively using the relative rooted mean squared error (RRMSE) between the neural network’s predicted wavefield $\hat{u}(x)$ and the true wavefield $u(x)$

Both breast and brain datasets are used to test all baseline methods. Considering the unequal number of breast and brain phantoms, we ensure comparable metrics by training and validating neural networks using a subset of breasts (2400 for training, 240 for validation) and the entire set of brains (1500 for training, 90 for validation).

The RRMSE results for NBSO and other baseline neural networks are presented in Table 1. In our experiments, we not only adhere to the original formulations (marked by superscripted star) of FNO, BFNO, and UNet, but also test their modified versions (without star) integrated with our scattering amplitude learning methodology. As demonstrated in the left half of Table 1, the scattering amplitude learning improves all frameworks in the fixed source scenario. Consequently, in the experiments with multiple sources, network performance is evaluated only with the scattering amplitude incorporated. As demonstrated in Table 1, NBSO surpasses all other baselines in both breast and brain datasets, regardless of the source configurations. Fig. 4 provides a visual comparison of the predicted wavefields for a representative breast phantom, as generated by various baselines, all trained on breast datasets and assuming multiple sources. NBSO captures high-frequency features with greater accuracy than other methods.

In this section, we evaluate the reconstruction quality of USCT FWI using the NBSO alongside various baseline methods. Our experimental setup simulates the USCT system with 256 transducers uniformly distributed around the ring, tasked with source activation and signal collection. As detailed in Sec. 4.3, image reconstruction for all methods is conducted using the adjoint method in conjunction with the LBFGS optimizer, and the quality of image reconstruction is assessed using SSIM and PSNR metrics. Given that our USCT FWI requires computing wavefields from multiple sources for gradient calculation, only neural networks trained in the multi-source setting are used in this section.

Table. 2 showcases the USCT FWI’s performance on 30 breast phantoms using NBSO, FNO, UNet, BFNO, and CBS solvers at different SNR levels (noise-free, 10dB, 5dB) for observed data. NBSO consistently outperforms all neural network-based methods in reconstruction quality, regardless of noise levels. This is further illustrated in Fig. 5, where we display the reconstruction results for breast phantoms of two different types, all-fatty and fibroglandular. The NBSO-based FWI effectively recovers key anatomical features in both cases, whereas other neural networks occasionally introduce significant artifacts. The superior performance of NBSO over other baselines is especially evident when applied to all-fatty breast tissues, where the higher contrast of sound speed indicates stronger scattering effects. Remarkably, the addition of measurement noise does not substantially degrades NBSO’s reconstructions, suggesting that the neural operator might possess some implicit prior knowledge about the target, acquired from training on anatomically realistic data. However, it is also noted that NBSO’s FWI results appear blurrier compared to the ground-truth and CBS-FWI, indicating NBSO’s potential shortfall in capturing high-frequency wavefield features.

While NBSO’s reconstruction quality may slightly lag behind that of CBS, its computational efficiency is substantially higher, as demonstrated in Table 3. We use Intel i7-13700KF and NVIDIA RTX4080ti for speed evaluations. NBSO surpasses GPU-accelerated CBS by an order of magnitude and CPU-based CBS by two orders of magnitude for a single source. Given the capability of neural networks for batch processing—predicting the wavefields of multiple inputs parallelly—the efficiency gap between NBSO and CBS becomes even more pronounced. In summary, NBSO presents a well-balanced solution, offering reasonable reconstruction quality alongside superior efficiency.

We further investigate the adaptability of NBSO-based FWI to out-of-distribution targets. As reported in Table 4 and illustrated in Fig. 6, we assess the reconstruction quality on both breast and brain phantoms, using NBSO trained on breast data, brain data, or a combination of both. NBSO demonstrates competent performance on targets that are within its training distribution. However, it shows a marked decrease in effectiveness for targets significantly divergent from its training set. Notably, training NBSO on a mixed dataset enhances its generalizability across datasets, yielding better image reconstruction quality than training on a single data type for brain imaging.

We refer to the Wavesim [31] for the details of the governing equations. We simulate the numerical solution on a square with side length of 24 centimeters. The step size of grid is 0.05cm thus the grid shape is $480\times 480$. For boundary condition, we use an absorbing boundary layer (type ARL in Wavesim) whose boundary width is $[75,75]$ grid outside the solving region. We choose the points per wavelenght $PPW=6$ for the absorbing boundary layer.

The training of neural network models, including the baselines (FNO, BFNO, UNet), is performed with the ADAM optimizer, with a learning rate $\eta=10^{-2}$, for 50 epochs in both single-source and multi-source settings. The training objective is to minimize the relative $L_{2}$ loss function, incorporating a weight decay rate of $10^{-5}$. We also implement a Cosine-Annealing learning rate scheduler with a step size of $10$ and a minimum learning rate of $\eta_{\min}=2\times 10^{-4}$. In the multi-source setting, models are trained using 4 GPUs in the DeepSpeed framework with a batch size of 10, while single-source models are trained on a single GPU with the same batch size.

Fig. 7 offers a visual comparison of wavefields predicted by NBSO for a representative brain phantom, alongside those generated by baseline models trained on the same brain datasets and assuming multiple sources. NBSO captures high-frequency features with greater accuracy than other methods.

Table. 6 showcases the USCT FWI’s performance on 30 brain phantoms using NBSO, FNO, UNet, BFNO, and CBS solvers at different SNR levels (noise-free, 10dB, 5dB) for observed data. NBSO consistently outperforms all neural network-based methods in reconstruction quality, regardless of noise levels. This is further illustrated in Fig. 8, where we display the reconstruction results for a brain phantom.