Learning spatially adaptive sparsity level maps for arbitrary convolutional dictionaries

TL;DR

This work builds on a recently proposed image reconstruction method, which is based on embedding data-driven information into a model-based convolutional dictionary regularization via neural network-inferred spatially adaptive sparsity level maps, and extends the method to achieve filter-permutation invariance as well as the possibility to change the convolutional dictionary at inference time.

Abstract

State-of-the-art learned reconstruction methods often rely on black-box modules that, despite their strong performance, raise questions about their interpretability and robustness. Here, we build on a recently proposed image reconstruction method, which is based on embedding data-driven information into a model-based convolutional dictionary regularization via neural network-inferred spatially adaptive sparsity level maps. By means of improved network design and dedicated training strategies, we extend the method to achieve filter-permutation invariance as well as the possibility to change the convolutional dictionary at inference time. We apply our method to low-field MRI and compare it to several other recent deep learning-based methods, also on in vivo data, in which the benefit for the use of a different dictionary is showcased. We further assess the method's robustness when tested on in- and out-of-distribution data. When tested on the latter, the proposed method suffers less from the data distribution shift compared to the other learned methods, which we attribute to its reduced reliance on training data due to its underlying model-based reconstruction component.

Figures (6)

• Figure 1: The reconstruction pipeline of the CDL-$\boldsymbol{\Lambda}$ method and the three different versions of $\mathrm{NET}_{\Theta}$ (\ref{['eq:NET_v1']}--\ref{['eq:NET_v3']}) (orange blocks) investigated in this study. First, a learned high-pass filtering step solving \ref{['eq:PL']} is performed (green block). Then, $\mathrm{NET}_{\Theta}$ estimates spatially varying sparsity level maps from an input (blue block). Last, a FISTA algorithm is unrolled to obtain an approximate solution of \ref{['eq:PR']}(ocher block). An estimate of the solution is obtained by adding the low-frequency component that was extracted in the first block. In contrast to V1 and V2, the proposed improved V3 is permutation invariant and allows for the use of dictionaries with a different number of filters $K$.
• Figure 2: MSE obtained by CDL-$\boldsymbol{\Lambda}$ with $\mathrm{NET}_{\Theta}$ V3 over the brain MR test set for 16 different choices of dictionaries. Note that, for training, we did not use the $K=128$-dictionaries.
• Figure 3: Eight out of $K=32$$\boldsymbol{\Lambda}$-maps with the largest variance, seen as indicative of the filter's importance in the image representation, for each $\mathrm{NET}_{\Theta}$ V1, V2, V3, and for filter size $k_f \times k_f = 11 \times 11$.
• Figure 4: A comparison of MoDL aggarwal2018modl, E2E VarNet sriram2020end, SRDEnseNet de2022deep and CDL-$\boldsymbol{\Lambda}$ method ($K=64, 11\times 11$-kernels) on in-distribution brain MR data (left, $\sigma^2=0.3$) and out-of-distribution knee MR data (right, $\sigma^2=0.2$).
• Figure 5: SSIM scores on the brain and knee MR test datasets for noise variances $\sigma^2=0.2$ and $\ \sigma^2=0.3$. The dashed lines denote the median SSIM value of the respective methods on the brain dataset. Note also the distribution of outliers for the knee test set, where, in some rare cases, E2E VarNet and SRDenseNet even seem to deteriorate the images compared to the adjoint.