Adaptive Sampling for Linear Sensing Systems via Langevin Dynamics

TL;DR

This paper introduces a Bayesian adaptive sampling method using Langevin dynamics to improve image quality and speed in sensing systems, demonstrated on MRI with significant quality gains.

Contribution

It presents a novel Bayesian adaptive sampling approach based on SGLD that generalizes well across different image priors and out-of-distribution data.

Findings

Improved MRI image quality by 2-3 dB PSNR.

Effective adaptive sampling enhances subtle detail restoration.

Method generalizes across analytical and neural network priors.

Abstract

Adaptive or dynamic signal sampling in sensing systems can adapt subsequent sampling strategies based on acquired signals, thereby potentially improving image quality and speed. This paper proposes a Bayesian method for adaptive sampling based on greedy variance reduction and stochastic gradient Langevin dynamics (SGLD). The image priors involved can be either analytical or neural network-based. Notably, the learned image priors generalize well to out-of-distribution test cases that have different statistics than the training dataset. As a real-world validation, the method is applied to accelerate the acquisition of magnetic resonance imaging (MRI). Compared to non-adaptive sampling, the proposed method effectively improved the image quality by 2-3 dB in PSNR, and improved the restoration of subtle details.

Equations

$$\mathbf{y}=\mathbf{A}\mathbf{x}+\mathbf{\epsilon },$$

$$l=\underset{n\in 1,2,\dots N}{\mathrm{argmax}}\mathrm{Var}\{[{\hat{\mathbf{y}}}_1^{(k)}{]}_n,\dots ,[{\hat{\mathbf{y}}}_{N_{\mathrm{sample}}}^{(k)}{]}_n\}.$$

$$\hat{\mathbf{x}}\sim p(\mathbf{x}|{\mathbf{y}}^{(k)})=p(\mathbf{x})p({\mathbf{y}}^{(k)}|\mathbf{x})/p({\mathbf{y}}^{(k)}),$$

$$\Delta {\mathbf{x}}_t={\eta }_t(\nabla \log p({\mathbf{y}}^{(k)}|{\mathbf{x}}_t)+\nabla \log p({\mathbf{x}}_t))+\sqrt{2{\eta }_t}\mathcal{N}(0,1),$$

$$\underset{\mathbf{\theta }}{\mathrm{argmin}}{\mathbb{E}}_{\mathbf{x}\in \mathcal{X}}\Vert \log p(\mathbf{x})-f_{\mathbf{\theta }}(\mathbf{x}){\Vert }_2^2$$

Full text

**Index Terms— ** adaptive sampling, diffusion model, score-based model, Bayesian experimental design, magnetic resonance imaging

1 Introduction

Many imaging systems acquire measurements sequentially. Reducing the number of measurements can accelerate the signal acquisition process and benefit modalities that require lower radiation, such as computed tomography (CT) and scanning electron microscopy (SEM). Nevertheless, this can result in an under-determined image reconstruction problem. To address this challenge, various reconstruction methods have been proposed, such as compressed sensing [1], to enable the recovery of an object from undersampled measurements.

Sampling strategy also plays a critical role in achieving high-quality images. For instance, many sub-Nyquist sampling patterns have been investigated in MRI, including analytical and data-driven designs [2]. However, predetermined strategies may not always be optimal for various imaging scenarios. To address this challenge, adaptive sampling or dynamic sampling techniques can select the next batch of ‘important’ data points based on existing observations. This approach enables better use of prior information from both signal statistics and observed signals, leading to improved image quality and acquisition speed. Relevant methods include Bayesian experimental design (BED) [3], neural network-based regression [4], and reinforcement learning [5]. These methods improved image quality in various applications. However, many neural network-based methods may lack generalization ability and explainability to out-of-distribution test sets and real-world applications.

This paper presents a model-based dynamic sampling approach that predicts new sampling locations by greedily minimizing the variance of posterior samples drawn from the posterior distribution [6]. The sampler uses stochastic gradient Langevin dynamics (SGLD) [7] and supports various image priors. We applied the proposed dynamic sampling to accelerate MRI acquisition. Across many experiment settings, the proposed approach significantly improved the image quality.

2 Methods

Consider a linear sensing system

[TABLE]

where $\bm{A}\in\mathbb{C}^{M\times N}$ denotes a sensing matrix, ${\mathbf{x}}\in\mathbb{C}^{N}$ denotes the object, and $\bm{y}\in\mathbb{C}^{M}$ denotes raw measurements. To accelerate the acquisition, we consider the ‘undersampled’ case where $\bm{y}$ has $L<N$ non-zero entries. Typically, the locations of non-zero entries in $\bm{y}$ follows pre-determined patterns. The proposed method, instead, dynamically chooses additional sample locations in a sequence of $K$ sampling iterations where the samples for iteration $k+1$ are based on the measurements $(\bm{y}_{1},\ldots,\bm{y}_{k})$ recorded in previous iterations.

Specifically, we apply a Bayesian approach [7]. At the $k$ iteration of additive sampling, based on the measurements acquired up until this point $\bm{y}^{(k)}=(\bm{y}_{1},\ldots,\bm{y}_{k})$, the first step draws samples from the posterior distribution $p({\mathbf{x}}|\bm{y}^{(k)})$, yielding a collection of reconstructed images denoted $\{\hat{{\mathbf{x}}}_{i}^{(k)}\}_{i=1}^{N_{\mathrm{sample}}}$. We use an SGLD sampler detailed below. The second step projects each estimate $\hat{{\mathbf{x}}}_{i}^{(k)}$ (typically in the image domain) back to the measurement domain using the sensing equation $\hat{\bm{y}}_{i}^{(k)}=\bm{A}\hat{{\mathbf{x}}}_{i}^{(k)}$. The third step selects the next sampling locations by greedily minimizing the variance of samples $\left\{\hat{\bm{y}}_{i}^{(k)}\right\}$ in the measurement domain. In detail, we select the next measurement location(s) $l$ for the $k+1$ iteration using the k-space locations having the maximum variance:

