Learning Sub-Sampling and Signal Recovery with Applications in Ultrasound Imaging

TL;DR

This paper introduces Deep Probabilistic Sub-sampling (DPS), a deep learning approach that learns task-specific sub-sampling patterns for ultrasound imaging, enabling efficient data acquisition and accurate reconstruction in real-time.

Contribution

The paper presents a novel deep learning framework that jointly learns sub-sampling patterns and reconstruction models, optimized for specific tasks in ultrasound imaging.

Findings

Effective in-silico sparse signal recovery from Fourier measurements.

Successful in-vivo reconstruction of anatomical images and tissue motion from sub-sampled ultrasound data.

Framework enables practical implementation via analog-to-digital conversion or sparse array design.

Abstract

Limitations on bandwidth and power consumption impose strict bounds on data rates of diagnostic imaging systems. Consequently, the design of suitable (i.e. task- and data-aware) compression and reconstruction techniques has attracted considerable attention in recent years. Compressed sensing emerged as a popular framework for sparse signal reconstruction from a small set of compressed measurements. However, typical compressed sensing designs measure a (non)linearly weighted combination of all input signal elements, which poses practical challenges. These designs are also not necessarily task-optimal. In addition, real-time recovery is hampered by the iterative and time-consuming nature of sparse recovery algorithms. Recently, deep learning methods have shown promise for fast recovery from compressed measurements, but the design of adequate and practical sensing strategies remains a challenge. Here, we propose a deep learning solution termed Deep Probabilistic Sub-sampling (DPS), that learns a task-driven sub-sampling pattern, while jointly training a subsequent task model. Once learned, the task-based sub-sampling patterns are fixed and straightforwardly implementable, e.g. by non-uniform analog-to-digital conversion, sparse array design, or slow-time ultrasound pulsing schemes. The effectiveness of our framework is demonstrated in-silico for sparse signal recovery from partial Fourier measurements, and in-vivo for both anatomical image and tissue-motion (Doppler) reconstruction from sub-sampled medical ultrasound imaging data.

Equations

$$\mathbf{y}$$

$$\mathbf{z}$$

$$\hat{\mathbf{z}}$$

$$r_m\sim \text{Cat}(N,{\mathbf{\pi }}_m),$$

$${\pi }_{m,n}=\frac{\exp {\varphi }_{m,n}}{{\sum }_{i=1}^N\exp {\varphi }_{m,i}},$$

$$\tilde{r}_{m}=\underset{n\in\{1,\ldots,N\}}{\mathrm{argmax}}\big{\{}w_{m-1,n}+\phi_{m,n}+e_{m,n}\big{\}}.$$

$$\displaystyle\mathbf{a}_{m}=\mathrm{one\_hot}_{N}\big{\{}\tilde{r}_{m}\big{\}}=$$

$$\displaystyle\mathrm{one\_hot}_{N}\Big{\{}\underset{n\in\{1,..,N\}}{\mathrm{argmax}}\big{\{}w_{m-1,n}+\phi_{m,n}+e_{m,n}\big{\}}\Big{\}}.$$

$${\nabla }_{{\mathbf{\varphi }}_m}{\mathbf{a}}_m:=$$

$$\displaystyle\nabla_{\boldsymbol{\phi}_{m}}\mathbb{E}_{\mathbf{e}_{m}}\big{[}\mathrm{softmax}_{\tau}(\mathbf{w}_{m-1}+\boldsymbol{\phi}_{m}+\mathbf{e}_{m})\big{]}=$$

$$\displaystyle\nabla_{\boldsymbol{\phi}_{m}}\mathbb{E}_{\mathbf{e}_{m}}\Bigg{[}\frac{\mathrm{exp}\{(\mathbf{w}_{m-1}+\boldsymbol{\phi}_{m}+\mathbf{e}_{m})\mathrm{/}\tau\}}{\sum_{i=1}^{N}\mathrm{exp}\{(w_{m-1,i}+\phi_{m,i}+e_{m,i})\mathrm{/}\tau\}}\Bigg{]},$$

$${\mathbf{a}}_m\triangleq \underset{\tau \to 0}{\lim }\mathrm{softma}{\mathrm{x}}_{\tau }({\mathbf{w}}_{m-1}+{\mathbf{\varphi }}_m+{\mathbf{e}}_m),\text{and}$$

$$\displaystyle\nabla_{\boldsymbol{\phi}_{m}}\mathbf{a}_{m}\triangleq\nabla_{\boldsymbol{\phi}_{m}}\mathbb{E}_{\mathbf{e}_{m}}\big{[}\mathrm{softmax}_{\tau}(\mathbf{w}_{m-1}+\boldsymbol{\phi}_{m}+\mathbf{e}_{m})\big{]},$$

$${\mathcal{L}}_{\mathcal{S}}=-\sum _{m=1}^M\sum _{n=1}^N{\pi }_{m,n}\log {\pi }_{m,n},$$

$$\hat{\theta },\hat{\mathbf{\Phi }}=\underset{\theta ,\mathbf{\Phi }}{\mathrm{argmin}}({\mathcal{L}}_{\text{mse}}+{\mathcal{L}}_{\text{pen}}),$$

$${\mathcal{L}}_{\text{mse}}={\mathbb{E}}_{(\mathbf{x},\mathbf{z})\sim p_{{}_{\mathcal{D}}}}\Vert \mathbf{z}-g_{\theta }({\mathbf{A}}_{\mathbf{\Phi }}\mathbf{x}){\Vert }_2^2,$$

$${\mathcal{L}}_{\text{pen}}=\lambda \Vert \theta {\Vert }_2+\mu {\mathcal{L}}_{\mathcal{S}},$$

$${\hat{\varphi }}_{m,n}=\alpha (n-\frac{N}{M}m{)}^4+\beta (n-\frac{N}{M}m{)}^2+{\gamma }_{m,n},$$

$$T(i,\zeta )=T_T(i,\zeta )+T_R(i,\zeta ),$$

$$T_T(i,\zeta )=\frac{\sqrt{f_d^2{\sin }^2\zeta -(r_f+f_d\cos \zeta {)}^2}-r_f}{c_0},$$

