Kernel Density Steering: Inference-Time Scaling via Mode Seeking for Image Restoration

TL;DR

Kernel Density Steering (KDS) enhances diffusion-based image restoration by guiding multiple samples towards high-density regions, reducing artifacts and improving image quality without retraining.

Contribution

Introduces KDS, a novel inference-time method that uses local mode-seeking with multiple diffusion samples to improve image restoration quality.

Findings

Significantly improves quantitative metrics on super-resolution and inpainting tasks.

Reduces artifacts and enhances image fidelity in experimental results.

Requires no retraining or external verifiers, easily integrating with existing diffusion models.

Abstract

Diffusion models show promise for image restoration, but existing methods often struggle with inconsistent fidelity and undesirable artifacts. To address this, we introduce Kernel Density Steering (KDS), a novel inference-time framework promoting robust, high-fidelity outputs through explicit local mode-seeking. KDS employs an $N$-particle ensemble of diffusion samples, computing patch-wise kernel density estimation gradients from their collective outputs. These gradients steer patches in each particle towards shared, higher-density regions identified within the ensemble. This collective local mode-seeking mechanism, acting as "collective wisdom", steers samples away from spurious modes prone to artifacts, arising from independent sampling or model imperfections, and towards more robust, high-fidelity structures. This allows us to obtain better quality samples at the expense of higher compute by simultaneously sampling multiple particles. As a plug-and-play framework, KDS requires no retraining or external verifiers, seamlessly integrating with various diffusion samplers. Extensive numerical validations demonstrate KDS substantially improves both quantitative and qualitative performance on challenging real-world super-resolution and image inpainting tasks.