Solving Inverse Problems with Score-Based Generative Priors learned from Noisy Data

I Introduction

Recent advances in score-based (diffusion) generative modeling have helped substantially improve the capabilities of solving ill-posed imaging inverse problems using fewer measurements and with higher reconstruction fidelity in various domains such as medical imaging , digital communications , image super-resolution , and more.

However, learning high-quality score-based generative models for distributions over real-world signals currently assumes a large database of fully-sampled and noise-free training samples are available . In many application domains, acquiring such a training set is impossible in practice, because noise is inherently present in the sensors used to acquire measurements. For example, noise is always present in single or multi-coil magnetic resonance imaging (MRI) due to thermal noise in the radio-frequency (RF) hardware as well as losses due to the body. Similarly, noise is present in communication transceivers, especially in wireless scenarios where mobile users operate devices under cost and energy constraints. Learning score-based generative models using only noisy samples is thus an important research problem. In this work, we propose a solution to this problem by combining denoising using Stein’s unbiased risk estimate (SURE) with the denoising score matching (DSM) objective used in diffusion model training , to achieve joint denoising and score learning with a single deep learning model.

Our framework connects the two already-similar objectives by leveraging Tweedie’s single-step denoising formula at a properly chosen noise level. We show that this objective is naturally compatible with the DSM formulation and allows for a simple modification which introduces two additional terms and a scalar weight in the training function, in the case of training data corrupted by independent and identically distributed (i.i.d.) Gaussian noise. We evaluate the effectiveness of our proposed approach by using the learned models as priors in two different applications: compressive multiple-input multiple-output (MIMO) channel estimation and accelerated MRI reconstruction. Our results show that the priors learned using the proposed formulation can be reliably used for solving linear inverse problems, even when the training dataset is captured in a poor signal-to-noise ratio (SNR) of dB.

There is an extensive body of prior work related to denoising of structured signals; see for a recent review paper. Most relevant to our work is the use of the self-supervised SURE objective , which introduces an unbiased estimate of the supervised denoising loss purely using noisy data. Using the SURE loss as a training objective for deep neural network denoisers has been previously investigated in a standalone denoising setting , and the connection to multi-level denoising score matching has only been recently explored in , where learning the score function on a restricted subset of noise levels is addressed.

Several other recent methods have used self-supervised learning for blind denoising without ground truth training data. One example is Noise2Noise , where multiple noisy measurements of the same sample are used to train a deep neural denoiser. The work in further introduces the Noise2Self framework, which requires only a single noisy measurement of each sample in the training set, similar to our proposed method. The recent work in introduces the Noise2Score framework for self-supervised learning of the score function in moderate and high SNR settings, which leaves the problem of learning score functions in low SNR settings open. More recently, the work in has investigated learning a score-based generative model from noisy diffusion MRI samples, still requiring access to multiple noisy scans of the same subject.

Pre-trained generative priors have also been used in solving linear inverse problems, surpassing classical compressed sensing approaches . Recently, a series of works have shown that score-based generative models produce competitive reconstruction performance when used to solve accelerated MRI in a supervised learning setting, and that the score can be learned implicitly from MMSE denoisers . The work in applies the same ideas to learning priors for MIMO channel estimation from a limited number of pilot measurements. The recent work of explores self-supervised learning of score models for MRI reconstruction by splitting the measurement data into distinct subsets, similar in approach to Noise2Self.

I-B Contributions

Our contributions in this work are the following:

We formulate SURE-Score: a framework for jointly learning a denoiser and score-based generative model using the same deep neural network. The loss connects SURE, DSM, and Tweedie’s rule in a single objective.

We evaluate SURE-Score by learning score-based generative models from noisy data for MIMO channels and MRI. Even though the distributional structure in the two domains is different, our results show that SURE-Score can learn accurate priors even at an SNR of dB.

Our simulation results show that MIMO channels acquired at dB can be used for compressed channel estimation with a performance drop of at most $8$ dB in normalized mean square error (NMSE) compared to fully supervised. We also show that self-supervised denoising and reconstruction at $5\times$ accelerated MRI is possible with about $15\%$ error increase when learning a prior from noisy samples acquired at $10$ dB.

