Title: Beta Diffusion
Venue: NeurIPS 2023
Accepted: yes
Keywords: Diffusion models

KL-divergence upper bounds

multiplicative transitions

scaled and shifted beta distributions
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

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.

1 Introduction

Diffusion-based deep generative models have been gaining traction recently. One representative example is Gaussian diffusion [57, 59, 23, 60, 35] that uses a Gaussian Markov chain to gradually diffuse images into Gaussian noise for training. The learned reverse diffusion process, defined by a Gaussian Markov chain in reverse order, iteratively refines noisy inputs towards clean photorealistic images. Gaussian diffusion can also be viewed from the lens of denoising score matching [28, 64, 59, 60] and stochastic differential equations [61]. They have shown remarkable success across a wide range of tasks, including but not limited to generating, restoring, and editing images [14, 24, 50, 53, 54, 70, 67], transforming 2D to 3D [49, 1], synthesizing audio [9, 37, 73], reinforcement learning [31, 66, 48], quantifying uncertainty [19], and designing drugs and proteins [56, 44, 32].

Constructing a diffusion-based generative model often follows a general recipe [57, 23, 35, 34]. The recipe involves three basic steps: First, defining a forward diffusion process that introduces noise into the data and corrupts it with decreasing signal-to-noise ratio (SNR) as time progresses from 0 to 1. Second, defining a reverse diffusion process that denoises the corrupted data as time reverses from 1 to 0. Third, discretizing the time interval from 0 to 1 into a finite number of intervals, and viewing the discretized forward and reverse processes as a fixed inference network and a learnable generator, respectively. Auto-encoding variational inference [36, 51] is then applied to optimize the parameters of the generator by minimizing a weighted negative ELBO that includes a Kullback-Leibler (KL) divergence-based loss term for each discretized reverse step.

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 as the analytic form of the conditional posterior for any earlier time given both a data observation and its corrupted version at the present time. The latter requirement, according to 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, 2, 18, 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 diffusion processes, this paper introduces beta diffusion as a novel addition to the family of diffusion-based generative models. Beta diffusion is specifically designed to generate data within bounded ranges. Its forward diffusion process is defined by the application of beta distributions in a 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 a data observation remains a beta distribution. We illustrate the forward beta diffusion process in Figure 1, which simultaneously 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 a 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 time-dependent KL loss term of the negative ELBO, examining it 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. In addition, there exists another Bregman divergence that upper bounds the KL divergence from the univariate marginal of a corrupted observation in the reverse chain to that in the forward chain. These two Bregman divergences, which can be derived either through KL 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 or counteracting accumulation errors over the course of reverse

diffusion. We further demonstrate that combining these two KLUBs presents a 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 a forward diffusion chain to that of its reverse. Either KLUB 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 KL 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 symmetric, resulting in KLUBs that are equivalent to (weighted) negative ELBOs.

Our main contributions are the introduction of beta diffusion as a 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 ELBOs. Additionally, we introduce the log-beta divergence, a Bregman divergence corresponding to the differentiable and strictly convex log-beta function, as a useful tool for analyzing KLUBs. These contributions enhance the existing family of diffusion-based generative models and provide a new perspective on optimizing them.

2 Beta Diffusion and Optimization via KLUBs

We begin by specifying the general requirements for constructing a diffusion-based generative model and establish the notation accordingly [57, 23, 35]. Let x 0 denote the observed data, and let z s and z t represent their corrupted versions at time s and time t , respectively, where 0 < s < t < 1 . In the forward diffusion process, we require access to random samples from the marginal distribution q ( z t | x 0 ) at any time t , as well as an analytical expression of the probability density function (PDF) of the conditional distribution q ( z s | z t , x 0 ) for any s < t .

The forward beta diffusion chain uses diffusion scheduling parameters α t to control the decay of its expected value over the course of forward diffusion, given by E [ z t | x 0 ] = α t x 0 , and a positive concentration parameter η to control the tightness of the diffusion process around its expected value. We typically set α t to approach 1 and 0 as t approaches 0 and 1, respectively, and satisfy the condition

Let Γ( · ) denote the gamma function and B ( · , · ) denote the beta function. The beta distribution Beta ( x ; a, b ) = B ( a, b ) -1 x a -1 (1 -x ) b -1 is a member of the exponential family [6, 65, 7]. Its log partition function is a log-beta function as ln B ( a, b ) = ln Γ( a ) + ln Γ( b ) -ln Γ( a + b ) , which is differentiable, and strictly convex on (0 , ∞ ) 2 as a function of two variables [15]. As a result, the KL divergence between two beta distributions can be expressed as the Bregman divergence associated with the log-beta function. Specifically, as in Appendix A, one can show by their definitions that

We refer to the above Bregman divergence as the log-beta divergence. Moreover, if ( α q , β q ) are random variables, applying Proposition 1 of Banerjee et al. [4], we can conclude that the optimal value of ( α p , β p ) that minimizes this log-beta divergence is ( α ∗ p , β ∗ p ) = E [( α q , β q )] .

Next, we introduce a conditional bivariate beta distribution, which given a 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 1 (Conditional Beta Bivariate Distribution) . Denote ( z s , z t ) as variables over a pair of time points ( s, t ) , with 0 < s < t < 1 . Given a random sample x 0 ∈ (0 , 1) from a probability-valued data distribution p data ( x 0 ) , we define a conditional bivariate beta distribution over ( z s , z t ) with PDF:

Marginals: Given x 0 , the two univariate marginals of this distribution are both beta distributed as

Conditionals: Given x 0 , a random sample ( z t , z s ) from this distribution can be either generated in forward order, by multiplying a beta variable from (3) with beta variable π s → t , as

or generated in reverse order, by combining a beta variable from (4) with beta variable p s ← t , as

The proof starts with applying change of variables to obtain two scaled and shifted beta distributions

