Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Practical Operator Sketching Framework for Accelerating Iterative Data-Driven Solutions in Inverse Problems

Junqi Tang, Guixian Xu, Subhadip Mukherjee, Carola-Bibiane Schönlieb

TL;DR

This work introduces an operator sketching framework to accelerate iterative data-driven reconstruction (IDR) methods for imaging inverse problems, focusing on plug-and-play and deep unrolling networks. By performing double-sided sketching in both image and data domains, the authors derive PnP-MS2G and sketched primal-dual networks (LSPD/SkLSPD), together with stochastic lazy denoisers (Lazy-PnP and Lazy-PnP-EQ) to dramatically reduce computation and memory costs. Theoretical recovery guarantees are established for the sketching framework, and extensive experiments on natural and tomographic imaging demonstrate substantial speedups with little loss in reconstruction quality. The methods extend to coarse-to-fine sketching strategies and equivariant priors, enabling practical acceleration for high-dimensional imaging tasks such as CT and MRI, with potential applicability to broader deep restoration and equilibrium-model frameworks.

Abstract

We propose a new operator-sketching paradigm for designing efficient iterative data-driven reconstruction (IDR) schemes, e.g. Plug-and-Play algorithms and deep unrolling networks. These IDR schemes are currently the state-of-the-art solutions for imaging inverse problems. However, for high-dimensional imaging tasks, especially X-ray CT and MRI imaging, these IDR schemes typically become inefficient both in terms of computation, due to the need of computing multiple times the high-dimensional forward and adjoint operators. In this work, we explore and propose a universal dimensionality reduction framework for accelerating IDR schemes in solving imaging inverse problems, based on leveraging the sketching techniques from stochastic optimization. Using this framework, we derive a number of accelerated IDR schemes, such as the plug-and-play multi-stage sketched gradient (PnP-MS2G) and sketching-based primal-dual (LSPD and Sk-LSPD) deep unrolling networks. Meanwhile, for fully accelerating PnP schemes when the denoisers are computationally expensive, we provide novel stochastic lazy denoising schemes (Lazy-PnP and Lazy-PnP-EQ), leveraging the ProxSkip scheme in optimization and equivariant image denoisers, which can massively accelerate the PnP algorithms with improved practicality. We provide theoretical analysis for recovery guarantees of instances of the proposed framework. Our numerical experiments on natural image processing and tomographic image reconstruction demonstrate the remarkable effectiveness of our sketched IDR schemes.

Practical Operator Sketching Framework for Accelerating Iterative Data-Driven Solutions in Inverse Problems

TL;DR

This work introduces an operator sketching framework to accelerate iterative data-driven reconstruction (IDR) methods for imaging inverse problems, focusing on plug-and-play and deep unrolling networks. By performing double-sided sketching in both image and data domains, the authors derive PnP-MS2G and sketched primal-dual networks (LSPD/SkLSPD), together with stochastic lazy denoisers (Lazy-PnP and Lazy-PnP-EQ) to dramatically reduce computation and memory costs. Theoretical recovery guarantees are established for the sketching framework, and extensive experiments on natural and tomographic imaging demonstrate substantial speedups with little loss in reconstruction quality. The methods extend to coarse-to-fine sketching strategies and equivariant priors, enabling practical acceleration for high-dimensional imaging tasks such as CT and MRI, with potential applicability to broader deep restoration and equilibrium-model frameworks.

Abstract