II System Model and Background

We model signals corrupted by additive white Gaussian noise as follows:

II-B Score-Based Generative Modeling

The score function $\psi_{X}$ of a probability distribution $p_{X}$ is defined as the gradient of its logarithm with respect to the sample :

Learning the score function of the probability distribution avoids the challenges of normalization and is useful to downstream tasks such as sampling from the learned distribution or solving inverse problems . The goal of score-based generative modeling is to learn the score function $\psi_{X}$, given a training dataset sampled from $p_{X}$.

When noiseless data are available for training and choosing the added noise $z$ as Gaussian i.i.d., the gradient of the conditional distribution can be expressed in closed form as $-z/\sigma^{2}$. Score-based generative models are trained using denoising score matching at multiple noise levels simultaneously . Following , we weigh the loss at each level by $\sigma^{2}$ to normalize the magnitude of the DSM loss across all noise levels. This yields the expression for the multilevel DSM training loss:

The distribution of training noise levels $p_{\Sigma}$ is a uniform distribution across geometrically distributed noise levels between $\sigma_{\textrm{min}}$ and $\sigma_{\textrm{max}}$. We use the guidelines in and implement the noise-conditional score network as $s_{\theta}(x;\sigma)=s_{\theta}/\sigma$, where $s_{\theta}(x)$ is a deep neural network with learnable weights $\theta$.

II-C Solving Inverse Problems using Annealed Langevin Dynamics

We use a modified version of annealed Langevin dynamics to sample from the posterior distribution. This is an iterative algorithm that starts from a random $x_{0}\sim N(0,I)$ and requires access to the noise-conditional score function $\psi_{X_{t}|Y}(x_{t}|y;\sigma)$ at each step. For brevity, we omit $\sigma_{t}$ from the notation in the following. The update rule is given by:

The conditional score function $\psi_{X_{t}|Y}$ can be further expanded via Bayes’ rule as:

where the denominator term vanishes due to not depending on $x$. The first term is related to the linear formulation of the inverse problem in (2). As $\sigma_{t}\rightarrow 0$, this term becomes:

where we choose $\gamma_{t}^{2}=\sigma_{t}^{2}$ in this work. The dependency of the second term in (8) on $\sigma_{t}$ is included in a pre-trained noise-conditional score-based model $s_{\theta}(x;\sigma_{t})$, which can readily approximate this term. The final update rule for sampling from the posterior distribution is given by:

We emphasize that our algorithm only approximates posterior sampling.

II-D Stein’s Unbiased Risk Estimate

Evaluating all partial derivatives in (13) requires a number of function evaluations proportional to the input dimension, which does not scale to real-world high-dimensional signals. Following , we approximate the divergence term using its Monte Carlo approximation with $n\sim\mathcal{N}(0,I)$ being i.i.d. Gaussian noise sampled for each sample in each batch, and the approximation given by:

thus it can be used as a training function for the model $g_{\theta}$, when only a noisy set of training samples is available , and a mini-batch of noisy samples is used. Note that the SURE loss can be applied when a single noisy version of each sample is available in the dataset, which matches our setting.

III Proposed Method: SURE-Score

Our proposed method learns the score of the perturbed distribution of minimum mean square error (MMSE) denoised samples. While this is distinct from learning the score of the true perturbed distribution of clean samples at arbitrary noise levels, our formulation couples SURE learning and denoising score matching via Tweedie’s rule for Gaussian corruptions such that the two objectives are consistent with each other at a specific noise level. We call our method SURE-Score. A block diagram of the training flow is shown in Figure 1.

Naively, learning the score of perturbed MMSE denoised samples could be done with two different functions, each with its own set of learnable parameters, and applied sequentially:

A function $g_{\phi}$ that learns the MMSE denoiser. This function can be learned purely from noisy data using the SURE training objective described in Section II-D.

A score-based generative model $s_{\theta}$ trained on the outputs of $g_{\phi}$. If sufficient training data are available, it could be split between the two stages, otherwise the same training set used to learn $g_{\phi}$ can be re-used in the second stage.