$$T_R(i,\zeta )=\frac{\sqrt{f_d^2{\cos }^2\zeta +{\{f_d\sin \zeta -\mathrm{\Delta x}(i-\frac{L-1}{2})\}}^2}}{c_0}.$$

$${\zeta }_g=\pm \arcsin (\frac{k\lambda }{\Delta x\cdot \frac{N}{M}}),$$

Full text

I Introduction

Advanced medical imaging techniques require transfer and storage of large amounts of data. Due to limited bandwidth, the raw sensor data must be compressed prior to its transfer to the backend system. Data compression, under-sampling, and subsequent reconstruction techniques have been an active area of research for medical imaging modalities such as computed tomography (CT) imaging [1, 2, 3, 4], ultrasound CT imaging [5], ultrasound (US) imaging [6, 7], and magnetic resonance imaging (MRI) [8, 9]. By carefully selecting when (i.e. in time) or where (e.g. in pixel-coordinate, across sensors, or k-space) to sample, not only the amount of data can be reduced, but also the acquisition time, power drain, and, particularly for x-ray CT, radiation exposure can be lowered. In this paper, we therefore propose a framework for learning a task-driven sub-sampling and reconstruction method that permits reduction of sensor data rates, while retaining the information required to perform a given (imaging) task.

Among diagnostic imaging options, US imaging is an increasingly used modality, owing to its portability, cost-effectiveness, excellent temporal resolution, minimal invasiveness, and radiation-free nature. Current US probes are increasingly compact, portable, and/or wireless [10, 11, 12]. Across many applications and form factors, data transfer is therefore subject to very limited capacity channels. For in-body imaging (where compactness is a necessity) data is transferred from the transducer to the back-end system via a thin low-bandwidth catheter. Likewise, the aforementioned wireless, or usb-connected, probes are also heavily bandwidth constrained. At the same time, emerging ultrafast 3D US imaging techniques [13, 14] cause data rates to drastically grow, which in turn poses even higher demands on the probe-to-system communication and processing rates. Given these challenges, US imaging serves as an excellent candidate for evaluating the effectiveness of the framework that we will introduce.

Commonly used techniques to reduce data rates in 2D and 3D echography applications are micro-beamforming [15, 16] and slow-time111In US imaging a distinction is made between slow-time and fast-time: slow-time refers to a sequence of snapshots (i.e., across multiple transmit/receive events), at the pulse repetition rate, whereas fast-time refers to samples along depth. multiplexing. The former compresses data from multiple (adjacent) transducer elements (i.e. channels) into a single focused line, thereby virtually reducing the number of receive channels. While effective, this impairs the attainable resolution and image quality. The latter only communicates a subset of the channel signals to the backend of the system for every slow-time transmission. This comes at the cost of reduced frame rates.

Compressed sensing (CS) permits low data rate sensing (below the Nyquist rate) with strong signal recovery guarantees under specific conditions [17, 18, 19, 20, 21]. In CS, a sparse signal $\mathbf{x}$ is to be recovered from measurements $\mathbf{y}$ that are taken at a sub-Nyquist rate through a sensing matrix $\mathbf{\Psi}$: $\mathbf{y}=\mathbf{\Psi}\mathbf{x}$, with $\mathbf{\Psi}$: \mathbb{R}^{N}$$\rightarrow$$\mathbb{R}^{M}, $M\ll N$. $\mathbf{\Psi}$ should preserve distance between distant signal vectors, i.e. it should satisfy the restricted isometry property (RIP) [22, 18].

Proven (RIP-compliant) designs for $\mathbf{\Psi}$ take randomly-weighted linear combinations of input vector elements [19, 21]. Unfortunately, such designs often impose challenges regarding practical implementability. For example, in US imaging, sensing weighted combinations of slow-time frames would require an, often infeasible and undesirably, large temporal signal support (including past and future values), and measuring linear combinations of channel signals imposes strong connectivity challenges. Alternatively, sampling a random subset of Fourier coefficients was also shown to be RIP-compliant [19, 21]. Whenever measuring in the Fourier domain is possible (e.g. in MRI), such partial Fourier measurements alleviate the above challenges.

After sensing, sparse signal recovery in CS is typically achieved through proximal gradient schemes, such as the Iterative Shrinkage and Thresholding algorithm (ISTA) [23]. Although proximal gradient schemes are effective tools for solving non-differentiable convex optimization problems, in practice, their time-consuming iterative nature makes them less suitable for real-time applications. Moreover, these methods require knowledge of an adequate sparsifying basis, which in practice is often difficult to formalize. This is particularly true for US imaging.

Recently, a number of deep learning approaches have been proposed for fast signal or image reconstruction in CS [24, 25], showing that deep neural networks can serve as a powerful alternative to conventional recovery techniques.

Inspired by both the challenge of finding adequate context-specific sensing matrices, and the given deep learning approaches for signal recovery, we present a deep learning solution that jointly learns a probabilistic generative model to form a context- and task-based sub-sampling pattern, and a corresponding downstream task model. We refer to the generative model as Deep Probabilistic Sub-sampling (DPS). Efficient joint learning of both models by error backpropagation is enabled through the adoption of the Gumbel-Softmax distribution [26, 27, 28], that circumvents the inherently non-differentiable nature of sampling. We demonstrate our framework for both reconstruction from partial Fourier measurements, and B-mode and tissue-motion (color Doppler) estimation from sub-sampled in-vivo US radio-frequency (RF) data.

The remainder of this paper is organized as follows, we start by providing some related work in Section II, followed by the general framework in Section III-A. Section III-B and III-C respectively elaborate on DPS and the task model. The training strategy is described in Section III-D. Section IV-A demonstrates our solution on a common Fourier domain sub-sampling problem. Its applications in US imaging are subsequently described in Section IV-B and IV-C. Results are given in Section V, which are discussed in Section VI. Final conclusions are drawn in Section VII.

Throughout the paper, bold letters denote vectors (lower case) and matrices (upper case). Matrices are indexed by row-first notation, i.e. $\mathbf{a}_{m}$ denotes the $m^{\text{th}}$ row-vector of matrix $\mathbf{A}$, of which $a_{m,n}$ is the scalar in the $n^{\text{th}}$ column.

II Related work

