Deep unrolling for learning optimal spatially varying regularisation parameters for Total Generalised Variation

TL;DR

The paper tackles learning spatially varying regularisation for Total Generalised Variation ($TGV$) in inverse imaging by coupling a CNN that predicts parameter maps $\Lambda_0,\Lambda_1$ with an unrolled PDHG solver for $TGV_{\Lambda_0,\Lambda_1}$. The end-to-end trainable framework, named U-TGV (and U-TV for comparison), demonstrates significant improvements in denoising and accelerated MRI reconstruction over scalar $TGV$, scalar TV, and unsupervised spatial methods, while revealing a structured triple-edge pattern near edges in the learned parameter maps. The key contributions include a practical deep unrolling approach to learn data-adaptive regularisation for a higher-order regulariser, and empirical evidence of interpretability through the learned maps and their edge behavior. This work advances data-driven regularisation in inverse imaging, offering both performance gains and directions for theoretical analysis of learned spatial maps.

Abstract

We extend a recently introduced deep unrolling framework for learning spatially varying regularisation parameters in inverse imaging problems to the case of Total Generalised Variation (TGV). The framework combines a deep convolutional neural network (CNN) inferring the two spatially varying TGV parameters with an unrolled algorithmic scheme that solves the corresponding variational problem. The two subnetworks are jointly trained end-to-end in a supervised fashion and as such the CNN learns to compute those parameters that drive the reconstructed images as close to the ground truth as possible. Numerical results in image denoising and MRI reconstruction show a significant qualitative and quantitative improvement compared to the best TGV scalar parameter case as well as to other approaches employing spatially varying parameters computed by unsupervised methods. We also observe that the inferred spatially varying parameter maps have a consistent structure near the image edges, asking for further theoretical investigations. In particular, the parameter that weighs the first-order TGV term has a triple-edge structure with alternating high-low-high values whereas the one that weighs the second-order term attains small values in a large neighbourhood around the edges.

Figures (5)

• Figure 1: Visualisation of the network $\mathcal{N}_{\theta}^{N}: f\mapsto (\mathrm{\Lambda}_{0}, \mathrm{\Lambda}_{1})\mapsto u$ of \ref{['unrolled_network']} (red arrows) and its training procedure \ref{['supervised_learning']} (black arrows) for the denoising case.
• Figure 2: Denoising results for the "turtle" image. The numbers is brackets show the [PSRN, SSIM] values. The colorbar is common among the visualisations of the parameter maps, except the $\mathrm{\Lambda}_0/\mathrm{\Lambda}_1$ ratio shown in logarithmic scale.
• Figure 3: Denoising results for the "parrot" image. The numbers in brackets show the [PSRN, SSIM] values. The colorbar is common among the visualisations of the parameter maps, except the $\mathrm{\Lambda}_0/\mathrm{\Lambda}_1$ ratio shown in logarithmic scale.
• Figure 4: Visualisation of the structure of the parameters $\mathrm{\Lambda}_0$ and $\mathrm{\Lambda}_1$ at the image edges. Both maps alternate between high-low-high values at the edges but $\mathrm{\Lambda}_0$ takes small values at a larger neighbourhood around the edges.
• Figure 5: MRI reconstruction results for a Brain image. The numbers in brackets show the [PSNR, SSIM] values.

Theorems & Definitions (3)

• Proposition 1
• proof
• Proposition 2: from Papafitsoros_Valkonen_2015