and then takes the products of them with their corresponding marginals, given by (3) and (4), 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 q ( z s , z t | x 0 ) defined in (2). The detailed proof is provided in Appendix B. To ensure numerical accuracy, we will calculate z s = z t +(1 -z t ) p s ← t in (6) in the logit space as

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 t for a given data observation x 0 , as shown in (4), stays as beta-distributed in the forward chain. For the beta distribution given by (4), we have

Thus when α t approaches 0 ( i.e. , t → 1 ), both z t and SNR t are shrunk towards 0, and if α 1 = 0 , we have z 1 ∼ Beta (0 , η ) , a degenerate beta distribution that becomes a unit point mass at 0.

Infinite Divisibility. We consider beta diffusion as a form of continuous diffusion, as its forward chain is infinitely divisible given x 0 . This means that for any time k ∈ ( s, t ) , we can perform forward diffusion from z s to z t by first setting z k = z s π s → k , where π s → k ∼ Beta ( ηα k x 0 , η ( α k -α s ) x 0 ) , and then setting z t = z k π k → t , where π k → t ∼ Beta ( ηα t x 0 , η ( α t -α k ) x 0 ) . The same approach can be used to show the infinite divisibility of reverse beta diffusion given x 0 .

Reverse Beta Diffusion. We follow Gaussian diffusion to use q ( z s | z t , x 0 ) to help define p θ ( z s | z t ) [58, 35, 72]. To construct a reverse beta diffusion chain, we will first need to learn how to reverse from z t to z s , where s < t . If we know the x 0 used to sample z t as in (3), then we can readily apply the conditional in (8) to sample z s . Since this information is unavailable during inference, we make a weaker assumption that we can exactly sample from z t ∼ q ( z t ) = E x 0 ∼ p data [ q ( z t | x 0 )] , which is the ' x 0 -free' univariate marginal at time t . It is straightforward to sample during training but must be approximated during inference. Under this weaker assumption on q ( z t ) , utilizing (8) but replacing its true x 0 with an approximation ˆ x 0 = f θ ( z t , t ) , where f θ denotes the learned generator parameterized by θ , we introduce our ' x 0 -free' and hence 'causal' time-reversal distribution as

2.3 Optimization via KLUBs and Log-Beta Divergence

KLUB for Time Reversal. The time-reversal distribution p θ ( z s | z t ) reaches its optimal when its product with q ( z t ) becomes equivalent to q ( z s , z t ) = E x 0 ∼ p data [ q ( z s , z t | x 0 )] , which is a marginal bivariate distribution that is ' x 0 -free.' Thus we propose to optimize θ by minimizing KL ( p θ ( z s | z t ) q ( z t ) || q ( z s , z t )) in theory but introduce a surrogate loss in practice:

Lemma 2 (KLUB (conditional)) . The KL divergence from q ( z s , z t ) to p θ ( z s | z t ) q ( z t ) , two ' x 0 -free' joint distributions defined by forward and reverse diffusions, respectively, can be upper bounded:

The proof in Appendix B utilizes the equation p θ ( z s | z t ) q ( z t ) = E x 0 ∼ p data [ p θ ( z s | z t ) q ( z t | x 0 )] and then applies the convexity of the KL divergence [11, 74] and the definition in (10).

Log-Beta Divergence. To find out the optimal solution under KLUB, following (1), we can express KLUB ( s, z t , x 0 ) given by (12) as a log-beta divergence as

We note KLUB s,t defined in (11) can also be written as E z t ∼ q ( z t ) E x 0 ∼ q ( x 0 | z t ) [ KLUB ( s, z t , x 0 )] , where the log-beta divergence for KLUB ( s, z t , x 0 ) , defined as in (13), includes x 0 ∼ q ( x 0 | z t ) in its first argument and the generator f θ ( z t , t ) in its second argument. Therefore, applying Proposition 1 of Banerjee et al. [4], we have the following Lemma.

Lemma 3. The objective KLUB s,t defined in (11) for any s < t is minimized when

Thus under the KLUB s,t -optimized θ ∗ , we have p θ ∗ ( z s | z t ) = q ( z s | z t , E [ x 0 | z t ]) , which is different from the optimal solution of the original KL loss in (11), which is p ∗ θ ( z s | z t ) = q ( z s , z t ) /q ( z t ) = q ( z s | z t ) = E x 0 ∼ q ( x 0 | z t ) [ q ( z s | z t , x 0 )] . 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 p θ ∗ ( z s | z t ) and q ( z s | z t ) and hence s needs to be close to t . Furthermore, the assumption of having access to unbiased samples from the true marginal q ( z t ) is also rarely met. Thus we need to discretize the time from 1 to t into sufficiently fine intervals and perform time-reversal sampling over these intervals. Specifically, we can start with z 1 = 0 and iterate (10) over these intervals to obtain an approximate sample from z t ∼ q ( z t ) . However, the error could accumulate along the way from z 1 to z t , to which we present a solution below.

KLUB for Error Accumulation Control. To counteract error accumulation during time reversal, we propose to approximate the true marginal q ( z ′ t ) using a 'distribution-cycle-consistency' approach. This involves feeding a random sample z t from q ( z t ) into the generator f θ , followed by the forward marginal q ( z ′ t | ˆ x 0 = f θ ( z t , t )) , with the aim of recovering the distribution q ( z ′ t ) itself. Specifically, we propose to approximate q ( z ′ t ) with p θ ( z ′ t ) := E z t ∼ q ( z t ) [ q ( z ′ t | ˆ x 0 = f θ ( z t , t ))] by minimizing KL ( p θ ( z ′ t ) || q ( z ′ t )) in theory, but introducing a surrogate loss in practice:

Lemma 4 (KLUB (marginal)) . The KL divergence KL ( p θ ( z ′ t ) || q ( z ′ t )) can be upper bounded:

The proof in Appendix B utilizes the fact that q ( z ′ t ) = E ( z t ,x 0 ) ∼ q ( z t | x 0 ) p data ( x 0 ) [ q ( z ′ t | x 0 )] .

Note that the mathematical definition of KLUB is reused throughout the paper and can refer to any of the equations (11), (12), (15), or (16) depending on the context. Similar to previous analysis, we have

