Mixture-of-Gaussians-Guided Schedule Design for Brownian Bridge Diffusion Models

2026-07-03Machine Learning

Machine LearningComputer Vision and Pattern Recognition
AI summary

The authors study Brownian Bridge Diffusion Models (BBDM), which improve image restoration by connecting the original clean image directly to the damaged one instead of to random noise. They develop a new mathematical framework to choose the best way to move between these images, using a special model called a Mixture-of-Gaussians. Their analysis shows there is a tradeoff between making restored images look better and making them more accurate. They also find general schedules that work well regardless of the image damage or prior knowledge. Tests on synthetic and real-world image tasks confirm their theoretical findings.

Brownian Bridge Diffusion Modelsimage restorationinverse problemsMixture of Gaussiansposterior distributionMMSE denoiserWasserstein distancemean squared errorimage inpaintingsuper-resolution
Authors
Ron Levi, Michael Elad
Abstract
Brownian Bridge Diffusion Models (BBDM) offer an appealing framework for image restoration and inverse problems by constructing a stochastic bridge from the clean signal directly to the degraded observation, rather than to pure noise. Despite their promise, the choice of bridge schedule is typically inherited from heuristics, and a principled analytical framework for schedule design has been lacking. In this work, we develop such a framework by offering a novel analysis of BBDM reverse dynamics under a Mixture-of-Gaussians (MoG) prior. This setting yields a closed-form ideal posterior and a corresponding MMSE denoiser, while the BBDM-induced reconstruction law is captured analytically through a tractable surrogate. Building on these expressions, we formulate two complementary schedule-design objectives: a Wasserstein criterion targeting perceptual quality and an MSE criterion targeting reconstruction fidelity. Our work exposes an inherent tradeoff between the two and proves the existence of universal schedules for both that are independent of the degradation and prior. Extensive experiments on controlled MoG settings confirm full alignment between theory and practice, and experiments on the FFHQ dataset across inpainting, deblurring, and super-resolution tasks validate the practical value of our schedule-design criteria.