Deep Diffusion Image Prior for Efficient OOD Adaptation in 3D Inverse Problems (2407.10641v1)

Abstract: Recent inverse problem solvers that leverage generative diffusion priors have garnered significant attention due to their exceptional quality. However, adaptation of the prior is necessary when there exists a discrepancy between the training and testing distributions. In this work, we propose deep diffusion image prior (DDIP), which generalizes the recent adaptation method of SCD by introducing a formal connection to the deep image prior. Under this framework, we propose an efficient adaptation method dubbed D3IP, specified for 3D measurements, which accelerates DDIP by orders of magnitude while achieving superior performance. D3IP enables seamless integration of 3D inverse solvers and thus leads to coherent 3D reconstruction. Moreover, we show that meta-learning techniques can also be applied to yield even better performance. We show that our method is capable of solving diverse 3D reconstructive tasks from the generative prior trained only with phantom images that are vastly different from the training set, opening up new opportunities of applying diffusion inverse solvers even when training with gold standard data is impossible. Code: https://github.com/HJ-harry/DDIP3D

• The paper presents the Deep Diffusion Image Prior (DDIP) framework that generalizes deep image prior to diffusion models for efficient OOD adaptation in 3D reconstruction.
• It introduces D3IP, a method that unifies parameter adaptation across 3D volumes, reducing computation from over six hours to just 40 minutes per 256³ volume.
• Meta-learning integration further refines the approach, enhancing performance in key medical imaging tasks such as sparse-view CT and MRI reconstructions.

Deep Diffusion Image Prior for Efficient OOD Adaptation in 3D Inverse Problems

The research paper titled "Deep Diffusion Image Prior for Efficient OOD Adaptation in 3D Inverse Problems" by Hyungjin Chung and Jong Chul Ye addresses a significant challenge in the field of inverse problem solvers using diffusion models. In particular, it focuses on the adaptation of generative priors in scenarios where there is a distribution mismatch between training and testing data, commonly referred to as an out-of-distribution (OOD) problem.

Overview of Key Contributions:

1. Introduction of Deep Diffusion Image Prior (DDIP): The authors propose the DDIP framework, which extends the concept of deep image prior (DIP) to the domain of diffusion models. This framework establishes a formal link between DIP and Steerable Conditional Diffusion (SCD), allowing for a more structured adaptation process during the reverse diffusion sampling.
2. Efficient 3D Adaptation Method - D3IP: To address the computational inefficiencies of adapting each 2D slice independently, the authors introduce D3IP. This method leverages a unified parameter adaptation across the entire 3D volume, significantly reducing the memory and computational complexity. D3IP integrates seamlessly with existing 3D inverse solvers to ensure spatial coherence in 3D reconstructions.
3. Incorporation of Meta-Learning: The research demonstrates that meta-learning techniques can further enhance the performance of the D3IP framework. By initializing the 3D adaptation model with meta-learned parameters and subsequently fine-tuning for specific tasks, the approach achieves superior results.

Methodological Insights:

The DDIP framework generalizes the adaptation process by considering multi-scale DIP on the probability-flow ODE path of diffusion models. This generalization results in a structured optimization process that refines the model's parameters progressively, from high noise levels to lower ones, thereby enhancing stability and robustness.

D3IP, as a variant of DDIP tailored for 3D inverse problems, shifts the adaptation paradigm from a slice-wise optimization to a volumetric approach. This transition dramatically improves computational efficiency, reducing the adaptation time for a 256³ volume from over six hours to just 40 minutes on a single GPU. The method also achieves better quality reconstructions by adapting parameters that capture the broader context within the entire volume, rather than isolated slices.

Experimental Validation:

The experimental section encompasses validation across three canonical inverse problems in medical imaging: sparse-view CT reconstruction, MRI reconstruction, and multi-coil MRI reconstruction. These experiments utilize diffusion priors trained solely on Ellipses phantom images. The paper demonstrates the capability of D3IP to adapt these priors effectively to real-world datasets, such as the AAPM grand challenge data for CT and the BRATS dataset for MRI, with significant improvements in PSNR, SSIM, and LPIPS metrics compared to prior state-of-the-art methods.

Practical and Theoretical Implications:

The practical implication of this research is profound, particularly in medical imaging, where obtaining high-quality ground truth data is often challenging, expensive, or impractical. The ability of D3IP to adapt to OOD data using priors trained on synthetic phantoms opens new avenues for practical deployment in clinical settings. Furthermore, the use of meta-learning within the D3IP framework hints at the potential for even more sophisticated adaptive systems that continually improve with exposure to new data.

The theoretical advancement provided by the paper lies in the robust connection between DIP and diffusion-based inverse problem solvers. This connection offers a new lens to view the adaptation process and sets the stage for future research to explore other forms of implicit priors within the diffusion framework.

Conclusion and Future Directions:

The paper presents a substantial step forward in the field of inverse problem solvers using diffusion models, addressing both computational and performance bottlenecks associated with OOD adaptation. Future research can build upon these findings by exploring more advanced meta-learning strategies, further optimizing the adaptation process, and extending the approach to other inverse problem domains beyond medical imaging. Additionally, integrating more sophisticated regularization techniques during adaptation could mitigate the issues related to catastrophic forgetting and enhance the model's generalizability to a wider range of OOD scenarios.