KLUB ( z t , x 0 ) = D ln B ( a,b ) [ ηα t x 0 , η (1 -α t x 0 )] , [ ηα t f θ ( z t , t ) , η (1 -α t f θ ( z t , t ))] , (17) and can conclude with the following Lemma.

Lemma 5. KLUB t 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 f θ via stochastic gradient descent (SGD). Specifically, denoting ω, π ∈ [0 , 1] as two weight coefficients, the loss term for the i th data observation x ( i ) 0 in a mini-batch can be computed as

2.4 Discretized Beta Diffusion for Generation of Range-Bounded Data

For generating range-bounded data, we discretize the beta diffusion chain. Denote z t 0 = 1 and let t j increase with j . Repeating (5) for T times, we define a discretized forward beta diffusion chain:

A notable feature of (19) is that the marginal at any discrete time step t j follows a beta distribution, similarly defined as in (4). We also note while q ( z t 1: T | x 0 ) defines a Markov chain, the marginal

in general does not. Unlike in beta diffusion, where the transitions between z t and z t -1 are applied multiplicatively, in Gaussian diffusion, the transitions between z t and z t -1 are related to each other additively and z t 1: T forms a Markov chain regardless of whether x 0 is marginalized out.

The discretized forward beta diffusion chain given by (19) is reversible assuming knowing x 0 . This means given x 0 , it can be equivalently sampled in reverse order by first sampling z t T ∼ q ( z t T | x 0 ) = Beta ( ηα t T x 0 , η (1 -α t T x 0 )) and then repeating (8) for t T , . . . , t 2 , with PDF q ( z t 1: T | x 0 ) = q ( z t T | x 0 ) ∏ T j =2 1 1 -z t j Beta ( z t j -1 -z t j 1 -z t j ; η ( α t j -1 -α t j ) x 0 , η (1 -α t j -1 x 0 ) ) . This non-causal chain, while not useful by itself, serves as a blueprint for approximate generation.

Specifically, we approximate the marginal given by (20) with a Markov chain in reverse order as

To start the reverse process, we choose to approximate q ( z t T ) = E x 0 ∼ p data ( x 0 ) [ q ( z t T | x 0 )] with p prior ( z t T ) = q ( z t T | E [ x 0 ]) , which means we let z t T ∼ Beta ( ηα t T E [ x 0 ] , η (1 -α t T E [ x 0 ])) . To sample z t 1: T -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 ' x 0 -free' joint distribution over all T steps of the discretized reverse beta diffusion chain, expressed as p θ ( z t 1: T ) , to approach that of the discretized forward chain, expressed as q ( z t 1: T ) . Thus an optimized θ is desired to minimize the KL divergence KL ( p θ ( z t 1: T ) || q ( z t 1: T )) . While this KL loss is in general intractable to compute, it can also be bounded using the KLUB shown as follows.

Lemma 6 (KLUB for discretized diffusion chain) . KL ( p θ ( z t 1: T ) || q ( z t 1: T )) is upper bounded by

We provide the proof in Appendix B. We note Lemma 6 is a general statement applicable for any diffusion models with a discrete forward chain q ( z t 1: T | x 0 ) = ∏ T j =1 q ( z t j | z t j -1 , x 0 ) and a discrete reverse chain p θ ( z t 1: T ) = p prior ( z t T ) ∏ T j =2 p θ ( z t j -1 | z t j ) . To estimate the KLUB in (22), however, during training, one would need to sample z t j ∼ p θ ( z t j | x 0 ) ∝ p ( x 0 | z t j ) p θ ( z t j ) , which is often infeasible. If we replace z t j ∼ p θ ( z t j | x 0 ) with z t j ∼ q ( z t j | x 0 ) , then ˜ KLUB s,t becomes the same as KLUB s,t given by (11), and KLUB t given by (15) can be considered to remedy the impact of approximating p θ ( z t j ) with q ( z t j ) . Therefore, we can consider the combination of KLUB s,t given by (11) and KLUB t given by (15) as a computationally viable solution to compute ˜ KLUB s,t , which hence justifies the use of the loss in (18) to optimize the discretized reverse beta diffusion chain. We summarize the training and sampling algorithms of beta diffusion in Algorithms 1 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), f θ can also be optimized using two argument-swapped KL divergences: KL ( q ( z t , z s ) || p θ ( z s | z t ) q ( z t )) and KL ( q ( z ′ t ) || p θ ( z ′ t )) . By the same analysis, these KL divergences can also be bounded by KLUBs 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 (14) in beta diffusion and are found to provide clearly inferior empirical performance.

Negative ELBO 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

where the first two terms are often ignored and the focus is placed on the remaining T -2 terms as

Lemma 7 ( -ELBO and argument-swapped KLUB) . Optimizing the generator f θ with -ELBO is equivalent to using an upper-bound for the augment-swapped KL divergence KL ( q ( z t 1: T ) || p θ ( z t 1: T )) .

The proof in Appendix B relies on the convex nature of both the KL divergence and the negative logarithmic function. We find for beta diffusion, optimizing KL ( p θ ( z t 1: T ) || q ( z t 1: T )) via the proposed KLUBs is clearly preferred to optimizing KL ( q ( z t 1: T ) || p θ ( z t 1: T )) via -ELBOs ( i.e., augmentswapped KLUBs) and leads to stable and satisfactory performance.

KLUBs and (Weighted) Negative ELBOs for Gaussian Diffusion. We note KLUBs are directly applicable to 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 argument-swapped 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 ELBO [23, 35, 20], 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 in both forward and reverse sampling. Gaussian diffusion starts at N (0 , 1) in its reverse process, while both Poisson and beta diffusion start at 0 . Beta diffusion's reverse sampling is a monotonically non-decreasing process, similar to Poisson diffusion, but while Poisson diffusion takes 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, 7].

