Beta Diffusion
Authors: Mingyuan Zhou
Tianqi Chen
Zhendong Wang
Huangjie Zheng
Venue: NeurIPS 2023 poster (2023)
Keywords: Diffusion models
KL-divergence upper bounds
multiplicative transitions
scaled and shifted beta distributions
TL;DR: We introduce beta diffusion for generative modeling of range-bounded data and propose KL-divergence upper bounds for optimizing diffusion models

Abstract: We introduce beta diffusion, a novel generative modeling method that integrates demasking and denoising to generate data within bounded ranges. Using scaled and shifted beta distributions, beta diffusion utilizes multiplicative transitions over time to create both forward and reverse diffusion processes, maintaining beta distributions in both the forward marginals and the reverse conditionals, given the data at any point in time. Unlike traditional diffusion-based generative models relying on additive Gaussian noise and reweighted evidence lower bounds (ELBOs), beta diffusion is multiplicative and optimized with KL-divergence upper bounds (KLUBs) derived from the convexity of the KL divergence. We demonstrate that the proposed KLUBs are more effective for optimizing beta diffusion compared to negative ELBOs, which can also be derived as the KLUBs of the same KL divergence with its two arguments swapped. The loss function of beta diffusion, expressed in terms of Bregman divergence, further supports the efficacy of KLUBs for optimization. Experimental results on both synthetic data and natural images demonstrate the unique capabilities of beta diffusion in generative modeling of range-bounded data and validate the effectiveness of KLUBs in optimizing diffusion models, thereby making them valuable additions to the family of diffusion-based generative models and the optimization techniques used to train them.