In this section we briefly list recent applications of conventional CS techniques for medical imaging that sub-sample the data. We then give promising applications of sparse arrays. These examples highlight the potential relevance for learning a task-driven sub-sampling pattern across a number of applications. The recent developments in deep learning for CS, that we discuss lastly, show state-of-the-art methods for learning-based data compression.

II-A Compressed sensing in medical imaging

Several CS approaches have been introduced for various medical imaging applications. In MRI, CS is applied by randomly sub-sampling the K-space [8, 9], i.e. the 2D spatial Fourier transform of the image. The authors of [29] extend this to sub-sampling in the K-time space, while preserving qualitative image reconstructions using their k-t BLAST and k-t SENSE algorithms for one coil and multiple coils, respectively. Likewise, CS has spurred low-dose X-ray CT through image reconstruction from sub-sampled projection measurements [3, 4]. For US imaging, in [7, 30], the authors leverage CS through Xampling, by passing the RF channel signals through analog sum-of-sinc filters, thereby permitting sampling of a partial set of Fourier coefficients. Related to this, we demonstrate how DPS permits learning of partial Fourier coefficients in Section V-A. In the discrete domain, the authors of [31] exploit compressed beamforming for image reconstruction from a randomly sub-sampled set of receive transducer elements, and the work in [6] proposes learned dictionaries for improved CS-based reconstruction from sub-sampled RF lines. In our work, we learn optimized discrete sub-sampling schemes from data in an end-to-end fashion.

II-B Sparse arrays

Significant research efforts have been invested in exploration of adequate sparse array designs [32]. Examples in medical US imaging are a non-uniform slow-time transmission scheme for spectral Doppler [33] and sparse arrays for reduction of the required number of channels for B-mode imaging222In US imaging, B-mode refers to “brightness mode”, a 2D intensity image at a single point in time., based on sparse periodic arrays [34] or sum coarrays [35]. In Section V-B and V-C, we show how DPS enables learning of these slow-time and array sampling patterns for US imaging in a task-based fashion.

II-C Deep learning for compressed sensing

Recently, a number of deep learning approaches have been proposed for fast signal or image reconstruction in CS [24, 25], showing that deep neural networks can serve as powerful signal or image recovery methods. The authors of [24, 36, 37, 38, 39, 40] extend learning beyond signal recovery, and simultaneously train signal compression methods. However, they all rely on taking (randomly weighted) (non)linear combinations of elements from the input vector, making them challenging to implement in hardware. Instead, DPS is based on sub-sampling, which is straightforwardly implementable and applicable across the applications given in Section II-A and II-B.

III Methods

III-A General framework

We consider a signal vector $\mathbf{x}\in\mathbb{C}^{N}$ that we wish to sub-sample through binary sub-sampling matrix $\mathbf{A}_{\boldsymbol{\Phi}}\in\{0,1\}^{M\times N}$ parametrized by $\boldsymbol{\Phi}\in$ $\mathbb{R}^{M\times N}$, to yield a measurement vector333$\mathbf{x}$ and $\mathbf{y}$ can also be higher dimensional. In that case all given formulas are applied on the dimension in which we want to sub-sample $\mathbf{x}$. $\mathbf{y}\in\mathbb{C}^{M}$ , with $M\ll N$:

[TABLE]

We subsequently aim to decode $\mathbf{y}$ into $\mathbf{z}$, some function of the original signal vector $\mathbf{x}$ in which we are interested (i.e. the task):

[TABLE]

To this end, we adopt a (potentially nonlinear) differentiable function approximator $g_{\theta}(\cdot)$ parametrized by a set $\theta$:

[TABLE]

where $\hat{\mathbf{z}}$ denotes the approximation of $\mathbf{z}$ from the sub-sampled measurements $\mathbf{y}$. The function $g_{\theta}(\cdot)$ may for instance be a neural network. Matrix $\mathbf{A}_{\boldsymbol{\Phi}}$ is constrained to have a row-wise $\ell_{0}$ norm equal to 1, i.e. every row contains exactly one non-zero element, and a column-wise $\ell_{0}$ norm that is either 0 or 1. As such, $\mathbf{A}_{\boldsymbol{\Phi}}$ selects a subset of $M$ (out of $N$) elements from input vector $\mathbf{x}$.

To permit joint learning of an adequate sub-sampling pattern for $\mathbf{x}$ and recovery of $\mathbf{z}$ through $g_{\theta}(\cdot)$ by backpropagation, we will introduce a probabilistic sampling strategy, on which we elaborate in the next section.

III-B DPS: Deep Probabilistic Sub-sampling

Each row $\mathbf{a}_{m}$ of $\mathbf{A}_{\boldsymbol{\Phi}}$, with $m\in\{1,..,M\}$, is defined as a one-hot encoding444The one-hot encoding of a scalar is defined as a binary vector with exactly one nonzero value, positioned at the index given by the scalar. We define operator $\mathrm{one\_hot}_{N}(\cdot)$, which returns a one-hot unit-vector of length $N$. of an independent categorical random variable

[TABLE]

where $\boldsymbol{\pi}_{m}\in$ $\mathbb{R}^{N}$ $=\{\pi_{m,1},\ldots,\pi_{m,N}\}$ is a vector containing $N$ class probabilities. Note that $\pi_{m,n}$ thus represents the probability of sampling the $n^{\text{th}}$ entry in $\mathbf{x}$ at the $m^{\text{th}}$ measurement $y_{m}$. We reparametrize $\pi_{m,n}$ using unnormalized log-probabilities (logits) $\phi_{m,n}$, such that

[TABLE]

where $\phi_{m,n}$ is the $n^{\text{th}}$ unnormalized logit of $r_{m}$.

To enable sampling from the categorical probability distribution, we leverage the Gumbel-max trick [41], i.e. sampling is reparametrized into a function of the distribution parameters and a Gumbel noise vector $\mathbf{e}_{m}\in$ $\mathbb{R}^{N}$, with $e_{m,n}\sim\mathrm{Gumbel}(0,1)$, $n\in\{1,\ldots,N\}$ i.i.d.. For $m\in\{1,\ldots,M\}$, a realization of $r_{m}$ is then defined as:

[TABLE]

