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

TL;DR

This paper introduces a universal operator-sketching framework to accelerate iterative data-driven reconstruction methods for high-dimensional imaging inverse problems, significantly improving efficiency and practicality.

Contribution

The paper proposes a novel operator-sketching paradigm and develops accelerated IDR schemes, including multi-stage sketched gradient and sketching-based primal-dual networks, with theoretical guarantees.

Findings

Significant speedup in high-dimensional imaging tasks.

Effective acceleration of PnP schemes with stochastic lazy denoising.

Demonstrated improved performance on natural and tomographic images.

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.