Score-Based Turbo Message Passing for Plug-and-Play Compressive Imaging

TL;DR

This paper introduces a novel message-passing algorithm called score-based turbo message passing (STMP) that leverages score-based generative models for improved compressive image recovery, especially under quantization, with fast convergence and better performance.

Contribution

The paper proposes STMP, integrating score-based MMSE denoisers into message passing for compressive imaging, and extends it to quantized measurements with Q-STMP, demonstrating superior performance and convergence.

Findings

STMP outperforms existing methods in performance-complexity tradeoff.

Q-STMP is robust under 1-bit quantization.

Both algorithms converge within 10 iterations.

Abstract

Message-passing algorithms have been adapted for compressive imaging by incorporating various off-the-shelf image denoisers. However, these denoisers rely largely on generic or hand-crafted priors and often fall short in accurately capturing the complex statistical structure of natural images. As a result, traditional plug-and-play (PnP) methods often lead to suboptimal reconstruction, especially in highly underdetermined regimes. Recently, score-based generative models have emerged as a powerful framework for accurately characterizing sophisticated image distribution. Yet, their direct use for posterior sampling typically incurs prohibitive computational complexity. In this paper, by exploiting the close connection between score-based generative modeling and empirical Bayes denoising, we devise a message-passing framework that integrates a score-based minimum mean-squared error (MMSE) denoiser for compressive image recovery. The resulting algorithm, named score-based turbo message passing (STMP), combines the fast convergence of message passing with the expressive power of score-based generative priors. For practical systems with quantized measurements, we further propose quantized STMP (Q-STMP), which augments STMP with a component-wise MMSE dequantization module. We demonstrate that the asymptotic performance of STMP and Q-STMP can be accurately predicted by a set of state-evolution (SE) equations. Experiments on the FFHQ dataset demonstrate that STMP strikes a significantly better performance-complexity tradeoff compared with competing baselines, and that Q-STMP remains robust even under 1-bit quantization. Remarkably, both STMP and Q-STMP typically converge within 10 iterations.