Previous realizations $\tilde{r}_{1},\cdots,\tilde{r}_{m-1}$ are masked with $w_{m-1,n}\in\{-\infty,0\}$ for $n\in\{1\ldots N\}$, by adding $-\infty$ to the logit of the previously selected category, enforcing sampling without replacement among the $M$ distributions.

Each row $\mathbf{a}_{m}\in\{1,\ldots,M\}$ can now be defined as:

[TABLE]

We define $\boldsymbol{\phi}_{m}\in$ \mathbb{R}^{N}$$=\{\phi_{m,1},\ldots,\phi_{m,N}\} as the $m^{\text{th}}$ row of a trainable matrix $\boldsymbol{\Phi}\in$ $\mathbb{R}^{M\times N}$that contains the unnormalized logits of all distributions. To permit optimization of $\boldsymbol{\Phi}$ by backpropagation, we require $\nabla_{\boldsymbol{\phi}_{m}}\mathbf{\mathbf{a}_{m}}$ to exist $\forall m\in\{1,\ldots,M\}$. Since $\mathrm{argmax}(\cdot)$ is a non-differentiable operator, we adopt the Straight-Through Gumbel Estimator [26, 27] as a surrogate for $\nabla_{\boldsymbol{\phi}_{m}}\mathbf{\mathbf{a}_{m}}$:

[TABLE]

with (row operator) $\mathrm{softmax}_{\tau}(\cdot)$ as a continuous differentiable approximation of the one-hot encoded $\mathrm{argmax}(\cdot)$ operation. We refer to sampling using the $\mathrm{softmax}_{\tau}(\cdot)$ function as soft sampling. Its temperature parameter $\tau$ serves as a gradient distributor over multiple entries (i.e. logits) in $\boldsymbol{\phi}_{m}$.

In the limit of $\tau\rightarrow 0$, soft sampling approaches the one-hot encoded $\mathrm{argmax}(\cdot)$ operator in (III-B) [26, 27], which results in the final trainable sub-sampling pattern of DPS:

[TABLE]

with $\tau>0$ and $m\in\{1,\ldots,M\}$.

III-C Task model

The task model is implemented using a neural network $g_{\theta}(\cdot)$, with trainable parameters $\theta$. Suitable network architectures are application-specific and therefore described per application in  Section IV.

III-D Training strategy

We train model parameters $\boldsymbol{\Phi}$ and $\theta$ by minimizing the mean squared error (MSE) between the model’s output $\hat{\mathbf{z}}$ and the target $\mathbf{z}$, assuming normally distributed prediction errors. To prevent overfitting and exploding gradients, the problem is regularized by adding an $\ell_{2}$ penalty on $\theta$. Besides, we promote training towards one-hot distributions by penalizing convergence towards high entropy distributions using:

[TABLE]

with $\pi_{m,n}$ defined as in (5).

The resulting optimization problem can be written as:

[TABLE]

with

[TABLE]

and

[TABLE]

where the input and target vectors, i.e. $\mathbf{x}$ and $\mathbf{z}$ respectively, follow data-generating distribution $p_{{}_{\mathcal{D}}}$. Penalty multipliers $\lambda$ and $\mu$ weigh the importance of the different penalties.

The Adam solver with hyperparameters $\beta_{1}=0.9$, $\beta_{2}~{}=~{}0.999$, and $\epsilon=1e{-7}$ [42] is used to stochastically optimize (12). In practice, we found that the appropriate learning rates for $\boldsymbol{\Phi}$ and $\theta$ were different. As such, two distinct learning rates were used, i.e. $\eta_{\boldsymbol{\Phi}}$ and $\eta_{\theta}$, by multiplying the update step for $\boldsymbol{\Phi}$ by a factor $\frac{\eta_{\boldsymbol{\Phi}}}{\eta_{\theta}}>1$. The adopted values are reported in Section IV, along with the values for the penalty multipliers $\lambda$ and $\mu$.

The temperature parameter $\tau$ in (10) is initialized at $5.0$ and gradually lowered to $0.5$ during training. This annealing scheme ensures that soft sampling approaches hard sampling towards the end of training, while confining the variance of the gradient updates at the beginning of training [26].

Logits matrix $\boldsymbol{\Phi}$ was initialized to have a prior on training towards a diagonal matrix, i.e. all elements $\phi_{m,n}$, with $m\in$ $\{1,\ldots,M\}$ and $n\in\{1,\ldots,N\}$ are initialized according to:

[TABLE]

with constants $\alpha=-2.73e{-7}$ and $\beta=-2.73e{-3}$, and $\gamma_{m,n}\sim\mathcal{N}(0,0.01)$ i.i.d..

A system-level overview of DPS followed by the task model is visualized in Fig. 1, and its pseudocode is given in Algorithm 1. It was implemented in Python using Keras [43] with a TensorFlow backend [44]. Training and inference were performed on a Titan XP (NVIDIA, Santa Clara, CA).

IV Validation methodology

IV-A Partial Fourier sampling of sparse signals

IV-A1 Data generation and task

Many practical CS applications require signal reconstruction from partial Fourier measurements [8, 9, 45], and we therefore first demonstrate DPS in such a scenario. To that end, we synthetically generate random K-sparse signal vectors $\mathbf{z}\in\mathbb{R}^{128}$, with $K=5$, which we subsequently Fourier-transform to yield the signal $\mathbf{x}\in\mathbb{C}^{128}$ that we aim to partially sample555For each experiment the length of the signal (in the dimension to be sub-sampled) was set to the integer multiple of the sub-sampling factor that is closest to 128, e.g. 126 for factor 6.. Training data is generated on-line, but we create a fixed hold-out test set of 1000 of such signals for model comparison.

In this experiment, the measurement $\mathbf{y}\in\mathbb{C}^{M}$, with $M\leq 128$, is a sub-sampled set of Fourier coefficients in $\mathbf{x}$, and the task is to recover the sparse signal $\mathbf{z}=f(\mathbf{x})=\mathcal{F}^{-1}(\mathbf{x})$, from measurement $\mathbf{y}$, where $\mathcal{F}^{-1}\in\mathbb{C}^{N\times N}$ is the inverse discrete Fourier transform (IDFT).

