Diff-Unfolding: A Model-Based Score Learning Framework for Inverse Problems (2505.11393v2)

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.

An Examination of Diff-Unfolding: Bridging Model-Based and Data-Driven Approaches for Inverse Problems

The paper "Diff-Unfolding: A Model-Based Score Learning Framework for Inverse Problems" introduces Diff-Unfolding, a novel architecture designed to tackle inverse problems in computational imaging. These problems involve reconstructing an original signal or image from degraded measurements, a task complicated by noise and incomplete data. Traditional methods rely on handcrafted priors, while recent advances utilize deep learning-based priors. The diffusion models (DMs), a class of deep generative models, have seen success in capturing intricate image distributions, making them suitable for addressing inverse problems.

Framework Design and Theoretical Justification

Diff-Unfolding distinguishes itself by marrying the strength of diffusion models with a model-based unrolling architecture, explicitly integrating the measurement operator into the network. This modular design departs from conventional conditional diffusion models (CDMs), where measurement models are often entangled with network parameters, hindering adaptability to different forward operators.

The work introduces a new theoretical framework, demonstrating that the computation of the posterior score function, essential for CDM-based recovery, can be formulated as a composite optimization problem. By disentangling this optimization into two terms—one with the measurement model and another with the prior—Diff-Unfolding facilitates a more interpretable and flexible approach. The paper rigorously proves this theoretical basis via complex mathematical derivations, ensuring that the approach remains grounded in solid analytical foundations.

Experimental Validation and Results

The paper presents substantial experimental evidence supporting the efficacy and efficiency of Diff-Unfolding across a spectrum of tasks, emphasizing image restoration and MRI reconstruction. In image restoration tasks encompassing Gaussian deblurring, super-resolution, and inpainting, Diff-Unfolding achieves consistently superior results compared to state-of-the-art methods, such as Conditionally-trained Diffusion Denoising Score Model (CDDB) and Diffusion Probabilistic Image Restoration (DiffPIR), often with improvements of 2 dB in PSNR and reductions in LPIPS by up to 22.7%.

In the context of MRI reconstruction, which holds significant clinical implications, Diff-Unfolding not only excels in metrics such as PSNR and SSIM but also significantly reduces inference time, proving practical for real-world deployments. The experiments demonstrate the framework's ability to maintain high fidelity and low perceptual loss, ensuring finer detail retention crucial for accurate diagnostic interpretation.

Implications and Future Directions

The introduction of Diff-Unfolding has several theoretical and practical implications. Theoretically, the method underscores a modular framework that allows dynamic substitution of forward operators across inverse problems, circumventing the need for retraining, as is customary in conventional CDM approaches. Practically, the reduced parameter size and faster inference time render Diff-Unfolding an appealing option for deployment in computationally constrained environments such as real-time medical imaging systems.

Future developments could extend Diff-Unfolding to accommodate nonlinear or non-Gaussian noise models, broadening its applicability to a wider range of real-world scenarios. Moreover, further exploration into non-linear parameterization of the prior and data fidelity terms could yield even more generalized adaptable architectures.

Conclusion

In summary, Diff-Unfolding advances the field of inverse problem-solving by showcasing how a synergy of model-based optimization and deep learning frameworks can lead to efficient and adaptable solutions. By incorporating the forward model directly into the unfolding network, this approach enables precise recovery of high-quality images across different application domains with reduced computational overhead. As such, this research represents an important step towards developing versatile, real-time systems capable of addressing diverse and complex inverse problems.