Several recent works have explored alternative diffusion processes closely related to Gaussian diffusion. Cold diffusion by Bansal et al. [5] builds models around arbitrary image transformations instead of Gaussian corruption, but it still relies on L 1 loss, resembling Gaussian diffusion's squared Euclidean distance. Rissanen et al. [52] propose an inverse heat dispersion-based diffusion process that reverses the heat equation using inductive biases in Gaussian diffusion-like models. Soft diffusion by Daras et al. [12] 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-based reverse diffusion for generation. By contrast, beta diffusion is distinct from all of them in forward diffusion, training loss, and reverse diffusion.

Aconcurrent work by Avdeyev et al. [3] utilizes the Jacobi diffusion process for discrete data diffusion models. Unlike Gaussian diffusion's SDE definition, the Jacobi diffusion process in Avdeyev et al. [3] is defined by the SDE dx = s 2 [ a (1 -x ) -bx ] dt + √ sx (1 -x ) dw , with x ∈ [0 , 1] and s, a, b > 0 . The stationary distribution is 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 (0 , η ) , a unit point mass at 0, not a Beta ( a, b ) random variable. The marginal distribution of beta diffusion at time t is expressed as q ( z t | x 0 ) ∼ Beta ( ηα t x 0 , η (1 -α t x 0 )) , 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.

Several recent studies have been actively exploring the adaptation of diffusion models to constrained scenarios [13, 40, 41, 17], where data is bounded within specific ranges or constrained to particular 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 hypercubic-constrained 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 [22, 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 f θ network, thereby eliminating the necessity for heuristic clipping.

To adapt Gaussian diffusion for high-resolution images, a 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 to process 1000 images of size 32 × 32 × 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]. It is worth exploring how to adapt these methods to 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, 78, 69], distilling the reverse diffusion chains [43, 55, 76], utilizing reinforcement learning [16], and transforming the SDE associated with Gaussian diffusion into an ODE, followed by fast ODE solvers [58, 39, 42, 75, 34]. Given these existing acceleration techniques for Gaussian diffusion, it is worth exploring their generalization to enhance the sampling efficiency of beta diffusion.

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 we can leverage these models for the betterment of society while mitigating any potential negative consequences.

4 Experiments

The training and sampling algorithms for beta diffusion are described in detail in Algorithms 1 and 2, respectively, in the Appendix. Our experiments, conducted on two synthetic data and the CIFAR10 images, primarily aim to showcase beta diffusion's effectiveness in generating range-bounded data. We also underscore the superiority of KLUBs over negative ELBOs as effective optimization objectives for optimizing beta diffusion. Additionally, we highlight the differences between beta and Gaussian diffusion, specifically in whether the data are generated through additive or multiplicative transforms and their ability to model the mixture of range-bounded distributions with disjoint supports.

We compare the performance of 'Gauss ELBO,' 'Beta ELBO,' and 'Beta KLUB,' which respectively correspond to a Gaussian diffusion model optimized with the SNR weighted negative ELBO [23, 35], a beta diffusion model optimized with the proposed KLUB loss defined in (18) but with the two arguments inside each KL term swapped, and a beta diffusion model optimized with the proposed KLUB loss defined in (18). On CIFAR-10, we also evaluate beta diffusion alongside a range of non-Gaussian or Gaussian-like diffusion models for comparison.

4.1 Synthetic Data

We consider a discrete distribution that consists of an equal mixture of five unit point masses located at x 0 ∈ D = { 1 / 7 , 2 / 7 , 3 / 7 , 4 / 7 , 5 / 7 } . We would like to highlight that a unit point mass can also be seen as an extreme case of range-bounded data, where the range is zero. Despite being simple,

Gauss_ELBO

Beta_ELBO

0.4

0.4

0.4

Frequency

E'0

Empirical Frequency

0.3

Empirical Frequency

0.3

Empirical F

0.2

0.2

0.2

0.1

0.1

0.1

0.0

0.0

0.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Bins

Bins

Bins

this data could be challenging to model by a continuous distribution, as it would require the generator to concentrate its continuous-valued generations on these five discrete points.

We follow previous works to choose the beta linear diffusion schedule as α t = e -1 2 β d t 2 -β min t , where β d = 19 . 9 and β min = 0 . 1 . This schedule, widely used by Gaussian diffusion [23, 61, 34], is applied consistently across all experiments conducted on synthetic data. We set η = 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 S cale = 1 and S hift = 0 . We use the same structured generator f θ for both Gaussian and beta diffusion. We choose 20-dimensional sinusoidal position embeddings [63], with the positions set as 1000 t . The network is an MLP structured as (21-256)-ReLU-(256-256)-ReLU-(256-1). We utilize the Adam optimizer with a learning rate of 5 e4 and a mini-batch size of 1000.

For data generation, we set NFE = 200 . We provide the generation results in Figure 3, which shows the true probability mass function (PMF) of the discrete distribution and the empirical PMFs over 100 equal-sized bins between 0 and 1. Each empirical PMF is computed based on 100k random data points generated by the model trained after 400k iterations. It is clear from Figure 3 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 E, 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 ω = 1 and ω = 0 , respectively, in the loss given by (18). Additionally, we evaluate beta diffusion on 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 α t = 1 / (1 + e -c 0 -( c 1 -c 0 ) t ) , which has been observed to offer greater flexibility than the beta linear schedule for image generation. This schedule bears resemblance to the sigmoid-based one introduced for Gaussian diffusion [35, 29]. We configure the parameters for beta diffusion as follows: c 0 = 10 , c 1 = -13 , S hift = 0 . 6 , S cale = 0 . 39 , η = 10000 , ω = 0 . 99 , and π = 0 . 95 . We utilize the Adam optimizer with a learning rate of 2 e4 . We use EDM'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.

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, η , and B . A broad spectrum of diffusion models is encompassed in Table 1, which shows that beta diffusion outperforms all non-Gaussian diffusion models on CIFAR10, including Cold Diffusion [5] and Inverse Heat Dispersion [52], as well as categorical and count-based diffusion models [2, 8, 10]. In comparison to Gaussian diffusion and Gaussian+blurring diffusion, beta diffusion surpasses VDM [35], 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 a competitive alternative that uses non-Gaussian based diffusion.