According to CS theory, the sensing matrix should be RIP-compliant in order to have strong signal reconstruction guarantees [22, 18]. Here, this sensing matrix is defined as $\Psi=\mathbf{A}_{\boldsymbol{\Phi}}\mathcal{F}$. Since the sparse basis, i.e. $\mathcal{F}$, is known, we can check whether $\Psi$ adheres to RIP for the learned sampling matrix $\mathbf{A}_{\boldsymbol{\Phi}}$. For our K-sparse signals, each $M\times K$ (with $M>K$) sub-matrix should be full rank, and all $M$ rows should be (close to) orthogonal to each other. Since $\Psi$ is a subset of the rows of the ortho-normal DFT matrix $\mathcal{F}$, we only need to check the rank of each of the sub-matrices.

IV-A2 Task model

We adopt a specific task model architecture $g_{\theta}(\cdot)$, inspired by the proximal gradient ISTA scheme [23]. It unrolls the iterative solution of ISTA as a 3-layer feedforward neural network with trainable parameters [46, 47]. To prevent dying gradients during backpropagation, we replace the conventional soft-thresholding operators in this learned ISTA (LISTA) method by a sigmoid-based soft-thresholding operator [48]. The thresholds are jointly trained as well.

We compare signal reconstruction by DPS to reconstruction from untrained fixed uniform or random sub-sampling patterns. The latter is typically adopted in CS [19, 21]. The adopted task model is jointly trained with the specified sub-sampling method. We refer to these methods as Uniform+LISTA, Random+LISTA, and DPS+LISTA, respectively. Additionally, we compare these task models to the ISTA reconstruction for (the same) random and learned sub-sampling realizations. These experiments are referred to as Random+ISTA and DPS+ISTA, respectively.

IV-A3 Training

We train for 96,000 iterations across randomly on-line generated mini-batches of 16 Fourier-transformed data vectors. We define one iteration as a parameter update using one mini-batch of data. The learning rates $\eta_{\boldsymbol{\Phi}}$ and $\eta_{\theta}$ are set at $5e{-3}$ and $1e{-3}$, and the penalty multipliers $\lambda$ and $\mu$ at 0.0 and $1e-8$, respectively.

IV-B Slow-time sub-sampling in ultrasound imaging

IV-B1 Data acquisition and pre-processing

Sequential (slow-time) echography data were acquired from two different porcine models. To that end, a 48-element linear array miniTEE s7-3t transducer with a pitch of 0.151 mm was used in combination with a Verasonics Vantrage system (Kirkland, WA). The center frequency for transmission and reception was 4.8 MHz and a 13-angle diverging wave scheme was used (with steering angles uniformly distributed between $-45\degree$ and $+45\degree$, and $F\text{-number}=-1.91$). The sampling rate of the received RF data was 19.2 MHz and coherently compounded beamformed frames (each with 68 scanlines) were collected at a frame rate of 474 Hz. For both the training and test set, all RF data frames were demodulated into their in-phase and quadrature (IQ) components, and subsequently normalized between -1 and +1.

Data belonging to either the training or test set differ in terms of the used porcine model, viewing angle, and location of the probe. The training set contains in-vivo open-chest extracardiac US measurements of the right atrium of the first porcine model. Two ablation surgeries took place, causing physical changes in the movement of the heart. Measurements with two different viewing angles, before, in between, and after the two ablations were included in this training set, leading to a total of 1,340 frames.

The test set contains in-vivo open-chest extracardiac US measurements from the second porcine model. This model had an infarct two months before the acquisition. The measurements were made at the left ventricular posterior wall with the long axis view. In total, 335 frames were recorded.

The experiments were carried out with prior approval from the Dutch Central Commission for Animal Studies (protocol nr. AVD/115002015205) (The Netherlands) and according to the European directives (2010/63/EU) (Belgium) and the Guidelines for the Care and Use of Laboratory Animals (National Institutes of Health, USA).

IV-B2 Tasks

Using the training dataset acquired according to the procedure described in the previous section, we employ DPS to learn a sparse slow-time transmit scheme666DPS learns to select $M_{\text{ST}}$ out of $N_{\text{ST}}$ slow-time frames, of which each is created by coherently compounding 13 diverging wave transmits. The group of diverging waves corresponding to a single slow-time frame is thus always sampled as a whole in this scenario., tailored towards two distinct imaging tasks. As a first task, we consider recovery of the envelope of the beamformed RF signal in order to produce anatomical B-mode US images. Here, $\mathbf{x}$ is (fully sampled) complex IQ data, $\mathbf{y}$ is a subset of slow-time transmits in $\mathbf{x}$, and the target $\mathbf{z}$ is the magnitude of $\mathbf{x}$. Second, we explore learning-based tissue-motion estimation (i.e. color Doppler [49, 50]). In this case, the target $\mathbf{z}$ is computed by applying the Kasai auto-correlator [51] on $\mathbf{x}$. We mask the Doppler targets such that Doppler information is only available at positions where tissue is present. We expect the two tasks to yield very distinct sampling patterns; while B-mode images are independently constructed per frame, Doppler shifts are obtained by measuring phase shifts across the slow-time sequence. Figure 2a illustrates the processing pipeline and associated data dimensions.

IV-B3 Task model

For prediction of $\hat{\mathbf{z}}$ from the subset of slow-time transmit events in $\mathbf{y}$, we employ a convolutional neural network [52]. The architecture of this task model differs slightly for the two defined tasks. For B-mode recovery, the model comprises a total of six convolutional layers, of which the first two are 1-dimensional and applied across the fast-time dimension (with respectively 256 and 128 features, and window length = 5). The last four convolutions are 2-dimensional and applied across both fast-time and slow-time, all having kernel sizes $5\times 5$ and respectively 32, 64, 32, and 1 feature(s). The receptive field of this task model across fast-time, sub-sampled slow-time, and radial scanlines was 25, $M_{\text{ST}}$ (i.e. all sampled slow-time frames), and 1, respectively.

