Unifying Bayesian Flow Networks and Diffusion Models through Stochastic Differential Equations

TL;DR

This work unifies Bayesian Flow Networks and diffusion models by identifying linear SDEs for BFN noise, showing that BFN regression aligns with denoising score matching, and interpreting BFN sampling as a first-order reverse-time SDE solver. By importing fast sampling techniques from score-based diffusion models, the authors design specialized BFN solvers (BFN-Solvers) that substantially speed up sampling with small NFEs on both image and text data. They extend the framework to discrete data through latent-variable SDEs and DSM-equivalence, achieving strong performance and efficiency without altering discrete data. The results offer a rigorous, scalable bridge between BFNs and DMs and provide practical tools for faster, high-quality generation across modalities, with code available openly.

Abstract

Bayesian flow networks (BFNs) iteratively refine the parameters, instead of the samples in diffusion models (DMs), of distributions at various noise levels through Bayesian inference. Owing to its differentiable nature, BFNs are promising in modeling both continuous and discrete data, while simultaneously maintaining fast sampling capabilities. This paper aims to understand and enhance BFNs by connecting them with DMs through stochastic differential equations (SDEs). We identify the linear SDEs corresponding to the noise-addition processes in BFNs, demonstrate that BFN's regression losses are aligned with denoise score matching, and validate the sampler in BFN as a first-order solver for the respective reverse-time SDE. Based on these findings and existing recipes of fast sampling in DMs, we propose specialized solvers for BFNs that markedly surpass the original BFN sampler in terms of sample quality with a limited number of function evaluations (e.g., 10) on both image and text datasets. Notably, our best sampler achieves an increase in speed of 5~20 times for free. Our code is available at https://github.com/ML-GSAI/BFN-Solver.

Figures (9)

• Figure 1: Illustration of BFN on continuous data and the corresponding SDEs. The SDEs are defined w.r.t. ${\bm{\mu}}$ on time $[0, 1-\eta]$.
• Figure 2: Illustration of BFN on discrete data and the corresponding SDEs. The SDEs are defined w.r.t. the latent variables ${\mathbf{z}}$, which are marginalized in BFN, on time $[0, 1-\eta]$.
• Figure 3: Fast sampling results on the continuous CIFAR-10 dataset. Sampling quality is measured by FID $\downarrow$, varying the NFE.
• Figure 4: Fast sampling results on the discrete text8 dataset. Sampling quality is measured by SA $\uparrow$, varying the NFE.
• Figure 5: User study results on the discrete text8 dataset with 10 NFE. We present the preference rates (with 95% confidence intervals) of BFN-Solver1 and BFN-Solver2 over BFN baseline.

Theorems & Definitions (12)

• Theorem 4.1: Proof in Appendix \ref{['proof:sde']}
• Proposition 4.2: Proof in Appendix \ref{['append:ancestral_sampling']}
• Theorem 5.1: Proof in Appendix \ref{['append:dbfn_sde']}
• Theorem 5.2: Proof in Appendix \ref{['app:discrete-dsm']}
• Proposition 5.3: Proof in Appendix \ref{['app:dis_bfn_sampler']}
• proof
• proof
• proof
• proof
• proof