Comprehensive Examination of Unrolled Networks for Solving Linear Inverse Problems

TL;DR

This work tackles the design complexity of unrolled networks for linear inverse problems by introducing Deep Memory Unrolled Networks (DeMUN), which incorporate gradients from all previous iterations via memory terms to let data determine the effective optimization path. Through extensive ablations, training with an unweighted intermediate loss $\ell_{i,1}$ and incorporating residual connections consistently yields superior reconstruction quality across Gaussian and DCT forward models, with PSNR improving with more projection steps before plateauing. The results demonstrate robustness to changes in the measurement matrix, additive noise, and image resolution, and provide practical guidelines for loss functions, projector capacity, and step count, showing that simple projector architectures can suffice. Overall, DeMUN offers a scalable, robust framework that reduces design choices and improves performance for solving linear inverse problems in imaging.

Abstract

Unrolled networks have become prevalent in various computer vision and imaging tasks. Although they have demonstrated remarkable efficacy in solving specific computer vision and computational imaging tasks, their adaptation to other applications presents considerable challenges. This is primarily due to the multitude of design decisions that practitioners working on new applications must navigate, each potentially affecting the network's overall performance. These decisions include selecting the optimization algorithm, defining the loss function, and determining the number of convolutional layers, among others. Compounding the issue, evaluating each design choice requires time-consuming simulations to train, fine-tune the neural network, and optimize for its performance. As a result, the process of exploring multiple options and identifying the optimal configuration becomes time-consuming and computationally demanding. The main objectives of this paper are (1) to unify some ideas and methodologies used in unrolled networks to reduce the number of design choices a user has to make, and (2) to report a comprehensive ablation study to discuss the impact of each of the choices involved in designing unrolled networks and present practical recommendations based on our findings. We anticipate that this study will help scientists and engineers design unrolled networks for their applications and diagnose problems within their networks efficiently.

Figures (13)

• Figure 1: Diagram of projected gradient descent. Starting with $x^0 = 0$, the $i^{\rm th}$ loss reducer unit performs the operation $\tilde{x}^i = x^i + \mu A^T (y-Ax^i)$, and the $i^{\rm th}$ projector unit performs $x^{i+1} = P_{\mathcal{C}} (\tilde{x}^i)$.
• Figure 2: An example of the memory terms combined into a single image.
• Figure 3: An example of the DnCNN architecture with $L = 3$ intermediate layers
• Figure 4: DeMUN (no residual connections) with loss $\ell_{ll}$. The networks are trained for $T=15$ (left) and $T=30$ (right), and the graph displays the PSNR after each intermediate projection.
• Figure 5: PGD (no residual connections) with loss $\ell_{ll}$. The networks are trained for $T=15$ (left) and $T=30$ (right), and the graph displays the PSNR after each intermediate projection.