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

2026-08-07

AI Summary

Existing diffusion-based image restoration methods struggle with spatially non-uniform degradation due to their reliance on fixed data constraints and uniform sampling steps, often leading to structural distortions, detail loss, and computational redundancy. This work proposes the LEADer framework, which introduces local epistemic uncertainty into diffusion-based image restoration for the first time. In the spatial domain, it dynamically modulates the strength of null-space priors based on pixel-wise uncertainty, enabling adaptive data consistency enforcement. In the temporal domain, it leverages the trace of uncertainty to prune the sampling trajectory, achieving efficient and adaptive inference. Evaluated across multiple state-of-the-art diffusion-based image restoration models, LEADer significantly improves restoration quality while substantially reducing sampling time, incurs negligible memory overhead, and guarantees strict data consistency along with deterministic error bounds.

Abstract

Diffusion models have demonstrated remarkable effectiveness in image restoration tasks. However, when guiding image reconstruction, existing Diffusion Model-based Image Restoration (DMIR) methods typically rely on fixed data constraints and uniform step sizes, thereby overlooking the dynamic nature of the generative process. Such rigid designs render the models vulnerable to spatially non-uniform degradations, thus resulting in structural distortions and loss of fine details. Meanwhile, uniform step sizes introduce computational redundancy, whereas naïve step reduction strategies tend to accumulate approximation errors. To address these limitations, we propose a Local Epistemic Uncertainty Guided Active Sampling framework (LEADer). In the spatial domain, LEADer leverages pixel-wise uncertainty to dynamically modulate the prior strength within the null space, which effectively balances detail preservation and artifact suppression. In the temporal domain, it quantifies sampling stability via the uncertainty trace to enable adaptive trajectory pruning, thereby accelerating convergence. Theoretical proofs demonstrate that our framework achieves strict data consistency, while the trajectory pruning strategy admits a deterministic error bound, thereby guaranteeing stable convergence under skip sampling. Notably, our plug-and-play method can be seamlessly integrated into various DMIR baselines. Extensive experiments show that LEADer improves the performance of multiple state-of-the-art DMIR methods, while significantly reducing sampling time with negligible memory overhead. Code is available at https://github.com/JiaqiZhang-Sengoku/LEADer.

Problem

• Diffusion Model
• Image Restoration
• Epistemic Uncertainty
• Active Sampling
• Non-uniform Degradation

Innovation

• epistemic uncertainty
• active sampling
• diffusion models
• adaptive trajectory pruning
• plug-and-play image restoration