Learning Image Priors through Patch-based Diffusion Models for Solving Inverse Problems

2024-06-04

Neural Information Processing Systems

AI Summary

Existing diffusion prior methods rely on full-image training, incurring prohibitive computational and data requirements that hinder scalability to high-dimensional or high-resolution images (e.g., 3D CT). To address this, we propose PaDIS—a position-aware diffusion prior learning framework based on local image patches. Our core innovation restricts diffusion model training to small, overlapping patches while incorporating learnable positional encodings to implicitly capture global structural priors—thereby circumventing memory and data bottlenecks inherent in full-image modeling. This paradigm is natively compatible with diffusion-based inverse solvers (DIS), enabling efficient solutions to inverse problems including CT reconstruction, deblurring, and super-resolution. Experiments demonstrate that, under limited training data, PaDIS significantly outperforms full-image baselines, achieving superior reconstruction fidelity and generalization across both natural and medical imaging domains.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search
• Machine Learning: Learning with Manifolds

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs

Abstract

Diffusion models can learn strong image priors from underlying data distribution and use them to solve inverse problems, but the training process is computationally expensive and requires lots of data. Such bottlenecks prevent most existing works from being feasible for high-dimensional and high-resolution data such as 3D images. This paper proposes a method to learn an efficient data prior for the entire image by training diffusion models only on patches of images. Specifically, we propose a patch-based position-aware diffusion inverse solver, called PaDIS, where we obtain the score function of the whole image through scores of patches and their positional encoding and utilize this as the prior for solving inverse problems. First of all, we show that this diffusion model achieves an improved memory efficiency and data efficiency while still maintaining the capability to generate entire images via positional encoding. Additionally, the proposed PaDIS model is highly flexible and can be plugged in with different diffusion inverse solvers (DIS). We demonstrate that the proposed PaDIS approach enables solving various inverse problems in both natural and medical image domains, including CT reconstruction, deblurring, and superresolution, given only patch-based priors. Notably, PaDIS outperforms previous DIS methods trained on entire image priors in the case of limited training data, demonstrating the data efficiency of our proposed approach by learning patch-based prior.

Problem

• Learning image priors from patches to avoid full-image training
• Solving inverse problems with limited data and computational resources
• Enabling high-resolution 3D image reconstruction using patch-based diffusion

Innovation

• Patch-based diffusion models for image priors
• Positional encoding enables whole image generation
• Plug-and-play compatibility with various solvers