Diff-Unfolding: A Model-Based Score Learning Framework for Inverse Problems

TL;DR

Diff-Unfolding introduces a modular, model-based score learning framework for inverse problems that generalizes across tasks by decoupling measurement models from learned priors, achieving state-of-the-art results in image restoration and MRI.

Contribution

The paper presents a novel unrolled optimization framework for learning posterior scores in conditional diffusion models, enabling flexible generalization across inverse problems without retraining.

Findings

Achieves up to 2 dB PSNR improvement in image restoration.

Reduces LPIPS by 22.7% in MRI reconstruction.

Operates efficiently with 0.63 seconds per image on high-resolution data.

Abstract

Diffusion models are extensively used for modeling image priors for inverse problems. We introduce \emph{Diff-Unfolding}, a principled framework for learning posterior score functions of \emph{conditional diffusion models} by explicitly incorporating the physical measurement operator into a modular network architecture. Diff-Unfolding formulates posterior score learning as the training of an unrolled optimization scheme, where the measurement model is decoupled from the learned image prior. This design allows our method to generalize across inverse problems at inference time by simply replacing the forward operator without retraining. We theoretically justify our unrolling approach by showing that the posterior score can be derived from a composite model-based optimization formulation. Extensive experiments on image restoration and accelerated MRI show that Diff-Unfolding achieves state-of-the-art performance, improving PSNR by up to 2 dB and reducing LPIPS by $22.7\%$, while being both compact (47M parameters) and efficient (0.72 seconds per $256 \times 256$ image). An optimized C++/LibTorch implementation further reduces inference time to 0.63 seconds, underscoring the practicality of our approach.