I$^2$SB: Image-to-Image Schrödinger Bridge (2302.05872v3)

Abstract: We propose Image-to-Image Schr\"odinger Bridge (I$^2$SB), a new class of conditional diffusion models that directly learn the nonlinear diffusion processes between two given distributions. These diffusion bridges are particularly useful for image restoration, as the degraded images are structurally informative priors for reconstructing the clean images. I$^2$SB belongs to a tractable class of Schr\"odinger bridge, the nonlinear extension to score-based models, whose marginal distributions can be computed analytically given boundary pairs. This results in a simulation-free framework for nonlinear diffusions, where the I$^2$SB training becomes scalable by adopting practical techniques used in standard diffusion models. We validate I$^2$SB in solving various image restoration tasks, including inpainting, super-resolution, deblurring, and JPEG restoration on ImageNet 256x256 and show that I$^2$SB surpasses standard conditional diffusion models with more interpretable generative processes. Moreover, I$^2$SB matches the performance of inverse methods that additionally require the knowledge of the corruption operators. Our work opens up new algorithmic opportunities for developing efficient nonlinear diffusion models on a large scale. scale. Project page and codes: https://i2sb.github.io/

• The paper introduces I²SB, a novel framework that modifies Schrödinger bridge systems to enable tractable image restoration via conditional diffusion models.
• It reformulates entropy-regularized optimal transport into efficient drift functions using linearly structured SDEs, bypassing intractable PDE complexities.
• Empirical validation shows I²SB improves sampling efficiency and restoration quality, outperforming standard SGMs on FID and CA across diverse tasks.

An Analysis of I$^2$SB: Image-to-Image Schrödinger Bridge

The paper presents a novel framework, titled Image-to-Image Schrödinger Bridge (I$^2$SB), for addressing image restoration problems by leveraging a specific class of conditional diffusion models. This research proposes an alternative to existing score-based generative models (SGMs) by introducing a tractable implementation of Schrödinger bridge (SB) systems, which builts diffusion paths directly between two distinct image domains: clean and degraded.

Methodological Advancements

I$^2$SB builds upon the theoretical foundations of SBs, which are classically associated with entropy-regularized optimal transport, by modifying the complex system of PDEs that usually hinder practical implementation due to intractability. Key to the framework is the recognition that traditional SB approaches can be reframed into a tractable format compatible with the computational approaches used in SGMs. This allows I$^2$SB to employ similar network architectures and sampling techniques, specifically those of the DDPM.

The researchers extend the SB framework by formulating drift functions as scored functions corresponding to linearly structured Stochastic Differential Equations (SDEs). Notably, the assumption that clean images can be represented as Dirac delta distributions facilitates solving these SDEs, resulting in tractable boundary conditions and rendering the computation feasible. It effectively bridges the gap between theoretical optimal transport solutions and practical, computationally efficient systems.

Empirical Validation and Results

Under the empirical exploration, I$^2$SB exhibits significant efficiency and interpretability improvements over standard conditional SGMs and existing SB methodologies across a range of high-dimensional image restoration tasks, specifically on ImageNet datasets. The tasks undertaken include 4× super-resolution, deblurring, JPEG restoration, and inpainting. Performance is measured using Frechet Inception Distance (FID) and Classifier Accuracy (CA), where I$^2$SB surpasses standard SGMs on multiple configurations and matches state-of-the-art results achieved by diffusion-based inverse models (DIMs) that necessitate more domain-specific knowledge during training and execution.

The tractability of I$^2$SB enables lower computational complexity and increased sampling efficiency, while its underlying models benefit from a significant reduction in performance drops when decreasing the number of function evaluations (NFEs) during sampling. This indicates a practical advantage and potential for substantial computational savings in real-world application scenarios.

Implications and Speculative Future Directions

I$^2$SB illustrates the vast potential of tailoring nonlinear diffusion models to exploit structured priors in image restoration tasks. The tractability facilitated by the proposed computational framework lays the groundwork for deploying advanced generative models across varied applications while maintaining operational efficiency in computational environments constrained by resources.

Looking ahead, the duality of theoretical soundness against tractable practicality could prompt further advancements into broader spectrum applications such as general image-to-image translations beyond restoration, adaptation to less structured domains, or incorporating additional modalities into diffusion processes, revealing latent cross-modal mappings. Moreover, extending this framework to unpaired data scenarios could widen its applicability significantly.

Conclusion

This research contributes a significant advancement in the practical application of theoretical models within conditional diffusion frameworks. By aligning the precision of Schrödinger bridge systems with the scalability and efficiency of SGM-inspired techniques, I$^2$SB presents a robust image restoration tool which balances computational efficiency, model interpretability, and empirical validity, carving a promising path in the landscape of machine learning and computer vision.