Energy-oriented Diffusion Bridge for Image Restoration with Foundational Diffusion Models

TL;DR

The paper introduces E-Bridge, a diffusion model framework that improves image restoration by reducing trajectory energy and enabling efficient, high-quality recovery with fewer steps.

Contribution

It proposes a novel energy-oriented diffusion bridge that shortens trajectories and uses a single-step mapping for efficient, adaptable image restoration.

Findings

Achieves state-of-the-art results across various restoration tasks.

Enables high-quality recovery with fewer sampling steps.

Provides a tunable trajectory length for task-specific balance.

Abstract

Diffusion bridge models have shown great promise in image restoration by explicitly connecting clean and degraded image distributions. However, they often rely on complex and high-cost trajectories, which limit both sampling efficiency and final restoration quality. To address this, we propose an Energy-oriented diffusion Bridge (E-Bridge) framework to approximate a set of low-cost manifold geodesic trajectories to boost the performance of the proposed method. We achieve this by designing a novel bridge process that evolves over a shorter time horizon and makes the reverse process start from an entropy-regularized point that mixes the degraded image and Gaussian noise, which theoretically reduces the required trajectory energy. To solve this process efficiently, we draw inspiration from consistency models to learn a single-step mapping function, optimized via a continuous-time consistency objective tailored for our trajectory, so as to analytically map any state on the trajectory to the target image. Notably, the trajectory length in our framework becomes a tunable task-adaptive knob, allowing the model to adaptively balance information preservation against generative power for tasks of varying degradation, such as denoising versus super-resolution. Extensive experiments demonstrate that our E-Bridge achieves state-of-the-art performance across various image restoration tasks while enabling high-quality recovery with a single or fewer sampling steps. Our project page is https://jinnh.github.io/E-Bridge/.