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

Diffusion models are extensively used for modeling image priors for inverse problems. We introduce Diff-Unfolding, a principled framework for learning posterior score functions of 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 × 256 image). An optimized C++/LibTorch implementation further reduces inference time to 0.63 seconds, underscoring the practicality of our approach.