Beta Diffusion
Mingyuan Zhou;
Tianqi Chen, Zhendong Wang, and Huangjie Zheng The University of Texas Austin Austin, TX 78712
Abstract
We introduce beta diffusion, a novel generative modeling method that integrates demasking and denoising generate data within bounded ranges. Using scaled and shifted beta distributions, beta diffusion utilizes multiplicative transitions over time create both forward and reverse diffusion processes, maintaining beta distributions in both the forward marginals and the reverse conditionals, given the data any point time. Unlike traditional diffusion-basee generative models relying on additive Gaussian noise and reweighted evidence lower bounds (ELBOs), beta diffusion is multiplicative and optimized with KL-divergence upper bounds (KLUBs) derived from the convexity of the KL. divergence. We demonstrate that the proposed KLUBs are more effective for optimizing beta diffusion compared to negative ELBOs, which can also be derived as the KLUBs of the same KL divergence with its two arguments swapped. The loss function of beta diffusion, expressed in terms of Bregman divergence, further supports the efficacy of KLUBs for optimization. Experimental results on both synthetic data and natural images demonstrate the unique capabilities of beta diffusion in generative modeling of range-bounded data and validate the effectiveness of KLUBs in optimizing diffusion models, thereby making them valuable additions to the family of diffusion- based generative models and the optimization techniques used to train them.
Introduction
Diffusion-basser deep generative models have been gaining traction recently. One representative example Gaussian diffusion (57 59 23 60 335 that uses Gaussian Markov chain to gradually diffuse images into Gaussian noise for training. The learned reverse diffusion process, defined by Gaussian Markov chain in reverse order, iteratively refines noisy inputs towards clean photo- realistic images. Gaussian diffusion can also be viewed from the lens of denoising score matching 59
60) and stochastic differential equations (6I). They have shown remarkable success acTOSS a wide range of tasks, including but not limited to generating, restoring, and editing images [14] 24 55
28 64
53
667, transforming 2D to 3D
49 I, synthesizing audio 9 37 773, reinforcement
learning 31
66
48 quantifying uncertainty 19 and designing drugs and proteins 56
44
32
The recipe involves three basic steps: First, defining a forward diffusion process that introduces noise into the data and corrupts with decreasing signal-to-noise ratio (SNR) as time progresses from 1. Second, defining reverse diffusion process that denoises the corrupted data as time reverses from to 0. Third, discretizing the time interval from 0 to 1 into finite number of intervals, and viewing the discretized forward and reverse processes respectively. Auto-encoding variational inference [36] 51 is then applied optimize the parameters of the generator by minimizing a weighted negative ELBO that includes Kullback-Leibler (KL) divergence- based loss term for each discretized reverse step.
Constructing diffusion-basee based generative model often follows general recipe (5
23
35
34
fixed inference network and a learnable generator,
"Corresponding to: mingyuan. zhou@mccombs www.uuteeeeeooo.. edu PyTorch code is available at: hhttps://g/iii http:///www...o/////ggggng/////a Diffusion
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

X
Figure I: Illustration of the beta forward diffusion process for two example images. The first column displays the original images, while the other 21 columns display the images noised and masked by beta diffusion time t 0, 0.05....... ... using yy 10000 and the sigmoid diffusion schedule with co 10 and c₁ 133
Figure 2: Illustration of reverse beta diffusion for two example generations. The time decreases from to 0 when moving from left to right. each image generation, the top row shows the trajectory 2, (rescaled for visualization), which has been demasked and denoised using reverse diffusion, whereas the bottom row shows &u = fe( t), whose theoretical optimal solution is equal to E[zo | æ.). See Appendix D Dooo more details.
Although the general diffusion-modeling recipe is simple in concept, it requires access to the corrupted data at any time during the forward diffusion process given a data observation, as well the analytic form of the conditional posterior for any earlier time given both data observation and its corrupted version at the present time. The latter requirement, according Bayes' rule, implies access to the analytical form of the distribution of a corrupted data observation at the present time given its value at any previous time. Linear operations of the Gaussian distributions naturally satisfy these requirements since they are conjugate to themselves with respect to their mean parameters. This means that the marginal form and the conditional distribution of the mean remain Gaussian when two Gaussian distributions are mixed. Similarly, the requirements can be met under the categorical distribution 26 [8] 27) and Poisson distribution [10]. However, few additional distributions are known to meet these requirements and it remains uncertain whether negative ELBO would be the preferred loss.
While previous works have primarily used Gaussian, categorical, or Poisson distribution- based diffu- sion processes, this paper introduces beta diffusion as a novel addition to the family of diffusion- based generative models. Beta diffusion specifically designed to generate data within bounded ranges. Its forward diffusion process is defined by the application of beta distributions in multiplicative manner, whereas its reverse diffusion process is characterized by the use of scaled and shifted beta distributions. Notably, the distribution at any point in time of the forward diffusion given data observation remains a beta distribution. We illustrate the forward beta diffusion process in Figure which simultancously adds noise to and masks the data, and the reverse one in Figure 2 which iteratively performs demasking and denoising for data generation. We provide the details on how these images are obtained in Appendix D
Since the KL divergence between two beta distributions is analytic, one can follow the general recipe to define negative ELBO to optimize beta diffusion. However, our experiments show that minimizing the negative ELBO of beta diffusion can fail to optimally estimate the parameters of the reverse diffusion process. For each individual demendepeddent KL loss term of the negative ELBO, examining in terms of Bregman divergence 4 reveals that the model parameters and corrupted data are placed in its first and second arguments, respectively. However, to ensure that the optimal solution under the Bregman divergence agrees with the expectation of the clean data given the corrupted data, the order of the two arguments must be swapped.
By swapping the Bregman divergence's two arguments, we obtain an upper bound on the KL divergence from the joint distribution of corrupted observations in the reverse chain to that in the forward chain. This bound arises from the convexity of the KL divergence. another Bregman divergence that upper bounds the KL divergence from the univariate marginal of a corrupted observation the reverse chain to that in the forward chain. These two Bregman divergences, which can be derived either through KI divergence upper bounds (KLUBs) or the logarithmically convex beta function, share the same optimal solution but have distinct roles in targeting reverse accuracy at each step counteracting accumulation errors over the course of reverse
addition, there exists

---

diffusion. We further demonstrate that combining these two KLUBs presents computationally viable substitution for a KLUB derived at the chain level, which upper bounds the KL divergence from the joint distribution of all latent variables in forward diffusion chain to that of its reverse. Either KI.UB works on its own, which is not unexpected as they both share the same optimal solutions in theory, but combining them could lead to the best overall performance. In beta diffusion, the KI divergence is asymmetric, which enables us to derive an alternative set of KLUBs by swapping its two arguments. We will demonstrate that these augment-swapped KLUBs, which will be referred to as AS-KLUBs, essentially reduce to negative ELBOs. In Gaussian diffusion, the KL divergence is often made symmetrie, resulting in KLUBs that are equivalent to (weighted) negative ELBOs.
Our main contributions are the introduction of beta diffusion as novel diffusion- based multiplicative generative model for range- bounded data, as well as the proposal of KLUBs as effective loss objectives for optimizing diffusion models, in place of (weighted) negative BL.BOs. Additionally, we introduce the log-beta divergence, a Bregman divergence corresponding to the differentiable and strictly convex log-bets function, as a useful tool for analyzing KLUBs. These contributions enhance the existing family of diffusion-based generative models and provide new perspective on optimizing them.
2 Beta Diffusion and Optimization via KLUBs
We begin by specifying the general requirements for constructing diffusion-based generative model and establish the notation accordingly [57 23 35 Let zo denote the observed data, and let 2, and 1. In
2, represent their corrupted versions time s and time t, respectively, where 0 the forward diffusion process, we require access to random samples from the marginal distribution ((2 | zg) at any time t. as well as an analytical expression of the probability density function (PDF) of the conditional distribution g(2 2,, To) for any < t.
The forward beta diffusion chain uses diffusion scheduling parameters a2 to control the decay of its expected value over the course of forward diffusion, given by TF [mo] = OcyTn, and a positive concentration parameter y) to control the tightness of the diffusion process around its expected value. We typically set an to approach and 0 as approaches O and 1, respectively, and satisfy the condition
2 0g > 0₁ > 0+ > an 2 0 for all & E (0, 1), / E (0, 1).
Let I()) denote the gamma function and B(.,) denote the beta function. The beta distribution Beta(æ; a, b) B(a, b)- (1 x) 6-1 is member of the exponential family 66 65 77. Its log partition function is a log-beta function as In B(a, 6) = Iaa) + lb) In I'(a + 6), which is differentiable, and strictly convex (0, )² as a function of two variables 115. As result, the KL divergence between two beta distributions can be expressed as the Bregman divergence associated with the log-beta function. Specifically, in Appendix A
one can show by their definitions that
(p)) (,)) In B(ag, 3a)
Va In B(ap: 3p) V In B B(pp: 3p)
(ag ap, By - 3p)
B(ap: 8p)
Din B(a.b) ((ag. 8q), (ap, Bp)).
We refer to the above Bregman divergence the log-beta divergence. Moreover, if (aigs 8gy are random variables, applying Proposition 1 of Banerjee et al. we can conclude that the optimal value of (ap, (p) that minimizes this log-beta divergence is (a (5) = (,))
Next, we introduce a conditional bivariate beta distribution, which given data observation has (scaled and shifted) beta distributions for not only its two marginals but also two conditionals. These properties are important for developing the proposed diffusion model with multiplicative transitions.
2.1 Conditional Bivariate Beta Distribution
We first present the conditional bivariate beta distribution in the following Lemma, which generalizes previous results on the distribution of the product of independent beta random variables (30 38 33
Lemma (Conditional Beta Bivariate Distribution), Denote (2a, 21) as variables over a pair of time points (s. t), with 0 < s < t <<> 1. Given a random sample 2o E (0, from a probability vaalued data
distribation Paata (20), we define a conditional bivariate beta distribution over (Za; 2t)) with PDF:
T(n)
z/ 90:00 (1 2g 2g)n(11 (11-0,00)--
(2,,, 21 (0)) = (oya)zt)[[(())) a, ( ( (1000)
(Zs - 21) -n a0
3

---

Marginals: Given 20, the two univariate marginals of this distribution are both beta distributed as g(2s s mo) Beta( 17(1 (rg:Ta),,
g(2x mo) = Beta( 17(1 (xz3a)).
(4)
Conditionals: Given To, a random sample (z,, z.) from this distribution can be either generated in forward order, by multiplying a beta variable from (3) with beta variable nT
24 = Tg As-tt ~ Beta(naa naa,, - az) 2s ~ g(2s mn))
(5)
or generated in reverse order, by combining beta variable from (4) with beta variable Pa-t- as
2, = 2t + (1 - Psssttt, Pset ~ Beta(n((s - o2) ry(1 agTo))) 21 ~ g( mo)- The proof starts with applying change of variables to obtain two scaled and shifted beta distributions
(6)
{ 2s, x0) = - Beta
;; nT(r, - -a))20)
( ze, xD) = 1 Beta
8 ( :naa a1)20. 17 - -a200))
(8)
and then takes the products of them with their corresponding marginals, given by and
respectively, to show that the PDF of the joint distribution defined in (5) and that defined in (6) are
both equal to the PDF of g 2 mo) defined in 2 The detailed proof is provided in Appendix B
To ensure numerical accuracy, we will calculate 2, 2 (1 Po in (6) in the logit space
logit(z.) (2)) = In (e (doai()))) (ogi((pa))) + login( (x)) lng)(()--)))
2.2 Continuous Beta Diffusion
Forward Beta Diffusion. We can use the conditional bivariate beta distribution to construct a forward beta diffusion chain, beginning with the beta distribution from (3) and proceeding with the scaled beta distribution from (7). The marginal at any time for a given data observation Zu, as shown in (4) stays beta- distribiiti in the forward chain. For the beta distribution given by (4). we have
El æo] = a, var[ze| xo] = (a.za)(1-aae))
SNR,
E|2, stdl Sz
7+1
L-a,20
Thus when o₁ approaches 0 (i.e., 1), both 2; and SNR, shrunk towards and if on O. have 21 ~ Beta(O, 77), a degenerate beta distribution that becomes a unit point mass at O.
Infinite Divisibility. We consider beta diffusion chain is infinitely divisible given xn. This means that for any time & E (s, t), we can perform forward diffusion from 2, to 2 by first setting z Ra where πz-
form of continuous diffusion, as its forward
T)(ck. )x))
and then setting ū = T= where TE-z ~ Beta(rja 1)(ckz - (gg) )=y) The same approach can be used to show the infinite divisibility of reverse beta diffusion given zo.
Reverse Beta Diffusion. We follow Gaussian diffusion use q 2e, xo) to help define pe(2, I =e)
58 35 72 To construct reverse beta diffusion chain, we will first need to learn how to reverse from 2/ to Σs, where s < f. If we know the 2g used to sample 2, as in (3) then we can readily apply
the conditional in (8) to sample 2.. Since this information is unavailable during inference, we make weaker assumption that we can exactly sample from 2; ~ q(2x) Ezg-paaee [q(2 [xo)] which is the "ro-free" univariate marginal at time t. It is straightforward to sample during training but must be approximated during inference. Under this weaker assumption on g(=+), utilizing (8) but replacing its
true zo with an approximation do fø(2, !), where fe denotes the learned generator parameterized by 8, we introduce our "ro-free" and hence "causal" reve-reversal distribution as
po(=5 2x) q(2, l 24, fo t))-
(10)
2.3 Optimization via KLUBs and Log- Log- Divergence
KLUB for Time Reversal. The time-reversal distribution (=x 2x) reaches its optimal when its product with q(>;) becomes equivalent to g(2., 2) Erorpaan (g(zs) #1 (ro)], which is marginal bivariate distribution that is "xe-free." Thus we propose to optimize 0 by minimizing (Lo 22) (=)((=))lg(()) (t)) in theory but introduce surrogate loss in practice:

---

Lemma 2 (KLUB (conditional)). The KL divergence from q(zs, 2x) to po(zs (2)q(2)), two "ro-free" joint distributions defined by forward and reverse diffusions, respectively, can be upper bounded:
( ())((=)) =,)) _ KL.BB,,t a,t = (0)) / Pa) [KLUB(s 24, 2o)], (11) KLUB(s, 2t, = (1-(( ze, 80 = Salz,, ()) žt, mo))]
(12)
The proof in Appendix B utiliize the equation (2(2a Ze. (+()) = Exorpaace [po( 2z 2 ) [o)]
and
then applies the convexity of the KL divergence [ (74) and the definition (T0)
To find out the optimal solution under KLUB, following we can express KLUB(s, (s. 24, Fo) given by (12 as a log-beta divergence as
Log-Beta Divergence.
Dn B(ab) { [mra, at)2o, 17(1 -a,0))] [m(a, ac) n( fo }.
(13)
We note KLUB,.t defined in (IT) can also be written as [e,,=((z.) 1 Begnnggzo |x) [KLUB(s, 2,, zo]],
where the log-bet divergence for (LUB(s, 24, 2o), defined as (13), includes 2g ( 2,) in its
first argument and the generator fo(=,, () in its second argument. Therefore, applying Proposition 1
of Banerjee et al. 44 we have the following Lemma.
Lemma 3. The objective KLUB,,1 defined in for any < is minimized when
fo- (zx, 1) = B(xol 2x] = E. R [x)] for all 2t ~ q(21).
(14)
Thus under the KLUB -oottmiidd 6" we have pa- (2g 21) = g(zs| 24. E| 21]), which is different
from the optimal solution of the original KL loss in (II,, which is pă z,) = g( 21) 21)/((z)) =
((2s | =x) = Bu sx) [g(=, ze, =oo]] It is interesting to note that they only differ on whether the expectation is carried out inside or outside the conditional posterior.
In practice, we need to control the gap between pe- (zg [2)) and g(2, |ze) and hence s needs to be close to f. Furthermore, the assumption of having access unbiased samples from the true marginal q(z)) is also rarely met. Thus we need to discretize the time from I to t into sufficiently fine intervals and perform time-reversal sampling over these intervals. Specifically, we can start with 21 = 0 and iterate (10) over these intervals to obtain an approximate sample from 2 g(2z). However, the error could accumulate along the way from 21 to 2t, to which we present solution below.
KLUB for Error Accumulation Control. To counteract error accumulation during time reversal, we propose to approximate the true marginal q(=() using "distribution serriburion-ycdeeconsitteee approach. This involves feeding random sample 2₁ from ((=x) into the generator fo. followed by the forward marginal q<² < &0 fe(ze, f)), with the aim of recovering the distribution q(х{) itself. Specifically, [9(])]
we propose to approximate g(x() with po(2() (2{) == -z,=g(=x)) [g(2]] f = fe(zz, (z. f)) by minimizing KL(po (1)po(/)]l(( (x())) in theory, but introducing a surrogate loss in practice:
Lemma (KLUB (marginal)). The KL divergence (L(po (=)))(()))((() can be upper bounded:
(a()))
(L((p(()))ll (x+))) < KLUB, = (e,, [oo (zo) [KLUB (2,, zo)], KLUB 2,, mn) = (L(((<1) fo(z,, ()) mo)).
(15)
(16)
The proof in Appendix B utilizes the fact that ((x()= = (=))0))=((= (mo)pacta (mo) [g(2) 20)].
Note that the mathematical definition of KLUB is reused throughout the paper and refer to any of the equations [T]. (12). (15), or (16) depending on the context. Similar to previous analysis, we have
KLUB(2,, Din B(a.b) { [M0,30 n7(1 acto)), [nouso(EE, (), n11 - orJo(E,,t))), (17)
and can conclude with the following Lemma.
Lemma 5. KLUB, in (15) is optimized when the same optimal solution given by 14 is met.
Optimization via KLUBs. With the two KLUBs for both time reversal and error accumulation control, whose optimal solutions are the same as in (14), we are ready to optimize the generator fo via stochastic gradient descent (SGD). Specifically, denoting w, æ E [0, 1] as two weight coefficients,
the loss term for the ith data observation xi in a mini-batch can be computed as
(i)
C wLLU 81, Zt.s (1)) + (1 (2K,
(18)
24: g((t (1))) Beta 17(1 a1, 2)))) s, nt. t Uni((((,
(6)
5

---

2.4 Discretized Beta Diffusion for Generation of Range-Bounded Data
For generating range-bounded data, we discretize the beta diffusion chain. Denote Stg 1 and let ty increase with j. Repeating 5 for T times, we define discretized forward beta diffusion chain:
q( ( (o) = [1 ((2,, [2,,,, mo) = IT Sty-1
Beta
Fty-1 na,, "onn at, ) (19)
notable feature of (19) that the marginal at any discrete time step i, follows beta distribution, similarly defined in We also note while q xo) defines Markov chain, the marginal
(Z = Ezzo-ppata(xo))(e [g( [=o]]
(20)
in general does not. Unlike in beta diffusion, where the transitions between 2, and 2 are applied multiplicatively, in Gaussian diffusion, the transitions between ž and 1 are related to each other additively and Ft.e forms Markov chain regardless of whether zo is marginalized out.
The discretized forward beta diffusion chain given by (19) is reversible assuming knowing ø0- This means given zo, it can be equivalently sampled in reverse order by first sampling Sty ( (z0) = Beta( To, 7711 Ot, To)) and then repeating (8) for (T, l2, with PDF
((turr [ = a( (r) I 1 Beta Fig - 1-2;,
; 04, )ro r) at mo)) This
1-
non-causal chain, while not useful by itself, serves as blueprint for approximate generation.
Specifically, we approximate the marginal given by 200 with a Markov chain in reverse order as
pe ( ) Pprior (Ztz) ) I[;=2 Po (2t, (2,)) Pprior II,= 1[, (( 20,, fo (1,),,,,))) )). (21)
To start the reverse process, we choose to approximate q(2,,) Esooppaec (ma) (g Stp x with
(
(ztp) 2( E[rra],, which means we let Sty Beta B[ro]] n7(1 Oep E [r )) To
sample 1 we use the remaining terms in (21) which are scaled and shifted beta distributions
that are specified as in (8) and can be sampled as in (6 and (9).
KLUB for Discretized Beta Diffusion. An optimized generator is expected to make the "xo-free" joint distribution over all T steps of the discretized reverse beta diffusion chain, expressed po (zt,.r ) to approach that of the discretized forward chain, expressed as q( ). Thus an optimized is desired to minimize the KL divergence (po((Z()))) (l44(2)))))))) )). While this KL loss is in general intractable to compute, it can also be bounded using the KLUB shown as follows.
Lemma (KLUB for discretized diffusion chain).
(pot. (l(((().))) upper bounded by
KLUB Re paptaa)) [KL (Prrior (Zty (2tr [xn) Σ KI.UB,,-44 where
(22)
KLUBB,, (ereo)))p so)p [KLUB(s, (1. mo)] KLUB 2, mb) KL =)]la(=, =.,))))
We provide the proof in Appendix B We note Lemma a general statement applicable for any
diffusion models with a discrete forward chain q( ro) = 11 a(2,, (0)) and a discrete
reverse chain po (2trr) ) = Pprior IT po Po(zt-1 1 m1y). To estimate the KLUB in (22) however, during training, one would need to sample 21, ~ po(22, |22)) x p(zol & (2(,), ). which is often infeasible. If we replace žt. ~ Po | wo) with žt, ~ q(zt, (0)). then KLUB becomes the same as KLUB,,1 given by (IT), and KLÜB, given by (15) can be considered to remedy the impact of approximating pe (2t,) with q(2i,). Therefore, we can consider the combination of KLUB given
by (IT) and KLUB, given by (15) as a computationally viable solution to compute KLUUB,. which hence justifies the use of the loss (18) to optimize the discretized reverse beta diffusion chain. We summarize the training and sampling algorithms of beta diffusion in Algorithms and 2 respectively.
2.5 Argument-swapped KLUBs and Negative ELBOs
We note that in theory. instead of the KL divergences in 11 and 15 fe can also be optimized using two argument-swapper KL divergences: (L. 2.)
Po 2g 2e 2x) and KL(g(=)) (<)))ppp===)))) (=f)). By the same analysis, these KI divergences can also be bounded by KL.BS and log-beta divergences that are equivalent to the previous ones, but with swapped arguments. The argument-swapped KLUBs and log-beta divergences will be shown to be closely related to optimizing a discretized beta reverse

---

diffusion chain via - -ELBO, but they do not guarantee an optimal solution that satisfies diffusion and are found to provide clearly inferior empirical performance.
14 in beta
Negative BL for Discretized Beta Diffusion. As an alternative to KLUB, one can also consider following the convention in diffusion modeling to minimize the negative ELBO, expressed as
pizo (t
- Pxp-paato (2r) In po (2o) < EEBOO= = Exor~pance (z0) Bu(eiir xo)
In
2(Z1u:r (m0)
= EE,, Byter (a) In p(zo ax,) + B,,, (L I ao) Pprioor (2t,)) )1 + Σ}= Ezo Bytey [2 (00) [L( 2, zo)l
where the first two terms are often ignored and the focus is placed on the remaining T terms as
L(t (1 z0) = (L( Rt, zo) -,, Re,, t0 = (Q(z),)t))))
B(obi { [n(ac,-- ar, )ol=,, 12). at, fo(,, [n(ex,- - atz, (ao.))o.n1111 ai,- 201] (23)
Lemma 7 (-ELBO and urgument-swapee KLUB). Optimizing the generator fo with - -ELBO is
equivalent to using an upper-bound for the augment- augment-swapped KL divergence KL(q(Zt,r))) )lppe (1po((2
(zt,:))) )).
The proof in Appendix relies on the convex nature of both the KL divergence and the negative logarithmic function. We find for beta diffusion, optimizing KI () (1(((,,))) via the proposed KLUBB is clearly preferred to optimizing (1 (ppol ([Ipo((t.z)) via - ELBOS (i.e., augment- swapped KLUBs) and leads to stable and satisfactory performance.
KLUBs and (Weighted) Negative ELBOs for Gaussian Difftssion. We note KLUBs are directly applicable Gaussian diffusion, but they may not result in new optimization algorithms for Gaussian diffusion that drastically differ from the weighted ELBO, which weighs the KL terms using the corresponding SNRs. Moreover, whether the default or aguuenntswappee KLUBs are used typically does not make any difference in Gaussian diffusion and would result in the same squared error-based Bregman divergence. We provide the derivation of the (weighted) ELBOs from the lens of KLUBs in Appendix C providing theoretical support for Gaussian diffusion to use the SNR weighted BLBO 23
35
which was often considered as a heuristic but crucial modification of ELBO.
3 Related Work, Limitations, and Future Directions
Various diffusion processes, including Gaussian, categorical, Poisson, and beta diffusions, employ specific distributions both forward and reverse sampling. Gaussian diffusion starts at (o 1) in its reverse process, while both Poisson and beta diffusion start at O. Beta diffusion's 's reverse sampling is a monotonically non-decreasing process, similar to Poisson diffusion, but while Poisson diffusion takes count-vat count-valued discrete jumps, beta diffusion takes probability- valued continuous jumps. A future
direction involves extending beta diffusion to encompass the exponential family (6 65 45
Several recent works have explored alternative diffusion processes closely related to Gaussian diffusion. Cold diffusion by Bansal al. (5] builds models around arbitrary image transformations instead of Gaussian corruption, but it still relies on L1 loss, resembling Gaussian diffusion's 's squared Euclidean distance. Rissanen al. [52) propose an inverse heat dispersion- basedd diffusion process that reverses the heat equation using inductive biases in Gaussian diffusion- models. Soft diffusion by Daras al. [2] uses linear corruption processes like Gaussian blur and masking. Blurring diffusion by Hoogeboom and Salimans (25 shows that blurring (or heat dissipation) can be equivalently defined using a Gaussian diffusion process with non-isotropic noise and proposes to incorporate blurring into Gaussian diffusion. These alternative diffusion processes share similarities with Gaussian diffusion in loss definition and the use of Gaussian -aased reverse diffusion for generation. By contrast, beta diffusion is distinct from all of them in forward diffusion, training loss, and reverse diffusion.
concurrent work by Avdeyev et 33 utilizes the Jacobi diffusion process for discrete data diffusion models. Unlike Gaussian diffusion's 's SDE definition, the Jacobi diffusion process in Avdeyev et al. 3
is defined by the SDE dx 2 [a(1 - x bx]dtt + ssr(11 z)dw, with x E [o. 1] and s, a, b >>>>> 0. The stationary distribution a univariate beta distribution Beta(a, b). Beta diffusion and the Jacobi process are related to the beta distribution, but they differ in several aspects: Beta diffusion ends its forward process at Beta 17), a unit point mass at O, not a Beta(a, b) random variable. The marginal distribution of beta diffusion at time / is expressed as g(2 (2)) ~ Beta (1)(), v/(1 azzoo),, ). while the Jacobi diffusion process involves an infinite sum. Potential connections between beta diffusion and the Jacobi process under specific parameterizations are worth further investigation.
7

---

Several recent studies have been actively exploring the adaptation of diffusion models to constrained scenarios [T3 40 I77. where data bounded within specific ranges or constrained to partic- ular manifolds. These approaches are all rooted in the framework of Gaussian diffusion, which involves the incorporation of additive Gaussian noise. In sharp contrast, beta diffusion introduces a distinct perspective by incorporating multiplicative noise, leading to the emergence of an inherently hypereubic-constriiee diffusion model that offers new development and exploration opportunities.
Classifier-free guidance (CFG), often used in conjunction with heuristic clipping, is a widely used technique to perform conditional generation with Gaussian diffusion [2] 41]. Beta diffusion offers the potential for a seamless integration of CFG by directly applying it to the logit space, the operating space of its fe network, thereby eliminating the necessity for heuristic clipping.
To adapt Gaussian diffusion for high-resolution images, common approach is to perform diffusion within the latent space of an auto-encoder [53 62] A promising avenue to explore is the incorporation of sigmoid or tanh activation functions in the encoder's final layer. This modification would establish a bounded latent space conducive to applying beta diffusion, ultimately leading to the development of latent beta diffusion tailored for high-resolution image generation.
One limitation of beta diffusion is that its training is computationally expensive and data-intensive, akin to Gaussian diffusion. Specifically, with four Nvidia RTX A5000 GPUs, beta diffusion and Gaussian diffusion (VP-EDM) both take approximately 1.46 seconds process 1000 images of size 32 32 x 3. Processing 200 million CIFAR- 10 images, the default number required to reproduce the results of VP-EDM, would thus take over 80 hours. Several recent works have explored different techniques to make Gaussian diffusion faster and/or more data efficient in training [68] 20 76 71..
is worth exploring how to adapt these methods enhance the training efficiency of beta diffusion.
Beta diffusion has comparable sampling costs to Gaussian diffusion with the same NFE. However, various methods have been developed to accelerate the generation of Gaussian diffusion, including combining it with VAEs, GANs, or conditional transport [77 for faster generation (72 47
69], distilling the reverse diffusion chains 43
53 76).
utilizing reinforcement learning [16], and
transforming the SDE associated with Gaussian diffusion into an ODE, followed by fast ODE solvers
75 34. Given these existing acceleration techniques for Gaussian diffusion, it is worth exploring their generalization to enhance the sampling efficiency of beta diffusion.
58
39
42
Beta diffusion raises concerns regarding potential negative societal impact when trained on image datasets curated with ill intentions. This issue is not exclusive to beta diffusion but applies to diffusion-based generative models as a whole. It is crucial to address how can leverage these models for the betterment of society while mitigating any potential negative consequences.
Experiments
The training and sampling algorithms for beta diffusion are described in detail in Algorithms [ and respectively, in the Appendix. Our experiments, conducted on two synthetic data and the CIFAR10 10 im- ages, primarily aim to showcase beta diffusion'' 's effectiveness in generating range-bounded data. We also underscore the superiority of KLUBs over negative ELBOs effective optimization objectives for optimizing beta diffusion. Additionally, we highlight the differences between beta and Gaussian diffusion, specifically in whether the data generated through additive multiplicative transforms distributions with disjoint supports.
and their ability to model the mixture of
We compare the performance of "Gauss ELBO," "Beta ELBO," and "Beta KLUB," which respectively correspond a Gaussian diffusion model optimized with the SNR weighted negative ELBO 23 35 beta diffusion model optimized with the proposed KLUB loss defined in (18) but with the two arguments inside each KL term swapped, and beta diffusion model optimized with the proposed KLUB loss defined in (8). On CIFAR-10, 10, we also evaluate beta diffusion alongside range of non-Gaussian or Gaussian-like diffusion models for comparison.
4.1 Synthetic Data
We consider discrete distribution that consists of an equal mixture of five unit point masses located at 3o E D = 3/77 We would like to highlight that unit point mass can also be seen an extreme case of range-bounded data, where the range is zero. Despite being simple,
8

---

Origital Priganee
03
R2
0.4
0.1
0.7
Figure 3: Comparison of the true PMF, marked in red square m. and the empirical PMFs of three different methods- Gauss ELBO, Beta ELBO, and Beta KL.U calculated over 100 equal-sized bins between and 1. Each empirical PMF is marked in solid dot ..
this data could be challenging to model by continuous distribution, as it would require the generator to concentrate its continuous-valuee generations on these five discrete points.
We follow previous works to choose the beta linear diffusion schedule as 0; = c 0at 2 mmit where 3a = 19.9 and 8mi 0.1. This schedule, widely used by Gaussian diffusion 23 consistently across all experiments conducted on synthetic data. We set 17 10000, π 0.95, and 0.5. As the data already falls within the range of 0 to 1, necessitating neither scaling nor shifting, we set Soate = I and Snaft = O. We use the same structured generator So for both Gaussian and beta diffusion. We choose 20-dimensional sinusoidal position embeddings (63] with the positions set as 1000t. The network is an MLP structured as (21-256)-Rel) the Adam optimizer with a learning rate of 5e-4 and a mini-batch size of 1000.
61 34 is applied
We utilize
For data generation, we set NFE = 200. We provide the generation results in Figure 3 which shows the probability mass function (PMF) of the discrete distribution and the empirical PMFs over 100 equal-sized bins between and 1. Each empirical PMF is computed based on 100k random data points generated by the model trained after 400k iterations. It clear from Figure B that "Gauss ELBO" is the worst in terms of mis-aligning the data supports and placing its data into zero-density regions; "Beta ELBO" is the worst in terms of systematically overestimating the density at smaller-valued supports; whereas "Beta KLUB" reaches the best compromise between accurately identifying the data supports and maintaining correct density ratios between different supports.
In Appendix [
we further provide quantitative performance comparisons between different diffusion models and conduct an ablation study between KLUB and its two variants for beta diffusion: "KLUB Conditional" and "KLUB Marginal," corresponding to w = and w O, respectively, in the loss given by 18 Additionally, we evaluate beta diffusion another synthetic data, which comes from a mixture of range-bounded continuous distributions and point masses supported on disjoint regions.
4.2 Experiments on Image Generation
We employ the CIFAR- 10 dataset and build upon VP-EDM 34 as the foundation of our codebase. Our initial foray into applying beta diffusion to generative modeling of natural images closely mirrors the settings of Gaussian diffusion, including the choice of the generator's network architecture. We introduce a sigmoid diffusion schedule defined as observed to offer greater flexibility than the beta linear schedule for image generation. This schedule bears resemblance to the sigmoid-base based one introduced for Gaussian diffusion [3] 29]. We configure the parameters for beta diffusion as follows: Cu 10, c = -13, Snaft 0.6, Seale 0.39, 7 10000, w 0.99, and æ 0.95. We utilize the Adam optimizer with learning rate of 2e-4. We use EDM's s data augmentation approach, but restrict augmented images to a 0-1 range before scaling and shifting. We use the beta diffusion model trained on 200M images to calculate the FID [21) Explanations regarding the intuition behind these parameter selections can be located in Appendix G
1/(1 + e -0 -co)t), which has been
Below we provide numerical comparison of beta diffusion with not only Gaussian diffusion models but also alternative non-Gaussian or Gaussian-like ones. We also conduct an ablation study to
compare KLUB and negative ELBO across different NFE, 17. and B. A broad spectrum of diffusion models is encompassed in Table D which shows that beta diffusion outperforms all non non-Gaussian diffusion models on CIFAR 10, including Cold Diffusion [5] and Inverse Heat Dispersion [52], as well
as categorical and count-based diffusion models g 10 comparison to Gaussian diffusion and
Caussian+blurring diffusion, beta diffusion surpasses VDM K5 Soft Diffusion 12] and Blurring Diffusion 25
While it may fall slightly short of DDPM 23
improved DDPM 46], TPDM+ (78
VP-EDM 34 it remains competitive alternative that uses non-Gaussian based diffusion.

---

Table 1: Comparison the FID scores of various diffusion models trained on CIFAR-I0.
Diflusion Space
HID (L)
DDPM 23
3.17
VDM
35
4.00
Gaussian
Improved DDPM 45
2.90
1DPM+
78
VP-EDM
14
1.97
Soft Diffusion
12
3.86
Gaussian+ Blissiin+BBuurin
Blurring Diffusion
25
3.17
80.08 (deblurring) (inpainting)
Cold Diffusion (image reconstruction)
[
Deterministic
Inverse Heat Dispersion
S2
18.96
Categorical
D3PM LousstLogistic 2
7.34
+LDR-10 10 S
3.74
Count
JUMP (Poisson Diffusion) 10
4.80
Runge-bounded
Beta Diffusion
3.06
Table 2:
Comparing FID scores for KLUB and negative EL opptiiidd Beta Diffusion on CIFAR-I0 10 with vary- ing NFE under = 10000 and two different mini-batch sizes B.
Loss 8
ELBO
-ELBO
KLUB 512
KLUB
512
288
288
20
16.04
16.10
17.06
16.09
50
6.82
6.82
6.48
5.96
200
4.55
4.84
3,69
3.31
500
4.39
4.65
3.45
3.10
(a) ELBO
(b) KLUB
1000
3,38
Figure 4: Uncurated randomly-genennnnn images by beta diffusion optimized with -ELBO or KLUB with n 10000 and B 288, The generation with NFE- 1000 starts from the same random seed.
2000
4.50
4.66
3.37
3.06
Table presents comparison between KLUB and negative ELBO-optimizer beta diffusion across different NFE under two different mini-batch sizes B. Table in the Appendix includes the results under several different combinations of y) and B. We also include Figure 4 to visually compare generated images under KLUB and negative ELBO. The findings presented in Tables and and Figure provide further validation of KLUB's efficacy in optimizing beta diffusion.
As each training run takes a long time and FID evaluation also time-consuming, we have not yet optimized the combination of these hyperparameters given the limit of our current computation resources. Thus the results reported in this paper, while demonstrating that beta diffusion can provide competitive image generation performance, do not yet reflect the full potential of beta diffusion. These results are likely to be further improved given an optimized hyperparameter setting or a network architecture that is tailored to beta diffusion. We leave these further investigations to our future work.
5 Conclusion
We introduce beta diffusion characterized by the following properties: 1) Analytic Marginal: Given
probability- data observation xo E (0, 1). the distribution at any time point E [0, 1] of the forward beta diffusion chain, expressed as 7( (o)). is a beta distribution. 2) Analytical Conditional: Conditioning on data zo and a borward-sampled latent variable 2; ~ q(2 22)), the forward beta diffusion chain can be reversed from time to the latent variable at any previous time E [0,t) by sampling from an analytic conditional posterior 2, ~ q(2, Zzz 2t) that follows scaled and shifted beta distribution. 3) KLUBs: We introduce the combination of two different Kullback-Leible Upper Bounds (KLUBs) for optimization and represent them under the log- Bregman divergence, showing that their optimal solutions of the generator are both achieved at fø- (22, t) E[2g| 2.). We also establish the connection between sugment-swapper KLUBs and (weighted) negative ELBOs for diffusion models. Our experimental results confirm the distinctive qualities of beta diffusion when applied to generative modeling of hounedd data spanning disjoint regions or residing in high-dimensional spaces, as well as the effectiveness of KLUBs for optimizing beta diffusion.
10

---

Acknowledgements
The authors acknowledge the support of NSF-IIS 2212418, NIH-R37 CA271186, the Fall 2022 McCombs REG award, the Texas Advanced Computing Center (TACC), and the NSF Al Institute for Foundations of Machine Learning (IFML).
References
[1] Mohammadreza Armandpour, Huangjie Zheng, Ali Sadeghian, Amir Sadeghian, and Mingyuan Zhou. Re-imagine the negative prompt algorithm: Transform 2D diffusion into 3D, alleviate Janus problem and beyond, arXiv preprint arXiv:2304, 23440,,,0, 2023. hhttpp::/pprppnaete github. io/
(the first three authors contributed equally).
[2] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981-17993 2021.
[3] Pavel Avdeyev, Chenlai Shi, Yuhao Tan, Kseniia Dudnyk, and Jian Zhou. Dirichlet diffusion score model for biological sequence generation. arXiv preprint arXiv:2305. 2023.
[4] Arindam Banerjee, Srujana Merugu, Inderjit Dhillon, Joydeep Ghosh, and John Lafferty. Clustering with Bregman divergences. Journal of machine learning research, 6(10), 2005.
[5] Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold diffusion: Inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208. 09392, 2022.
[6] O. Sonndorff-Niesenn Information and Exponential Families In Statistical Theory. Wiley Series in Probability and Statistics Ser. John Wiley & Sons, Incorporated, New York, 2nd ed. edition, 2014. ISBN 1881118857373.
[7] David M Blei, Alp Kucukelbir, and Jon D McAuliffe, Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518)859--77, 2017.
[8] Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Informa- fion Processing Systems, 2022. http:ttts///:pterreett httppperreveee..eeeerrereeeeeeee seet/foorumiiddDdTa622YY
[9] Nanxin Chen, Yu Zhang, Heiga Zen, Ron Weiss, Mohammad Norouzi, and William Chan. WaveGrad: Estimating gradients for waveform generation. In International Con- ference on Learning Representations, 2021. URL http::://popeeeeeeiee...o wtw.:hwwwwwetffouummmmmm NaMLjcFa080
[10] Tianqi Chen and Mingyuan Zhou. Learning Jump: Thinning and thickening latent counts for generative modeling. In ICML 2023: International Conference on Machine Learning, July 2023. URL CTW htttp://axxiiii org/ aba/2305 18375
[11] Thomas M Cover and Joy A Thomas. Elements of information theory. 1991.
[12] Giannis Daras, Mauricio Delbracio. Hossein Talebi, Alex Dimakis, and Peyman Milanfar. Soft diffusion: Score matching with general corruptions. Transactions on Machine Learning Re- search, 2023. ISSN 2835-8856. (RRp:htts:///peereeeee
net www.::eeeeoooooeeee www.mmiidd-9rrrbbxiuu
[13] Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thomton, Yee Whye Teh, and Arnaud Doucet. Riemannian acore-based generative modelling. Advances in Neural Information Processing Systems, 15:2406-2422. 2022.
[14] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan, editors, Advances Neural Information Processing Systems, 2021. htllhhttps://penreete wtwpeerrriiieeeeeefrrrrer net/forum?
Ad=

---

[15] Sever Silvestru Dragomir, Ravi P. Agarwal, and N. S. Barnett. Inequalities for beta and gamma functions via some classical and new integral inequalities. Journal of Inequalities and Applications, 2000:504054, 2000.
[16] Ying Fan and Kangwook Lee. Optimizing DDPM sampling with shortcut fine-tuning. arXiv preprint arXiv:2301. 13362, 2023.
[17] Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson. Diffusion models for constrained domains. arXiv preprint arXiv:2304.0555:0 304.0036464 2023.
[18) Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10696-10706, 2022.
[19] Xizewen Han, Huangjie Zheng. and Mingyuan Zhou. CARD: Classification and regression diffusion models. In Advances in Neural Information Processing Systems, 2022.
[20] Tianka Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via Min-SNR weighting strategy. arXiv preprint arXiv:2303.094 09556, 2023.
[21] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to a local Nash equilibrium. In Advances Neural Information Processing Systems, pages
2017.
[22] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.125000 12598, 2022.
[23] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[24] Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res., 20,
[25] Emiel Hoogeboom and Tim Salimans. Blurring diffusion models. The Eleventh International Conference on Learning Representations, 2023. URL
(26) Emiel Hoogehoom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 14:12444-1244 2021.
[27] Minghui Hu, Yujie Wang, Tat-Jen Cham, Jianfei Yang, and Ponnuthurai N Suganthan. Global context with discrete diffusion in vector quantised modelling for image generation. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11502-11511, 2022.
[28] Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(24):695-709, 2005. http:ht:://mmll org/ www.:/tpr..////ooeeeeo pappps//esss / hyvarinen05a. html
[29] Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. arXiv preprint arXiv:2212.119722 2022.
[30] M. V. Jambunathan. Some properties of beta and gamma distributions. The Annals of mathe- matical statistics, 25(2):401-455, 1954. ISSN 0003-4851.
[31] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.099: 229999991, 2022.
12

---

[32] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi S. Jaakkola. Torsional diffusion for molecular conformer generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL wtttthttts:://. http:///ppppweee.e/eets/ee/ee/et wttp://www.ee//foomdd-ddddd. j2r622 __
[33] N. L. Johnson, A. W. Kemp, and S. Kotz. Univariate Discrete Distributions. John Wiley & Sons, 2005.
[34] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openrevies: htte//oorrum.iccc///UT/MMMO
[35] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-21707, 2021.
[36] Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. arXiv preprint 1rXiv:1312.61 14, 2013.
[37] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DiffWave: A versatile diffusion model for audio synthesis. In Intermational Conference on Learning Representations, 2021. httpstt:::/ //ppe//ee/e/..eeeeeeeoe ftrp:////eeeee.....eeee.orrrrrrr dat.ruummmmdaaa---8mmammmm
[38] Kotlarski. On groups of independent random variables whose product follows the beta distribution. Colloquium Mathematicum, 9(2):325-332) 0 1962.
[39] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations, 2022. URL attps://openrevier.
[40] Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations, 2023. URL www.:eeeewondieeeeeeee nett/forumIdd-WHHyCCOO.....
[41] Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04740. 2023.
[42] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-solver: A fast ODE solver for diffusson probabilistic model sampling in around 10 steps. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022. URLLhtttp://opeerrvviie wttp:oviwwi.ntttftoeeoee www./eeuunnmeem wewforomidd-uuaaaa.....
_
[43] Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.023888 2021.
[44] Shitong Luo, Yufeng Su, Xingang Peng. Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen- specific antibody design and optimization with diffusion- generative models for protein structures. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://opeerevieee net/ forum?idd=JSorG22Ti.
[45] Kevin P Murphy, Machine Learning: A Probabilistic Perspective. MIT Press, 2012.
[46] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[47] Kushagra Pandey, Avideep Mukherjee, Piyush Rai, and Abhishek Kumar. Diffusevac: Effi- cient, controllable and high-fidelity generation from low-dimensional latents. arXiv preprint arXiv:2201.00300. 2022.
[48] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et Imitating human behaviour with diffusion models. arXiv preprint 13X1:23011.1667 2023.
13