[TABLE]

To compute a collection of reconstructions or estimates $\{\hat{{\mathbf{x}}}_{i}\}$, we sample from the posterior

[TABLE]

where $p({\mathbf{x}})$ denotes the prior and $p(\bm{y}^{(k)}|{\mathbf{x}})$ denotes the likelihood. In contrast, a typical iterative image reconstruction algorithm gives a point estimate, such as the MAP estimator. SGLD [7] samples from the posterior distribution using the update

[TABLE]

where $\eta_{t}$ denotes the time-dependent step size [8, 9]. Intuitively, SGLD explores the solution space by injecting Gaussian noise similar to the Langevin Monte Carlo sampler.

In applications where the noise $\bm{\varepsilon}$ is Gaussian, the gradient of likelihood has the closed-form solution $\nabla\log p(\bm{y}|{\mathbf{x}})=-\bm{A}^{\prime}(\bm{A}{\mathbf{x}}-\bm{y}).$ The prior term $p({\mathbf{x}})$, or the score function $\nabla\log p({\mathbf{x}})$ can take various forms. For example, a simple prior that penalizes first-order roughness has the form $p({\mathbf{x}})=\mathrm{e}^{-\lambda\|\bm{T}{\mathbf{x}}\|_{2}^{2}/2}$, where $\bm{T}$ is the first-order finite difference transform; its corresponding score function is $\nabla\log p({\mathbf{x}})=-\lambda\bm{T}^{\prime}\bm{T}{\mathbf{x}}$. Analytical priors may not be informative and many studies propose to learn score functions from datasets. Score matching approximates the score function with a learnable function $f_{\bm{\theta}}({\mathbf{x}})$ and learns from a training set $\mathcal{X}$:

[TABLE]

Recent improvements in score matching, such as sliced score matching and denoising score matching [8, 9], have extended the method’s effectiveness and made it more applicable to large datasets [10, 11] To demonstrate the adaptability of our algorithm, we tested both analytical priors and score functions based on neural networks. Alg. 1 details the proposed approach.

3 Experiments

We applied the proposed dynamic sampling method to MRI data that reside in the Fourier domain (k-space). For our experiment with Cartesian sampling, the sensing matrix $\bm{A}$ contained both FFT and coil sensitivity (calculated by methods described [12]). The score functions included both a simple analytical one $f({\mathbf{x}})=-\lambda\bm{T}^{\prime}\bm{T}{\mathbf{x}}$ and a learned U-Net-based model. We evaluated the analytical priors on multiple MRI datasets [13, 14, 15], using both 1D and 2D sampling patterns. We compared the dynamic sampling patterns with well-received fixed sampling patterns, such as Poisson-disk, for $N_{\mathrm{add}}$ = 50 and $N_{\mathrm{step}}$ = 200.

We used the same U-Net-based architecture (NCSN++) and configurations as in [11] to train the learned prior on the fastMRI brain dataset. The complex-valued image was formulated as two input channels. To demonstrate the generalization ability, we tested it on test sets that contained different anatomies and sequences than the fastMRI database, including an MP-RAGE sequence of human brains [15] and a GRE sequence of mouse brains, without any fine-tuning. For the mouse brain dynamic contrast-enhanced (DCE) data, we learned the sampling pattern from a ‘pilot’ frame and then applied it to subsequent frames. We used $N_{\mathrm{add}}$ = 30 and $N_{\mathrm{step}}$ = 100 and the accelerated sampler described in [16]. The sequence $\bm{\eta}$ used the same configuration as described in [11].

4 Results

For the analytical prior, Fig. 1 and Fig. 2 show dynamic sampling patterns and corresponding reconstruction examples. Compared to predetermined sampling patterns, the proposed method reduced aliasing artifacts across multiple anatomies and contrasts.

For the learned prior (NSCN++), Fig. 3 shows an out-of-distribution example, using GRE sequences of the human brain. With the proposed adaptive sampling, the fine details and tissue contrast in the reconstructed images were improved compared to predetermined sampling patterns. Fig. 4 shows another out-of-distribution case, mouse brain DCE imaging. The adaptive sampling scheme was optimized for the first frame and applied to subsequent frames. Adaptive sampling led to less blurred structures and improved SNR.

5 Discussion

The posterior sampling processes can be computationally expensive, determined by both the system matrix $\bm{A}$ and the score function $\nabla\log p({\mathbf{x}})$. Simpler analytical priors may accelerate the sampling. The sampling is embarrassingly parallel and can benefit from parallel computing and hardware improvements. In its current form, the proposed dynamic sampling is particularly useful for dynamic imaging applications such as fMRI and DCE-MRI where a ‘pilot’ scan is available to design tailored sampling patterns for subsequent frames and avoid the long computation time that may compromise the benefits of dynamic sampling.

The sampling from the posterior distribution may benefit from faster samplers [17]. Some ‘single-shot’ samplers based on neural network methods can sample faster than SGLD [18] however, they are trained on a certain dataset and may lack the ability to generalize to out-of-distribution applications.

The proposed dynamic sampling method has demonstrated decent robustness in simulated experiments and analytical priors worked well for different test cases. The learned priors were trained on a fastMRI brain dataset but generalized well to different anatomies, vendors, sequences, and field strengths. Future work will include a systematic comparison with prior arts and prospective in-vivo experiments.