Beta Diffusion

Summary

In this submission, the authors introduce a diffusion model for range-bounded data. To 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. They 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$. They then propose an upper bound on the KL divergence between the forward process and this backward process. They showcase their model on a range of synthetic 1d datasets.

Strengths

- The proposed approach enables modelling distribution of data which is supported on collection of bounded ranges. - The beta diffusion process leads to closed form conditionals. - Since the Kullback-Leibler divergence between Beta distributions can be available in closed-form, the training loss can be computed efficiently.

Weaknesses

- I found this submission not so easy to follow at times. For instance, why is a separate Section 2.4 needed? - 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? - 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]. Dirichlet diffusion score model for biological sequence generation, Avdeyev, Pavel and Shi, Chenlai and Tan, Yuhao and Dudnyk, Kseniia and Zhou, Jian, 2023.

Questions

- line 123: This seems to hint at the fact that the proposed diffusion models could be written as the discretisation of an SDE which would admits as stationary distribution a Beta distribution. I'd be curious to know whether the Jacobi process used in [Pavel et al. 2023] is that SDE. - 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. - 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? - 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. - line 192: What are the consequences of the multiplicative vs additive noise? -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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

N/A

Summary

This work introduces a new type of diffusion model based on the beta distribution. The end of this is to defined a diffusion model on data that lives inside a specified range, i.e. on an interval [a,b]. This 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. They 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. The methodology is demonstrated on a couple of simple 1D settings that show the method works well in these cases. I 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.

Strengths

- The work is cleanly presented and easy to understand. - There is an extensive exploration of the options available for training these models. - The work opens up a discussion for other non-additive noise diffusion models, potentially using other infinitly divisible exponential families. - The experiments are sufficient to convince me that this method does indeed work, and could potentially be practically useful. - The experimental set up is well explained and looks highly reproducible.

Weaknesses

- I would have liked to see evaluation on some real world data settings to demonstrate the methods utility (in the main paper). - [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. - 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?) [1] Lou, Aaron, and Stefano Ermon. "Reflected diffusion models." arXiv preprint arXiv:2304.04740 (2023). [2] Fishman, Nic, et al. "Diffusion Models for Constrained Domains." arXiv preprint arXiv:2304.05364 (2023).

Questions

- Equation (9) is the specified short-time reversal distribution for this diffusion, akin to the Normal distribution in a typical discrete time diffusion model. In that setting the small time reverse diffusion isn't actually a normal, but one can show it converges to a normal with a specific form as the time-step goes to zero, and with a reasonable convergence rate, justifying this choice of variational distribution. Do you have a way of justifying the pick of a beta distribution with the specified form as the small time reverse distribution? For large time steps this obviously cannot be true (e.g. $p(z_0|z_1)$ is certainly not beta distributed unless the data itself is). - How well does this approach scale with dimension? - Have the image experiments been finished, and compared properly to other generative models for images? [1] Lou, Aaron, and Stefano Ermon. "Reflected diffusion models." arXiv preprint arXiv:2304.04740 (2023). [2] Fishman, Nic, et al. "Diffusion Models for Constrained Domains." arXiv preprint arXiv:2304.05364 (2023).

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors do not really discuss the limitations of their work, and I would like to see this better discussed, particularly with respect to other methods on the same topic. For example unlike score-based SDE models one cannot get a likelihood out of these models.

Summary

This paper addresses a prevalent assumption found in deep generative diffusion models, which is the Gaussian assumption in both the forward and reverse processes. In this study, the authors explore the utilization of the beta distribution in these processes and establish several key properties. Firstly, they analytically compute the marginal distribution within the beta distribution framework. Additionally, they develop an analytical computation for the conditional distribution, given both the data point and a noisy sample, in the form of a scaled and shifted beta distribution. To optimize the beta diffusion process, the authors propose a combination of two distinct KL upper bounds, which is better than using the standard ELBO. The experimental results on a synthetic setup demonstrate the ability of the beta distribution to effectively model range-bounded data distributed across disjoint regions. Furthermore, they show the efficacy of the proposed objective functions. The supplementary material also includes preliminary experiments conducted on image generation tasks.

Strengths

[Uniqueness of the methodology] Unlike 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. [Presentation] The manuscript is overall well-written and exhibits clarity in its presentation. Particularly, the introduction and related works sections provide concise and comprehensive summaries of prior research, effectively highlighting the distinction between the Gaussian diffusion process and the proposed Beta distribution process. To further enhance the motivation behind this work, it would be beneficial to include a practical scenario that illustrates the necessity of constructing the Beta diffusion process. This addition would greatly strengthen the motivation behind the study.

Weaknesses

[Absence of real-world experiments] The 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. [KLUB] I 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.