For Doppler reconstruction, $\mathbf{y}$ is first zero-filled, i.e. zeros are added at non-sampled indices, after which six 2D convolutional layers with $5\times 5$ kernels and respectively 32, 64, 128, 64, 32, and 1 feature(s) are leveraged. Applying zero-filling after sub-sampling retains the original time-index of the sampled frames, which we found beneficial for Doppler recovery. The receptive field for fast-time, sub-sampled and zero-filled slow-time, and scanlines is 25, 25, and 1, respectively. In both task models, all convolutional layers, except the last, have leaky ReLU activations ($\alpha=0.1$) [53].

Next, for B-mode recovery, the model contains four 2D convolutions layers (over fast-time and slow-time) having kernel sizes $5\times 5$ and respectively 32, 64, 32, and 1 feature(s). The relatively more difficult task of Doppler recovery required a slightly deeper model, containing six 2D convolutional layers with kernel sizes $5\times 5$ and respectively 32, 64, 128, 64, 32, and 1 feature(s). We use leaky rectified linear unit (leaky ReLU) activation functions ($\alpha=0.1$) across all convolution layers, except the last, which has no activation function [53].

IV-B4 Training

For both tasks, the models are stochastically optimized using the Adam solver, with settings as described in Section III-D, and learning rates $\eta_{\boldsymbol{\Phi}}=2e{-3}$ and $\eta_{\theta}=1e{-4}$. We train for 640,000 iterations with mini-batches consisting of 16 randomly selected patches. Each patch contains 128 sequential slow-time samples of 256 fast-time IQ samples for a single radial scanline. The logits in $\boldsymbol{\Phi}$ are initialized according to (15). Penalty multipliers $\lambda$ and $\mu$ are set at $1e{-5}$ and $1e{-7}$, respectively, and $\mu$ linearly increases with $1e{-8}$ per epoch (1 epoch consists of 64 mini-batch updates) to more strictly promote low-entropy distributions towards the end of training.

IV-C Channel sub-sampling in ultrasound imaging

IV-C1 Data acquisition and pre-processing

The imaging setup and division into training and test set as described in Section IV-B1, was also used to demonstrate DPS for sub-sampling across the 48-channel array, prior to receive beamforming from one diverging wave transmit. To facilitate the subsequent receive beamforming stage, we first pre-delay the channel signals for 68 different scanlines (with steering angles $\zeta$ in $[-\frac{\pi}{4},\frac{\pi}{4}]$) [54]. Taking into account the transmit delay $T_{T}$ (i.e. the time-of-flight (TOF) between the virtual point source behind the array and the focus point in our diverging wave transmission scheme), and the receive delay $T_{R}$ (i.e. the TOF of the back-scattered wave between the focus point and the array element location, indexed by $i$), the total delay function $T(i,\zeta)$ for the central wave transmit is defined as [54]:

[TABLE]

in which

[TABLE]

and

[TABLE]

Focal depth is denoted by $f_{d}$, $r_{f}$ is the distance between the surface of the transducer array and the virtual point source behind the array, $\Delta x$ and $L$ are respectively the pitch and total number of channels of the array, and $c_{0}$ denotes the speed of sound in soft tissue. The adopted values for these parameters are: $f_{d}=40$ mm, $r_{f}=13.5$ mm, ${\Delta x=0.151}$ mm, $L=48$, and $c_{0}=1540$ m/s.

After computing 68 delayed signals per channel, we obtain a 4D dataset spanning slow-time frames, fast-time samples, channels, and radial scanlines. Note that pre-computing these delays is merely only done to accelerate training, and can in practice be performed post-sampling, as shown in Fig. 2b. Finally, the pre-delayed RF channel signals are demodulated into their IQ components, and thereafter normalized between -1 and +1.

IV-C2 Tasks

We again distinguish two tasks, B-mode reconstruction and color Doppler estimation. For both tasks, targets are generated using fully sampled DAS-beamformed channel data from one wave transmit. These beamformed IQ-demodulated RF signals are processed as described in Section IV-B2 to create B-mode, respectively Doppler, targets.

IV-C3 Task model

For recovery of $\hat{\mathbf{z}}$ from the sub-sampled channel data in $\mathbf{y}$, we leverage a convolutional neural network having the same architecture for both tasks. After sub-sampling, we first apply zero-filling for the non-sampled channels, after which four 2D convolutional layers, having $5\times 5$ kernels and respectively 64, 128, 64, and 48 features, are adopted. Convolutions take place across the fast- and slow-time dimension, i.e. the channels are fully connected. Each convolutional layer is followed by a leaky ReLU activation function ($\alpha=0.1$) [53]. The network’s last layer is a fully connected layer across the (sub-sampled) channel dimension, and therefore acts as an array apodization, as used in typical DAS beamforming [55]. The receptive field of this model for respectively fast-time, slow-time, sub-sampled channels, and scanlines was 17, 17, $M_{\text{CH}}$ (i.e. all selected channels), and 1.

IV-C4 Training

For both tasks, the networks are stochastically optimized using the Adam optimizer, with its settings as described in Section III-D. The learning rates $\eta_{\boldsymbol{\Phi}}$ and $\eta_{\theta}$ are set at $2e{-3}$ and $1e{-4}$ respectively, and we train for 192,000 mini-batch updates. Randomly selected mini-batches consist of 16 patches spanning 32 slow-time frames, 64 fast-time samples, 48 channels, and one radial scanline. The trainable matrix $\boldsymbol{\Phi}$ is initialized according to (15) and hyperparameter $\lambda$ is set at $1e{-5}$. The entropy penalty multiplier $\mu$ is set at $1e{-8}$ and linearly increases with $1e{-11}$ per epoch (1 epoch consisted of 64 mini-batch updates), to more strictly promote one-hot distributions towards the end of training.

V Results

V-A Partial Fourier sampling of sparse signals

Figure 3 displays sparse signal recovery from partial Fourier measurements for an example signal from the test set. The Fourier space is sub-sampled (with factor $\frac{N}{M}=4$) uniformly (a), randomly (b), or using a realization from the learned sub-sampling distributions by DPS (c). A quantitative evaluation of the signal recoveries for different sub-sampling and reconstruction methods is given in Fig. 4a. Uniform sub-sampling performed poorly due to aliasing, resulting in a repeated prediction pattern (Fig. 3a-right). However, (CS-inspired) random sampling (Fig. 3b-right) shows on par results with learned sampling by DPS (Fig. 3c-right). Interestingly, the learned pattern also exhibits (pseudo-random) irregular sampling (Fig. 3c-left), and its corresponding sensing matrix $\boldsymbol{\Phi}$ showed to be RIP-compliant.