We propose a new operator-sketching paradigm for designing efficient iterative data-driven reconstruction (IDR) schemes, e.g. Plug-and-Play algorithms and deep unrolling networks. These IDR schemes are currently the state-of-the-art solutions for imaging inverse problems. However, for high-dimensional imaging tasks, especially X-ray CT and MRI imaging, these IDR schemes typically become inefficient both in terms of computation, due to the need of computing multiple times the high-dimensional forward and adjoint operators. In this work, we explore and propose a universal dimensionality reduction framework for accelerating IDR schemes in solving imaging inverse problems, based on leveraging the sketching techniques from stochastic optimization. Using this framework, we derive a number of accelerated IDR schemes, such as the plug-and-play multi-stage sketched gradient (PnP-MS2G) and sketching-based primal-dual (LSPD and Sk-LSPD) deep unrolling networks. Meanwhile, for fully accelerating PnP schemes when the denoisers are computationally expensive, we provide novel stochastic lazy denoising schemes (Lazy-PnP and Lazy-PnP-EQ), leveraging the ProxSkip scheme in optimization and equivariant image denoisers, which can massively accelerate the PnP algorithms with improved practicality. We provide theoretical analysis for recovery guarantees of instances of the proposed framework. Our numerical experiments on natural image processing and tomographic image reconstruction demonstrate the remarkable effectiveness of our sketched IDR schemes.
Paper Structure (29 sections, 2 theorems, 61 equations, 8 figures, 2 tables, 3 algorithms)

This paper contains 29 sections, 2 theorems, 61 equations, 8 figures, 2 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Contributions of this work
  3. Background
  4. Imaging Inverse Problems
  5. Iterative Data-Driven Reconstruction
  6. Plug-and-Play algorithms
  7. Deep unrolling
  8. Iterative Operator Sketching Framework
  9. Doubly-Sketched PnP
  10. Algorithmic Framework
  11. Stochastic Lazy Denoisers
  12. Lazy-PnP
  13. Equivariant Lazy PnP
  14. Accelerating the Deep Unrolling Schemes via Operator Sketching
  15. One side sketching: Subset Approximation
  16. ...and 14 more sections

Key Result

Theorem 1

(Upper bound) Assuming A.1-3, let $\eta = \frac{1}{qL_s}$ and $b = Ax^\dagger + w$, the output of $k$-th stage of PnP-MS2G has the following guarantee for the estimation of $x^\dagger$: where $x_{\mathrm{init}}$ denotes the initial point of $k$-th stage of PnP-MS2G, $\alpha = \kappa(1 - \frac{\mu_c}{ L_s})$, $\kappa=1$ if $\mathcal{M}$ is convex, $\kappa=2$ if $\mathcal{M}$ is non-convex, $\varep

Figures (8)

  • Figure 1: One simple example of the practical choices for the building blocks of one layer of our LSPD network. Both dual and primal subnetworks are consist of 3 convolutional layers. The dual subnet has 3 input channels concatenating $[M_ib, M_iAx_k, y_k]$, while the primal subnet has 2 input channels for $[(M_iA)^Ty_{k+1}, x_k]$
  • Figure 2: Example for applying MS2G (minibatched) in X-ray CT reconstruction, comparing to the PnP stochastic gradient descent (PnP-SGD) proposed by Sun et al sun2019online. Note that each iteration of PnP-MS2G is more computationally efficient than PnP-SGD due to the dimensionality reduction by operator sketching. Surprisingly, even in terms of iteration-number the PnP-MS2G can provide better convergence rate comparing to PnP-SGD.
  • Figure 3: (Lazy-Denoiser)Example for applying our Lazy-Denoiser scheme with equivariant-PnP for image superresolution ($4\times$). The denoiser network we choose here is DnCNN. Here we show the reconstructed image at 2000-th iteration.
  • Figure 4: (Lazy-Denoiser)Example for applying our Lazy-Denoiser scheme with equivariant-PnP for image superresolution ($4\times$). The denoiser network we choose here is DnCNN. Here we show the reconstructed image at 2000-th iteration.
  • Figure 5: (Lazy-Denoiser)Example for applying our Lazy-Denoiser scheme with equivariant-PnP for image superresolution ($4\times$). The denoiser network we choose here is DnCNN. Here we show the reconstructed image at 2000-th iteration.
  • ...and 3 more figures

Theorems & Definitions (2)

  • Theorem 1
  • Theorem 2