Equivariant Sampling for Improving Diffusion Model-based Image Restoration

TL;DR

This work identifies two core limitations in problem-agnostic diffusion-model-based image restoration: reliance on a single sampling trajectory and non-optimized time-step scheduling. It introduces EquS, which injects equivariant information through a dual-trajectory sampling framework using an equivariant inverse mapping, and further enhances efficiency with a Timestep-Aware Schedule (TAS), yielding EquS$^+$. The approach achieves consistent, state-of-the-art improvements across multiple IR tasks (CS, inpainting, SR, deblurring, colorization) on ImageNet and CelebA-HQ without added computational cost, and demonstrates robustness to noise and masking. By integrating seamlessly with existing DMIR methods, EquS provides a practical, generalizable boost to diffusion-prior based image restoration in zero-shot settings.

Abstract

Recent advances in generative models, especially diffusion models, have significantly improved image restoration (IR) performance. However, existing problem-agnostic diffusion model-based image restoration (DMIR) methods face challenges in fully leveraging diffusion priors, resulting in suboptimal performance. In this paper, we address the limitations of current problem-agnostic DMIR methods by analyzing their sampling process and providing effective solutions. We introduce EquS, a DMIR method that imposes equivariant information through dual sampling trajectories. To further boost EquS, we propose the Timestep-Aware Schedule (TAS) and introduce EquS$^+$. TAS prioritizes deterministic steps to enhance certainty and sampling efficiency. Extensive experiments on benchmarks demonstrate that our method is compatible with previous problem-agnostic DMIR methods and significantly boosts their performance without increasing computational costs. Our code is available at https://github.com/FouierL/EquS.

Figures (9)

• Figure 1: Our method offers (a) superior quantitative performance, (b) improved qualitative results. It is (c) adaptable to various IR applications, (d) robust to different scales, and (e) resilient to different noise levels. $\mathbf{y}$ represents the degraded image, $\mathbf{x}_0$ denotes the sampling result, SR represents super-resolution and CS represents compressed-sensing.
• Figure 2: (a) Sampling process of diffusion model-based IR. (b) Previous vs. Ours. (c) Effect of our methods.
• Figure 3: Conceptual illustration of the trajectories of two different sampling processes. $\mathcal{H}$ represents the contours of the data distribution. EquS enables dual-trajectory sampling, determining $\mathbf{x}_0$ by considering bidirectional information.
• Figure 4: The sampling process of the Equivariant Sampling (EquS).
• Figure 5: Deconstruction of the diffusion reverse process: stochastic steps and deterministic steps. $\mathbf{x}_{0|t}$ is more certain in the deterministic steps.