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

TL;DR

This work addresses the challenge of adapting diffusion-model priors to out-of-distribution data in 3D inverse problems, where training data may be unavailable or non-representative. It introduces Deep Diffusion Image Prior (DDIP), linking DIP with steerable diffusion to enable stable, multi-scale prior adaptation along the PF-ODE trajectory, and extends this to 3D with D3IP for efficient volume-wide adaptation. D3IP achieves substantial speedups and often superior reconstruction quality by jointly adapting a single parameter set across a volume, and can integrate 3D diffusion solvers like DiffusionMBIR and meta-learning-based initializations to further boost performance. The results demonstrate strong OOD performance on three canonical 3D medical-imaging tasks while keeping training entirely unsupervised (priors learned from phantoms), highlighting practical potential for real-world, data-scarce applications in biomedical imaging and astronomy, with code released for reproducibility.

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

Figures (9)

• Figure 1: OOD inverse problem setting. Pre-trained diffusion model learns $p_\theta({\boldsymbol x})$, but at test time we only have ${\boldsymbol y}_{\rm out}$ obtained from unknown OOD distributions, and aim to sample from $p_{\rm out}({\boldsymbol x}|{\boldsymbol y}_{\rm out})$.
• Figure 2: OOD adaptation schemes in DIS. (a) DDIP/SCD performs independent adaptation across slices and requires ${\mathcal{O}}(N)$ compute & memory. (b) D3IP (base) performs joint adaptation with stochastic gradients from MC sampling (blue dotted line) and requires ${\mathcal{O}}(1)$ compute & memory. (c) $\theta_{\rm vol}$ adapted from D3IP (base) can be used as a meta-parameter to be further adapted to specific slices.
• Figure 3: Representative results on 3 different tasks. (row 1-2): 3D SV-CT, (row 3-4): 3D MRI, (row 5-6): CS-MRI. Comparison against DPS chung2023diffusion, DDS chung2024decomposed, and SCD barbano2023steerable. Ours: D3IP (base). Cyan arrows indicate regions of remaining artifacts even after adaptation with SCD. Green boxes illustrate the acquisition scheme of the measurement (acquisition angle, sub-sampling pattern).
• Figure 4: 3D-MRI reconstruction with DDS chung2024decomposed, DDIP, D3IP (mbir). Cyan and red arrows indicate artifacts from prior mismatch and slice-wise independent reconstruction, respectively. 1-4$^{\rm th}$ row: $xy, yz, xz$ slice, and 3D rendering.
• Figure 5: Comparison of reconstructions with D3IP (base) and D3IP (meta).