Comparing signal reconstruction using either ISTA (300 iterations) or LISTA (3 folds), Fig. 4a shows that for high sub-sampling factors, LISTA outperforms ISTA, while enabling much faster (speed-up factor $>1000$) inference (Fig. 4b). This finding confirms the results provided by the authors of [47], who show that an unfolded scheme can achieve better performance for few unfoldings compared to the iterative scheme for the same amount of iterations.

V-B Slow-time sub-sampling in ultrasound imaging

Figure 5 demonstrates anatomical B-mode (a-e) and color Doppler (f-j) reconstructions for both fixed uniform and DPS-driven sub-sampled slow-time transmits. For uniform sub-sampling, the DPS block in Fig. 2a is replaced by a (non-trainable) uniform sub-sampling operation. The adopted sampling patterns (with selected frames in black) are visualized below the reconstructions. Note that for DPS, these sampling patterns correspond to a single stochastic realization from the trained probability distributions, on which we detail in Appendix A. To demonstrate both the individual and joint contributions of DPS and the task-model to performance, we moreover compare our trained neural network task model to a non-learned cubic-interpolation-based reconstruction method for both sampling strategies.

Interestingly, DPS’ learned sampling patterns for both tasks showed to be very distinct. For B-mode reconstruction, the learned pattern exhibited an almost perfectly uniform sampling pattern (Fig. 5c and d), whereas For Doppler reconstruction, the sub-sampling pattern found by DPS exhibited an ‘ensemble’-type of sampling (Fig. 5h and i). For both tasks, similar patterns were clearly visible for the other tested sub-sampling factors as well. The learned ‘ensemble’-style sub-sampling pattern for Doppler recovery efficiently captures high frequency slow-time signals due to tissue displacements (Doppler shifts) within ensembles, and relatively low frequency information (changes in Doppler shifts over time) among these ensembles. Given the (nearly) uniform sampling pattern learned by DPS for B-mode reconstruction, the resulting reconstructions using fixed uniform versus learned sub-sampling (Fig. 5a, b and 5c, d, respectively), were found to be similar. For increasing factors, reconstructed images displayed increased blurring for all sampling-reconstruction combinations. The MSEs for different sub-sampling factors, both for uniform and learned sampling by DPS combined with task model $g_{\theta}(\cdot)$, are visualized in Fig. 6a. It confirms the similarity of reconstructed B-mode images for uniform and learned sub-sampling patterns.

The need for accompanying DPS with a jointly trained nonlinear task model becomes particularly evident when comparing the Doppler reconstructions in Fig. 5h and i. Simply interpolating frames and using a standard Doppler estimator (Kasai) leads to significant aliasing in Fig. 5h. For anatomical B-mode imaging, which does not rely on inter-frame processing, the difference between reconstruction methods mainly manifests itself by the inherent denoising priors of the chosen convolutional architecture.

That joint training of sampling and a dedicated task model yields best results can also be seen in Fig. 6b, where we quantitatively compare fixed versus learned sampling.

Using the trained network for inference on 256 slow-time frames, containing 68 scanlines and 2048 fast-time samples from the test set, took on average 1.29 s ($\sigma=29.2$ ms). Accordingly, the reconstruction network allows a reconstruction speed of 198 sub-sampled frames per second.

V-C Channel sub-sampling in ultrasound imaging

Figure 7 displays the anatomical B-mode and Doppler reconstructions for learned channel sub-sampling by DPS, compared to uniform channel sub-sampling. Interestingly, for both tasks, the abundant sampling of center channels was prioritized by DPS.

Assessing the respective contributions of sampling and subsequent learned processing to the overall performance of these sparse arrays, we find that using adequate jointly learned processing (here in the form of a deep network) is crucial (Fig. 7d and i). Compared to naive DAS/Kasai processing (Fig. 7c and h), we indeed observe boosted lateral resolution, contrast, and mitigation of aliasing artifacts. Yet, despite the power of these learned deep networks, they do not recover the severely aliased signals obtained by the (non-learned) uniformly sampled sparse array - strong grating lobes dominate the image in that case (Fig. 7b and g).

For phased array US probes, it is common practice to adopt a pitch which is half the signal’s wavelength in order to prevent grating lobes in the field of view [55]. Increasing the pitch between channels by uniformly sub-sampling the channel array thus caused grating lobes to appear in the B-mode images, both in grayscale (Fig. 7a and b) and with Doppler overlay (Fig. 7f and g), which we indicated by the white dashed lines in Fig. 7a and f-top. For B-mode recovery, we also show a slice of the provided B-mode frames (middle plots in Fig. 7a-e) at an axial depth indicated by the red arrow in Fig. 7e. The relative angle $\zeta_{g}$ of the $k^{\text{th}}$ grating lobes (with respect to the main beam) can be calculated as [55]:

[TABLE]

where $\Delta x$ ($=0.151~{}$mm) is the (original) pitch of the array, and $\lambda$ ($=0.3~{}$mm) is the wavelength of the signal. For $\frac{N}{M}=4$, the angle of the first ($k=1$) grating lobe is $\zeta_{g}=29.8\degree$, which is in accordance with the visually present grating lobes.

Figure 8a and b show quantitative results for both B-mode and Doppler reconstruction using different sub-sampling factors and the two different sampling strategies (uniform versus DPS), while adopting the dedicated trained task model for reconstruction. Here we again observe that the joint training regime of sampling and reconstruction outperforms fixed sampling with a trained reconstruction model.

