Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration
ACM MM 2026 Oral!
1Jiangsu University, 2Xi'an University of Technology
*Equal contribution. †Corresponding author.
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.
- We propose a novel perspective on the diffusion reverse process based on local epistemic uncertainty, which reveals the spatiotemporal limitations of the sampling strategies in existing DMIR methods.
- We propose an active diffusion sampling framework that proactively resolves local distortion and sampling redundancy via uncertainty-calibrated prior modulation and state-aware trajectory pruning.
- Extensive experimental results demonstrate that our proposed plug-and-play framework effectively enhances image restoration quality and performance while improving inference efficiency.
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.
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.
Ablation Analysis
As shown in Table 2, UCPM mainly improves restoration fidelity by making prior modulation spatially adaptive, while SATP reduces redundant computation by adjusting the effective number of sampling steps according to the uncertainty trace. Their combination gives the strongest trade-off between quality and speed, as the improved sampling stability from UCPM makes trajectory pruning safer. 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 corresponding sensitivity results in Table 3 and Figure 4 show that 0.01 provides a balanced default setting.
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{LEADer,
title={Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration},
author={Jiaqi Zhang and
Zheng Pang and
Rongrong Gao and
Qiyuan Zhang and
Yang Yang},
journal={arXiv preprint arXiv:2608.06981},
year={2026}
}