---

[49] Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. DreamFusion: Text-to-3D using
2D diffusion. In The Eleventh International Conference on Learning Representations, 2023.
ht..ttpps:/ htttts:://pperreieiett
[50] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprini urXiv:2204.06/25, 2022.
[51] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In International conference on machine learning, pages 1278-1286. PMLR, 2014.
[52] Severi Rissanen, Markus Heinonen, and Arn Solin. Generative modelling with inverse heat dissipation. In The Eleventh International Conference on Learning Representations, 2023. URL
[53] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[54] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding, Advances in Neural Information Processing Systems, 35: (6479-364994 2022.
[55] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022. URL https://pppeereeeioo
net//foruaPia-TTaiiIppboo
[56] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference Machine Learning, pages 9558-9568. PMLR, 2021.
[57] Jascha Sohl-Dickstein Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[58] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2020.
[59] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, pages 11918-11999 2019.
[60] Yang Song and Stefano Ermon. Improved techniques for training score-basee generative models. Advances neural information processing systems, 23:12438-12448, 2020.
[61] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. URL httttt///openrree/ het//formmP:d:PxII11122H
[62] Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space.
In Advances in neural information processing systems, 2021.
[63] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998 0038, 2017.
[64] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674 2011.
[65] Martin Wainwright and Michael Irwin Jordan, Graphical models, exponential families, and variational inference. Now Publishers Inc, 2008.
14

---

[66] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL http://pws.pppppeeereeeeeeeeeeea
[67] Zhendong Wang, Yifan Jiang. Yadong Lu, Yelong Shen, Pengcheng He, Weizhu Chen, Zhangyang Wang, and Mingyuan Zhou. In-context learning unlocked for diffusion models. arXiv preprint arXiv: 01005001110 2023.
[68] Zhendong Wang. Yifan Jiang, Huangjie Zheng, Peihao Wang, Pengcheng He, Zhangyang Wang, Weizhu Chen, and Mingyuan Zhou. Patch diffusion: Faster and more data-efficient training of diffusion models. arXiv preprint arXiv:2304. 12526, 2023. URL httpttts/aaaii..oooo aba/
2304. 12526
[69] Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion-GANe GAN: Training GANs with diffusion. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: /stps:///ppereeveew net/forum?ii HZf7UbpWHuA
[70] Zhixin Wang, Xiaoyun Zhang, Ziying Zhang, Huangjie Zheng, Mingyuan Zhou, Ya Zhang, and Yanfeng Wang. DR2: Diffusion-basee robust degradation remover for blind face restoration. In CVPR 2023: Conference on Computer Vision and Pattern Recognition, 2023. URL https //openreview d-t.:///wwweeoooo/dieeaaaaaaa HHnDUUUPEGHGO
[71] Zike Wu, Pan Zhou, Kenji Kawaguchi, and Hanwang Zhang. Fast diffusion model. arXiv preprint arXiv:2306.0699) 2023.
[72] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion GANs. In International Conference on Learning Representations, 2022. URL /ttp::tts///ppeereviii
net//orrumidd
[73] Dongchao Yang, Jianwei Yu, Helin Wang, Wen Wang, Chao Weng, Yuexian Zou, and Dong Yu. Diffsound: Discrete diffusion model for text-to-sound generation. JEEE/ACM Transactions on Audio, Speech, and Language Processing, 31:1720-1733, 2023. doi: 10...P9/
3268730.
[74] Mingzhang Yin and Mingyuan Zhou Semi-implicit variational inference. In International Conference Machine Learning, pages 5660-5669, 2018.
[75] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2023, URL nttps://openreree hew///oouuuriddtto////ttttb
[76] Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference Machine Learning, pages 42400. PMLR, 2023.
[77] Huangjie Zheng and Mingyuan Zhou. Exploiting chain rule and Bayes' theorem to compare probability distributions. Advances in Neural Information Processing Systems, 34:14993-15000 15006. 2021.
[78] Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Truncated diffusion prob- abilistic models and diffusion-based adversarial auto-encoders. In The Eleventh International Conference on Learning Representations, 2023. URL hnttps::/opeereeeeeeeee www.rrrviieeneeeffrummm net/forum? id=HDxgaakk9561
15