Questions

[About the formulation] #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. #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. #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. [About the experiments] #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. #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. #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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The paper does not explicitly outline its limitations. My suggestions to improve this work are described in the sections above. As this work is primarily theoretical in nature, it is unlikely to have any significant negative societal impacts.

Summary

The presented work proposes to design the forward and reverse process of diffusion models based on beta distributions (as opposed to Gaussian noise). It also showcases a new loss term based on KL-divergence upper bounds (KLUB) that can aid performance over the conventional ELBO-based diffusion model objective when used to optimize the proposed beta diffusion model. The paper shows that the proposed beta diffusion model can better model/generate synthetic distributions, especially point masses, compared to Gaussian diffusion, but lags behind it when used for CIFAR-10 image generation.

Strengths

- The paper is well-written - Alternatives to Gaussian diffusion processes are an interesting and relevant direction - The synthetic experiments give a good intuition on the advantages of beta diffusion over Gaussian diffusion (mostly in terms of modeling point masses)

Weaknesses

- Comparison and discussion with related work is lacking. As a result, the novelty of the presented work is hard to assess. While not having a related work section may be okay in some cases, discussing relevant related work thoroughly is very important. In particular, the presented paper completely misses to discuss or even cite any closely related papers on alternative diffusion processes of which at least [1, 2, 3, 4] come to my mind. - 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). - 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. - 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: 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). 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). 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). 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. [1] https://arxiv.org/abs/2208.09392 [2] https://research.google/pubs/pub52492/ (TMLR 23) [3] https://openreview.net/forum?id=OjDkC57x5sz (ICLR 23) [4] https://openreview.net/forum?id=4PJUBT9f2Ol (ICLR 23) [5] https://arxiv.org/abs/2305.10699 (ICML 23)

Questions

- What does \beta_max in the appendix correspond to in the main text? - How many generated samples are used to compute the Cifar10 FID? - More direct comparison to Gaussian diffusion on CIFAR 10 would be valuable (even if on a subset of Cifar10) - In terms of computational efficiency, are there any differences between beta and Gaussian diffusion? - 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). ************************************************* After 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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The paper would benefit from a more thorough discussion of the limitations of the presented work. As far as I can tell, the only limitation mentioned is that no hyperparameter tuning was performed, which is totally okay, but is not a limitation of the method itself (which would be much more valuable and interesting to read about).

Maintain original Score; Recommend acceptance

I am grateful for the authors' comprehensive response to my comments. They have effectively addressed the main concerns I raised in the initial reviews. Furthermore, it is encouraging to learn that the authors are actively exploring potential extensions of the beta diffusion framework. This includes its application to the exponential family as well as its integration with latent diffusion models. With these positive developments in mind, I am confident in maintaining my original score and recommend the acceptance of this work.

Raise score and recommend acceptance

Thank you for the convincing rebuttal, especially regarding the comparison to [1-4] in terms of both methodological and performance differences. I will raise my score and recommend acceptance as a result. Optimally, it would have been beneficial if the related work discussion and the Cifar 10 experiments had been ready for the main paper submission deadline. Please include those in the revised paper, incl. your verbal comparison with refs [1-5] (I think that discussing [5] in the paper would be great, even if it was published after your first submission for Neurips). > Indeed, KLUB-based objectives are applicable to diffusion models in general, including both beta and Gaussian diffusions. We believe this aspect was carefully addressed without any self-contradiction (Line 219-225) or overclaiming. For further clarification, please refer to Appendix D, showing how KLUB-based objectives relate to (weighted) -ELBO for Gaussian diffusion, and also our response to “W2” from Reviewer FX5S. I understand that KLUB can be used as an objective for Gaussian diffusion too. But, as you note in multiple places (e.g. appendix D), this seems to lead to the same existing objective (-ELBO). What I am complaining about in the referenced lines (e.g 78, 238, 346) is that they make it sound to me as though KLUB is a new way of optimizing diffusion models *in general*, even though practically it only leads to a new (and more effective, indeed) way of optimizing beta diffusion. As such, in my opinion it would be better to only claim having shown that KLUBs are an effective objective for *beta* diffusion models. > For toy data, we set the concentration parameter as, while for images, to distinguish between 256 pixel values, we choose a larger value (...) Please include this ablation of the concentration parameter on Cifar 10 to the revised paper.

response

Thanks for the detailed response to my comments! This does address my concerns. I'd suggest clarifying / stressing the following points in the manuscript - closed-form forward (no need to discretise) - convergence to $Beta(0, \eta)$ (which is a delta mass) - forward vs reverse KL <=> KLUB vs -ELBO and discussion thereof - connection with [Pavel et al. 2023] I've updated my score accordingly