Realizing that a bigger aperture yields higher lateral resolution, it is remarkable that DPS did not assign high probability to sampling (some of the) outer channel elements of the array. To check whether DPS could have found a more optimal sub-sampling pattern, possibly including more outer channel elements, we also compare DPS to a fixed sparse array (using half of the elements777The theoretical lower bound on the amount of necessary channels (allowing non-integer channel positions) is given at $2\sqrt{C}$, where the full array contains $2C-1$ elements [35]. This theoretical lower bound is thus 9.9 channels for our 48-channel array. However, we can only sample integer channel positions (i.e. physically existing channels). As such, we could only create a design for sub-sampling factor $2$.) that is designed based on the notion that its sum coarray should be full [35]. For both sampling strategies, a dedicated task model $g_{\theta}(\cdot)$ was trained again. Notably, the learned sparse array by DPS (Fig. 9b), also has a sum coarray that is full, and yields a performance that is up to par with that of the hand-designed sum coarray (Fig. 9).

Computationally, running inference on patches from the test set showed an average reconstruction time of 36.7 ms ($\sigma=1.65$ ms) for IQ data from 12 channels, steered towards 68 scanlines with 2048 fast-time samples at one point in (slow-)time, implying a frame reconstruction rate of 27 frames per second.

VI Discussion

Recent technological trends in medical imaging have spurred the demand for imaging pipelines that rely on less data without compromising image quality, temporal resolution, or more generally, diagnostics. We here consider the notion of task-adaptive sampling, in which sampling schemes are optimized not to recover the sensor signals themselves, but to fulfill a specific imaging task.

To that end, in this paper we proposed a sampling framework, DPS, that permits joint learning of a context- and task-specific sub-sampling pattern, and an adequate task model. We demonstrated that this paradigm yields improved performance compared to non-trained sampling and reconstruction models across a variety of ultrasound imaging experiments.

Unlike other recently introduced learned compressed sensing techniques, DPS learns to sub-sample rather than to take full linear measurements, which face practical implementation challenges. Sub-sampling permits straightforward implementation of the learned sampling pattern into sensing applications, with examples being sparse array design, slow-time US pulsing schemes, (non-uniform) analog-to-digital converters (ADC) and partial Fourier measurements.

In US imaging, we specifically applied DPS for slow-time pulse scheme design and the array channel selection problem. Besides data reduction, the former reduces the amount of transmit events, which has the additional advantage of drastically reducing power consumption. Reduced power consumption also benefits battery life for wireless applications, and reduces heat generation of ADCs, which is particularly relevant for in-body applications. The applications, or tasks, that we considered within the US imaging domain were B-mode imaging and tissue-motion (color Doppler) imaging. DPS yielded distinct sampling patterns for each task, with e.g. tissue-motion estimation spurring a pattern that uses compact groups of slow-time frames with a short inter-pulse time.

Several points of discussion are important to note here. First, since DPS performs sampling without replacement among the distributions, duplicate sampling does not occur. Yet, in particular for channel sub-sampling, we observed that in some cases several of the $M$ distributions trained towards assigning a high probability for the same channel (Appendix A, Fig. 10b and d). The entropy penalty multiplier $\mu$ and the ratio between the learning rates $\eta_{\boldsymbol{\Phi}}$ and $\eta_{\theta}$ were found to have predominant influence on this training process. Although extensive fine-tuning of hyperparameters may therefore further improve the results, it was out of the scope of this research. To promote convergence towards different distributions, one might also consider adding an additional column-wise penalty on logits matrix $\boldsymbol{\Phi}$.

Second, the experiments for channel sub-sampling were performed using a driving scheme comprising only a single diverging wave transmit, causing relatively low lateral target resolution. This may have steered the optimization of sub-sampling patterns by DPS dominantly towards the central channels. Importantly, we would like to stress that the presented sub-sampling patterns are not generic across data types or imaging setups, but rather fully context and task-adapted in an end-to-end optimized data-driven fashion. Adaptation to a new context, data distribution, and imaging setup is done based on new training data.

Third, the focus of this work was on the development of a framework that permits backpropagation-based joint learning of (discrete) sampling with a task model, and we did not heavily optimize, nor fine-tune, for different task model architectures. Comparing Fig. 5b and d to Fig. 5a and c, respectively, we notice that the convolutional nature of our task model leads to inherent denoising and loss of speckle, likely due to the deep image priors [56]. Therefore, if desired, in future work, one could resort to more US-specific adaptive beamforming deep networks such as ABLE [57, 58] to better preserve speckle in B-mode images.

We expect that other US imaging applications, such as super-resolution US localization microscopy (ULM), can similarly benefit from learned and dedicated sampling schemes. In ULM, millions of highly sparse point-scatterers (intravascular microbubbles) are to be detected and localized across thousands of frames at ultrafast imaging rates [59]. Consequently, data rates are extremely high. Recently, deep neural networks have been proposed for fast ULM recovery [60, 50], and one can envisage the use of DPS to learn adequate sampling patterns that reduce data rates in this context.

Beyond the US applications considered here, future work may include learning sub-sampling and reconstruction for compressed sensing MRI [8], where measurements are inherently performed by sampling the spatial Fourier domain. MRI thus shares strong similarities with signal reconstruction from partial Fourier measurements (shown in Section V-A), making it an excellent candidate for DPS. Also investigating the use of DPS for sparse view CT imaging is of interest, potentially permitting reduction of the amount of transmit events, and therewith exposure to harmful radiation.

VII Conclusions

In this paper we have presented Deep Probabilistic Sub-sampling (DPS), a probabilistic deep learning framework that permits joint optimization of a task-based sub-sampling scheme and a downstream task model. We have demonstrated its effectiveness for sensing partial Fourier coefficients of sparse signals and a number of US imaging applications, showing that the proposed method indeed learns sampling schemes that are dedicated to a given task. As such, DPS opens up a wide range of new opportunities; beyond US imaging, we foresee its application in other medical imaging domains (e.g. MRI and CT) and, more generally, in compressed sensing problems.

Appendix A Trained probability distributions by DPS

DPS relies on a generative probabilistic sampling model, in which samples are selected stochastically based on learned probabilities for the entire set of candidate samples. For each selected sample, a full categorical distribution across all $N$ candidates is trained, leading to $M$ probability distributions that express the model’s belief regarding the relevance of particular candidate samples at any stage of the training process. These probabilities often converged to near-one-hot representations, with high beliefs for particular candidate samples. Interestingly this did not occur for channel selection with the task of color Doppler recovery. Figure 10 shows the trained probability distributions across tasks and sampling strategies, along with a particular realization of these random variables in the form of a sampling scheme.