(a)

-

ELBO

Table 2 presents a comparison between KLUB and negative ELBO-optimized beta diffusion across different NFE under two different mini-batch sizes B . Table 3 in the Appendix includes the results under several different combinations of η and B . We also include Figure 4 to visually compare generated images under KLUB and negative ELBO. The findings presented in Tables 2 and 3 and Figure 4 provide further validation of KLUB's efficacy in optimizing beta diffusion.

As each training run takes a long time and FID evaluation is 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 a probability-valued data observation x 0 ∈ (0 , 1) , the distribution at any time point t ∈ [0 , 1] of the forward beta diffusion chain, expressed as q ( z t | x 0 ) , is a beta distribution. 2) Analytical Conditional: Conditioning on a data x 0 and a forward-sampled latent variable z t ∼ q ( z t | x 0 ) , the forward beta diffusion chain can be reversed from time t to the latent variable at any previous time s ∈ [0 , t ) by sampling from an analytic conditional posterior z s ∼ q ( z s | z t , x 0 ) that follows a scaled and shifted beta distribution. 3) KLUBs: We introduce the combination of two different Kullback-Leibler Upper Bounds (KLUBs) for optimization and represent them under the log-beta Bregman divergence, showing that their optimal solutions of the generator are both achieved at f θ ∗ ( z t , t ) = E [ x 0 | z t ] . We also establish the connection between augment-swapped KLUBs and (weighted) negative ELBOs for diffusion models. Our experimental results confirm the distinctive qualities of beta diffusion when applied to generative modeling of range-bounded data spanning disjoint regions or residing in high-dimensional spaces, as well as the effectiveness of KLUBs for optimizing beta diffusion.

References

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.04968 , 2023. URL https://perp-neg. github.io/ . (the first three authors contributed equally).

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.

Pavel Avdeyev, Chenlai Shi, Yuhao Tan, Kseniia Dudnyk, and Jian Zhou. Dirichlet diffusion score model for biological sequence generation. arXiv preprint arXiv:2305.10699 , 2023.

Arindam Banerjee, Srujana Merugu, Inderjit S Dhillon, Joydeep Ghosh, and John Lafferty. Clustering with Bregman divergences. Journal of machine learning research , 6(10), 2005.

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.

O. Barndorff-Nielsen. 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 9781118857373.

David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 112(518):859-877, 2017.

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 Information Processing Systems , 2022. URL https://openreview.net/forum?id=DmT862YAieY .

Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan. WaveGrad: Estimating gradients for waveform generation. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id= NsMLjcFaO8O .

Tianqi Chen and Mingyuan Zhou. Learning to Jump: Thinning and thickening latent counts for generative modeling. In ICML 2023: International Conference on Machine Learning , July 2023. URL http://arxiv.org/abs/2305.18375 .

Thomas M Cover and Joy A Thomas. Elements of information theory. 1991.

Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alex Dimakis, and Peyman Milanfar. Soft diffusion: Score matching with general corruptions. Transactions on Machine Learning Research , 2023. ISSN 2835-8856. URL https://openreview.net/forum?id=W98rebBxlQ .

Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. Advances in Neural Information Processing Systems , 35:2406-2422, 2022.

Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https://openreview.net/forum? id=AAWuCvzaVt .

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.

Ying Fan and Kangwook Lee. Optimizing DDPM sampling with shortcut fine-tuning. arXiv preprint arXiv:2301.13362 , 2023.

Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson. Diffusion models for constrained domains. arXiv preprint arXiv:2304.05364 , 2023.

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.

Xizewen Han, Huangjie Zheng, and Mingyuan Zhou. CARD: Classification and regression diffusion models. In Advances in Neural Information Processing Systems , 2022.

Tiankai 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.09556 , 2023.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In Advances in Neural Information Processing Systems , pages 6626-6637, 2017.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , 2020.

Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res. , 23:47-1, 2022.

Emiel Hoogeboom and Tim Salimans. Blurring diffusion models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum? id=OjDkC57x5sz .

Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems , 34:12454-12465, 2021.

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 Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11502-11511, 2022.

Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(24):695-709, 2005. URL http://jmlr.org/papers/v6/ hyvarinen05a.html .

Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. arXiv preprint arXiv:2212.11972 , 2022.

M. V. Jambunathan. Some properties of beta and gamma distributions. The Annals of mathematical statistics , 25(2):401-405, 1954. ISSN 0003-4851.

Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

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 https://openreview.net/forum?id=w6fj2r62r_H .

N. L. Johnson, A. W. Kemp, and S. Kotz. Univariate Discrete Distributions . John Wiley & Sons, 2005.

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://openreview.net/forum?id=k7FuTOWMOc7 .

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems , 34:21696-21707, 2021.

Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114 , 2013.

Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DiffWave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=a-xFK8Ymz5J .

I. Kotlarski. On groups of n independent random variables whose product follows the beta distribution. Colloquium Mathematicum , 9(2):325-332, 0 1962.

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 https://openreview.net/forum?id=PlKWVd2yBkY .

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 https://openreview.net/forum?id=WH1yCa0TbB .

Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04740 , 2023.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-solver: A fast ODEsolver for diffusion probabilistic model sampling in around 10 steps. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=2uAaGwlP_V .

Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388 , 2021.

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific antibody design and optimization with diffusion-based generative models for protein structures. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/ forum?id=jSorGn2Tjg .

Kevin P Murphy. Machine Learning: A Probabilistic Perspective . MIT Press, 2012.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Kushagra Pandey, Avideep Mukherjee, Piyush Rai, and Abhishek Kumar. Diffusevae: Efficient, controllable and high-fidelity generation from low-dimensional latents. arXiv preprint arXiv:2201.00308 , 2022.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301.10677 , 2023.

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. URL https://openreview.net/forum?id=FjNys5c7VyY .

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning , pages 1278-1286. PMLR, 2014.

