Exploring the Design Space of Diffusion Bridge Models via Stochasticity Control

2024-10-28

arXiv.org

AI Summary

Existing diffusion bridge and stochastic interpolation models for pixel-space image-to-image translation suffer from technical fragmentation due to incompatible mathematical assumptions and neglect the insufficient diversity under fixed sampling budgets. This paper proposes the Stochasticity-Controlled Diffusion Bridge (SDB), the first framework to jointly regulate three sources of stochasticity—sampling SDE dynamics, transition kernels, and base distributions—along the noise-source dimension. SDB avoids training and sampling singularities and introduces a differentiable diversity metric. Built upon extended diffusion bridge theory and SDE-based modeling, SDB integrates controllable noise injection and joint FID/diversity evaluation. Empirically, SDB achieves state-of-the-art performance: it maintains high visual fidelity while accelerating sampling by 5× over baselines, significantly reducing FID, and substantially improving generation diversity.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Stochastic Optimization
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

Diffusion bridge models effectively facilitate image-to-image (I2I) translation by connecting two distributions. However, existing methods overlook the impact of noise in sampling SDEs, transition kernel, and the base distribution on sampling efficiency, image quality and diversity. To address this gap, we propose the Stochasticity-controlled Diffusion Bridge (SDB), a novel theoretical framework that extends the design space of diffusion bridges, and provides strategies to mitigate singularities during both training and sampling. By controlling stochasticity in the sampling SDEs, our sampler achieves speeds up to 5 times faster than the baseline, while also producing lower FID scores. After training, SDB sets new benchmarks in image quality and sampling efficiency via managing stochasticity within the transition kernel. Furthermore, introducing stochasticity into the base distribution significantly improves image diversity, as quantified by a newly introduced metric.

Problem

• Unifying diffusion bridge models for image translation
• Enhancing sampling efficiency and image quality
• Addressing low sample diversity in fixed conditions

Innovation

• Extend Stochastic Interpolants with preconditioning
• Optimize sampling algorithm for efficiency
• Modify base distribution for diversity