Piecewise-Guided Diffusion for Inverse Problems

Faster posterior sampling for diffusion-based inverse problems — ~25% speedup with negligible quality loss.

Official implementation of
"Diffusion Models for Solving Inverse Problems via Posterior Sampling with Piecewise Guidance"
TMLR 2026

---

🧠 Overview

Diffusion models provide a powerful framework for solving inverse problems via posterior sampling.
However, existing approaches (e.g., ΠGDM) require computationally expensive operations at every diffusion step, limiting their practicality.

We propose a piecewise guidance framework that:

• ⚡ Reduces inference time by ~25%
• 🎯 Maintains comparable reconstruction quality (PSNR, SSIM)
• 🔁 Works with a single pretrained diffusion model (no retraining required)
• 📉 Eliminates expensive operations (e.g., Jacobians, matrix inverses) in early steps

📄 Source: OpenReview version

---

🚀 Key Idea

The key observation is:

Diffusion noise is negligible at early → exact guidance is unnecessary

We exploit this by using a piecewise approximation of the guidance term:

• Early stage (low t)
  → Use a closed-form, efficient approximation
  → No Jacobians, no matrix inverses
  → No backpropagation through the denoising model

• Late stage (high t)
  → Use accurate guidance (ΠGDM-style)

This yields a speed–accuracy tradeoff controlled by a single parameter T₀.

---

✨ Contributions

• 🧩 Piecewise guidance framework for diffusion-based inverse problems
• ⚙️ Eliminates costly operations (vector-Jacobian products, matrix inversions) in early steps
• 📊 Theoretical analysis via KL divergence bounds for approximation quality
• 🎯 Problem-agnostic: works across multiple inverse problems
• 📉 Achieves:
  • 25% faster inference for inpainting
  • 23–24% faster for super-resolution
  • With negligible loss in PSNR/SSIM

📄 Verified in experiments across ImageNet-based tasks.

---

🔬 Method Overview

The guidance term at each diffusion step is:

$$ \nabla_{x_t} \log p(x_t \mid y) = \nabla_{x_t} \log p(x_t) + \nabla_{x_t} \log p(y \mid x_t) $$

🔹 Piecewise Guidance

We approximate:

• For early steps (t < T₀):

  • Closed-form Gaussian likelihood
  • Efficient computation

• For later steps (t ≥ T₀):

  • One-step denoising approximation (ΠGDM-style)

This leads to:

• ⚡ Faster sampling
• 🎯 Accurate reconstruction
• 🔄 Tunable tradeoff via T₀

---

📊 Results

Quantitative Performance

Task	PSNR ↑	SSIM ↑	LPIPS ↓	Time ↓
Inpainting (center)	18.45	0.67	0.31	53s
Inpainting (random)	26.57	0.78	0.19	53s
SR ×4	23.07	0.64	0.24	57s
SR ×8	20.00	0.47	0.30	53s

Compared to ΠGDM:

• ✅ Similar PSNR / SSIM
• ⚡ ~25% faster inference

📄 From Table 2 — OpenReview version

---

🎥 Visual Results

🪶 Reconstruction Example

Qualitative results across inverse problems:

• Inpainting (center mask)
• Inpainting (random mask)
• Super-resolution (×4, ×8)

---

📉 Perceptual Quality vs Guidance Threshold

👉 Increasing T₀:

• Improves speed
• Slightly worsens perceptual quality (LPIPS)
• Minimal impact on PSNR/SSIM

---

⏱️ Inference Time vs T₀

👉 Increasing T₀ leads to consistent runtime reduction across all tasks.

📄 Based on Figures 2–5 and 8–9 — OpenReview version

---

💻 Code Usage & Explanation Guide

👉 Open Full Guide

---

📚 Citation

If this work helped you with your research, please consider citing our paper and give our repo a ⭐!

@article{
mohseni-sehdeh2026diffusion,
title={Diffusion Models for Solving Inverse Problems via Posterior Sampling with Piecewise Guidance},
author={Saeed Mohseni-Sehdeh and Walid Saad and Kei Sakaguchi and Tao Yu},
journal={Transactions on Machine Learning Research},
issn={2835-8856},
year={2026},
url={https://openreview.net/forum?id=nvw3XfvBi7},
note={}
}