The previous identities suggest that a score-based generative model for a specific noise level and the MMSE denoiser for the same noise level are interchangeable. Thus, we propose to learn them jointly using the following loss function at a single noise level $\sigma$, with training data corrupted at noise level $\sigma_{w}$:

where $\lambda>0$ is a hyper-parameter controlling the relative weighting between the two losses. Note that in the above, the DSM term uses the output of the denoiser $g_{\theta}$ – this leads to a recurrent call of the same underlying network. Each of the two terms in the above corresponds to one of two objectives: the first term learns the score-based generative model for the perturbed distribution of MMSE denoised samples, while the second term learns the MMSE denoiser itself.

The Tweedie re-parameterization introduces a consistency between the two objectives at noise level $\sigma_{w}$, that is when $\sigma=\sigma_{w}$ in (20): the more accurately the true score function is approximated at noise level $\sigma_{w}$, the more accurately denoising becomes, and vice versa. This can be viewed as a coarse version of the recently introduced consistency objective in , where we ensure that removing and adding back noise at noise level $\sigma_{w}$ leads to consistently learned gradients of the perturbed score function. In practice, both terms at $\sigma=\sigma_{w}$ are bounded and non-zero due to the non-zero residuals of MMSE denoising. Additionally, the above consistency is only introduced for $\sigma=\sigma_{w}$, but in practice we find that this does not hinder performance at other noise levels, where extrapolation is sufficiently accurate.

Expanding the expected versions of the two losses in (20) yields the expression in (16). Replacing $g_{\theta}$ with its Tweedie re-parameterization yields our SURE-Score loss function in (17), which only depends on the score network $s_{\theta}$ and regularization $\lambda$. The training loss we use in practice is a mini-batch version of (17), with expectations identical to those in (5).

Finally, our training and reconstruction algorithms require the selection of several hyperparameters without access to ground-truth for tuning. In particular, we inherit all hyperparameters related to both denoising score matching and to SURE; notable among them is the step size schedule in DSM given by $p_{\Sigma}$, which is known to depend on the dataset , and the Monte Carlo approximation $\epsilon$, which is known to depend on the noise level $\sigma_{w}$ . We also must choose $\lambda$ to balance the two losses in (17).

IV Experimental Results and Discussion

We test our approach on two linear inverse problems: estimating the MIMO wireless channel matrix from compressed pilot measurements, and multi-coil MRI reconstruction from limited k-space measurements. We compare our approach to following baselines: (i) supervised score-model training, (ii) naive DSM training on the noisy data, (iii) BM3D, (iv) Noise2Self, and (v) Noise2Score. In (iii-iv), we apply DSM training on top of the denoised dataset.

We qualitatively evaluate score models learned from training data at different SNR levels (given by SNRw). We also quantitatively evaluate posterior reconstruction performance. Following convention, we report normalized mean squared error (NMSE) in dB for channel estimation and normalized root mean squared error (NRMSE) for MRI, each at particular sub-sampling ratios and for different values of $\sigma_{w}^{2}$ and $\sigma_{n}^{2}$.

IV-B MIMO Channel Modeling and Estimation

IV-B2 Posterior Reconstruction

IV-C Multi-Coil MRI Reconstruction

IV-C2 Posterior Reconstruction

We use annealed Langevin dynamics as previously described in (7). We simulate an acceleration factor of five using a vertical sampling mask with fully sampled central k-space and uniform random sub-sampling elsewhere. We assume high SNR and set $\sigma_{n}=0$ during inference.

V Conclusion

Through our results, we can observe the impact of SURE-based denoising on prior sampling as well as for solving inverse problems in the wireless and MR domain while learning from datasets corrupted with a noise-to-signal ratio of up to dB. We also observe that self-supervised techniques like SURE-Score can match supervised denoising performance. While our algorithm matches the convergence speed of supervised score models in terms of epochs, we find the runtime per iteration increases due to an additional pass through the network at each step. Another challenge, like in other self-supervised methods, is choosing hyper-parameters without access to ground truth data. Finally, our work currently assumes white Gaussian noise corruption but could be extended to arbitrary exponential families using Generalized SURE .

References