Severi Rissanen, Markus Heinonen, and Arno Solin. Generative modelling with inverse heat dissipation. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=4PJUBT9f2Ol .

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

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: 36479-36494, 2022.

Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations , 2022. URL https://openreview. net/forum?id=TIdIXIpzhoI .

Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning , pages 9558-9568. PMLR, 2021.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations , 2020.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems , pages 11918-11930, 2019.

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems , 33:12438-12448, 2020.

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 https://openreview. net/forum?id=PxTIG12RRHS .

Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In Advances in neural information processing systems , 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems , pages 5998-6008, 2017.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Martin J Wainwright and Michael Irwin Jordan. Graphical models, exponential families, and variational inference . Now Publishers Inc, 2008.

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=AHvFDPi-FA .

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:2305.01115 , 2023.

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 https://arxiv.org/abs/ 2304.12526 .

Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion-GAN: Training GANs with diffusion. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id= HZf7UbpWHuA .

Zhixin Wang, Xiaoyun Zhang, Ziying Zhang, Huangjie Zheng, Mingyuan Zhou, Ya Zhang, and Yanfeng Wang. DR2: Diffusion-based robust degradation remover for blind face restoration. In CVPR 2023: Conference on Computer Vision and Pattern Recognition , 2023. URL https: //openreview.net/forum?id=HmcDUPEGA0 .

Zike Wu, Pan Zhou, Kenji Kawaguchi, and Hanwang Zhang. Fast diffusion model. arXiv preprint arXiv:2306.06991 , 2023.

Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion GANs. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=JprM0p-q0Co .

Dongchao Yang, Jianwei Yu, Helin Wang, Wen Wang, Chao Weng, Yuexian Zou, and Dong Yu. Diffsound: Discrete diffusion model for text-to-sound generation. IEEE/ACM Transactions on Audio, Speech, and Language Processing , 31:1720-1733, 2023. doi: 10.1109/TASLP.2023. 3268730.

Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. In International Conference on Machine Learning , pages 5660-5669, 2018.

Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=Loek7hfb46P .

Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference on Machine Learning , pages 42390-42402. PMLR, 2023.

Huangjie Zheng and Mingyuan Zhou. Exploiting chain rule and Bayes' theorem to compare probability distributions. Advances in Neural Information Processing Systems , 34:14993-15006, 2021.

Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Truncated diffusion probabilistic models and diffusion-based adversarial auto-encoders. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum? id=HDxgaKk956l .

Beta Diffusion: Appendix

Algorithm 1 Training of Beta Diffusion

Require: Dataset D whose values are bounded from 0 to 1, Mini-batch size B , concentration parameter η = 10000 , data shifting parameter S hift = 0 . 6 , data scaling parameter S cale = 0 . 39 , generator f θ , loss balance coefficient ω = 0 . 99 , time reversal coefficient π = 0 . 95 , and a decreasing function that returns scheduling parameter α t ∈ (0 , 1) given t ∈ [0 , 1] , such as a beta linear schedule defined by α t = e -0 . 5( β max -β min ) t 2 -β min t , where β max = 20 and β min = 0 . 1 , and a sigmoid schedule defined by α t = 1 / (1 + e -c 0 -( c 1 -c 0 ) t ) , where c 0 = 10 and c 1 = -13 .

1: repeat

2: Draw a mini-batch X 0 = { x ( i ) } B from D

10:

0 i =1 3: for i = 1 to B do ▷ can be run in parallel 4: t i ∼ Unif (1 e5 , 1) 5: s i = πt i 6: Compute α s i and α t i 7: x ( i ) 0 = x ( i ) 0 ∗ S cale + S hift 8: z t i ∼ Beta ( ηα t i x ( i ) 0 , η (1 -α t i x ( i ) 0 )) 9: ˆ x ( i ) 0 = f θ ( z t i , t i ) ∗ S cale + S hift Using (18) to compute the loss as L

or swap the two arguments of both log-beta Bregman divergences in (24) if the loss is (weighted) -ELBO.

11: end for

13: until converge

12: Perform SGD with 1 B ∇ θ ∑ B i =1 L i

Algorithm 2 Sampling of Beta Diffusion

Require: Number of function evaluations (NFE) J = 200 , generator f θ , and timesteps { t j } J j =0 :

3:

else

5: end if

7:

z

t

∼

Beta

(

ηα

t

ˆ

x

0

, η

(1

-

α

t

J

ˆ

x

0

))

▷ which is implemented in the logit space as

▷ which is implemented in the logit space as

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 ln B ( a, b ) = ln Γ( a ) + ln Γ( b ) -ln Γ( a + b ) , which is differentiable, and strictly convex on (0 , ∞ ) 2 as a function of a and b , can be expressed as

which is equivalent to the analytic expression of

B Proof

Proof of Lemma 1. The joint distribution of z t and z s in (5) can be expressed as

The joint distribution of z t and z s in (6) can be expressed as

The joint distribution shown in (25) is the same as that in (26).

Proof of Lemma 2. Since we can re-express q ( z t ) and q ( z s , z t ) as

and q ( z s , z t | x 0 ) = q ( z s | z t , x 0 ) q ( z t | x 0 ) , using the convexity of the KL divergence, we have

Proof of Lemma 4. Since we can re-express p θ ( z ′ t ) and q ( z ′ t ) as

using the convexity of the KL divergence, we have

KL ( p θ ( z ′ t ) || q ( z ′ t ))

Proof of Lemma 6. The proof utilizes the convexity of the KL divergence to show that

We note p θ ( z t j | x 0 ) ∝ p ( x 0 | z t j ) p θ ( z t j ) , which in general is intractable to sample from, motivating us to replace p θ ( z t j | x 0 ) with q ( z t j | x 0 ) for tractable computation.

Proof of Lemma 7. Utilizing the convexity of the KL divergence, we present an upper-bound of the augmented-swapped KL divergence KL ( q ( z t 1: T ) || p θ ( z t 1: T )) , referred to as AS-KLUB, as

Utilizing the convexity of the negative logarithmic function, we have

Thus, when we disregard its first term, AS-KLUB is equivalent to -ELBO without its first two terms. Since these three terms typically do not affect optimization, optimizing the generator f θ with AS-KLUB is equivalent to using -ELBO.

