Direct Diffusion Bridge using Data Consistency for Inverse Problems (2305.19809v2)

Abstract: Diffusion model-based inverse problem solvers have shown impressive performance, but are limited in speed, mostly as they require reverse diffusion sampling starting from noise. Several recent works have tried to alleviate this problem by building a diffusion process, directly bridging the clean and the corrupted for specific inverse problems. In this paper, we first unify these existing works under the name Direct Diffusion Bridges (DDB), showing that while motivated by different theories, the resulting algorithms only differ in the choice of parameters. Then, we highlight a critical limitation of the current DDB framework, namely that it does not ensure data consistency. To address this problem, we propose a modified inference procedure that imposes data consistency without the need for fine-tuning. We term the resulting method data Consistent DDB (CDDB), which outperforms its inconsistent counterpart in terms of both perception and distortion metrics, thereby effectively pushing the Pareto-frontier toward the optimum. Our proposed method achieves state-of-the-art results on both evaluation criteria, showcasing its superiority over existing methods. Code is available at https://github.com/HJ-harry/CDDB

• The paper proposes a data consistent extension to diffusion bridges (CDDB) that enhances image reconstruction fidelity and accelerates inference.
• It utilizes a time-conditional neural network over all timesteps to flexibly balance the trade-off between perceptual quality and distortion.
• Empirical results demonstrate state-of-the-art performance, achieving up to 50 times faster inference without additional training.

Introduction

Research in solving inverse problems using diffusion models has recently garnered significant attention. Diffusion model-based Inverse Problem Solvers (DIS) have demonstrated remarkable performance, capitalizing on the capabilities of diffusion models as generative priors. However, these models conventionally suffer from slow inference times. This fundamental limitation is rooted in the need to start the sampling process from noise. To alleviate this issue, recent works have developed a diffusion process bridging the gap between clean and corrupted forms for specific inverse problems, termed Direct Diffusion Bridges (DDB).

Direct Diffusion Bridges (DDB) Framework

The unification of existing works into the Direct Diffusion Bridges (DDB) framework stands as a pivotal contribution. The DDB methodology defines the diffusion process from an image to a measurement distribution as a continuum facilitated by a time-conditional neural network trained across all timesteps. The selection of neural function evaluations (NFEs) provides a flexible approach to manage the trade-off between the desired level of perceptual quality and distortion.

However, a salient limitation in the DDB framework is its lack of data consistency. It's critical that reconstructions adhere closely to the given measurements, which DDB methods have traditionally overlooked. To that end, the research proposes a Data Consistent Direct Diffusion Bridge (CDDB) method that enforces data consistency without the necessity of fine-tuning the pre-trained model. The CDDB method is shown to enhance both perception and distortion metrics, effectively shifting the Pareto-frontier toward more optimal reconstructions.

CDDB: Enhancing Data Consistency

The CDDB develops upon existing DDB algorithms by incorporating a modification that maintains data consistency throughout the sampling process. The proposed method employs a modified inference procedure that corrects predicted measurements to satisfy data constraints, thus greatly improving the sample quality. This modification is achieved without additional training, utilizing a pre-trained model as its foundation.

Theoretical underpinnings reveal that CDDB is an extension of previously established methods, such as DDS and DPS, within the DDB landscape. This positions CDDB as a versatile tool that scales the existing methods' performance, especially in terms of speed and stability.

Empirical Validation and Implications

Empirical evaluation validates that CDDB achieves state-of-the-art results across diverse tasks, substantiating the method's effectiveness. CDDB not only pushes the envelope in terms of better sample quality but does so at an accelerated rate, exceeding 50 times faster in certain cases. It's shown that even for tasks where DDB could not previously improve performance due to certain constraints, CDDB-deep, a variant of the core approach, rises to the challenge by using deeper gradients for the correction steps.

The societal and practical implications of CDDB are significant especially within domains that demand both high fidelity and perceptual quality in image reconstructions. While the CDDB's reliance on learned priors from data distribution could potentially perpetuate biases within that data, it's a consideration that will have to be managed with careful deployment.

In conclusion, the CDDB framework emerges as a powerful, consistent, and efficient methodology for solving inverse problems within imaging, leveraging the strengths of diffusion models and enhancing them with the critical aspect of data consistency. Its remarkable versatility and robust performance are likely to make it a seminal development in the field of generative AI.