VIPaint: Image Inpainting with Pre-Trained Diffusion Models via Variational Inference

1. Problem Statement

Image inpainting involves reconstructing missing or corrupted regions of an image. While Diffusion Probabilistic Models (DPMs) have achieved state-of-the-art results in image generation, applying them to inpainting (an inverse problem) presents significant challenges:

• Large Mask Regions: Existing methods often fail to maintain global coherence when large portions of an image are masked, producing inconsistent or unrealistic structures.
• Latent Diffusion Models (LDMs): Most successful inpainting algorithms are designed for pixel-based diffusion models. They struggle to adapt to LDMs (like Stable Diffusion), which operate in a compressed latent space and are computationally more efficient for high-resolution images.
• Training Requirements: Many current approaches require specialized training on specific mask distributions or fine-tuning of the generative model for every new task, limiting their "zero-shot" applicability.
• Posterior Approximation: Methods like REDdiff attempt to use Variational Inference (VI) but often rely on simple Gaussian approximations at the noise-free level ($t=0$). This leads to optimization instability, an inability to capture uncertainty (producing a single mode rather than diverse samples), and poor performance on LDMs.

2. Methodology: VIPaint

The authors propose VIPaint, a hierarchical variational inference algorithm that optimizes a non-Gaussian Markov approximation of the true diffusion posterior. It allows for "zero-shot" inpainting using pre-trained DPMs (both pixel-based and latent).

Core Architecture

Instead of optimizing a single image at $t=0$, VIPaint defines a hierarchical posterior over a coarse sequence of intermediate latent timesteps (keypoints) $\{h(1), \dots, h(K)\}$, where $h(K)$ is high noise and $h(1)$ is low noise.

1. Variational Posterior Formulation:
   The posterior $q_\lambda(z_{h(1):h(K)})$ is defined as a Markov chain:
   $$q_\lambda(z_{h(1):h(K)}) = q_\lambda(z_{h(K)}) \prod_{i=1}^{K-1} q_\lambda(z_{h(i)} | z_{h(i+1)})$$

   • High Noise ($h(K)$): Modeled as a factorized Gaussian.
   • Intermediate/Low Noise: Modeled as a Gaussian where the mean is a convex combination of the diffusion prior prediction ($\hat{x}_\theta$) and a learnable, image-specific variational parameter ($\mu$). The variance is also learned, allowing the model to dynamically increase uncertainty in masked regions.
   • Non-Gaussian Nature: Although the variational distribution is Gaussian at each step, the overall approximation is non-Gaussian because it incorporates the non-linear denoising network $\hat{x}_\theta(\cdot)$.

2. Optimization Objective (Variational Bound):
   VIPaint minimizes a variational lower bound on the marginal likelihood of the observation $y$, composed of three terms:

   • Reconstruction Loss: Ensures samples at the lowest noise level $z_{h(1)}$ are consistent with the observed pixels $y$. It uses a one-step approximation to avoid expensive backpropagation through the full generation chain.
   • Prior Loss: Regularizes the highest noise level $z_{h(K)}$ to align with the high-level semantics encoded by the diffusion prior in the range $[h(K), T]$.
   • Hierarchical Loss: Ensures consistency between adjacent keypoints in the hierarchy, capturing mid-level details in the critical noise range $[h(1), h(K)]$.

3. Sampling Strategy:

   • Optimization: Variational parameters $\lambda$ are optimized for each specific image query using gradient descent (no amortization).
   • Ancestral Sampling: Once optimized, diverse samples are generated by ancestral sampling from $h(K)$ down to $h(1)$.
   • Refinement: For $t < h(1)$, the model uses gradient-based refinement (similar to Langevin dynamics) to generate fine-grained details consistent with the observation.

3. Key Contributions

• Hierarchical Non-Gaussian Approximation: VIPaint introduces a novel VI framework that approximates the diffusion posterior over a sequence of latent steps rather than a single point. This captures both global semantics (high noise) and local details (low noise) simultaneously.
• Zero-Shot Applicability: The method works directly with pre-trained models (including Stable Diffusion 3.5) without requiring retraining, fine-tuning, or specialized datasets for specific mask shapes.
• Stability and Uncertainty: By avoiding direct optimization at the noise-free level ($t=0$) and learning variances, VIPaint achieves stable optimization and naturally captures posterior uncertainty, enabling the generation of multiple diverse, plausible inpaintings.
• LDM Compatibility: The method is explicitly designed to work with Latent Diffusion Models, overcoming the limitations of previous pixel-based VI approaches.

4. Experimental Results

The authors evaluated VIPaint on ImageNet64, ImageNet256, LSUN-Churches, and Stable Diffusion 3.5 (1024x1024).

• Quantitative Performance:

  • Large Masks: VIPaint significantly outperforms baselines (DPS, RePaint, CoPaint, REDdiff) in PSNR, SSIM, and LPIPS, particularly for large masked regions (40-80% of the image).
  • LDMs: On Stable Diffusion 3.5, VIPaint achieves superior results compared to DPS and specialized models like ControlNet and SD-Inpainting, despite not being trained for the task.
  • Other Inverse Problems: The method also demonstrates strong performance in Gaussian deblurring and super-resolution tasks.

• Qualitative Performance:

  • Global Coherence: Unlike baselines that often produce blurry or structurally inconsistent results in large masks, VIPaint preserves global structure and semantic consistency.
  • Diversity: VIPaint generates multiple distinct, high-quality completions for the same masked image, whereas many baselines collapse to a single mode or produce inconsistent variations.
  • Artifact Reduction: It produces sharper results with fewer artifacts at mask boundaries compared to gradient-based sampling methods like DPS.

• Efficiency:

  • While optimization is required per image, VIPaint-2 (with 2 keypoints) is computationally competitive with sampling-based methods and significantly faster than methods requiring extensive fine-tuning.

5. Significance

VIPaint represents a paradigm shift in applying diffusion models to inverse problems. By framing inpainting as a variational inference problem over a hierarchical latent trajectory, it bridges the gap between the generative power of foundation models and the constraints of partial observations.

Its ability to handle large masks and latent diffusion models without retraining makes it a robust, general-purpose tool for image restoration. This approach suggests that future inverse problem solvers can leverage the full capacity of pre-trained foundation models through principled probabilistic inference rather than heuristic guidance or task-specific training.