C Derivation of (Weighted) ELBOs of Gaussian Diffusion from KLUB

Let us denote α 0 = 1 and z 0 = x 0 . Following the definition in Ho et al. [23] and Song et al. [58], we define a Gaussian diffusion-based generative model as

where the diffusion scheduling parameters α t in this paper is related to β t in Ho et al. [23] as

The same as the derivations in Ho et al. [23], we can express the forward marginal distribution at time t as

Assuming E [ x 0 ] = 0 and var [ x 0 ] = 1 , the signal-to-noise ratio at time t is defined as

Since √ α t -1 α t x t ∼ N ( z t -1 , α t -1 α t -1) and z t -1 ∼ N ( √ α t -1 z 0 , 1 -α t -1 ) , using the conjugacy of the Gaussian distributions with respect to their mean, we can express the conditional posterior as q ( z t -1 | x 0 , z t )

The forward Gaussian diffusion chain

can be equivalently sampled in reverse order as

As this reverse chain is non-causal and non-Markovian, we need to approximate it with a Markov chain during inference, which is expressed as

The usual strategy [58, 35, 72] is to utilize the conditional posterior to define

In what follows, we redefine the time t as a continuous variable between 0 and 1. We let t 1: T be a set of T discrete time points within that interval, with 0 ≤ t 1 < t 2 < . . . < t T ≤ 1 . The corrupted data observations over these T time points are defined as z t 1: T .

C.1 Optimization of Gaussian Diffusion via Negative ELBO

Viewing p θ ( x 0 | z 1: T ) as the decoder, p θ ( z 1: T ) as the prior, and q ( z 1: T | x 0 ) as the inference network, we can optimize the generator parameter via the negative ELBO as

where the first two terms are often ignored and the focus is placed on the remaining T -2 terms, defined as

where 0 < s < t < 1 . We choose s = max( t -1 /T, 0) during training.

Since z t = √ α t x 0 + √ 1 -α t ϵ t , ϵ t ∼ N (0 , 1) , it is true that

Instead of directly predicting x 0 given x t , we can equivalently predict ˆ ϵ t = ϵ θ ( z t , t ) given z t and let

Thus (27) can also be written as

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 KLUB (Conditional), for Gaussian diffusion we have

For two Gaussian distributions q 1 and q 2 that have the same variance, we have KL ( q 1 || q 2 ) = KL ( q 2 || q 1 ) and hence

Therefore, over the same set of time points t 1: T , optimizing via KLUB conditional is identical to optimizing via -ELBO. As analyzed before, optimizing via KLUB conditional ( i.e. , -ELBO) may not be able to directly counteract the error accumulated over the course of diffusion, and hence could lead to slow convergence.

C.3 Optimization of Gaussian Diffusion via KLUB Marginal and SNR-weighted Negative ELBO

Following the definition of KLUB (Marginal), for Gaussian diffusion we have

It is surprising to find out that the KLUB Marginal is identical to the SNR weighted -ELBO first introduced in Ho et al. [23] and further discussed in Hang et al. [20].

C.4 Optimization of Gaussian Diffusion via KLUB

Following beta diffusion, we can also combine two KLUBs as

Since α s 1 -α s 1 -α t α t is in general close to 1 when s is not too far from t , when ω is not too close to 1, a combination of these two KLUBs does not result in an algorithm that clearly differs from the use of an SNR-weighted negative ELBO.

D Illustration of Forward and Reverse Beta Diffusion

Illustration of Forward Beta Diffusion. We first visualize the beta forward diffusion process by displaying a true image x 0 and its noise-corrupted versions over the course of the forward diffusion process. Specifically, we display the noise corrupted and masked image at time t = 0 , 0 . 05 , 0 . 1 . . . , 1 as

It is clear from Figure 1 that with multiplicative beta-distributed 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 as 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 at a 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, as illustrated below.

Illustration of Reverse Beta Diffusion. We further illustrate the reverse process of beta diffusion by displaying

where z t j -1 = z t j +(1 -z t j ) p ( t j -1 ← t j ) is iteratively computed as in Algorithm 2. We also display

where the approximation would become more and more accurate when θ approaches its theoretical optimal θ ∗ , under which we have f θ ∗ ( z t j , t j ) = E [ x 0 | z t j ] .

As shown in Figure 2, starting from a random image drawn from

most of whose pixels would be completely dark, beta diffusion gradually demasks and denoises the image towards a clean image through multiplicative transforms, as shown in Algorithm 2.

Statisticaldistancebetweenthetrueandgenerateddatadistributions

Jensen-Shannondivergence

0.6

Wassersteinldistance

GauSS -ELBO

0.8

0.100

Beta-ELBO

Hellinger distance

Beta-KLUB

0.4

0.6

0.075

0.050

0.4

0.2

0.025

0.2

0.000

0.0

0.0

103

104

105

103

104

105

103

104

105

Iterations

lterations

Iterations

Statisticaldistancebetweenthetrueandgenerateddatadistributions

0.08

Jensen-Shannondivergence

0.6

Wassersteinldistance

0.06

Beta-KLuB_marginal

Beta-KLUB_conditional

Hellinger distance

0.8

BetaKLUB

0.4

0.6

0.04

0.4

0.2

0.02

0.2

0.00

0.0

0.0

103

104

105

103

104

105

103

104

105

Iterations

Iterations

Iterations

E Additional Results for Synthetic Data

Quantitative Evaluation Metrics. We consider Wasserstein-1, Jensen-Shannon divergence (JSD), and Hellinger distance as three complimentary evaluation metrics. Wasserstein-1 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 are no data. By contrast, both JSD and Hellinger can well reflect the misalignments between the supports of high data density regions.

Wasserstein-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 at random from a discrete set D , then we compute the Wasserstein-1 distance by sorting a set of 10k generated samples from small to large and calculating the mean of their absolute element-wise differences with a non-decreasing vector of the same size, which consists of an equal number of copies for each unique value in D .

