Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration

ACM MM 2026

Jiaqi Zhang^1,*, Zheng Pang^1,*, Rongrong Gao¹, Qiyuan Zhang², Yang Yang^1,†

¹Jiangsu University, ²Xi'an University of Technology

*Equal contribution. †Corresponding author.

Paper

Code

Poster

Conference

Input

ProDiff

+ LEADer

Abstract

Diffusion models have demonstrated remarkable effectiveness in image restoration tasks. However, existing Diffusion Model-based Image Restoration (DMIR) methods typically rely on fixed data constraints and uniform step sizes, overlooking the dynamic nature of the generative process. These rigid designs are vulnerable to spatially non-uniform degradations and introduce computational redundancy. We propose Local Epistemic Uncertainty Guided Active Sampling, termed LEADer. In the spatial domain, LEADer leverages pixel-wise uncertainty to dynamically modulate the prior strength within the null space, balancing detail preservation and artifact suppression. In the temporal domain, it quantifies sampling stability via the uncertainty trace to enable adaptive trajectory pruning. LEADer can be seamlessly integrated into various DMIR baselines and improves restoration quality while significantly reducing sampling time.

Method Overview

LEADer revisits plug-and-play diffusion restoration from a state-aware sampling perspective. Instead of applying the same restoration constraint everywhere and at every denoising step, it asks how reliable the diffusion prior is for each local region and how stable the current sampling state is. This uncertainty signal then controls both the spatial correction strength and the temporal sampling schedule.

Architecture of the LEADer framework. — Figure 1. Architecture of LEADer. Local epistemic uncertainty guides reverse sampling in both spatial and temporal dimensions: UCPM adaptively modulates prior constraints, while SATP selects adaptive step sizes to skip redundant iterations.

Local Epistemic Uncertainty Quantification: LEADer estimates local uncertainty from the diffusion model's state, treating it as a spatiotemporal signal that reflects where the restoration is confident and where ambiguity remains.
Uncertainty-Calibrated Prior Modulation: UCPM uses pixel-level uncertainty to adaptively modulate null-space prior strength, preserving details in confident regions while suppressing artifacts in uncertain areas.
State-Aware Trajectory Pruning: SATP uses the uncertainty trace to select adaptive step sizes, skipping redundant reverse-diffusion updates while keeping the approximation error bounded.
Plug-and-play Integration: The framework can be inserted into multiple zero-shot DMIR baselines without retraining the diffusion model.

This design keeps the strict data-consistency benefit of null-space projection while making the reverse process more flexible. Confident regions avoid unnecessary over-regularization, ambiguous regions receive stronger prior guidance, and stable trajectory segments can be skipped to save inference time.

Comparison Analysis

LEADer improves both visual quality and efficiency across restoration tasks, including super-resolution, Gaussian deblurring, motion deblurring, and compressed sensing.

The comparison highlights the core benefit of uncertainty-guided restoration: baseline methods with fixed sampling rules often lose local details or introduce artifacts under non-uniform degradations, while LEADer adaptively allocates correction strength according to the current restoration difficulty.

Qualitative comparisons on five restoration tasks. — Figure 2. Qualitative comparisons on five image restoration tasks.

Quantitative results across five image restoration tasks on CelebA-HQ 1K and ImageNet 1K. — Table 1. Quantitative results for five image restoration tasks on CelebA-HQ 1K (top) and ImageNet 1K (bottom). The average change reports performance improvements and time reductions over each baseline.

Sampling Trajectory Analysis

Compared with fixed-step baselines, LEADer dynamically adjusts the sampling pace and converges to clearer local structures at lower numbers of function evaluations.

The trajectory visualization shows that DDNM improves details gradually through dense and uniform updates. In contrast, LEADer identifies stable states earlier and prunes redundant steps, while preserving the important denoising updates needed for local texture recovery.

Comparison between DDNM and LEADer at different NFEs. — Figure 3. Comparison of restored states between DDNM and LEADer at different NFEs.

Ablation Analysis

Ablation studies indicate that UCPM improves restoration fidelity, SATP reduces redundant computation, and their combination gives the strongest trade-off between quality and speed.

UCPM mainly contributes to perceptual and distortion improvements by making the prior modulation spatially adaptive. SATP contributes to efficiency by adjusting the effective number of sampling steps according to the uncertainty trace. When used together, the improved sampling stability from UCPM makes trajectory pruning safer.

Ablation study results for UCPM and SATP. — Table 2. Ablation study of the UCPM and SATP components on three representative restoration tasks.

The information-loss budget controls the balance between restoration quality and acceleration. A conservative budget permits fewer skipped steps, while an excessively large budget reduces runtime at the cost of noticeable quality degradation. The experiments use 0.01 as the default balanced setting.

Hyperparameter sensitivity analysis for the information-loss budget. — Table 3 and Figure 4. Hyperparameter sensitivity analysis on the information-loss budget B, together with ablation results across different sampling steps on CelebA-HQ 1K.

Ablation study on different sampling steps. — Table 3 and Figure 4. Hyperparameter sensitivity analysis on the information-loss budget B, together with ablation results across different sampling steps on CelebA-HQ 1K.

Memory Consumption

LEADer introduces only marginal GPU memory overhead when integrated into existing DMIR methods. Across five ImageNet restoration tasks, the additional memory consumption remains between 0.07% and 0.28%, showing that the uncertainty-guided modules improve quality and efficiency without materially increasing deployment cost.

BibTeX

@article{zhang2026leader,
  title={Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration},
  author={Zhang, Jiaqi and Pang, Zheng and Gao, Rongrong and Zhang, Qiyuan and Yang, Yang},
  year={2026}
}