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

2025-02-23

AI Summary

For inverse imaging problems such as image denoising and MRI reconstruction, this paper proposes the first end-to-end deep unrolling framework that jointly learns two spatially varying regularization parameters in Total Generalized Variation (TGV). The method integrates convolutional neural networks (CNNs) with an unrolled iterative structure derived from TGV variational optimization, enabling supervised joint training. Its key contribution is the first interpretable, co-learned estimation of spatially varying dual TGV parameters—revealing structured patterns in the parameter maps (e.g., a triple-edge response for the first-order parameter at object boundaries), thereby offering new insights for TGV theoretical modeling. Extensive experiments demonstrate significant improvements over scalar TGV and various unsupervised spatially varying methods across multiple benchmarks, with consistent gains in PSNR and SSIM, as well as superior qualitative and quantitative reconstruction fidelity.

Technology Category

• Computer Vision: Learning & Optimization for CV
• Machine Learning: Structured Learning
• Search and Optimization: Non-convex Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Bridging structured and unstructured data

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.

Problem

• Learning spatially varying regularization parameters
• Improving image reconstruction quality
• Exploring parameter structures near image edges

Innovation

• Deep unrolling framework
• Spatially varying TGV parameters
• Supervised CNN training