JSD and Hellinger distance: We discretize the 100k generated data into 100 equal-sized bins between 0 and 1 and compute the frequency of the data in each bin, which provides an empirical probability mass function (PMF). We then compute both the JSD 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, it 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-swapped version leads to good 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 a much better recovery of the true underlying density in comparison to its argument-swapped version, as confirmed by examining the first subplot in Figure 5 and comparing the second and third subplots in Figure 3.

Gauss_ELBO

Beta_ELBO

0.25

0.25

0.25

Frequency

0.20

Frequency

0.20

Frequency

0.20

0.15

0.15

0.15

ical

ical

ical

0.10

0.10

0.10

Em

品

0.05

Em

0.05

id

E

0.05

0.00

0.00

0.00

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Bins

Bins

Bins

Statisticaldistancebetweenthetrueandgenerateddatadistributions

Jensen-Shannondivergence

Wasserstein1distance

0.10

GauSs-ELBO

0.4

Beta-ELBO

Hellinger distance

0.08

Beta-KLUB

0.3

0.6

0.06

0.2

0.4

0.04

0.1

0.02

0.2

0.00

0.0

103

104

105

103

104

105

103

104

105

Iterations

lterations

Iterations

Statisticaldistancebetweenthetrueandgenerateddatadistributions

0.10

Jensen-Shannondivergence

Wasserstein1 distance

Beta-KLUBconditional

0.4

0.08

Beta-KLUB_marginal

Hellinger distance

0.6

Beta-KLUB

0.3

0.06

0.2

0.4

0.04

0.02

0.1

0.2

0.00

0.0

103

104

105

103

104

105

103

104

105

Iterations

Iterations

Iterations

We further conduct an ablation study between KLUB and its two variants: 'KLUB Conditional' and 'KLUB Marginal,' corresponding to ω = 1 and ω = 0 , respectively, in the loss given by (18). The results suggest 'KLUB Conditional' and 'KLUB Marginal' have comparable performance, which is not that surprising considering that in 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' in aligning the generation with the data supports, but eventually delivers comparable performance, and combining them leads to the 'KLUB' that has a 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 is consistently better than or comparable to the other methods across various combinations of model parameters. While the toy data is 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 an equal mixture of three range-bounded distributions and two unit point masses, including a Uniform distribution between 0.1 and 0.2, a Beta(1,5) distribution first scaled by 0.1 and then shifted by 0.3, a unit point mass at x 0 = 0 . 5 , a Beta(0.5,0.5) distribution first scaled by 0.1 and

then shifted by 0.6, and a unit point mass at x 0 = 0 . 8 . A random sample is generated as

We follow the same experimental protocol in Section 4.1. The results are shown in Figures 7 and 8, from which we again observe that 'Gauss ELBO' does a poor job in aligning its generation with the true data supports and places a 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 a good practice to precondition the input of a neural network. For beta diffusion, most of our computation is operated in the logit space and hence we will consider how to precondition logit ( z t ) before feeding it as the input to the generator f θ . Specifically, since we have

we can draw the logit of z t as

where

We further estimate the variance of logit ( z t ) as

where

Now we are ready to precondition logit ( z t ) as

We use g ( z t ) as the input to the generator f θ and add a skip connection layer to add g ( z t ) into the output of the generator. Specifically, following the notation of EDM, we define the network as

where for simplicity, we set c skip ( t ) = c out ( t ) = c in ( t ) = 1 and c noise ( t ) = -logit ( α t ) / 8 . A more sophisticated selection of these parameters has the potential to enhance the performance of beta diffusion. However, due to our current limitations in computing resources, we defer this investigation to future studies.

G Parameter Settings of Beta Diffusion on CIFAR10

We present the intuition on how we set the model parameters, including S hift , S cale , η, ω, π, c 0 , and c 1 . Given the limitation of our computation resources, we leave the careful tuning of these model parameters to future work.

We set ω = 0 . 99 and π = 0 . 95 across all experiments conducted on CIFAR10 with beta diffusion.

We pre-process the range-bounded data by linearly scaling and shifting them to lie between [ S hift , S hift + S cale ] , where 0 < S hift < S hift + S cale < 1 . This linear transformation serves two purposes: firstly, it enables us to use beta diffusion to model any range-bounded data, and secondly, it 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 CIFAR10 images, as beta diffusion is diffusing pixels towards Beta (0 , η ) , our intuition is to set both S cale and S hift large enough to differentiate the diffusion trajectories of different pixel values. However, S cale + S hift needs to be smaller than 1, motivating us to choose S cale = 0 . 39 and S hift = 0 . 60 .

For the diffusion concentration parameter η , our intuition is that a larger η 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 η = 10000 and the mini-batch size as B = 288 by default, but also perform an ablation study on CIFAR10 with several different values, as shown in Table 3.

We use a sigmoid-based diffusion schedule as

where we set c 0 = 10 and c 1 = -13 by default. We note a similar schedule had been introduced in Kingma et al. [35] and Jabri et al. [29] for Gaussian diffusion.

For the CIFAR-10 dataset 2 , we utilize the parameterization of EDM 3 [34] as the code base. We replace the variance preserving (VP) loss and the corresponding network parameterization implemented in EDM, which is 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 c noise = -logit ( α t ) / 8 and simplify the other parameters as c skip = 1 , c out = 1 , and c in = 1 . The image pixel values between 0 and 255 are divided by 255 and then scaled by S cale and shifted by S hift , and hence we have x 0 ∈ [ S hift , S hift + S cale ] . The two inputs to the generator network f θ ( · , · ) , originally designed for VP-EDM and adopted directly for beta diffusion, are g ( logit ( z t )) and logit ( α t ) , where g ( · ) is a preconditioning function described in detail in Appendix F. The output of the generator network is first transformed by a sigmoid function and then scaled by S cale and shifted by S hift , and hence ˆ x 0 = f θ ( z t , t ) ∈ [ S hift , S hift + S cale ] .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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).\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n[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.\n\nWEAKNESSES:\n[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\nQUESTIONS:\n[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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe 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).", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}