---

Beta Diffusion: Appendix
Algorithm Training of Beta Diffusion
Require:
Dataset whose values are bounded from 1, Mini-batch size B, concentration
parameter 10000, data shifting parameter Shift = 0.6, data scaling parameter Seale = 0.39, generator Sa. loss balance coefficient w - 0.99, time reversal coefficient я 0.95, and decreasing function that returns scheduling parameter E (0, 1) given E [0, 1J. such as beta linear schedule defined by ar 0.5( Bmt.) Bm where 3max 20 and Bmin 0.1, and sigmoid schedule defined by a₁ 1/(1 - -o--00 (o)-)))))) where Co = 10 and C1 - 13.
1: repeat
2: Draw mini-batch Xo = (2o Ji=1 from D
(0)
3: for = to B do
can be run parallel
4:
4 Unif(1e-5,
5:
8; = πt;
6:
Compute as, and at
(1 = i0
(i) Seate + Shift
7:
8:
2t, Beta(nor, 20
(6) (1 01, (0))
(.1 Ta , 4;) Seale + Stift
9:
10:
Using (8) compute the loss
C, WKLUB(Bs; Zt,, 200 + (1 w)KLUB( (2L, 20. (0))
Din B(atti (msss -an)o 7111 as. 200
[n(an. ar, )40 , 1011-a,210 ax)
+ (1 w) Dm B(a.)) |non, 20 n11 (n )
a .1(1 a1.
(24)
or swap the two arguments of both log- beta Bregman divergences in (24) if the loss is (weighted)
-ELBO.
11: end for
Perform SGD with Vo ΣB EE14.
B
13: until converge
Table 3: Comparing FID scores for KLUB and negative EL.BO-optimized Beta Diffusion on the CIFAR-10 image dataset with varying NFE under several different combinations of concentration parameter and mini-batch size B. We train the model with 200M images and use it to calculate FID. We compute FID one time each experiment. The other model parameters
as Scate 0.39,
Shift = 0.60, π 0.95, 0.99, Ir 2e-4, co 10, and = 13.
Loss
7)
10
NFE 10
20
50
200
500
1000
2000
ELBO
1
512
37.64
16.04
6.82
4.55
4.39
4.41
4.50
ELBO
1
288
37.54
16.10
6.82
4.84
4.65
4.61
4.66
KLUB
1
512
39.05
17.06
6.48
3.69
3.45
3.38
3.37
KLUB
0.1
512
41.28
20.38
9.72
5.85
4.98
4.90
4.88
KLUB
0.3
512
36.70
16.47
7.03
4.10
3.65
3.67
3.66
KLUB
1
288
37.67
16.09
5.96
3.31
3.10
3.08
3.06
KLUB
0.8
288
36.46
15.58
5.98
3.49
3.22
3.21
3.23
KLUB
1.2
288
38.49
16.59
6.31
3.68
3.36
3.37
3.24
KLUB
1
128
38.16
16.47
6.29
3.74
3.44
3.40
3,47
KLUB
2
128
39.80
17.09
6.58
3.93
3.76
3.75
3.65
16

---

Algorithm 2 Sampling of Beta Diffusion
Require: Number of function evaluations (NFE) J 200, generator fo, and timesteps {6))=00
t (1- le-5) (J-J)/U--1 for j 1..... and 0
1: if NFE > 350 then
2:
Oit; 1/(1 -co-(cr-c)))
3: else
4:
5: end if
6: Initialize to Em(o] Seale + Shift
7: Σt, ~ Beta Fo, 1111 Ozt, mo))
8: for
do
9: 3o = (a at,) )+ Scale + Shift
10:
P(ty 1t) Beta (n(a,,,, 0z, )ap 7/1 -t)-
which implemented the logit space as
logit(p(()-1-4,)) In Inv, u Gamma (n( - 0it 100 1), Gamma(n - S0) 1)
11:
2t-1 2; (1 2t, (()))
> which is implemented the logit space as
(ogit(2r,_11) In (ogit((e,)) > + - + lugulzz, (iliir,)Hooggt Htopill)()
end for 13: return (ãn Snift)/Senl or (2to/as Shift)/Sonts
A Log-beta Divergence and KL Divergence between Beta Distributions
By the definition of Bregman divergence, the log-beta divergence corresponding to the log-beta function In B(a, b) In (a) In (bb) - In [(a +bb. b), which is differentiable, and strictly convex (0, )² as function of and b, can be expressed as
Din B(a,)) ((ag. Bq), (p))
In B( Bg)
V., In B(am 3p)
(0t Ogg 0g.33g -p)
B(ap- Bp)
V,, In B(ap Sp)
In B(ag, Bg)
(ai Op. Bq 3p) (b((ap) <bapp + 3p)) (b(Bp) - <laap + 1p))
B(aps 3p)
In B(ag. (g)
(0r app)((p))))))))))))))))))))))) (Bg (p)) (ag 0pp B, - 3p)s(a, 8p),
B(app 8p)
which is equivalent to the analytic expression of
KL (Letaat(aap) [ (oeta(ag.3 8g)).
17

---

B Proof
Proof of Lemma !
The joint distribution 24 and 2g in (5) can be expressed as
((2,; 2e mo)
=<(=e 22)g(2a (o)
1
Beta
- 17(c= at)20)
Beta(zai 711 a0)
2a
1
[(no
ma120-1
(1
(1-3 2z
a(a a(a,-a:)a0-1
2. ((narto)T(n(a a1 (1)ro))
2s
[(n)
x ((na,zo)[[ril (n(1-_2z0)) ag2o))
20,500 -1 (1 -1-26]
2z)
[(n)
T(garzo) ( 04 (0)) [nn - a,
x = (zo 22)0(a 22)0000-0))0011( (1 2z) (1100000)))
(25)
The joint distribution of 24 and 24 in (6) can be expressed as
q(2,, 2e 220) = a(2x 2,, ro ( xo)
2.
2t n1a, - at) n711 - a,20) narto a,))))
= 1 - Beta
1 - 24
T(n(1 - (250))
1- [(n(1 (aatt) ac))
1(1-a,x0)-1
2. 2t
2. -1 21
(1-
21
21
[(n)
narm0 (1 (22
r(maazo) ( - -1))
4
[(n)
- [ a
X (Z+ - pla = (1 22) )2)0(1 (111aa ma) 10,00000011
(26)
The joint distribution shown in (25 is the same as that in 226
Proof of Lemma
2
Since we re-express ((2)) and o(2,, 21) as
(2t) Ezg Paore (mm (g(2 mu)],
((2,, 21) Exoxpenca (mo) (0(2,, z 2o))
and g(=,, 2e xo) = (2s 22, 2o (2 zo)) using the convexity of the KL divergence, we have
(I.(po ==)g(==)))(()) =))
(B (=x 21) )() (o) (
(2v) (g(==s #t (n))
Exorpaanc (po 2. 24)((22 æa) (19 21 [x))]]
EEzo-panca (z0) (2, (=))) (m)) In po (=a Ze )(( (a))
(2, E4s 2o ( (x))
pe (Ze (2))
-Ezo-panta (zo) Ez mo)
(220 2. Fts 20)
E( solpaars (2n)
KKL [KL 20)ll((22 Re, xo))
E (egg (2o)
[L 2e, to = fols,, 1)114(2, [1, |
18

---

Proof of Lemma
Since we can re-express po (2/) and g(2/) as
po(z)) (2)0o(=)))((44 80 fe( t))] E (= ( ( (mo) (g((2// (mo))
fo(z,, t))]
((2/) (=z)=a)>=o(z mal [g(24 mo)],
using the convexity of the KL divergence, we have
KL(po(z))la(()))
(L(((aa, ma)rauce (mo) [g(2) 8g fo(e,,
(ea,o) to(e,, (onpro)) (o mn)))
I ((z (o)macce (g (2,, [x)))) zo))].
Proof of Lemma 0
The proof utilizes the convexity of the KL divergence to show that
KL(po (l KL(Banapaaan (mo)
[po 21:11
[20)] [ | (#o) (o(0))l((1111))
|20))
< Exg (mo) KL (Z 1o)) 1 ( [o))
= -Bro-panca E ppo (m)) In Po ( 28)
q ( 20)
(roopessea(ao)
KL (pprior (Zer [xo)lla(Zer xo))
T Bxy, rpp(=s, (
122)) In po po(2) (2,)) ( Reg 20o)
J=2
We note po (a( (oo) x p(aro ze, )o((z,), which in general is intractable sample from, motivating us to replace pa (Zt, 200) with q(zt, 22)) for tractable computation.
Proof of Lemmal 7 Utilizing the convexity of the KL divergence, we present an upper-bound of the pogmented-swappeee KL divergence (L(g(=t, (2), referred to AS-KLUB,
(L(( )lpo (Ero (mo) (4( St,rr [o)]
(2t, [x))))
S AS-KLUB Eezorpaaca Paara (ze) (I(((Zt1r mo)llpo (2 xo))
(Z 2o)
Ezoopanco oor og lao)
Pa (Ztz.z x0))
EEzo
KL(g(Zer [o)llpprior (ztz (x)))
T
Es, ,a(se, 1x)) Bes,. -el lx, mo) In g(z,,-, 22, 20) po (2t5 1 2zz)
j=2
Utilizing the convexity of the negative logarithmic function, we have
Eng (mo) In po E Peooreaea (mo) In T (p [p(20) serr
p(xo 2t (por (2
EEro- (2o) In Eg(E Io)
a(Zero x0)
< EEBOO = -- -Engrppaaa (mo) Buleneer mo)) [In p(oal se, )po(k))))) a(Sesor 220)
-Eoopiana(zo))
Eyten (m)) In 22,) + KL(g(Eer [o) (2tr))
+ Σ E2, maaa, mm Ba ~a(s, (0)) In 9 21, I0) 2=2
Po (2,,-1 /2,)
Thus, when we disregard its first term, AS-KLUB is equivalent to -ELBO without its first two terms. Since these three terms typically not affect optimization, optimizing the generator fo with AS-KLUB is equivalent to using -ELBO.
19

---

C Derivation of (Weighted) ELBOs of Gaussian Diffusion from KLUB
Let denote og 1 and 2 dо- Following the definition in Ho et al. 23 and Song et al. [58] define a Gaussian diffusion-based generative model as
21 (N -a); 21,
22 N (V
02
Q2
a1
a1
2x N
2/- (/ar1
at
04-1
2T N
(V
VT-1"
27
Of
QT-1
where the diffusion scheduling parameters at in this paper is related to 3g in Ho et al. 23
at
3 =
The same the derivations in Ho et al. time as
23 express the forward marginal distribution at
(VVor at)-
Assuming E[z]] = and var [ar[ro] = 1, the signal-to-noise ratio at time /. is defined as
SNR, = Exo
std[22 22 xol
[2 mol
27
= 1 - at
at
E[z
at 1 ag
Since 0t-1
Et (2-11, Ox-1 1) and 2 ~ (W - 04-1); ) using the conjugacy of
the Gaussian distributions with respect to their mean, we can express the conditional posterior as
( 2o, 2x)
1
at
/a6-120
yai-1a,
(1- 04-1
at
=N
+
1 0/t-
0t-1
0t
04-1
04-1 0't
at
Ote
=
01-1 (1
at
)ro
+
1 - a4-1
at
21.
1 - ar-1
(1
at
))
1 aq
Cr-1
1 - 04
V a--1
1 a4
a
The forward Gaussian diffusion chain
T ( mо) = II
p(z
2-11)
(=1
can be equivalently sampled in reverse order as
( mo) q( ro)
II ( 24,
As this reverse chain is non-causal and Marroiionn,,,, we need to approximate it with Markov chain during inference, which is expressed as
p(21:) Pprior (2T) II (2 (2).
The usual strategy
158 35 72
to utilize the conditional posterior to define
Pe (0 zx) = q 2,, fng- = (e t)).
In what follows, we redefine the time / as continuous variable between 0 and 1. We let 4 be set of T discrete time points within that interval, with S 4 < la < < lr 1. The corrupted data observations over these I' time points are defined as 2t..r
20

---

C.1 Optimization of Gaussian Diffusion via Negative ELBO
Viewing Pe(zy 21:T) as the decoder, po(21:T) as the prior, and ((1:2) Fo) as the inference network, we can optimize the generator parameter via the negative ELBO as
EEz In (xoo) S ELBO Ban Ry(e,, or ( In p( | st 1p(())
( 126)
Ex,Byr (m) In p(zo Rt, + BBx,, ( (tt xo) Pprioor (#e,)) Σラ= Bry Bulty amo) [[[t)-1, xo
where the first two terms are often ignored and the focus is placed on the remaining T 2 terms, defined as
21,
KL(g(
24, 20)
ze,
fo(2), t)))
1- (1- 0tt
va.
0s )
||:00 Sa(z,, 1)1
2 2-o0
(1
0t)
1 0gs
at
||:00 - Sa( 1)122
(27)
2
1-0
02s
1-
where 0 < <t 1. We choose s = max(t 1/T, 0) during training.
Since Va,zo + vI Oytt, NO0, 1), it is true that
2t /1 CxfCe
Vat
Instead of directly predicting zo given 2t, we can equivalently predict & = co(2,. f) given and let
vI CnEE
in So(2),
Vot
Thus (7 can also be written as
L 20)
1
02
11:00 fo( t)
2
02s
1 04
1
(SNR, - SNR:)
-f((z, 1)13
2
1
as
1- 0:1 )
1
||ce (2),,))11
2
1 0s8
a₁
1
SNR,
1) ||e, E0((z, 1)12
(28)
2
SNR,
which agrees with Equations 13 and 14 in Kingma et al.
35
C.2
Optimization of Gaussian Diffusion via KLUB Conditional
Following the definition of KIJUB (Conditional), for Gaussian diffusion we have
(sUB(s, 2t, mo) KL(g(== ( fo(z,, (), 2<)1(g(za J zx))
For two Gaussian distributions Q1 and 82 that have the same variance, have KL(gill02)
(L(gzllq)) and hence
KLUB(s, 24, No) = (L(g(ze = fo(z,. t), 21)19((2 zo. 2()) KL(g(=, Ixv 21) 11)2((z, ft fo f), z))
=(ss, TT))
Therefore, over the same set of time points 4.TT, optimizing via KLUB conditional identical to optimizing via -E.BO. As analyzed before, optimizing via KLUB conditional (i.e., -ELBO) may be able directly counteract the error accumulated over the course of diffusion, and hence could lead to slow convergence.
21

---

C.3
Optimization of Gaussian Diffusion via KLUB Marginal and wNR-weighted Negative
ELBO
Following the definition of KLUB (Marginal), for Gaussian diffusion we have
(IU mn) = KL.(N Varto, (v))l((aaaf(e,))) (), o4)
0t
1 |100 f=
SNN, f(2,, t)11
llso cz(2,, 1)122
It is surprising to find out that the KLUB Marginal is identical to the SNR weighted - -ELBO first introduced in Ho al. [23] and further discussed in Hang al. [20].
C.4
Optimization of Gaussian Diffusion via KLUB
Following beta diffusion, we can also combine two KLUBs as
w KKLUB Since Ca 1-0t is in general close to 1 when s not too far from f, when not too close 1, I-o, combination of these two KLUBs does not result in an algorithm that clearly differs from the use of an SNR- -eiighed negative ELBO.
8, Zt, Fo) + (1 (I,
0gg 1
w) ||co co f)
2
as Ct
D Illustration of Forward and Reverse Beta Diffusion
Illustration of Forward Beta Diffusion. We first visualize the beta forward diffusion process by displaying true image zo and its noise-corrupted versions over the course of the forward diffusion process. Specifically, we display the noise corrupted and masked image at time 0, 0.05, 0.1 as
& = max (min
St Snift) Snift).)).))
Beta( 77( - (c,tt)).
Scate
It is clear from Figure that with multiplicative beta-distribuce noises, the image becomes both noisier and sparser as time increases, and eventually becomes almost completely dark. Thus the forward process of beta diffusion can be considered simultaneously noising and masking the pixels. This clearly differs beta diffusion from Gaussian diffusion, whose forward diffusion gradually applies additive random noise and eventually ends Gaussian random noise. In addition, the reverse diffusion process of Gaussian diffusion can be considered a denoising process, whereas that of beta diffusion is simultaneously demasking and denoising the data, illustrated below.
Illustration of Reverse Beta Diffusion. by displaying
We further illustrate the reverse process of beta diffusion
-1,11 max min Scale
01.-1
where =t,-1 22, + (1 zt, )p (+())1+(5)) is iteratively computed as in Algorithm We also display
so fo(2t, (1)
Ra E[zo b
where the approximation would become more and more accurate when approaches its theoretical optimal 0*, under which we have fe- (2,, E[xo 21,1.
As shown in Figure 2 starting from random image drawn from
24,3
Beta(naa,, 2o. 17(1 at Ro)), 2g E(rol.
most of whose pixels would be completely dark, beta diffusion gradually demasks and denoises the image towards clean image through multiplicative transforms, shown in Algorithm 2
22

---

Statish distance between the true and genairtted data cistributions
Datxx E-50
Beta EL50
0.075
Beta KLUB
0.6
0.050
07
Hellegge
11.4
0.025
0.000
0.0
10
107
10
107
20F
10"
20
Figure 5: Comparison of the statistical distances between the true and generated data distributions over course of training. The blue. green, and orange curves are for "Gauss ELBO," "Beta ELBO," and "Beta KLUB," respectively. From the left to right are the plots for Wasserstein-1 distance, Hellinger distance, respectively.
divergence, and
Statist distance
trua and generated cistributions
0.08
0.6
asbince
conditrraa
0.0
0.00
Beta- KLUS marginal
Bota
KLVB
04-
Jeron Cramen
Mrn
Hellinge
0.4
0.2
0.0
101
10"
10
10
101
20"
10'
10
20
Eerations
Binia tons.
Eerations
Figure 6: Analogy to Figure 5 for comparing "KLUB Conditional," "KLUB Marginal," and "KLUB."
E
Additional Results for Synthetic Data
Quantitative Evaluation Metrics. We consider Wasserstein- Jensen Jensen-Shannon divergence (ISD), and Hellinger distance as three complimentary evaluation metrics. Wasserstein- has a high tolerance of misalignment between the supports of the true data density and these of the generated data. In other words, it is not sensitive to both misaligning the modes of the true density and these of the generated density and placing density in the regions where there and Hellinger can well reflect the misalignments between the supports of high data density regions.
data. By contrast, both JSD
Wusserstein- 1: We monitor the performance of different algorithms during training by generating 100k data points using the trained model and drawing 100k data points from the data distribution, and computing the Wasserstein-1 distance between their empirical distributions, which can be done by sorting them and taking the mean of their element-wise absolute differences. If the true distribution is discrete in that it takes values uniformly random from discrete set 2, then Wasserstein-1 distance by sorting a set of 10k generated samples from small to large and calculating the mean of their absolute element-wite wise differences with non-decreasing vector of the same size, which consists an equal number of copies for each unique value in D.
compute the
JSD and Hellinger distance: We discretize the 100k generated data into 100 equal-sized bins between and and compute the frequency of the data each bin, which provides an empirical probability mass function (PMF). We then compute both the ISD and Hellinger distance between the empirical PMF of the generated data and the empirical (true) PMF of the true (discrete) data.
Additional Results. The observations in Figure 3 are further confirmed by Figure 5 which shows that while Gaussian diffusion has good performance measured by the Wasserstein-1 distance, is considerably worse than beta diffusion in aligning the supports between the true and generated data distributions. Within beta diffusion, either KLUB or its argument-swapp performance measured by both the JSD and Hellinger distance, suggesting their excellent ability to concentrate their generations around the true data supports. However, the proposed KLUB has much better recovery of the true underlying density in comparison to its argument-swappen as confirmed by examining the first subplot in Figure and comparing the second and third subplots in Figure
version leads good
version,
23

---

Gause ELBO ELBO
Bata ELBO
Beta, Beta_KLJ8 KLJB
0.25
0.2
0.29
Emprrcal Preqgrky
0.20
Empiticall Inesurrcy
0.23
0.28
0.30
0.13
0.05
0.03
0.00
0.00
0.03
0.2
0.4
0.8
0.4
0.8
0.4
Figure 7: Analogy Figure 3 for the mixture of continuous range-bounded data and point masses.
Statistical distoncee between the true ond generated data cistributions
0.30
Gaias ELSO KLUB
Hellinger distenec
0.3
0.04
0.1
0.02
101
10
107
10a
10
20
107
10f
20a
tersuons
IteraJonns
teretions
Figure 8: Analogy to Figure 5 for the mixture of continuous range- boendedoddd data and point masses.
Statist
cistributions
assance
Beta KLUS conditional Beta KUJS marginal
0
Hellegge divtaanc
Bota KLUB
0.6
0.04
02
12.4
0.02
02-
0.08
0.0
10a
10*
10F
10*
10a
Figure Analogy Figure for comparing "KLUB Conditional." "KLUB Marginal," and "KLUB."
We further conduct an ablation study between KLUB and its two variants: "KLUB Conditional" and "KLUB Marginal," corresponding to w 1 and w O, respectively, in the loss given by (18). The results suggest "KLUB Conditional" and "KI.UB Marginal" have comparable performance, which is not that surprising considering that theory, they both share the same optimal solution given by (14). More closely examining the plots suggests that "KLUB Marginal" converges faster than "KLUB Conditional" aligning the generation with the data supports, but eventually delivers comparable performance, and combining them leads to the "KLUB" that has good overall performance.
We note we have conducted additional experiments by varying model parameters and adopting different diffusion schedules. Our observation is that while each algorithm could do well by carefully tuning these parameters, "Beta KLUB" is the least sensitive to parameter variations and consistently better than or comparable to the other methods across various combinations of model parameters. While the toy data valuable for showcasing the unique properties of beta diffusion, the performance of "Beta KLUB" is not that sensitive to model parameters, and the observed patterns may be disrupted by image-specific settings. Hence, their utility in tuning beta diffusion for image generation, where data dimension and model size/architecture are much larger and more complex, is limited.
E.1 Mixture of Continuous Range-bounded Distributions and Point Masses
We consider equal mixture of three range-bounded distributions and two unit point masses, including a Uniform distribution between 0.1 and 0.2, a Beta (,5) distribution first scaled by 0.1 and then shifted by 0.3, unit point mass at 2o 0.5, (eta(0.5,0.5) distribution first scaled by 0.1 and

---

then shifted by 0.6, and unit point mass at 2u = 0.8. A random sample is generated as
yjk, ht ~ Discrete 2, 4, 5)), y1 ~ 0..,
0.1 Beta( U2 02-0.3, 1,5,,
0.1
Vs 0.5, U4 0.1 Beta 0a-0-6; 0.5, 0.5), 5/5 0.9.
(29)
We follow the same experimental protocol in Section 4.1 The results are shown in Figures] 7 and g from which again observe that "Gauss ELBOO does poor job in aligning its generation with the true data supports and places proportion of its generation to form smooth transitions between the boundaries of high-density regions; "Beta ELBO" well captures the data supports and density shapes, but has a tendency to overestimate the density of smaller-valued data; whereas "Beta KLUB" very well aligns its generated data with the true data supports and captures the overall density shapes and the sharp transitions between these five range- bounded density regions, largely avoiding placing generated data into zero-density regions.
F Preconditioning of Beta Diffusion for Image Generation
It is often good practice precondition the input of neural network. For beta diffusion, most of our computation operated in the logit space and hence we will consider how to precondition logit(z:) before feeding it as the input to the generator so. Specifically, since we have
2t ~ Beta(njazz 17(1 -az0)),
can draw the logit of 2, as
logit(zs) 22) = In In Us, ut ~ Gamma(naacoo 1), Vs -ammm))((1aa 1).
Assume Fg Unif m. where Fiii = Shift and Fmax = Scate + Shift, we have E[logit(z+)] EiIn uz] Blln Rez [v(()))0)) F.. (vh(1(1 - a(oo))),
where
1 no (Fmax Fminn
Mac dln (z))
Ezo (naraao)
nas
1
[In T (()oeTmmxx - In [nnaemmiu))
mos (Tmax Fmin
1
(11 (1-aaammn)))
Exo (w1(111 ac200))
dln[(z)
nou (Tmax Xmin)
n(1 ((1=a(mmmx))
1
[n [nn((n(1 ax.min)) In T(7(1 0x4mmax))
(Tmaxxxx Fmin
further estimate the variance of logit(z;) as var[logit(z:)) var uz] var(Inuu) wx]
= -E (Ma)r)))) + Exg (x) (n7(1 (xTo))) + vara, [p( + var [+( ax20))]
where
Ez Exgly (nanno)) mart (Fmax Fminn)
(bb (1(xstmax) - (1]eeemmn)]
Exo (ss (n(1 cx200)) noat (Xmax Fmin)
[0b(1(1 - axxtmin)) (1((1 axx
varao [[((1100000] = Eng (1)o02)))) - (Ezo (1((1)))2)))))
100
2/ (nJss (Imin + 1 (#max - 2min))) (Ez, (v( 0). (1(==0)1((=100)
100
2s max
100
25

---

varzo (r.oo<b(((111 - anto))l = Ezo Ezol (n(1 -axx))))ll (Eo, ac
100
2
(n or (Fmin + 100 (Fmax #miul]]) (Bzol ( o)
max
100
21(==0)+1(==
Now we are ready to precondition logit(z)) as
g(ze) logit(z,) [[)gi((zz]]
V [=]]
We use g(ze) the input the generator fo and add skip connection layer to add g(z,) into the output of the generator. Specifically, following the notation of EDM, we define the network as
fo(=1, Cakip (t)g Coul (f) Fo (Cin (f)g)()),))) Croise (t)),
where for simplicity, we set Cskip(t) = Cour (tur(t) = Cin(t) = 1 and (t)ise = -ooiit (a)) more sophisticated selection of these parameters has the potential enhance the performance of beta diffusion. However, due to our current limitations in computing resources, we defer this investigation future studies.
G Parameter Settings of Beta Diffusion on CIFAR10
We present the intuition on how the model parameters, including Stift, Seates " w, n, co, and C1. Given the limitation of our computation resources, we leave the careful tuning of these model parameters future work.
We set w = 0.99 and π = 0.95 across all experiments conducted on CIFARIO with beta diffusion.
We pre-process the range- bounded data by linearly scaling and shifting them to lie between Shaft + Seaiel, where purposes: firstly, it enables us to use beta diffusion to model any range-bounded data, and secondly, helps avoid numerical challenges in computing the log-beta divergence when its arguments are small. When evaluating the performance, we linearly transform the generated data back to their original scale by reversing the scaling and shifting operations applied during the pre-processing step. For CIFARI0 images, beta diffusion diffusing pixels towards Beta (0, 17), intuition is to set both Seale and Shafe large enough to differentiate the diffusion trajectories of different pixel values. However, Seate Spift needs to be smaller than 1, motivating to choose Scale 0.39 and Shift = 0.60.
Shift Shift + Scale 1. This linear transformation serves two
For the diffusion concentration parameter y), our intuition is that larger 1) provides a higher ability to differentiate different pixel values and allows a finer discretization of the reverse diffusion process, but leads to slower training and demands more discretized steps during sampling. We set v) 10000 and the mini-batch size as B = 288 by default, but also perform an ablation study on CIFARIO with several different values, as shown in Table
We use a sigmoid-based diffusion schedule as
Qt 1/(1 t ee-
where we set Cg = 10 and C1 = 133 by default. We note a similar schedule had been introduced in Kingma et al. 35 and Jabri al. [29] for Gaussian diffusion.
For the CIFAR- datase we utilize the parameterization of EDM as the code base. We replace the variance preserving (VP) loss and the corresponding network parameterization implemented in EDM, which the weighted negative ELBO of Gaussian diffusion, with the KLUB loss of beta diffusion that is given by 24 We keep the generator network the same as that of VP-EDM, except that we set Cnoise = logit( 1/88 and simplify the other parameters as Cskip = 1, Cout = 1, and Gn = I. The image pixel values between 0 and 255 are divided by 255 and then scaled by Smate and shifted by Spaft. and hence we have 3o E [Suafe, Shaft + Scate]. The two inputs to the generator network )) originally designed for VP-EDM and adopted directly for beta diffusion, glogit( (2)) and logit(a)). where g(.) is preconditioning function described in detail in Appendix F The output of the generator network is first transformed by sigmoid function and then scaled by Seale and shifted by Shaft, and hence to = fo(2, t) E [Snafts Shift + Seate]-
hhttps://wwet cs. toronto............ wwwpantttdddo//iiiii/a/ httpp::/ htt.tth...../aab///eeeeoo
html
26


Reviews:
[{"review_id": "ib5Xt0ZsvA", "reviewer": "Reviewer_n1uL", "summary": "The presented work proposes to design the forward and reverse process of diffusion models based on beta distributions (as opposed to Gaussian noise). It also showcases a new loss term based on KL-divergence upper bounds (KLUB) that can aid performance over the conventional ELBO-based diffusion model objective when used to optimize the proposed beta diffusion model. The paper shows that the proposed beta diffusion model can better model/generate synthetic distributions, especially point masses, compared to Gaussian diffusion, but lags behind it when used for CIFAR-10 image generation.", "questions": "- What does \\beta_max in the appendix correspond to in the main text? \n- How many generated samples are used to compute the Cifar10 FID?\n- More direct comparison to Gaussian diffusion on CIFAR 10 would be valuable (even if on a subset of Cifar10)\n- In terms of computational efficiency, are there any differences between beta and Gaussian diffusion?\n- The authors mention that numerical instabilities can arise in fp32 for beta diffusion, and it would be valuable to quantify this (e.g. show plot that shows which hyperparameter ranges may cause that).\n\n*************************************************\nAfter the rebuttal, I would like to raise my score to 6 and recommend acceptance. I'm looking forward to seeing the new experiments, related work discussion & benchmarking, and ablations in the revised paper. ", "limitations": "The paper would benefit from a more thorough discussion of the limitations of the presented work. As far as I can tell, the only limitation mentioned is that no hyperparameter tuning was performed, which is totally okay, but is not a limitation of the method itself (which would be much more valuable and interesting to read about).\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper is well-written\n- Alternatives to Gaussian diffusion processes are an interesting and relevant direction\n- The synthetic experiments give a good intuition on the advantages of beta diffusion over Gaussian diffusion (mostly in terms of modeling point masses)\n", "weaknesses": "- Comparison and discussion with related work is lacking. As a result, the novelty of the presented work is hard to assess. While not having a related work section may be okay in some cases, discussing relevant related work thoroughly is very important. In particular, the presented paper completely misses to discuss or even cite any closely related papers on alternative diffusion processes of which at least [1, 2, 3, 4] come to my mind. \n- In addition, a quick literature search on my end popped out [5] (ICML 23), which to me seems like  extremely related work. [5] also proposes to use Beta distributions to design a novel diffusion process (and even go beyond that by extending it to a Dirichlet distribution-based diffusion model). While to be fair [5] can be deemed concurrent work to the one presented (and was uploaded to Arxiv only in May), I think that it is absolutely key to discuss this paper, outline how the presented work is different (or not) from it, and potentially even benchmark against it (in [5] they show that they outperform D3PM, which is a baseline in this paper too).\n- There are multiple passages in the draft (e.g. line 78, 238, 346) that make it sound like KLUB-based objectives can be used for diffusion models in general, but there is no experiment backing this claim and it directly contradicts the authors themselves (lines 219-225). This seems like overclaiming to me.\n- Experiments are not entirely convincing. While I see the value of the presented synthetic experiments, I believe that a more comprehensive evaluation is required. Some concerns that I have are: \n    1) Both synthetic experiments seem to show to me that beta diffusion can model point masses better than Gaussian diffusion. This is great, but makes me question the value of having two synthetic toy datasets in the main paper that lead to the same conclusion, especially compared to including more challenging datasets (Cifar10 experiments are only shown in the appendix and not mentioned in the main text).  \n   2) The lack of other convincing experiments makes it unclear how beta diffusion would perform on non-toy datasets (basically anything except for point masses and beta distributions), and it is totally nebulous in which cases, if any, beta diffusion could be a strongly performing alternative to other diffusion models (on cifar10, beta diffusion FID scores are worse than the ones from Gaussian diffusion or [2, 3]). The paper would strongly benefit from a more realistic example that showcases where beta diffusion may have an edge over other diffusion models on any relevant metric (does not need to be state-of-the-art performance). \n   3) The authors can do a better job of describing how and why anyone should use a beta diffusion model. For example, the hyperparameter values for the synthetic experiments are quite different from the ones used for Cifar10 (some beta diffusion-specific ones include the diffusion concentration parameter and how to bound the data) , and there is no discussion on how to choose them in practice (or even how the authors arrived at their particular choice of values).\n   4) The synthetic datasets would seem to be a great opportunity for a detailed exploration and/or ablation of the design space for beta diffusion. However, there are only few ablative studies presented, and for the ones presented it is unclear how statistically significant they are, since no error bars are shown (esp. Figures 2,3,5 should have them) . One ablative study that would be important to include is ablating \\beta_max from the appendix – To my understanding, when \\beta_max is chosen to be too large, the end distribution (of z_1) will be just a black image, in which case there would be no stochasticity at all coming from the initial conditions during the reverse sampling process, which implies that a trade-off exists that is important to understand. Additionally, it would be valuable to ablate using ELBO vs KLUBs as beta diffusion objective on a non-toy dataset (e.g. Cifar10), as well as ablating NFE (i.e. #iterations during inference) for any trained (fixed) beta diffusion model.\n\n[1] https://arxiv.org/abs/2208.09392\n\n[2] https://research.google/pubs/pub52492/  (TMLR 23)\n\n[3] https://openreview.net/forum?id=OjDkC57x5sz (ICLR 23)\n\n[4] https://openreview.net/forum?id=4PJUBT9f2Ol  (ICLR 23)\n\n[5] https://arxiv.org/abs/2305.10699 (ICML 23)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "The reviewer found our paper to be well-written and the exploration of beta diffusion as an interesting and relevant direction, but was very concerned about the lack of discussion and comparison with related works. We aim to decisively alleviate these concerns with our response and revision.\n\n> *W1: Related work.* \n\nIn Introduction, we carefully reviewed key references on Gaussian diffusion, effectively highlighting how beta diffusion is distinct (kindly refer to Reviewer FX5S). \n\nWe appreciate your mention of [1-4] as related works on alternative diffusion processes. In **Response to All**, we have included FID comparison to [1-4] on CIFAR10, showing that beta diffusion now outperforms all except [4]. Below, we further discuss how beta diffusion is distinct from [1-4] to emphasize its novelty:\n\n- [1]: Cold Diffusion builds models around arbitrary transformations (e.g., blurring) instead of Gaussian corruption. [2]: Soft Diffusion uses linear corruption processes (e.g., Gaussian blur and masking).  [4]: Inverse Heat Dispersion explicitly reverses the heat equation by exploring inductive biases in Gaussian diffusion-like models. \n\n- [3]: Extending [4], Blurring Diffusion adds blurring (or heat dispersion) into Gaussian diffusion. \n\n- These alternative diffusion processes in [1-4], Blurring Diffusion in particular, still resemble Gaussian diffusion in not only how the squared Euclidean distance is used to define the training loss but also their use of Gaussian-based reverse diffusion for generation.\n\n- Beta diffusion, on the other hand, is distinct from [1-4] in forward diffusion, training loss, and reverse diffusion. Its loss can be expressed as a log-beta Bregman divergence, which theoretically guarantees that $f_{\\theta^*}(z_t,t)=\\mathbb E[x_0|z_t]$. \n\n> *W2: [5] is extremely related but not cited.* \n\nDue to [5 (Pavel Avdeyev et al., ICML 2023)] appearing on ArXiv after the May 17 submission deadline, we were unaware of its existence. \n\nTherefore, we kindly request you to reconsider its relevance in evaluating beta diffusion. Furthermore, in our reading of [5], we noted significant differences, as elaborated in **Response to All**.\n\n> *W3: Can KLUBs be used for diffusion models in general?*\n\nIndeed, KLUB-based objectives are applicable to diffusion models in general, including both beta and Gaussian diffusions. We believe this aspect was carefully addressed without any self-contradiction (Line 219-225) or overclaiming. For further clarification, please refer to Appendix D, showing how KLUB-based objectives relate to (weighted) -ELBO for Gaussian diffusion, and also our response to “*W2*” from Reviewer FX5S.\n\n> *W4-1:*\n\nThe 2nd toy data differs from the 1st in that it also includes range-bounded continuous variables. We will utilize the extra page available in Camera Ready to include the updated CIFAR10 results in the main paper.\n\n>*W4-2:*\n\nAs in **Response to All**, beta diffusion now outperforms non-Gaussian diffusion models, including [1,2,4], on CIFAR10, and is competitive against Gaussian diffusion with or without adding blurring [3]. \n\n>*W4-3:*\n\nWe will release the code along with detailed instructions to enable domain experts to explore beta diffusion on their own data and applications.\n\n- For toy data, we set the concentration parameter as $\\eta=1000$, while for images, to distinguish between 256 pixel values, we choose a larger value of $\\eta=10000$. A larger $\\eta$ might result in slower training and the need for a finer discretization of reverse diffusion, with FIDs on CIFAR10 differing by approximately 1 between $\\eta=10000$ and $\\eta=1000$.\n\n- As beta diffusion is diffusing pixels $x_0$ towards $Beta(0, \\eta)$, our intuition suggests setting both Scale and Shift large enough to differentiate the diffusion trajectories of different pixel values. With this in mind, [Shift=0.5, Scale=0.4] provides a good compromise.\n\n- It is common practice to choose data and network-architecture specific parameters for diffusion models. While ideally, we would perform a grid search of these parameters, we haven't done so for beta diffusion yet due to limited computing resources.\n\n>*W4-4:* \n\nWe appreciate your suggestion regarding exploring/ablating the design space on toy datasets. While they are valuable for code debugging and showcasing the unique properties of beta diffusion, their performance is not that sensitive to model parameters, and the observed patterns may be disrupted by image specific settings. Hence, their utility in tuning beta diffusion for image generation, where data dimension and model size/architecture are much larger and more complex, is limited.\n\n- For toy data, we did not include error bars as the performance gaps between KLUB and -ELBO were large. We are glad to add them into Camera Ready. \n\n- For both Gaussian and beta diffusions, selecting an appropriate $\\beta_{\\max}$ is desirable. In Gaussian diffusion, an overly large $\\beta_{\\max}$ can lead to many wasted diffusion steps at $x_t\\sim N(0,1)$. Similarly, in beta diffusion that starts its reverse from $x_1\\sim Beta(0,\\eta)$, a point mass at 0, if $\\beta_{\\max}$ is too large, many training and reverse steps for $t$ close to 1 will be wasted at $x_t=$ MIN = torch.finfo(torch.float32).tiny= 1.1754943508222875e-38, the minimum threshold in sampling Beta in PyTorch Float32 precision.\n\n- **Response to All** now includes -ELBO vs KLUB on CIFAR10 and FIDs under different NFEs.\n\n> *Q 1-5:*\n\n- The main paper employs the sigmoid schedule and, therefore, does not include $\\beta_{\\max}$, which is a parameter of the beta linear schedule used by DDPM (see Algorithm 1).\n\n- As per convention, we used 50,000 generated samples to compute FID.\n\n- Please see **Response to All**.\n\n- No clear difference. See Line 542 in Appendix F.\n\n- Numerical instabilities may arise when Shift approaches zero. This is not surprising, as KL[Beta($a,b$)||Beta($c,d$)] exhibits numerical challenges when either $a$ or $c$ is close to zero."}, {"review_id": "ES2qNdTB7q", "reviewer": "Reviewer_FX5S", "summary": "This paper addresses a prevalent assumption found in deep generative diffusion models, which is the Gaussian assumption in both the forward and reverse processes. In this study, the authors explore the utilization of the beta distribution in these processes and establish several key properties. Firstly, they analytically compute the marginal distribution within the beta distribution framework. Additionally, they develop an analytical computation for the conditional distribution, given both the data point and a noisy sample, in the form of a scaled and shifted beta distribution. To optimize the beta diffusion process, the authors propose a combination of two distinct KL upper bounds, which is better than using the standard ELBO. The experimental results on a synthetic setup demonstrate the ability of the beta distribution to effectively model range-bounded data distributed across disjoint regions. Furthermore, they show the efficacy of the proposed objective functions. The supplementary material also includes preliminary experiments conducted on image generation tasks.", "questions": "[About the formulation]\n\n#1. One area where the real benefits of Beta diffusion could be highlighted is in guided sampling. In classifier-free guidance, as well as classifier-guided sampling, the presence of outliers can lead to low-quality generated samples, particularly in cases with high guidance scale. This typically leads to the generated samples of high-contrast. This issue is commonly observed in many open-source diffusion models such as SD, Karlo, and IF. The Imagen paper by Google addresses this problem by employing a dynamic threshold heuristic. However, I believe that Beta diffusion could serve as a fundamental solution to mitigate this drawback in guided sampling. Due to the design of the diffusion process, the data range is naturally constrained, making Beta diffusion well-suited to address this challenge.\n\n#2. Is it possible to extend this framework to the exponential family? If so, it would allow for a unified perspective encompassing Gaussian, Categorical, and Beta diffusion processes. This extension could contribute to a more comprehensive understanding of diffusion models.\n\n#3. In image generation tasks, various approaches, including Stable Diffusion, train diffusion models on the latent space learned by an autoencoder. I believe that Beta diffusion could be readily applied to the latent space with relative ease, but I would be interested to hear the authors' opinion on this matter. It is worth noting that in the latent dimension, the range is no longer limited to [0, 1], but it can be regulated by applying appropriate regularization techniques to the stage 1 model, such as the autoencoder in LDM.\n\n[About the experiments]\n\n#1. Have you considered comparing the simplified denoising loss, which discards the weights in the ELBO, to Beta KLUB? The DDPM paper has shown that this simplified loss outperforms the reweighted loss in image generation tasks. I'm curious if you observed similar results in your experiments.\n\n#2. Did you clip the values during the sampling steps in Gaussian diffusion to the pre-defined range ([0,1] or [-1, +1])? In practice, clipping the values during the reverse steps can improve the quality of the samples, especially when dealing with high CFG scale.\n\n#3. It would be helpful to provide more practical scenarios where the Beta diffusion process is more suitable compared to the vanilla Gaussian diffusion. While the supplementary material includes a preliminary experiment on image generation tasks, in my opinion, the Gaussian diffusion process already works quite well in the image domain. Therefore, showcasing more compelling applications of the Beta diffusion process would significantly strengthen the motivation behind your work.\n", "limitations": "The paper does not explicitly outline its limitations. My suggestions to improve this work are described in the sections above. As this work is primarily theoretical in nature, it is unlikely to have any significant negative societal impacts.\n", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "[Uniqueness of the methodology]\n\nUnlike previous works on diffusion-based generative models that heavily rely on the Gaussian assumption for both the forward and reverse processes, this work stands out as the first attempt to introduce Beta distributions into these processes. This unique approach sets it apart from existing literature. Furthermore, the authors have devised an objective function specifically tailored to optimize the beta diffusion process, further distinguishing their methodology from previous approaches.\n\n[Presentation]\n\nThe manuscript is overall well-written and exhibits clarity in its presentation. Particularly, the introduction and related works sections provide concise and comprehensive summaries of prior research, effectively highlighting the distinction between the Gaussian diffusion process and the proposed Beta distribution process. To further enhance the motivation behind this work, it would be beneficial to include a practical scenario that illustrates the necessity of constructing the Beta diffusion process. This addition would greatly strengthen the motivation behind the study.", "weaknesses": "[Absence of real-world experiments] \n\nThe main paper primarily presents experiments conducted on a synthetic dataset. While strong theoretical analysis can support the use of synthetic experiments, including more real-world experiments would further reinforce the proposed methods. It would be beneficial to provide practical scenarios that demonstrate why the beta diffusion process should be considered over the well-established Gaussian diffusion. In the supplementary material, preliminary experiments on image generation tasks are included, indicating potential in that direction. However, it is not entirely clear why the image generation task specifically requires a beta diffusion process. One could argue that since images are typically bounded within the range [0, 1], the beta distribution is more suitable. Theoretical reasoning supports this argument. However, in practice, simply clipping the values during the reverse steps is often sufficient to remove outliers during sampling.\n\n[KLUB] \n\nI remain unconvinced about the superiority of KLUB over the negative ELBO in optimizing the beta diffusion process. Could you provide further elaboration in the rebuttal? Additionally, I am curious if this argument can be generalized to the standard Gaussian diffusion process as well.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your excellent summary of the key contributions of our paper and your recognition of the uniqueness of the proposed methodologies. We are delighted that you found our paper to be well-written and presented with clarity. Below please see our point-by-point response.\n\n> *W1: Absense of real-world experiments in the main paper.*\n\nIn **Response to All**, we demonstrate that beta diffusion now surpasses many existing non-Gaussian diffusion models in generating CIFAR10 images, and approaches Gaussian ones.\n\nRegarding the heuristic of \"clipping\" in reverse Gaussian diffusion, we acknowledge its widespread use, especially under classifier-free guidance (CFG). However, the reference [Lou & Ermon, April 2023], pointed out by Reviewer Npra, aims to eliminate the need for this heuristic in Gaussian diffusion. Beta diffusion offers a principled alternative to avoid this heuristic by performing CFG in the logit space that our Unet operates on, which we plan to investigate as part of our future work.\n\nWe also recognize the value of evaluating beta diffusion in various tasks where its multiplicative construction and hypercubic-constrained nature could provide unique advantages. We eagerly anticipate collaborating with domain experts in different applications, such as modeling high-dimensional biological data, which often exhibit range-bounded, sparse, skewed, and heavy-tailed characteristics.\n\n> *W2: Is KLUB superior over -ELBO?*\n\nYes, please see Table 2 in **Response to All**. We acknowledge the need for clarifying the four different types of KLUBs:\n\n- KLUB-conditional: A KLUB of $KL[p_{\\theta}(z_s|z_t)q(z_t)||q(z_s,z_t)]$, as shown in Eq.10-11.\n- KLUB-marginal: A KLUB of $KL[p_{\\theta}(z’_t)||q(z’_t)]$, as shown in Eq.14-15\n- KLUB-conditional-AS (AS stands for augment-swapped): A KLUB of $KL[q(z_s,z_t)||p_{\\theta}(z_s|z_t)q(z_t)]$, which is the same as Eq.11 except for swapping the two arguments inside the KL.\n- KLUB-marginal-AS: A KLUB of $KL[q(z_t')||p_{\\theta}(z_t')]$, which is the same as Eq.15 except for swapping the two arguments inside the KL.\n\nIt is important to note that optimizing with KLUB-conditional-AS is equivalent to optimizing with -ELBO. All four KLUB types can be expressed using the log-beta Bregman divergence. However, only the first two KLUBs guarantee the optimal solution as shown in Eq.13. Using $w$KLUB-conditional-AS+$(1-w)$ KLUB-marginal-AS or -ELBO does not provide such a guarantee. This explains why the proposed KLUB (i.e., $w$ KLUB-conditional + $(1-w)$ KLUB-marginal) is found to be better than KLUB-AS or -ELBO for optimizing beta diffusion.\n\nFor Gaussian diffusion, we find interesting theoretical justifications for the SNR weighted -ELBO, which is often considered a heuristic but crucial modification of -ELBO:\n\n- As shown in Appendices D.1-D.2, in Gaussian diffusion, KLUB-conditional = KLUB-conditional-AS = -ELBO.\n- As shown in Appendix D.3, KLUB-marginal = KLUB-marginal-AS = SNR weighted -ELBO.\n\nThis result arises from the symmetric nature of the squared Euclidean distance, which is the Bregman divergence under $KL[p||q]=KL[q||p]$, where $p$ and $q$ are univariate Gaussians with equal variance. The findings shed light on the justified use of SNR weighted -ELBO in Gaussian diffusion.\n\n\n> *F1: Potenial of beta diffusion for guided sampling.*\n\nThank you for highlighting the potential benefit of beta diffusion. In our future work, we plan to explore beta diffusion with label/text guidance, with/without CFG, and compare its performance to the principled solution for guided Gaussian diffusion proposed by [Lou & Ermon, April 2023].\n\n>*F2: Is it possible to extend this framework to the exponential family?*\n\nYes, we are actively working on the extension to infinitely divisible exponential family distributions. Notably, we have already addressed the technical aspects for both the exponential and Gamma distributions, two additional members of the exponential family, and are in the process of conducting experiments with them.\n\n> *F3: Latent beta diffusion?*\n\nThanks for the suggestion! We'll explore using sigmoid/tanh activation in the encoder's last layer to provide a bounded latent space for beta diffusion to work on. This will be valuable for scaling beta diffusion to high-resolution images in our future work.\n\n> *E1: Simplified loss for beta KLUB?*\n\nEq. 1 demonstrates that dropping the $\\alpha_t$-related weight coefficients in beta ELBO/KLUB is not straightforward. For Gaussian diffusion, we default to using the SNR weighted -ELBO (i.e., simplified -ELBO). As explained in our response to \"*W2*\", in Gaussian diffusion, KLUB-conditional = -ELBO and KLUB-marginal = simplified -ELBO. Therefore, we can interpret the comparison between KLUB-conditional and KLUB-marginal as analogous to comparing the original loss with the simplified loss.\n\nOur toy experiments indicate that KLUB-marginal often converges faster, and combining both KLUB-conditional and KLUB-marginal yields the best overall performance.\n\n> *E2: Clipping or not?*\n\nWe did not apply clipping in the toy data experiments. When quoting the unconditional CIFAR10 image generation results of Gaussian diffusion from their original papers, we considered the absence of CFG. Additionally, we assumed that if clipping was important, the original models had already applied it.\n\n>*E3: Where beta diffusion are more suitable than Gaussian diffusion?*\n\nGaussian diffusion was initially introduced in 2015 but only gained prominence with DDPM due to its optimized combination of architecture, loss modification, and noise scheduling. We believe that beta diffusion holds promise for image generation, particularly with proper customization and tuning of its network architecture and hyperparameters, and can remove the need for heuristic clipping under CFG. In scenarios where Gaussian-like distributions are unsuitable for modeling, such as many biological datasets, we anticipate that beta diffusion could prove highly advantageous.\n"}, {"review_id": "8xb8fv67Jy", "reviewer": "Reviewer_Npra", "summary": "This work introduces a new type of diffusion model based on the beta distribution. The end of this is to defined a diffusion model on data that lives inside a specified range, i.e. on an interval [a,b].\n\nThis is achieved by first defining a forward noising process conditioned on the initial data sample $x_0$, $q(z_t, z_s | x_0)$ as a bi-variate beta distribution. The marginals of this distribution are also beta distributions, and the conditional $q(z_t | z_s, x_0)$ is also a beta distribution. The distributions are parameterised in such a way that $q(z_0 | x_0) = \\delta_{x_0}$ and $q(z_1 | x_0) = \\delta_0$. Evolving $t$ from $0 \\to 1$ therefore maps between the data distribution and a delta at 0. This process can be seen as a form of continuous diffusion where the noise applied is multiplicative, rather than additive like in Gaussian based diffusions.\n\nThey then propose to reverse this distribution by learning the distribution $p(z_s | z_t) = \\mathbb{E}_{x_0\\sim p_0}[q(z_s | z_t, x_0)] \\approx q(z_s | z_t, f_\\theta(z_t, \\alpha_t))$ and reversing this with small time steps. The learning of this distribution is done by matching $q(z_s | z_t, f_\\theta(z_t, \\alpha_t)) q(z_t)$ to $ \\mathbb{E}[q(z_t, z_s | x_0)] $. This is done via an upper bound on the KL, that can be shown to be minimised when $f_\\theta(z_t, \\alpha_t) = \\mathbb{E}[x_0 | z_t]$. Various other options for KL upper bounds are also proposed to minimise this distance.\n\nThe methodology is demonstrated on a couple of simple 1D settings that show the method works well in these cases. \n\nI think this is a very good paper. I would give it a higher score if the preliminary image experiments were finished and added to the main text. I think they would make the paper much stronger - not because they are images per-se but because they show the method works well in high dimension. Other high dimension tasks would be similarly convincing.\n\n", "questions": "- Equation (9) is the specified short-time reversal distribution for this diffusion, akin to the Normal distribution in a typical discrete time diffusion model. In that setting the small time reverse diffusion isn't actually a normal, but one can show it converges to a normal with a specific form as the time-step goes to zero, and with a reasonable convergence rate, justifying this choice of variational distribution. Do you have a way of justifying the pick of a beta distribution with the specified form as the small time reverse distribution? For large time steps this obviously cannot be true (e.g. $p(z_0|z_1)$ is certainly not beta distributed unless the data itself is).\n- How well does this approach scale with dimension?\n- Have the image experiments been finished, and compared properly to other generative models for images? \n\n[1] Lou, Aaron, and Stefano Ermon. \"Reflected diffusion models.\" arXiv preprint arXiv:2304.04740 (2023).\n\n[2] Fishman, Nic, et al. \"Diffusion Models for Constrained Domains.\" arXiv preprint arXiv:2304.05364 (2023).", "limitations": "The authors do not really discuss the limitations of their work, and I would like to see this better discussed, particularly with respect to other methods on the same topic. For example unlike score-based SDE models one cannot get a likelihood out of these models.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The work is cleanly presented and easy to understand.\n- There is an extensive exploration of the options available for training these models.\n- The work opens up a discussion for other non-additive noise diffusion models, potentially using other infinitly divisible exponential families.\n- The experiments are sufficient to convince me that this method does indeed work, and could potentially be practically useful.\n- The experimental set up is well explained and looks highly reproducible.", "weaknesses": "- I would have liked to see evaluation on some real world data settings to demonstrate the methods utility (in the main paper).\n- [1,2] were publicly available about a month before the submission deadline. I appreciate that is close, but I would expect to see at least a related work discussion to these works, which is not present, as they cover the same range-bound data setting as this work. Ideally we would see a comparison to some of the experimental settings in these works.\n- I was very surprised to discover image experiments in the appendix after nearly having finished writing this review - these should certainly be mentioned in the main paper (were these completed after the submission but before the supplementary deadline?)\n\n[1] Lou, Aaron, and Stefano Ermon. \"Reflected diffusion models.\" arXiv preprint arXiv:2304.04740 (2023).\n\n[2] Fishman, Nic, et al. \"Diffusion Models for Constrained Domains.\" arXiv preprint arXiv:2304.05364 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for providing a comprehensive summary of our paper. We deeply appreciate your positive feedback on the technical details, overall quality, and recognition of the potential for future work. In what follows, we will address any concerns you may have, and we are confident that our response will further emphasize the value and significance of the proposed beta diffusion.\n\n> *W1: No demsontration on real world datasets in the main paper.*\n\nAs the camera-ready version permits an extra page, we will relocate the CIFAR10 results from the Appendix to the main paper. Kindly refer to **Response to All** for the latest updates on CIFAR10.\n\n> *W2: Missing discussion on two recent works by [Lou & Ermon, April 2023] and [Fishman et al., April 2023].*\n\nThank you for bringing these two recent works to our attention. Both of them appeared on ArXiv about a month before the submission deadline. In the camera-ready version, we will include discussions on both papers, as suggested.\n\nLooking ahead, we plan to compare our future extensions of beta diffusion with these works. Specifically, we aim to explore the incorporation of classifier-free guidance (CFG) into logit($z_t$) under beta diffusion and assess its performance in guided image generation, in comparison with [Lou & Ermon, April 2023]. Additionally, we seek to extend the hypercubic constraint, intrinsically satisfied by beta diffusion, to encompass more general manifold constraints, as studied in [Fishman et al., April 2023]. These exciting future endeavors hold great promise.\n\nUpon carefully analyzing both papers, we have noticed several distinctions:\n\n- Both papers attempt to adapt unconstrained Gaussian diffusion to constrained settings, whereas beta diffusion, by design, inherently satisfies the hypercubic constraint.\n\n- While [Lou & Ermon, April 2023] proposes principled methods to handle range-bounded data, and [Fishman et al., April 2023] explores more general manifold constraints (e.g., convex polytope), they are still rooted in the framework of Gaussian diffusion with additive Gaussian noise. By contrast, beta diffusion, with its multiplicative noise, introduces a new foundational hypercubic-constrained diffusion model with potential for further development, such as CFG.\n\n- As of now, there is no evidence of the scalability of [Fishman et al., April 2023] to high dimensions.\n\nFor future research, it would be intriguing to investigate how the techniques from [Lou & Ermon, April 2023] and [Fishman et al., April 2023] can be leveraged to enhance beta diffusion for specific applications or broaden its applicability to more general constraints. Such explorations could yield valuable insights and further advancements in the field.\n\n> *W3: Image experiments were not mentioned in the main paper.*\n\nWe acknowledge that reviewing supplementary material is at the discretion of the reviewers and sincerely appreciate that you took the time to read our supplementary material before submitting your review.\n\nAs indicated in Line 247, more experiments were planned. We had the image generation results on CIFAR10 ready after the main paper deadline. Since these results provided valuable information that directly supported the content of our submission, we believed including them in the Appendix was permissible according to the guidelines outlined in the \"NeurIPS 2023 Call for Papers,\" and therefore, we decided to include them.\n\nSince the initial submission, we have been continuously improving our code by learning from the techniques used by EDM to enhance the training of Gaussian diffusion. Some of these techniques have proven beneficial for beta diffusion as well, resulting in a clear improvement in its FID on CIFAR10. Our latest FID achieved is 3.56. For further details, please refer to **Response to All**.\n\n> *Q1: Justification for Eq.9 under large-step time reversal.*\n\nWe would like to address your question from two perspectives:\n\n- First, the justification of Eq. 9 can be found in Line 150-155, demonstrating that it provides an approximation to the true time reversal distribution, given by $q(z_s | z_t) = \\mathbb E_{x_0\\sim q(x_0|z_t)}[q(z_s|z_t,x_0)]$. The optimization of $\\theta$ results in Eq. 9 becoming equivalent to $p_{\\theta^*}(z_s|z_t) = q(z_s|z_t,\\mathbb E[x_0|z_t])$.\n\n- Second, while Gaussian diffusion uses a small-step time reversal analysis for its justification, we have noted that the specific parameterization utilized in reverse Gaussian diffusion can also be justified by letting $p(z_s|z_t) = q(z_s| z_t, \\hat x_0 = f(z_t,t) )$, as shown in Eq. 9 of [Kingma et al. 2021, Variational Diffusion Models]. This reference, which was already cited in our paper, will be further highlighted to support the use of Eq. 9 in reverse beta diffusion.\n\n> *Q2: How well does this approach scale with dimension?*\n\nOur image experiments indicate that beta diffusion scales with dimension similar to Gaussian diffusion. While we have already tested beta diffusion on CIFAR10, we are also planning to conduct experiments with higher-resolution images. As detailed in Line 541-543 in Appendix F, using a machine equipped with four Nvidia RTX A5000 GPUs, both beta diffusion and Gaussian diffusion (VP-EDM) take approximately 1.46 seconds to process 1000 images of size 32×32×3 during training.\n\n> *Q3: Have the image experiments been finished, and compared properly to other generative models for images?*\n\nWe have continually worked on improving our training code by incorporating specific techniques used by EDM to enhance the training of Gaussian diffusion (DDPM, VP-EDM). However, we have yet to thoroughly explore methods to further enhance the current beta reverse diffusion sampler, which we aim to investigate in future research. For our most recent FID results on CIFAR10 and a comparison with representative non-Gaussian and Gaussian-like diffusion models, kindly refer to **Response to All**."}, {"review_id": "f6HPlj5gNi", "reviewer": "Reviewer_buWN", "summary": "In this submission, the authors introduce a diffusion model for range-bounded data.\nTo do so they introduce a noising process that is multiplicative (and not additive) that leads to conditionals that are beta distributed and converges to a beta distribution which is independent of the original datapoint.\nThey propose to approximate the reverse diffusion by learning the denoised value of $x_0$ given the noised data $z_t$ and the level of noise $\\alpha_t$.\nThey then propose an upper bound on the KL divergence between the forward process and this backward process.\nThey showcase their model on a range of synthetic 1d datasets.", "questions": "- line 123: This seems to hint at the fact that the proposed diffusion models could be written as the discretisation of an SDE which would admits as stationary distribution a Beta distribution. I'd be curious to know whether the Jacobi process used in [Pavel et al. 2023] is that SDE.\n- line 154: That's because of the specific choice of parametrisation of $p_\\theta(z_s|z_t) = q(z_s|z_t, x_0 = f_\\theta(z_t, \\alpha_t))$ right? It's quite an interesting remark indeed.\n- line 157-158: How samples from marginals $q(z_t)$ are obtained? Does this require discretising the forward process? When is this required? If so what is the added computational cost of this requirement during training?\n- line 160: Why starting with $z_1 = 0$? Isn't $z_1 \\sim B(0, \\eta)$? Does this converge to $\\delta_0$ as $\\eta \\rightarrow 0$? It's worth adding some discussion on this in Section 2.2.\n- line 192: What are the consequences of the multiplicative vs additive noise?\n-line 227: The difference is solely whether one looks at KL[p|q] or KL[q|p] right? I suppose the theoretical advantage of the latter being that the expectation is taken over the forward process which is easy to sample from. I am not sure to understand the claim that 'they can agree if the first argument of the KL divergence bounded by the KLUB is the PDF of a diffused data distribution'. Any idea why this ELBO loss seems to perform slightly worse? Is it because it's mode seeking?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The proposed approach enables modelling distribution of data which is supported on collection of bounded ranges.\n- The beta diffusion process leads to closed form conditionals.\n- Since the Kullback-Leibler divergence between Beta distributions can be available in closed-form, the training loss can be computed efficiently.", "weaknesses": "- I found this submission not so easy to follow at times. For instance, why is a separate Section 2.4 needed?\n- The need to discretised the process to get accurate samples of $q(z_t)$ is a potential drawback of the approach. Is the neural network $f_\\theta$ evaluated at each discretisation step during training?\n- The experiment are limited to synthetic datasets. It would be interesting to tackle binary label data, or categorical labels via stick-breaking construction as in [Pavel et al. 2023].\n\nDirichlet diffusion score model for biological sequence generation, Avdeyev, Pavel and Shi, Chenlai and Tan, Yuhao and Dudnyk, Kseniia and Zhou, Jian, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable questions. We will now provide detailed clarifications.\n\n>*W1: Why Section 2.4 is needed?*\n\nTo implement Algorithm 1 for training beta diffusion, Section 2.4 is not necessary. However, it is needed to demonstrate how discretization functions under beta diffusion. Specifically, Lemma 3 and its proof provide insights into why minimizing Eq. 17 can be seen as an approximate method for minimizing $KL(p_{\\theta}(z_{t_{1:T}}) || q(z_{t_{1:T}}))$. Additionally, Eq. 21 illustrates how the trained beta diffusion is used to generate new data with a Markov chain of length $T$.\n\n> *W2: Is there a need to discretize the process to sample from $q(z_t)$? Is $f_{\\theta}$ evaluated at each discretisation step during training?*\n\nBeta and Gaussian diffusions share similarities in how their forward and reverse diffusion processes are constructed. As illustrated in Algorithms 1 and 2 in Appendix:\n\n- For the forward process, no discretization is required. We only evaluate $f_\\theta$ once at each randomly selected $t$.\n\n- For the reverse process, we perform ancestral sampling to reverse from $t=1$ to $t=0$ in a discrete number of steps, where we need to evaluate $f_\\theta$ at each discrete step.\n\n> *W3: The experiment are limited to synthetic data. Comparison with [Pavel et al. 2023].*\n\nWe included preliminary experimental results on natural images in the Appendix. While we understand that reviewers were not required to read the Appendix, we appreciate your attention to these results and **Response to All**, where we shared new updates on experiments. \n\nWe note that [Pavel Avdeyev et al. 2023] appeared on ArXiv **after** the May 17 submission deadline. We have now carefully reviewed [Pavel Avdeyev et al. 2023] and found distinct differences, kindly refer to **Response to All**.\n\n>*Q1: SDE for beta diffusion?*\n\nWe don't believe that the Jacobi process used in [Pavel Avdeyev et al. 2023] corresponds to our beta diffusion for the following reasons:\n\n- First, in beta diffusion, the forward univariate marginal $q(z_t | x_0)$ follows a beta distribution, while Eq. 6 in [Pavel Avdeyev et al. 2023] demonstrates that $q(z_t | x_0)$ in the Jacobi process involves an infinite sum and appears much more complex.\n\n- Second, the endpoint of beta diffusion is $Beta(0, \\eta)$, a degenerate beta distribution, rather than $Beta(a, b)$ as seen in  [Pavel Avdeyev et al. 2023].\n\nNonetheless, there could be potential connections between the proposed beta diffusion and the Jacobi process under specific parameterizations. If you identify clearer connections, we would appreciate your insights and feedback.\n\n>*Q2: Is the analysis in Line 152-154 due to the specific choice shown in Eq. 9?*\n\nThe specific choice made in Eq. 9 is justifiable by the analysis in Line 152-154. In essence, optimizing $f_{\\theta}(z_t,t)$ aims to make it approach $\\mathbb E[x_0|z_t]$. This rationale supports using the particular formulation in Eq. 9.\n\n>*Q3: line 157-158: does forward process require discretization?*\n\nWhile the forward process is infinitely divisible, as discussed in Line 123-127, we avoid discretization when training. Instead, we can directly transition from $x_0$ to $z_t$ by sampling from a beta distribution, as shown in Eq. 17 and Step 7 in Algorithm 1. Therefore, no discretization is required during training, eliminating any additional discretization cost.\n\nThe same as Gaussian diffusion, discretization is only needed during generation.\n\n>*Q4: Why starting from $z_1=0$?*\n\nFollowing your suggestion, we will add more discussion in Section 2.2:\n\n- By construction, as $t$ approaches 1, $\\alpha_t$ approaches 0, leading to $z_t$ approaching $Beta(0,\\eta)$, a degenerate beta distribution with a point mass at zero.\n\n- In practice, when the shape of beta distribution becomes small, beta random samples will likely take the smallest positive constant supported by the Float32 precision in PyTorch (i.e., MIN = torch.finfo(torch.float32).tiny=1.1754943508222875e-38). Thus, starting from $z_1$= MIN is common when $\\alpha_{t_T}$ is set to be sufficiently close to 0.\n\n>*Q5: Questions on multiplicative vs additive noise,  KL[p||q] vs KL[q||p], KLUB vs ELBO, and why ELBO performs worse for beta diffusion.*\n\nThe forward process of beta diffusion, shown in Eqs. 5 and 19, is multiplicative, while Gaussian diffusion's forward process, illustrated in Eq. 31 in Appendix D, is additive. Our toy data experiments favor the multiplicative construction for modeling distributions with sharp transitions between density regions.\n\nThe key difference lies in whether we use $KL[p||q]$ or $KL[q||p]$: minimizing the KLUB for $KL( q(z_{t_{1:T}}) || p_{\\theta}(z_{t_{1:T}}) )$ corresponds to minimizing -ELBO, while we minimize the KLUB of $KL( p_{\\theta}(z_{t_{1:T}}) || q(z_{t_{1:T}}) )$ for beta diffusion. Notably, this choice doesn't impact Gaussian diffusion (Line 220 and Appendix D).\n\nYou are correct that $KL( q(z_{t_{1:T}}) || p_{\\theta}(z_{t_{1:T}}) )$ is easy to evaluate by sampling from $q(z_{t_{1:T}})$. To evaluate $KL( p_{\\theta}(z_{t_{1:T}}) || q(z_{t_{1:T}}) )$ in beta diffusion, we approximate the sampling of $p_{\\theta}(z_{t_j} | x_0)$ with $q(z_{t_j} | x_0)$ (Lemma 3 proof).\n\nWe note $q(z_{t_{1:T}})$ is the PDF of a diffused data distribution. The statement in Line 232 is trying to say that the KLUB for $KL( q(z_{t_{1:T}}) || p_{\\theta}(z_{t_{1:T}}) )$ in beta diffusion becomes equivalent to terms used in -ELBO for optimization.\n\nThe inferior performance of -ELBO in beta diffusion arises from the lack of a theoretical guarantee that minimizing it results in $f_{\\theta^*}(z_t,t) = \\mathbb E[x_0 | z_t]$ (Eq. 13). This relates to the Bregman divergence property. A comparison of Eqs. 12 and 24 highlights distinctions between the proposed KLUB and -ELBO under the log-beta Bregman divergence. Gaussian diffusion is unaffected by using the KLUBs of $KL[p||q]$ or $KL[q||p]$, kindly refer to our response to \"*W2*\" of Reviewer FX5S."}]

Author rebuttal:
# Response to All

We express our gratitude to all four reviewers for their valuable insights and suggestions. Their comments have been instrumental in identifying relevant recent works on alternative diffusion processes, refining the paper's positioning within the literature, elucidating the properties of KLUBs in relation to (weighted) ELBOs, and enhancing comparison studies. We will diligently incorporate these enhancements into the final version of the paper.

The feedback from the reviewers reflects their overall satisfaction with the technical aspects of the paper. Nevertheless, the absence of certain related works concerning alternative diffusion models has prompted legitimate concerns. This absence has hindered a comprehensive evaluation of the distinct contributions that our proposed beta diffusion model offers to the rapidly evolving landscape of diffusion models. Although we elucidated the differences between beta diffusion and Gaussian diffusion in our paper, we acknowledge that a more comprehensive discussion and comparison with alternative diffusion models would enhance the contextualization of our work.

To address these concerns effectively, we have taken care to discuss and compare to related works on alternative diffusion processes and provide thorough responses to each reviewer. Furthermore, we extend our explanations below to provide additional clarity, helping the reviewers in their assessment of the beta diffusion model's novelty, significance, and potential for generalization. We also emphasize the adaptability of the optimization techniques associated with the beta diffusion model, suggesting their potential applicability to a broader spectrum of exponential family distributions.

1. Following submission, we have refined our code by incorporating design strategies from EDM that enhance the training of Gaussian diffusion. Specifically, we pre-condition logit($z_t$) with its mean and standard deviation before feeding it into the Unet $f_{\theta}$ for CIFAR10 image experiments. This refinement has notably improved the FID score. Our initial reported FID of 6.17 in the Appendix has been lowered to **3.56** due to these improvements. Additionally, we are keen on exploring EDM's techniques to enhance our reverse beta diffusion sampler in future research endeavors.

2. Our analysis in Table 1 within the attached PDF encompasses a broad spectrum of diffusion models, both Gaussian and non-Gaussian. We highlight that beta diffusion now outperforms all non-Gaussian diffusion models on CIFAR10, including Cold Diffusion and Inverse Heat Dispersion, as well as categorical and count-based diffusion models. In comparison to Gaussian diffusion and Gaussian+blurring diffusion, beta diffusion surpasses Soft Diffusion. While it may fall short of DDPM, VP-EDM, and Blurring Diffusion, it remains a competitive alternative that uses non-Gaussian based diffusion.

3. Table 2 within the attached PDF has been introduced to compare KLUB and negative ELBO-optimized beta diffusion across various NFEs. We also include a Figure to visually compare generated images under KLUB and ELBO. The findings presented in Table 2 and the Figure provide further validation of KLUB's efficacy in optimizing beta diffusion.

4. We acknowledge the mention of [Pavel Avdeyev et al., 2023] by two reviewers. As this paper emerged on ArXiv post the May 17 submission deadline, we deem it concurrent work. We kindly request reviewers to reconsider its relevance in evaluating our paper. Furthermore, our analysis of the paper reveals key distinctions:

    - [Pavel Avdeyev et al., 2023] employs the Jacobi diffusion process for discrete data diffusion models. In contrast to Gaussian diffusion's SDE definition, the Jacobi diffusion process follows $dx = \frac{s}{2}[a(1-x)-bx]dt+\sqrt{sx(1-x)}dw$, with $x\in[0,1]$ and $s,a,b>0$. Its stationary distribution is a univariate beta distribution $\mbox{Beta}(a,b)$.

   - While beta diffusion and the Jacobi process are connected by the beta distribution, they differ in various aspects:

       - Beta diffusion concludes its forward process at $\mbox{Beta}(0,\eta)$, a unit point mass at 0, not a $\mbox{Beta}(a,b)$ random variable.

        - The marginal distribution of beta diffusion at time $t$ is denoted as $q(z_t| x_0)\sim \mbox{Beta}(\eta\alpha_t x_0,\eta(1-\alpha_t x_0))$, distinct from the infinite sum in the Jacobi diffusion process. The potential links between beta diffusion and the Jacobi process, under specific parameterizations, warrant further exploration.

    - The Jacobi diffusion process and its stick-breaking-based multivariate extension primarily address discrete data (binary and categorical) in relatively low dimensions. Their adaptability to range-bounded continuous data and scalability to high dimensions remain unestablished. In contrast, beta diffusion effectively models range-bounded continuous data and scales comparably to Gaussian diffusion in higher dimensions.

5.  Two of the reviewers highlighted the promising prospect of integrating classifier-free guidance (CFG) into beta diffusion, without resorting to heuristic clipping. We share the reviewers' enthusiasm regarding this potential advancement, and we believe that beta diffusion is poised to seamlessly integrate CFG by applying it to the logit space that our Unet operates on. We extend our gratitude to the reviewers for pointing out this future research direction, which we are eager to explore and further develop.

Meta-review: The authors develop a method for diffusions that model data in bounded support. Their approach is based on multiplicative transitions derived from the beta distribution. They develop an upper bound on the KL divergence to train this model. The approach is fresh and well-described. All of the reviewers are in agreement that the paper should be accepted.  In the final version, the authors should update their related work especially to include the concurrent piece of work (Dirichlet diffusion score model for biological sequence generation, Avdeyev et al., ICML 2023) that was discussed.
Decision: Accept (poster)