Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation
Authors: Diederik P Kingma
Ruiqi Gao
Venue: NeurIPS 2023 oral (2023)
Keywords: Diffusion Model
Evidence Lower Bound
Maximum Likelihood
TL;DR: We develop a theoretical understanding of the training objective of diffusion models as a the ELBOs subject to data augmentation.

Abstract: To achieve the highest perceptual quality, state-of-the-art diffusion models are optimized with objectives that typically look very different from the maximum likelihood and the Evidence Lower Bound (ELBO) objectives. In this work, we reveal that diffusion model objectives are actually closely related to the ELBO.

Specifically, we show that all commonly used diffusion model objectives equate to a weighted integral of ELBOs over different noise levels, where the weighting depends on the specific objective used. Under the condition of monotonic weighting, the connection is even closer: the diffusion objective then equals the ELBO, combined with simple data augmentation, namely Gaussian noise perturbation. We show that this condition holds for a number of state-of-the-art diffusion models. 

In experiments, we explore new monotonic weightings and demonstrate their effectiveness, achieving state-of-the-art FID scores on the high-resolution ImageNet benchmark.

Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation
Diederik P. Kingma Google DeepMind durke google com
Ruiqi Gao Google DeepMind ruiqigegoogle.com com con
Abstract
To achieve the highest perceptual quality, state-of-the-arr diffusion models are optimized with objectives that typically look very different from the maximum likelihood and the Evidence Lower Bound (ELBO) objectives. In this work, we reveal that diffusion model objectives are actually closely related to the ELBO.
Specifically, show that all commonly used diffusion model objectives equate to a weighted integral of ELBOs over different noise levels, where the weighting depends the specific objective used. Under the condition of monotonic weight- ing, the connection is even closer: the diffusion objective then equals the ELBO, combined with simple data augmentation, namely Gaussian noise perturbation. We show that this condition holds for number of state-of-the-art diffusion models. In experiments, we explore new monotonic weightings and demonstrate their effectiveness, achieving state-of-the-ar FID scores on the high-resolution ImageNet benchmark.
Introduction
Diffusion-based generative models, diffusion models in short, first introduced by Sohl- Dickstein et al. [2015]. After years of relative obscurity, this class of models suddenly rose to prominence with the work of Song and Ermon [2019] and Ho al. [2020] who demonstrated that, with further refinements in model architectures and objective functions, diffusion models can perform state-of-the-art image generation.
37th Conference on Neural Information Processing Systems (NeurIPS 2023).
(a) 512 512
(b) 256 256
(c) 128 128
Figure 1: Samples generated from our VDM++ diffusion models trained on the ImageNet dataset; Section 5 for details and Appendix M for more samples.

---

The diffusion model framework has since been successfully applied to text-to-image generation [Rombach et al., 2022b, Nichol al., 2021, Ramesh al., 2022, Saharia et al., 2022b, Yu al.. 2022, Nichol and Dhariwal, 2021, Ho et al., 2022, Dhariwal and Nichol, 2022, Ding et al., 2021], image-to-imaae generation [Saharia et al., 2022c,a, Whang et al., 2022], 3D synthesis [Poole et al., 2022, Watson et al., 2022], text-to-speech [Chen et al., 2021a, Kong et al., 2021, Chen et al., 2021b], and density estimation [Kingma et al., 2021, Song et al., 2021a].
Diffusion models can be interpreted as a special case of deep variational autoencoders (VAEs) [Kingma and Welling, 2013, Rezende et al., 2014] with a particular choice of inference model and generative model. Just like VAEs, the original diffusion models optimized by maximizing the variational lower bound of the log- likelihood of the data, also called the evidence lower bound, or ELBO for short. It was shown by Variational Diffussio Models (VDM) [Kingma et al., 2021] and [Song et al., 2021a] how to optimize continuous-time diffusion models with the ELBO objective, achieving state-of-the-ee likelihoods on image density estimation benchmarks.
et al., 2015] were
However, the best results in terms of sample quality metries such as FID scores were achieved with other objectives, for example a denoising score matching objective [Song and Ermon, 2019] or simple voise-prediction objective [Ho et al.. 2020]. These now-popular objective functions look, on the face of it, very different from the traditionally popular maximum likelihood and ELBO objectives. Through the analysis in this paper, we reveal that all training objectives used in state-of-thh-aa diffusion models are actually closely related to the ELBO objective.
This paper is structured as follows:
In Section 2 we introduce the broad diffusion model family under consideration.
In Section 3, we show how the various diffusion model objectives in the literature can be understood as special cases of weighted loss [Kingma et al., 2021, Song et al., 2021a], with different choices of weighting. The weighting function specifies the weight per noise level. In Section 3.2 we show that during training, the noise schedule acts importance sampling distribution for estimating the loss, and is thus important for efficient optimization. Based on this insight we propose a simple adaptive noise schedule.
In Section 4, we present our main result: that if the weighting function is monotonic function of time, then the weighted loss corresponds to maximizing the ELBO with data augmentation, namely Gaussian noise perturbation. This holds for, for example, the v- prediction loss of [Salimans and Ho, 2022] and flow matching with the optimal transport path [Lipman et al., 2022].
In Section 5 we perform experiments with various new monotonic weights on the ImageNet dataset, and find that our proposed monotonic weighting produces models with sample quality that are competitive with the best published results, achieving state-of-art FID and scores on high resolution ImageNet generation.
1.1 Related work
The main sections reference much of the related work. Earlier work [Kingma et al., 2021, Song et al., 2021a, Huang et al., 2021, Vahdat al., 2021]. including Variational Diffusion Models [Kingma et al., 2021], showed how to optimize continous-tin time diffusion models towards the ELBO objective. We generalize these earlier results by showing that any diffusion objective that corresponds with monotonic weighting corresponds to the ELBO, combined with a form of DistAug [Child et al., 2019]. DistAug a method of training data distribution augmentation for generative models where the model is conditioned on the data augmentation parameter at training time, and conditioned on 'no augmentation' inference time. The type of data augmentation under consideration in this paper, namely additive Gaussian noise, is also form of data distribution smoothing, which has been shown to improve sample quality in autoregressive models by Meng et al. [2021].
Kingma et al. [2021] showed how the ELBO is invariant to the choice of noise schedule, except for the endpoints. We generalize this result by showing that the invariance holds for any weighting function.
2

---

Noise schedules
Noise schedules
10
Cosine
EDM training
0.150
Cosine
EDM training
Ar (log SNR)
5-
EDM sampling
0.125
EDM sampling
0.100
0.075
0.050
0.025
0.000
0.6
10
c
A (log SNR)
Figure 2: Left: Noise schedules used in our experiments: cosine [Nichol and Dhariwal, 2021] and EDM [Karras et al., 2022] training and sampling schedules. Right: The same noise schedules, expressed as probability densities p()) ddt/dA. See Section 2.1 and Appendix E.3 for details.
Model
Suppose we have a dataset of datapoints drawn from g(x). We wish to learn a generative model (x) that approximates g(x). We'll use shorthand notation p := po-
The observed variable x might be the output of pre-trained encoder, as in latent diffusion models [Vahdat al., 2021, Rombach al., 2022a], on which the popular Stable Diffusion model is based. Our theoretical analysis also applies to this type of model.
In addition the observed variable x, we have series of latent variables z, for timesteps t [0. 1]: Zo.....1 = Zo: .... Z1. The model consists of two parts: a forward process forming a conditional joint distribution q( (zoo..... 1 x), and generative model forming joint distribution p(
2.1 Forward process and noise schedule
The forward process is a Gaussian diffusion process, giving rise to a conditional distribution (Zzo..... (x)) see Appendix E.1 for details. For every t [0, 1], the marginal distribution q(z, xx is given by:
(1) In case of the often-used variance preserving (VP) forward process, a3 = (Xgmoid(Ae) and of = sigmoid(-Ar), but other choices possible; our results are agnostic to this choice. The log signal-to-noise ratio (log-SNR) for timestep is given by X = log(a3/o3).
Z, axx OxE where N(..I).
The noise schedule is strictly monotonically decreasing function SA that maps from the time variable t (0, 1] to the corresponding log-SNR X: À (t). We sometimes denote the log-SNR as ) to emphasize that it is function of t. The endpoints of the noise schedule are given by Amax := fx (0) and Amin == (1). See Figure 2 for a visualization of commonly used noise schedules in the literature, and Appendix E.3 for more details.
Due to its monotonicity, /x is invertible: 1 t= /- (A). Given this bijection, we can do a change of variables: function of the value can be equivalently written function of the corresponding value X, and vice versa, which we'll make use of in this work.
U(00 1), then compute À = ft).). This results in a distribution over noise levels p(A) = ddt/ dt/dA (t) (see Section E.3), which we also plot in Figure? 2.
During model training, we sample time t uniformly:
Sometimes is best use different noise schedule for sampling from the model than for training. During sampling, the density p(A) gives the relative amount of time the sampler spends at different noise levels.
3

---

2.2 Generative model
The data x D, with density q(x), plus the forward model defines a joint distribution (Zzu, ...,,.) ((zo, .... Zı xx (x)((x)xxx with marginals q: (z) == (Zz).. The generative model defines a correspond- ing joint distribution over latent variables: p(zo. .... Z1).
For large enough Amux. Zn is almost identical to x, so learning a model p(zo) is practically equivalent to learning model p(x). For small enough Amin: 21 holds almost no information about x, such that there exists distribution p(z1) satisfying [ (x)])((z)) ss 0. Usually we can use p(zi) (0, I).
Let (z( A) denote a score model, which is a neural network that we let approximate Vz log or (z) through methods introduced in the next sections. If (z; A) V log (z),, then the forward process can be exactly reversed; see Appendix E.4.
If D (q ) [p(z))) = 0 and (z; A) 2 Vx log qe (z), then we have good generative model in which achieves our goal. So, our generative modeling task is reduced to learning score network so (z; A) that approximates V. log (z).
the sense that DK (q(zo..... (l((.......)))) ()) Ro which implies that (g((((z)) ((l))((lll())))))))
Sampling from the generative model can be performed by sampling ~ p(zi). ). then (approximately) solving the reverse SDE using the estimated So (z; A). Recent diffusion models have used increasingly sophisticated procedures for approximating the reverse SDE; see Appendix E.4. In experiments we use the DDPM sampler from Ho al. [2020] and the stochastic sampler with Heun's 's second order method proposed by Karras et al. [2022].
3 Diffusion Model Objectives
Denoising score matching. Above, we saw that we need to learn a score network sp(z; At) that approximates Vz log (z), for all noise levels X It was shown by [Vincent, 2011, Song and Ermon, 2019] that this can be achieved by minimizing denoising score matching objective over all noise scales and all datapoints
D:
LUSM (x) -. [m(t) l|so Ar) -zz, log ((z x
(2)
where 2, = =xx + OXE.
The c-prediction objective. Most contemporary diffusion models are optimized towards noise- prediction loss introduced by [Ho et al., 2020]. In this case, the score network is typically parame- terized through noise-prediction (e-prediction) model: sz A) éé A)/ox. (Other options, such as x- x-prediction, v-prediction, and EDM parameterizations, are explained in Appendix E.2.) The noise-prediction loss is:
1
(x( = = - ((co(zz;\)) -
(3)
Since [|so Ae) Xx, log [x)BB ox (zz) 24) el22. this is simply a version of the denoising score matching objective in Equation 2, but where w(t) of. Ho et al. [2020] showed that this noise-prediction objective can result in high-quality samples. Dhariwal and Nichol [2022] later improved upon these results by switching from 'linear' 'cosine' noise schedule X (see Figure 2). This noise-prediction loss with the cosine schedule is currently broadly used.
The ELBO objective. was shown by [Kingma al., 2021, Song al., 2021a] that the evidence lower bound (ELBO) of continuous-time time diffusion models simplifies to:
ELBO(x) = B
dt 1160(2;; Ar) - +
(4)
where e constant w.r.t. the score network parameters.

---

Loss function
Implied weighting w(A)
Monotonic?
ELBO [Kingma al., 2021, Song 202
IDDPM (e-prediction "cusine" schedule) [Nichol Dhariwal, 2021]
sech(A/2)
[Karras et al., 2022] (Appendix D.1)
(1;;244,244)) (e-
v-prediction with 'cosine' schedule [Salimans Ho, 2022] (Appeadix D.2)
A/2
Flow Matching with OT path (FM.OT) OT) [L-ipman et al., 2022] (Appendix D.3)
A/2
InDI (Delhracio and Milanfar, 2023] (Appendix D.4) F2 weighting with "cosine' schedule (Chos et al., 2022] (Appendix D.5)
(2/4)
Min-SNR-y [Hang et al., 2023] (Appendix D.6)
soch(A/2) min((1,3e
Table 1: Diffusion model objectives in the literature are special cases of the weighted loss with a weighting function w(A) given in this table. See Section 3.1 and Appendix D for more details and derivations. Most existing weighting functions are non-monotonic, except for the ELBO objective and the v-prediction objective with 'cosine' schedule.
Non-monotonic weighting functions
Monotonic weighting functions
1.0
IDDPM
1.0
EDM
prediction and FM OT
W(A) (normalizeer allizdd
w(A) (norma (normalized) lized]
EDM monotonic (Ours)
min-SNR-y
sigmoid(-X (Ours)
0.6
0.6
0.4
0.2
0.2
2.5
5.0
10.0
-5.0
-2.5
0.0
7.5
10.0
A (log SNR)
À (log SNR)
Figure 3: Diffusion model objectives in the literature are special cases of the weighted loss with non-monotonic (left) or monotonic (right) weighting functions. Each weighting function is scaled such that the maximum is over the plotted range. See Table and Appendix D.
3.1 The weighted loss
The different objective functions used in practice, including the ones above, can be viewed as special cases of the weighted loss introduced by Kingma al. [2021]' with different choices of weighting function w (Ae):
dd dt
L (x) 11.. ((()))
-|ée(zzz <11
(5)
See Appendix D for derivation of the implied weighting functions for all popular diffusion losses. Results are compiled Table 1, and visualized in Figure 3.
The ELBO objective (Equation 4) corresponds to uniform weighting. i.c. w(>e) = 1.
The popular noise-prediction objective (Equation 3) corresponds w(\e) - This is more compactly expressed as w(A,) p(Ac), i.e., the PDF of the implied distribution over noise levels during training. Typically, the noise prise-prediction objective used with the cosine schedule A which implies (>)) x sech(Ar/2) See Section E.3 for the expression of p(>e) for various noise schedules.
3.2 Invariance of the weighted loss to the noise schedule X
In Kingma et al. [2021], I, it was shown that the ELBO objective (Equation 4) is invariant to the choice of noise schedule, except for its endpoints Amia and Amax- This result can be extended to the much
Mooe specifically, Kingma et al. [2021] expressed the weighted diffusion loss in terms of x-prediction, which is equivalent the expression above due the relationship / ||e éz( X)ll dA f (2;; >) dd ||x (0( 2) de, where e² equals the signal-to-nol -oise ratio (SNR).
5

---

more general weighted diffusion loss of Equation 5, since with change of variables from to A, it can be rewritten
1 Amas
(x)
u(A)EE-NV(O.)) [Zx A) -c1 dd
Amin
Note that this integral does depend the noise schedule fx (the mapping between and A), except for its endpoints Amax and Amin- The shape of the function fa between Amin and Xmax does not affect the loss; only the weighting function w(A) does. Given a chosen weighting function w(A), the loss is invariant to the noise schedule A between / = 0 and = 1. This is important, since means that the only real difference between diffusion objectives their difference in weighting w(A).
This invariance does not hold for the Monte Carlo estimator of the loss that we use in training. based random samples t (/(0, 1), (O. The noise schedule still affects the variance of this Monte Carlo estimator and its gradients; therefore, the noise schedule affects the efficiency of optimization. More specifically, we'll Il show that the noise schedule acts as importance sampling distribution for estimating the loss integral of Equation Note that p()) -1/(dX/dt) We can therefore rewrite the weighted loss the following, which clarifies the role of p(A) as an importance sampling distribution:
(())
(x) =
2 E. 2e
ll ZA: - p(A)
Using this insight, we implemented an adaptive noise schedule, detailed in Appendix F. We find that by lowering the variance of the loss estimator, this often significantly speeds op optimization.
4 The weighted loss as the ELBO with data augmentation
We prove the following main result:
Theorem If the weighting w=\x) is monotonic, then the weighted diffusion objective of Equation 5 is equivalent to the ELBO with data augmentation (additive noise).
With monotonic w():)) we mean that w is a monotonically increasing function of f, and therefore a monotonically decreasing function of X.
We'll use shorthand notation L(t: x) for the KL divergence between the joint distributions of the forward process ( (x) and the reverse model p(z p(z,,...) for the subset of timesteps from t to 1:
x(t;x) = Dxx(g(z.....)) (x)
(8)
In Appendix A.1, prove that?:
d
dA
C(t:)) = E.-N(O [lle éz(zx;A)]
(9)
di
2 dt
As shown in Appendix A.1, this allows us to rewrite the weighted loss of Equation 5 as simply:
(x) = S.' d
dt L(t:x) w(Ae) dt
(10)
In Appendix A.2, we prove that using integration by parts, the weighted loss can then be rewritten as: (x) = S/' d dt =w()s) L(t;x) dt + w(Amax) L(0; x) + constant
(11)
*Interestingly, this reveals new relationship between KL. and Fisher divergences: Dr(((((( (p(z,))- See Appendix G for a derivation, a
2X Dxi(g(Ze (xx))//.......))))(/z...... discussion, and comparison with similar result by Lyu [2012].

---

Now, assume that w())) is a monotonically increasing function of t E [o, 1). Also, without loss of generality, assume that (x(e)) normalized such that w(11) 1. We can then further simplify the weighted loss to an expected KI. divergence:
L (x) = (t) [L(t; x)] + constant
(12)
where (t) is probability distribution determined by the weighting function, namely (t) := (d/dt w(),)), with support on t [0, 1]. The probability distribution (to has Dirac delta peak of typically very small mass (Amma)) at = 0.
Note that:
C(l:)) = DKI(g(Zt,... (x) [p((e.....))
(13)
2 Dxz(((Zeh ())()(())) = - -Bgen (x) [log [log p(z.)] + constant.
(14)
More specifically, [[t; equals the expected negative ELBO of noise noise-perturbed data, plus a constant; Section C for a detailed derivation.
This concludes our proof of Theorem
This result provides us the new insight that any of the objectives with (implied) monotonic weighting, listed in Table 1, can be understood as equivalent to the ELBO with simple data augmentation, namely additive noise. Specifically, this is a form of Distribution Augmentation (DistAug), where the model is conditioned on the augmentation indicator during training, and conditioned on 'no augmentation' during sampling.
Monotonicity of w(A) holds for number of models with state-of-the- perceptual quality, including VoiceBox for speech generation [Le et al., 2023), and Simple Diffusion for image generation [Hooge- boom et al., 2023].
4.1 Understanding the effect of weighting functions perceptual quality
In this subsection, we propose possible explanations of the reason that the weighted objective, or equivalently the ELBO with data augmentation assuming monotonic weightings. can effectively improve perceptual quality.
Connection to low-bit training. Earlier work [Kingma and Dhariwal, 2018] found that training on 5-bit data, removing the three least significant bits from the original 8-bit data, can lead to higher fidelity models. A likely reason is the more significant bits also more important to human perception, and removing the three least significant bits from the data allows the model to spend more capacity modeling the more significant bits. In [Kingma and Dhariwal, 2018], training 5-bit data was performed by adding uniform noise (with the appropriate scale) to the data before feeding it to the model, but it found that adding Gaussian noise had similar effect. Adding Gaussian noise with a single noise level is special case of the weighted objective where the weighting function is a step function. The effect of uniform noise can be approximated with a sigmoidal weighting function; in Appendix I we dive into more details.
Fourier analysis. To better understand why additive Gaussian noise can improve perceptual quality, consider the Fourier transform (FT) of a clean natural image perturhed witjh additive Gaussian noise. Since the FT is linear transformation, the FT of a the sum of natural image and Gaussian noise, equals the sum of the FT of the clean image and the FT of the Gaussian noise. For natural images, is known that the power spectrum, i.c. the average magnitude of the FT as a function of spatial frequency, decreases fast as the frequency increases [Burton and Moorhead, 1987], which means that the lowest frequencies have by far the highest magnitudes. On the other hand, the power spectrum of a Gaussian white noise is roughly constant. When adding noise to a clean natural image, the signal-to-noise ratio (i.e. the ratio of the power spectra of the clean and Gaussian noise images) in high frequency regions is lower. As result, additive Gaussian noise effectively destroys high-frequency information the data more quickly than the low-frequency information, pushing the model focus more on the low frequency components of the data, which often correspond to high-level content and global structure that is more crucial to perception. See Appendix J for a more detailed discussion.
7

---

DDPM sampler
EDM sampler
Model parameterization
Training schedule
Weighting function
Monotonic?
FID 4
18
FID 4
18
e-parametrizatoon
Cosine
soch(AA2) (Baseline)
1.85
54.1
+ 0.79
1.55 59.2 0.78
Cosine
sigmoid, A
1.75
55.3
1 1.23
Coxine
sigmaid( -
1.8
56.8
1: 0.85
1.46
60.4 1: 0.86
Cosine
sigmoid( -À
1.73
Cosine
sigmoidd +)
sigmoid( - ()
1.94
53.5
1.12
Adaptive
sigmoid( - 2)
1.70
1.44
+
Adaptive
EDM momotonic
1.67
56.8
1. 0.90
1.44
61.1 1: 1.80
[Karras al., 2022] (training)
(Baseline)
repriduction)
(training)
(Raseline)
+
Adaptive
EDM
1.43
63.2 + 1.76
Adaptive
sigmoid( - + 2)
1.55
63.7
+ 1.14
Adaptive
EDM-monotonic
1.43
63.7
+ 1.88
v. parametrizatioon
Adaptive
exp(- (Baseline)
1:
Adaptive
sigmoid( -> + 2)
1.51
64.4 +
Adaptive
EDM-monotonic
1.45
64.6 + 1.35
Table 2: ImageNet 64x64 results. See Section 5.2.
5
Experiments
Inspired by the theoretical results in Section 4, in this section we propose several monotonic weighting functions, and present experiments that test the effectiveness of the monotonic weighting functions compared to baseline (non-monotonic) weighting functions. In addition, we test the adaptive noise schedule (Section 3.2). For brevity we had to pick a name for our models. Since we build on earlier results on Variational Diffusion Models (VDM) [Kingma al., 2021], and our objective is equivalent to the VDM objective combined with data augmentation, we name our models VDM++.
5.1 Weighting functions
We experimented with several new monotonic weighting functions. First, a sigmoidal weighting of form sigmoid( + k)), inspired by the sigmoidal weighting corresponding to low-bit training (see Section 4.1 and Appendix for details), where k is hyper-parameter. In addition, inspired by the EDM weighting function from Karras et [2022] (non-monotonic, Table 1), we test a weighting function indicated with EDM-monotonic', which is identical to the EDM weighting, except that it is made monotonic by letting wAA) = max, (AA) for À < arg max, w (A). where (iA) indicates the original EDM weighting function. Put simply, 'EDM-monotonic" equals the EDM weighting, except that it is constant to the left of its peak; see Figure 3.
5.2 ImageNet 64x64
All experiments on ImageNet 64x64 were done with the U-Net diffusion model architecture from [Nichol and Dhariwal, 2021]. We carried out extensive ablation studies over several design choices of diffusion models, namely model parametrization. training noise schedules, weighting functions and samplers. For e-parametrization model, we took iDDPM [Nichol and Dhariwal, 202 I] as the baseline, which utilized cosine noise schedule and non-monotonie sech( A/2) weighting. For EDM parmetrization, we recruited the setting in Karras et al. [2022] as the baseline, with EDM training noise schedule (Figure 2) and mon-oonoto EDM weighting (Figure 3). For parammtrrization model, we followed Salimans and Ho [2022] taking the v-prediction loss as the baseline, which leads monotonic exp(-X/2) weighting. We always adaptive noise schedule under this setting. We started by searching the optimal hyperparameter for the sigmoidal weighting under the r-parametrization model, and then applied the best sigmoidal weighting and EDM-monotonic weighting under other settings. We evaluated models with two samplers, DDPM and EDM samplers, with their corresponding sampling noise schedules. Table summarizes the FID [Heusel al., 2017] and Inception scores [Salimans et al., 2016] across different settings. From the table we made several observations follows.

---

First, our adaptive noise schedule (Appendix F) works equally well as the hand-tuned fixed noise schedules. We verified this under e-parametrization with sigmoid(- + 2) weighting and EDM parametrization with EDM weighting. Under both settings, the adaptive noise schedule results in slightly better FID and Inception scores with the EDM sampler, and slightly worse scores with the DDPM sampler for the former setting. Although the adaptive noise schedule did not significantly affect the end result, it allows us to experiment with new weighting functions while avoiding hand- tuning of the noise schedule. In addition, it sometimes leads to faster training; see Appendix F.
Second, is possible modify an existing non- monotonic weighting function to a monotomic one with the minimal change, and results in on on-par sample quality. Specifically, under EDM parameterization and adaptive noise schedule, we change EDM weighting to mom-moooooiic weighting with the very straightforward approach stated in Section 5.1. It results in nearly the same FID score and slightly better Inception score.
Third, our proposed new weighting functions, sigmoid( - À + 2) and EDM monotonic weightings, work better than the baseline settings across different model parameterization and samplers in most cases. Larger performance gain has been observed under e-parameterization and v-parameterization settings. Dur-reimplementation of the EDM parameterization baseline setting could not exactly reproduce their reported FID number (1.36), but comes close (1.45). We observed less significant improvement under this setting. possibly because their already exhaustive search of the design space.
Finally, the above observations remain consistent across different model parameterization and sam- plers, indicating the generalizability of our proposed weighting functions and noise schedules.
5.3 High resolution ImageNet
In our final experiments, we tested whether the weighting functions that resulted in the best scores on ImageNet 64x64, namely sigmoid(-X+ 2) and 'EDM-monotonic' also results in competitive scores on high-resolution generation. As baseline we use the Simple Diffusion model from Hoogeboom et al. [2023], which reported the best FID scores to date high-resolution ImageNet without sampling modifications ( guidance).
We recruited the large U-VIT model from Simple Diffusion [Hoogeboom et al., 2023], and changed the training noise schedule and weighting function to our proposed ones. Note that for higher- resolution models, Hoogeboom et al. [2023] proposed a shifted version of the cosine noise schedule (Table 4), that leads to a shifted version of the weighting function wA)). Similarly, we extended our proposed sigmoidal and 'EDM-monotonic" weightings to their shifted versions (see Appendix D.2.1 for details). For fair comparison, we adopted the same vanilla DDPM sampler as Simple Diffusion, without other advanced sampling techniques such second-ordd As shown in Table 3, with our adaptive noise schedule for training, the two weighting functions we proposed led to either better or comparable FID and IS scores ImageNet 128x 128, compared to the bascline Simple Diffusion approach.
sampling or rejection sampling.
Next, we test our approach on ImageNet generation of multiple high resolutions (i.e., resolutions 1 128. 256 and 512), and compare with existing methods in the literature. See Table 3 for the summary of quantitative evaluations and Figure 1 for some generated samples by our approach. With the shifted version of IEDM-monotonic weighting, we achieved state-of-the-arr FID and IS scores on all three resolutions of ImageNet generation among all approaches without guidance. With classifier-free guidance (CFG) [Ho and Salimans, 2022], our method outperforms all diffusion- approaches on resolutions 128 and 512. On resolution 256, our method only falls bit behind Gao et al. [2023] and Hang et al. [2023], both of which were build upon the latent space of a pretrained auto-encoder from latent diffusion models [Rombach et al., 2022a] that was trained on much larger image datasets than ImageNet, while our model was trained on the ImageNet dataset only. It is worth noting that we achieve significant improvements compared to Simple Diffusion which serves as the backbone of our method, on all resolutions, with and without guidance. is possible to apply our proposed weighting functions and adaptive noise schedules to other diffusion-based approaches such Gao et al. [2023] to further improve their performance, which we shall leave to the future work.

---

Without gnidance
With guidance
FID
Method
Irain
train
128 128 resolution
ADM [Dhariwal and Nichol, 2022]
2,97
CDM [Ho al.,
3.76
128.8
2.5
2022]
Simple Diffusion (U-Net) [Hoogehoom al,, 2023]
2.88
137.3 2.0
Simple Diffusion (U-VIT, L) [Hoogeboom al,, 2023]
3.23
171.9
2.05 3.57
189.9 +
VDM++ (Ours), ar) = sigmoid( - + 2)
1.91
3.41
183.1
2.2
VDM++ (Ours), EDM-meoctenic weighting
1,75
2,88
171.1
2.7
1,78
3.16
190.5 +
resolution
BigGAN-deepp (no truncation! [Brock et al., 2018]
6.9
17144 2.0
MaskGIT [Chang et al., 2022]
182.1
ADM [Dhoriwal and Nichol, 2022]
10.94
3.94
215.9
CDM [Ho ot al., 2022]
4.88
4.63
158.7 2.3
RIN [Jabri al., 2022]
3.42
182.0
Simple Diffusion (U-Net) [Hoogehooo al., 2023]
3.76
3.71
171.6 -e 3.1
Simple Diffusion (U-VT, L) HHoogeboom al., 2023]
2.77
3,75
211.8
+
2,44
4,08
256.3 +
VDM+ (Ours), monotonno weighting
2.40
3.36
225.3
3.2
2.12
3.69
267.7 1: 4.9
Latent diffusion with pretrained VAE:
DET XL/2 [Peebles and Xie, 2022]
9.62
121.5
2.27
278.2
U. VIT [Bao al., 2023]
3.40
Min-SNR-~ [Hang al., 2023]
2.06
MDT [Gao et al., 2023]
6.23
143.0
1.79
283.0
512 512 resolution
MaskGIT [Chang et sl., 2022]
7.32
156.0
ADM [Dhariwal and Nichol, 2022]
23.24
3.85
231.7
RIN [Jabri al., 2022]
3.95
216.0
Simple Diffusion (U-Net) [Hoogebooo 2023]
4,30
4.28
171.0
3.0
Simple Diffusion (U-VIT [Hoogeboom al., 2023]
4.53
205.3
3.02
+
VDM++ (Ours). menotenic weighting
4.09
2.65
Latent diffission with pretrained VAE:
DIT-XU2 [Peehles and Xis, 2022]
12.03
105.3
3.04
240.8
LDM [Rambach 2022a]
10.56
103.5
1.2
3.60
247.7 1. 5.6
Table Comparison approaches in the literature for high-resolution ImageNet generation. 'With guidance' indicates that the method was combined with classifier-free guidance [Ho and Salimans, 2022]. Models under 'Latent diffusion with pretrained VAE' use the pre-trained VAE from Stable Diffusion [Rombach et al., 2022a], which used much larger training corpus than the other models in this table.
6 Conclusion and Discussion
In summary, we have shown that the weighted diffusion loss, which generalizes diffusion objectives in the literature, has an interpretation weighted integral of ELBO objectives, with ELBO per noise level. If the weighting function is monotonic, then we show that the objective has an interpretation the ELBO objective with data augmentation, where the augmentation is additive Gaussian noise, with a distribution of noise levels.
Our results open up exciting new directions for future work. The newfound equivalence between monotonically weighted diffusion objectives and the ELBO with data augmentation allows for direct apples-to-apples comparison of diffusion models with other likelihood example, it allows one to optimize other saclihood-basee models, such as autoregressive models, towards the same objective monotonically weighted diffusion models. This would shine light on whether diffusion models are better or worse than other model types, as measured in terms of their held-out objectives as opposed to FID scores. We leave such interesting experiments to future work.
models. For
10

---

Acknowledgments
We'd like to thank Alex Alemi and Ben Poole for fruitful discussions and feedback on early drafts. We thank Emiel Hoogeboom for advice and help on the implementation of Simple Diffusion.
References
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, (233):3333266 1982.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
2023.
Andrew Brock, Jeff Donahue, and Karen Simonyan, Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv: 1096, 2018.
Geoffrey J Burton and lan Moorhead. Color and spatial structure in natural scenes. Applied optics, 26(1):157-170, 1987.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11315-11322 2022.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan. WaveG- rad: Estimating Gradients for Waveform Generation. In ICLR, 2021a.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, Najim Dehak, and William Chan. WaveGrad Iterative Refinement for Text-to-Speech Synthesis.. In INTERSPEECH, 2021b.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon. Perception prioritized training of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11472-11481, 2022.
Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arXiv preprint arXiv:2303. 11435, 2023.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. In NeurIPS, 2022.
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffusion for categorical data, arXiv preprint arXiv:2211.15088 2022.
Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, and Jie Tang. Cogview: Mastering text-to-image generation via transformers, 2021. URL https: //aaxvvvvvv/./////////aa/aaaaav. oowoor//aao... abs / 2205..111111 13290.
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is strong image synthesizer. arXiv preprini arXiv:2303. 14389, 2023.
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.0953 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Berharrd Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. arXiv preprint 1oXiv:1706.08500, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint 2xXx:22077.15590 2022.

---

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilisticc models. arXiv preprint vrXiv:2006.11255 2020.
Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. JMLR, 2022.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint arXiv:2301.11093, 2023.
Chin-Wei Huang, Jae Hyun Lim, and Aaron Courville. A variational perspective on diffusion-based generative models and score matching. arXiv preprint arXiv:2106.02880 2021.
Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. arXiv preprint arXiv:2212. 2022.
Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. arXiv preprint 2211:22006000364,, 2022.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-217077 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Diederik P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions. arXiv preprini arXiv: 0007003039, 2018.
Diederik Kingma and Max Welling. Auto-encoding variational Bayes. International Conference on Learning Representations, 2013.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DiffWave: A Versatile Diffusion Model for Audio Synthesis. In ICLR, 2021.
Matthew Le, Apoorv Vyas, Bowen Shi, Brian Karrer, Leda Sari, Rashel Moritz, Mary Williamson, Vimal Manohar, Yossi Adi, Jay Mahadeokar, al. Voicebox: Text-guided multilingual universal speech generation at scale. arXiv preprint 1.XX:::23661668.8 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint 2rX:v:2210002477, 2022.
Siwei Lyu. Interpretation and generalization of score matching. arXiv preprint arXiv:1205.2629, 2012.
Chenlin Meng, Jiaming Song, Yang Song, Shengjia Zhao, and Stefano Ermon. Improved autoregres- sive modeling with distribution smoothing. arXiv preprint arXiv:2103. 15089, 2021.
Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv preprint arXiv:2 2021.
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Bob McGrew Pamela Mishkin, Ilya Sutskever, and Mark Chen. GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models. In arXhv:2112.107411 2021.
William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09978, 2022.
Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. DreamFusion: Text-to-3D using 2D Diffusion. arXiv, 2022.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical Text- Conditional Image Generation with CLIP Latents. In arXiv, 2022.
Danilo J Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. International Conference on Machine Learning, pages 1278-1286 2014.
12

---

Robin Rombach, Andreas Blattmann, Dominik Lorenz. Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pages 10684-10695, 2022a.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- Resolution Image Synthesis with Latent Diffusson Models. In CVPR, 2022b.
Chitwan Saharia, William Chan, Huiwen Chang, Chris A. Lee, Jonathan Ho, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Palette: swage-to-Image Diffusion Models. In SIGGRAPH, 2022a.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang. Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding. In NeurlPS, 2022b.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions Pattern Analysis and Machine Intelligence, 2022c.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005: 2022.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference Machine Learning, pages 2256-2265, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, pages 11895-11907, 2019.
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. arXiv e-prints, pages arXiv-2101. 2021a.
Yang Song, Jascha
Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling Through Stochastic Differential Equations. In International Conference on Learning Representations, 2021b.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. arXiv preprint 10X.v:2/0660099: 2021.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674 2011.
Daniel Watson, Ricardo Chan, William Martin-Brunn Brualla, Jonathan Ho, Andrea Tagliasacchi, and Mohammad Norouzi. Novel View Synthesis with Diffusion Models. arXiv, 2022.
Jay Whang, Mauricio Delbracio, Hossein Talebi, Chitwan Saharia, Alexandros G. Dimakis, and Peyman Milanfar. Deblurring via Stochastic Refinement. CVPR, 2022.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, Ben Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, and Yonghui Wu Jason Baldridge. Scaling Autoregressive Models for Content-Rich 2022.
Text-to-Image Generation. In
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Improved techniques for maximum likelihood estimation for diffusion odes. arXiv preprint arXiv:2305. 03935, 2023.
13

---

A
Main proof
Here we'll provide a proof of Equation 11.
Note that like in the main text, we use shorthand notation:
x(1;)) = Dxi(g(Z (x)(,
(15)
A.1 Time derivative of D k
Let dt denote an infinitesimal change in time. Note that L(t - dt; x) can be decomposed as the sum of KL divergence and an expected KL divergence:
L(t dt; x) [(t:x) + Bg(x, (x) [D (q(z)- Zz, p))
(16)
Due to this identity, the time derivative d/ dl L(I; x) can be expressed
d di [(::))
1 dl (C(t;x) L(t dt;x))
(17)
(2()))x) Dr (g aa (2,, x) lp((z de [2)]
(18)
dt
In Appendix E [Kingma et al., 2021], it shown that our model, this equals:
dt [(t;x) = S dt) - SNR(t) xx (x(( A:)112
d
(19)
dt -2 SN SN L0 1
(20)
where 2, = axx + xxe, and R(I) == exp(A) in our model, and SN) = d/d e² dd/dl, so terms of our definition of X, this is:
_ ((:x))
1 dd -e0 BEwNro.D [||x (0 Ax)
(21)
In terms of e-prediction (see Section E.2), because ||e é0l13 = e> c^|[x= &oll3 this simplifies to:
1 dd
dt [(t:x)
2 dt EC-N(O.)) - Éol 2e11
(22)
where Σy = axx + OxE. This can be easily translated to other parameterizztionss, see E.2.
This allows to rewrite the weighted loss of Equation 5 as:
dA ||60(zz;;)) A) <11
(x) 2 E, 11-11(0.))
wx\e) dt dA w(Ac) 2 dt EenN(o.)) [Ilea(z(:] c11
(23)
(24)
d [(t:x) w(>e)
Em-r(011) dt
(25)
S/' d dt w())) dt
(26)
A.2 Integration by parts Integration by parts a basic identity, which tells us that:
(()) /" ('())()) + Ma))((a) - ((b)g(b)
14

---

This allows to further rewrite the expression of the weighted loss in Equation 26 as:
(x) L ['"
d
w((e) dt [():)
(27)
di
Lu ae mr())) L(t;x) dt w (mmax )((;; ( x11;x)
d
(28)
where
-w( [1;;) -w(mmmn)) DEL (g(z) (x)((((()))
(29)
constant w.r.t. 0, since does not involve the score function, and typically very small, since D DKL ((lzii))) is typically small by design.
The term ((Amax) ( (q(o.......... (I(zo.....)) (xx small (see Figure 3).
typically small, since w (maax) typically very
This concludes our proof of Equation 11.
B Visualization
We tried to create helpful visualization of the result from Section A.2. Note that we can rewrite:
L d di [(lx)) dt w(Xe) df(t;x)
-==
[();))
[(1;x)
(30)
The relationship in Equation 28 can be rewritten
r=0
[ C(l;x) duo()e) f
((miia) L(1)
|
(())))
vo()aax C()) +
(31)
The first LHS term equals a weighted prior loss term, and the second LHS term equals the weighted diffusion loss. From a geometrie perspective, the LHS terms together define an area that equals the area given by the right term, as illustrated in the figure below.
m((1)
w(Ar)
0
w(A.) [mmil
-((:x)
wxh)) x((t:))
w((o)) [(0:x)
x(::))
[(t:x)
On the left, we have a rectangular area that equals a weighted prior loss w [(L), [(1), plus a curved area equal to the weighted diffusion loss f w(Ac) dd(l:x). This integral can be intuitively understood as Riemann sum over many tiny rectangles going from left (t = 1) to right = 0), each with height wAx) and width dC(t:x). On the right, have the same total area, but divided up into two different subareas: a rectangular area (mmax C(0)
C(0) and a curved area that equals the integral
t=1 [(I; x) (>:) going upwards from / = 0 to / 1, which can also be intuitively understood as another Riemann sum, with each tiny rectangle having width L(t; x) and height (Ae).. The area of each of those tiny rectangles on the right can understood as the ELBO at each noise level, [(t:x). times the weight of the ELBO at each noise level, (w().).
C Relationship between DKL (g((gg(z....
and the ELBO
First, note that:
C(t:x) = DKI (gcllglz,.... (Dx:(((l......./ Dro(g(2,)<)lll(<z))
(32)
15

---

More precisely, the joint KL divergence D (g(zz..... (Z(.......) is the expected negative ELBO of noise-perturbed data, plus a constant entropy term:
C[I; D (g(z,..... (x)(((((......))) = - B (2, (L.OO,2.)) H(((Zz/x))
(33)
constant
where the ELBO of noise-perturbed data is:
ELBO,,z:) == (x plg ži) log ((z)(z))
(34)
logp(zz)
(35)
where ă; = dt-
So, [(t:x) is the expected negative ELBO of noise-perturbee data 2,:
x(t;x) (xg [ELBO,(z,) + constant 2 -xx(( (log p (z:) + constant
(36)
Therefore, the expression of the weighted loss in Equation 12 be rewritten
(x) = Epp (1) [C(4:)) + constant
(37)
= pi (f): (x) (z,) + constant
(38)
ELBO of poise-pertureed data
2 (2) p(z))
+ constant
(39)
Log-likelihood of noise-perturbed data
where w(Amin) constant w.r.t. the diffusion model parameters. Therefore, minimizing C (x) is equivalent to maximizing this expected ELBO of noise-pentt passe-prrurred data.
D Derivation of weighting functions for previous works
The loss function used in previous works equivalent the weighted loss with certain choice of noise schedule and weighting function wAA). In this section, we derive these weighting functions wA))
D.1 'Elucidating Diffusion Models' (EDM) [Karras et al., 2022]
Karras et [2022] proposes the following training objective:
Leam (x) Eg=xt(t), [w(ō)][x Ro (22;5))11)
(40)
where p(ō) and (ō) are defined as:
p(log ō) NNoo Pmeany Pard):
(41)
(( (ō2 + o(atta)(("" odsts),
(42)
where à is equivalent to the standard deviation of noise added to the clean sample x the VP SDE
case, Ø² - Karras al. [2022] used hyperparameters P Pmean -1.2, Patd 1.2 and
Odata = 0.5.
This can be rewritten in terms of À = -2 2og as: Cedm(x) (x) Ep(A). -
(43)
where p(A) and (A) are defined as
p\)) = ((
(44)
(AA) = (e + 0.53) 0.52).
(45)
Comparing Equation 45 with the weighted loss expressed terms of x-prediction parameterization (Section E.2), we see that (ignoring the constant scaling factor 1/2) the EDM objective is special case of the weighted loss with weighting function:
w ()) = p(A)e= (A)
(46)
= (A)((e>> + 0..5)
(47)
where the divison by 2 can be ignored since it's constant. This leads to:
w(A) 2,,4 0.5?)
(48)
16

---

D.2
The v-prediction loss / 'SNR+1' "SNR--"-weihhhinng [Salimans and Ho, 2022]
Salimans and Ho [2022] introduced the y-parameterriztion, with v-prediction model v, where:
axÊ dxx
(49)
== axē - O2X
(50)
They propose to minimize a v-prediction loss, Eillv 4/12]. Note that é = (z) -xx8//x. For our general family, this implies:
||v - 111 = - + 1) x12
(51)
= =a + 1) E
(52)
In the special case of the variance preserving (VP) SDE, this simplifies to:
||v - +12 = (e^ + 1) x - x 12
(53)
(e- + 1)]l 211
(54)
Since the e-prediction loss corresponds to minimizing the weighted loss with w(A) = p(A), the v-prediction loss corresponds to minimizing the weighted loss with w(A) = (cc=A11)p(A)
Note that Salimans and Ho [2022] view the loss from the x-prediction viewpoint, instead of our ELBO viewpoint; so in their view, minimizing simply ||x X||2 means no weighting. Note that e is
the Sggaaltt-noise ratio (SNR). Since ||e -- éll e^||x x they call the e-prediction loss 'SNR weighting', and since ||v v112- = (e* + 1)] - x12 they call this 'SNR+1 wwiighting.
Salimans and Ho [2022] propose use optimize VP SDE with cosine schedule p())
sech()/))((27) = 1/(2m cosh(-A/2) and the v-prediction loss: Billv - v v1 This corresponds minimizing the weighted loss with:
w(() (c + 1)p(A)
(55)
= (c - + 1)/(2)/))))))) (2 cosh(-A/2))
(56)
ne X/2
(57)
The factor π can be ignored since it's constant, so we can equivalently use:
w(A) = A/2
(58)
D.2.1 With shifted cosine schedule
(Hoogeboom et al., 2023] extended the cosine schedule to a shifted version: p(A) = sech $)/(2m) where log 1og(64/d) (6//), where 64 is the base resolution and d is the model resolution (e.g.
128, 256, 512, etc.). In this case the weighting is:
((A) = ( 1)p(A)
(59)
-e
(60)
Since (2/m)e-* constant w.nt X, the weighting equivalent to the weighting for the unshifted cosine schedule.
D.3 Flow Matching with the Optimal Transport flow path (FM-OT) [Lipman et al., 2022]
Flow Matching [Lipman et al., 2022] with the Optimal Transport flow path can be seen as special case of Gaussian diffusion with the weighted loss.
D.3.1 Noise schedule
Note that in [Lipman et al., 2022], time goes from we go forward in time. Here, we'll let time go from 0 to 1 as we go forward in time, consistent with the rest of this paper. We'll Wllll also assume
17

---

To = 0, for which we can later correct by truncation (see Section E.3.1). In this model, the forward process g(zz(x) defined by:
Ze = 0+X + o,E
(61)
(1 1)x + LE
(62)
This implies that the log-SNR is given by:
x fx(t) = log(a (a////
(63)
2 log(( - -1)/t)
(64)
Its inverse is given by:
1 = = 222)
(65)
sigmoid(-1/2)
(66)
The derivative, as a function of t, is:
dA
Ji 2/(1))
(67)
di
This derivative of its inverse, as function of X, is:
dt d
d da (a)
d sigmoid(-2/2) sechh sSech((A/4)/
(68)
dd
dd
The corresponding density
d d ()) sech sec42))))
p(A)
(69)
which is a Logistic distribution; see also Table 4.
D.3.2 Score function parameterization and loss function
Lipman et al. [2022] then propose the following generative model ODE:
dz Ó(z,, t)dt
(70)
The model is then optimized with the Conditional flow matching (CFM) loss:
(x) E=..:oo),..-oooolloooo 6
(71)
where they use the parameterization:
== X -
(72)
D.3.3 Weighting function
What is the weighting function w(X) corresponding to this loss? Note that this parameterization means that:
= (1 - t)x + tc
(73)
= t)o + E
(74)
0 = [z+ - e)/(1 - t)
(75)
Since t t= 1/(1 + eA/2), we have that 1/(1 - t) =1 + e -/2, parameterized as function of X
we have:
o= (z) - cc)(1 + -
(76)
Likewise, we parameterize o o-prediction in terms of e-prediction:
ō(z,, A) = (ZΛ é(z), (1) + - 1/22)
(77)
We translate the o o-prediction loss to e-prediction loss:
||o ó(2, | = (1 + - E(Zx: 2)112
(78)
18

---

Therefore, combining the derivations above, the CFM loss, formulated in terms of the À parame- terization instead of t, and in terms of the e-prediction parameterization instead the o-prediction parameterization, is:
(x) = - ((z)) 1
(79)
-ocooollo-((zz, -0(z, L - 1) dt
(80)
dt di ( dA
(81)
Ania
Anna
(secF(AA///S)Beevvanl// Ö(22, A)11 dl
(82)
Annn
Arnas
+e (/P/
E(2x, 1)12 da
(83)
Amna
WAA)E - -(zx, 2)1121 da
(84)
2 Aman
(()) ((A)
E Bc-NNO.))
||e - E(Zx. 2)1l
(85)
2
where ÿ is any distribution with full support on DAmin Amux) and where:
w ()= (A44)))()
(86)
= e A/2
(87)
Therefore, this weighting is equivalent to the weighting for the v-prediction loss with cosine schedule (Section D.2): the CFM loss is equivalent to the v-prediction loss with cosine schedule.
D.4 Inversion by Direct Iteration (InDI) [Delbracio and Milanfar, 2023]
Delbracio and Milanfar [2023] propose Inversion by Direct Iteration (InDI). Their forward process is identical to the forward process of FM-OT [Lipman et al., 2022] introduced in Section D.3:
2: = (1 - f)x +
(88)
As derived in Section D.3 above, this means that the distribution over log-SNR À is the Logistic distribution: p(A) sech² (A/4((AA)))))) (A/4)/8. The proposed loss function is the x- x-prediction loss:
LIDI (x)
|X x(z,, 1
(89)
Since ||x - x11 =e - é E and the e-prediction loss corresponds to minimizing the weighted
loss with w(A) p(A), the x-prediction loss above corresponds to minimizing the weighted loss with:
w = e P(A)
(90)
= = sech 2 (A/4)/8
(91)
Which is slightly different weighting then the FM-OT weighting, giving a bit more weighting to lower noise levels.
D.5 Perception prioritized weighting (P2 weighting) [Choi et al., 2022]
Choi al. [2022] proposed new weighting function: -dt/d>
P(A)
w()) (k: e> (I: + +2)²
(92)
where empirically they set A: = 1 and 7 either 0.5 or 1. Compared to the e-prediction objective, where w(A) dt/di = p()), this objective put more emphasis on the middle regime of the whole noise schedule, which Choi et al. [2022] hypothesized to be the most important regime for creating content that is sensitive to visual perception. When combined with the most commonly used cosine noise schedule [Nichol and Dhariwal, 2021], the weighting function becomes w(X) sech(2/2)/CC e²)7.
19

---

D.6 Min-SNR-y weighting [Hang et al., 2023]
Hang al. [2023] proposed the following training objective:
MinsNR (x) = [E,-20001),,
[min(e" . ||x - x(zi; A)1
(93)
min(1,,e" 7e x E(zii A)1
(94)
dt
dd ||e é(z;; 1)
-
min
(95)
di
dt
dd dt || E(2i; A)11
-
min }p(())
(96)
Therefore, corresponds to ((A) ye The motivation of the work avoid the model focusing too much on small noise levels, since it shares similar hypothesis to [Choi et al., 2022] that small noise levels are responsible for cleaning up details that may not be perceptible. cosine noise schedule is then combined with the proposed weighting function, leading to w(A) >e
sech(A/2) min(1, 7e J. is set as 5 empirically.
E
Useful Equations
E.1
SDEs
The forward process is a Gaussian diffusion process, whose time evolution is described by stochastic differential equation (SDE):
dz f(2, dt drift
g(t) dw
(97)
diffusion
For derivations of diffusion SDEs, Appendix B of [Song et al., 2021b]. Their B(t) equals
dr log(1 + e>e our formulation, and their Jo 3(x)ds equals log(1 + e->x), where they assume that À 00 at 8 0.
E.1.1 Variance-preserving (VP) SDE
common choice the poriance-preserving (VP) SDE, which generalizes denoising diffusion models [Ho et al., 2020] to continuous time [Song et al., 2021b, Kingma et al., 2021]. In the VP case:
1 d
log(1 + +
f(2,1) = 2
(98)
dt
d
g())² = log(1 + e edx) dt
(99)
a = sigmoid(A) (A)
(100)
of sigmoid(-)) X)
(101)
p(z1) = I),
(102)
E.1.2 Variance-exploding (VE) SDE
Another common choice of the variance-exploding (VE) SDE. In the VE case:
f(z,t) 0
(103)
s(t)² d dt
log(1 + e ~>o)
(104)
0f =1
(105)
of A
(106)
p(z1) N(00 e -
(107)
20

---

E.2 Possible parameterizations of the score network
There are various ways parameterizing the score network:
sz A)
(With the gradient of an energy baagy model)
(108)
(With a noise prediction model) (109)
= ar (o;;)(;
(With data prediction model) (110)
We can let neural network output any of se A), ég or &g, and we can convert the variables to each other using the equalities above.
The chosen relationship between 2,, &, é and so( A) above, are due the following relationships between ZA, x and E:
2 = 0AX + O2E
(111)
" (z) 026)
(112)
0 (z) a1x)
(113)
And:
22 log (og (x) Ve, - -112A 01x11//(02))
(114)
-- a1x)
(115)
-0x2(axx +OC 02x)
(116)
=</00>
(117)
In addition, there's the v-prediction parameterization (v := 0'E TAX) explained in D.2, and the o-prediction parameterization (o : x <) explained D.3.
Karras al. [2022] proposed specific F- F-paramctrization, with an F-prediction model Fg. In the special case of variance explosion (VE) SDE, it is formulated as:
Ha 22
e- EXl2 F
(118)
Ve-* data
data
where data = 0.5. Generalizing this to our more general family with arbitrary drift, this corresponds to:
e 2/22
ax data
X c + σ2
2A +
(119)
Ve + 02 data
data
So that we have:
Ve e + data
Tdata
F
A/22daata
x
e- e- data
(120)
Ve-x+ data
c/2 + oz Jasta
Z
(121)
Tdata
Ve-x F
In summary, given these different parameterizations, the e-prediction loss can be written terms of other parameterizations as follows:
|| é0113 e x x0l13
(e-prediction and x-prediction error)
(122)
VRl/Ve, log (zx)(z) se
(score prediction)
(123)
= 2 + +1 1) voll
(v-prediction, general)
(124)
-e- + 11 vol
(v-prediction with VP SDE)
(125)
-c->/00 + +1) 1 Fol1
(F-prediction)
(126)
Interestingly, we set 22 2 1, the training objectives of F-prediction and v-prediction are the
same.
21

---

Noine schedude name
A = fx(0) E...
t= (A)==
p(A) (A)=
Cosine
2log(ta((st/22)
(2/n) arctan(e A/2)
sccb(A/2)((x))
Shifted cosine
2log(tan(mt/2)) +
(2/ aretan(e -A/2-1)
sech(A/2 ()/
EDM (training)
FX' (1:444 2.49)
2(4;244,2444)
EDM (samplingg
+01-000 1/e (m
2e 2olo 1/e
1- -A((8)) 1/c
1/0 OoiA
Re(e -
172 1/P)
>7(2pp)
How Matching with OT (sec D.3)
2 1og((1 6)/4)
1/(1 -0----3 1/3
sech"(A/4)3
Table 4: Noise schedules used in our experiments: cosine [Nichol and Dhariwal, 2021], shifted cosine [Hoogeboom et al., 2023], and EDM [Karras et al., 2022] training and sampling schedules. Note that these are the noise schedules before truncation (Section E.3.1).
E.3
Noise schedules
During model training, we sample time / uniformly: ~ U(0,) 1), then compute À fx((). This results in a distribution over noise levels p()), whose cumulative density function (CDF) is given by 1 1-55'(A). For E [Amin Amux] the probability density function (PDF) is the derivative of the CDF, which is p(A)) = - (d/dA) (A) = dl/d/ = - 1//1(() Outside of [Aiinn Amax] the probability density
In table 4 we provide some popular noise schedules: cosine [Nichol and Dhariwal, 2021], shifted cosine [Hoogeboom et al., 2023], and EDM [Karras et al., 2022]. We do not list the 'linear' schedule by [Ho al., 2020], fx(t) log(ee 1) which has fallen out of use. Note that:
In the shifted cosine schedule, s = log(64/d), /d), where 64 the base resolution and d the used resolution (e.g. 128, 256, 512, etc.).
In the EDM training schedule, the function Fxr(A; μr, o²) is the Normal distribution CDF, and N (N;; μu, o is its PDF.
In the EDM sampling schedule, p == 7, Juin 0.002, Fmax 80. The density function p(A) in the table has support
log " log m. Outside this range, p()) = O.
Note that the table gives the noise schedules before truncation, resulting in [FA(0), Fx()] = [oo, - -oo]. The truncation procedure is given in E.
E.3.1 Truncation
The noise schedules above are truncated, resulting a noise schedule (0) whose endpoints have desired values [f)(0), (1)) = Amaax Amin]:
(x) == ft + - (t))t)
(127)
where: to = fx (Amax)
(128)
41 == S (Amia)
(129)
Its inverse is:
-(0))(11))66) (o) -10)
(130)
And the corresponding probability density:
d
d
if Amin S Xmax p(A) dx (A)
(A)= =-
(2)
dx (2)/((1 -t) - pp(2)/(( -ta
(131)
clse: p(A) 0
(132)
E,4 Sampling
Anderson [1982] showed that if so (z; A) Vz log qt (z), then the forward SDE exactly reversed by the following SDE:
dz = [((z t) " (z: 1) + g(t)dw
(133)
22

---

Fixed vs adaptive schedules
Fixed vs adaptive schedules
0.090
Fixed noise schedule (Cosine)
Fixed noise schedule (EDM)
0.089
Adaptive noise schedule
15.7
Adaptive noise schedule (Ours)
0.088
15.6
Lo55
0.087
15.5
0.086
0.085
15.4
0.084
50000
100000
150000
200000
250000
50000
100000
150000
200000
250000
Training iteration
Training Iteration
(a)
(b)
Figure 4: Our proposed adaptive noise schedule allows to freely change the weighting function without needing to handtune a corresponding noise schedule. For some models it leads to slightly worse convergence speed (left) compared to a static noise schedule, probably because the noise schedule was already well-tuned for the weighting function, while in other cases it led to faster convergence (right). The adaptive noise schedule did not significantly affect the end result, but it did allow us to more freely experiment with weighting functions.
Recent diffusion models have used increasingly sophisticated samplers. As an alternative to solving the SDE, [Song et al., 2021b] showed that sampling from the model can alternatively be done by solving the following probability flow ODE:
dz = [f(. 2 { A)]dt
(134)
which, under the assumption that Sg is conservative vector field, will result in the same marginals p(z:) as the SDE of Equation 133 for every / E [0, 1J, and therefore also the same marginal p(x).
Note that due the continuous-time nature of the model, any sampling method is necessarily approximate, with the discretization eπror depending on various factors including the choice of noise schedule. For sampling we can therefore typically use a different noise schedule fx for sampling than for training, and we can change the SDE drift term; as long as we appropriately rescale the input to the score network, this would still result in correct samples.
F Adaptive noise schedule
The invariance shown in Section 3.2 holds for the loss L (x), but not for the Monte Carlo estimator of the loss that we use in training, based on random samples ~ ((00) 1), E ~ (V I). The noise schedule still affects the variance of this Monte Carlo estimator and its gradients; therefore, the noise schedule affects the efficiency of optimization.
In fact, the noise schedule acts as an importance sampling distribution for estimating the loss integral of Equation 6. Specifically, note that p()) -1/(dX/dt), We can therefore rewrite the weighted loss as the following, which clarifies the role of p(A) as an importance sampling distribution:
(A)
L (x) = 1 (... 2
|l (εo(z)))) -
=11
(135)
In order to avoid having to hand-tune the noise schedule for different weighting functions, we implemented an adaptive noise schedule. The noise schedule Xe is updated online, where we let p()) x Ex= ( (A)) (0(z;; A) - =l|ž]. This noise schedule ensures that the loss is spread evenly over time, i.e. that the magnitude of the loss ExxD.e>>N(0 [(v(())/(AA]ll A) < <[]]
approximately invariant to À or t. We find that this often significantly speeds op optimization.
We implemented an adaptive noise schedule p()). where:
p(A) x E Ex Px-DDe>>N000 ((...) (A)lIC- (oolic-eo(zxx)
Co(2x; 2)1
(136)
23

---

In practice we approximate this by dividing the range Dmiin Amax] into 100 evenly spaced bins, and during training keep an exponential moving average (EMA) of w(A)] é éz A)ll within each bin. From these EMAs we construct a piecewise linear function (x(() such that Equation 136 is approximately satisfied. The EMAs and corresponding noise schedule p(A) are updated at each training iteration. Similar noise schedules have been proposed in Zheng et al. [2023], Nichol and
Dhariwal [2021], Dieleman et al. [2022].
In experiments we measure the effect of changing the fixed noise schedule of existing modules with an adaptive schedule. We found that this lead to approximately equal FID scores. In half of the experiments, optimization was approximately as fast as with the original noise schedule, while in the other half the adaptive noise schedule lead to faster optimization (see Figure 4). The end results were not significantly altered.
G Relationship between the KL divergence and Fisher divergence
We'll use the following definition of the Fisher divergence [Lyu, 2012]:
(x(((x))lp(x)) = [x) log g(x) Vx lo logg((x))1
(137)
Theorem 2. Assume a model in the family specified in Section 2, and assume the score network encodes a conservative vector field: sa(zn, Ae) Va, log p(z)) (not assumed by the other theorems). Then:
/x0_l/ = J dA
DKu(g(Z
(138)
-2"
Proof of Theorem 2. Note that (see Equation 117):
22, log g(z)(x) = - c/Ox
(139)
And assume the score network encodes conservative vector field:
Vx, logp(z)) = so (2,,, Ae) = - (z( 0)/ox
(140)
So the time derivative of Equation 22 can be expressed as:
d
Dr dx
- (2)1))))))) [l|V log (ogg(2x()) - VzL (og)(2))
(141)
Equation 138 follows from the definition of the Fisher divergence.
G.1 Comparison with Theorem by Lyu [2012]
Lyu [2012] prove a similar result in their Theorem 1. We'll translate their result into notation. In particular, let the forward process as our family, such that ((z//x) = (zz axx, I. The marginal (data) distribution is qxx), such that ((z)) q(z,) (x))xx Similarly, let the generative model have marginal p(x), and p(z,) J p(z, x (xx))xxxxx.)..))))x))))))))) So far the assumptions are the same in our family.
They assume that q(z x) = WZz,, x), which corresponds to a variance exploding (VE) diffusion process, with o. so À ltg Importantly, they make the assumption that p(z/(x) = xx), ((,,)x), i.e. that the forward process for p equals the forward process for Given these assumptions, Lyu [2012] show that:
d D = Pr(g(z))ll(())))
(142)
dt
Which, given the noise schedule X log(t), can be rewritten
D Dxz(((2))llx(z)) D
(143)
dA
which looks a lot like our Equation 138. One difference are that in Equation 138, the left-hand- side distributions are joint distributions, and q conditions on x, while Equation 143 is about the
24

---

unconditional q. Another key difference is that for Equation 138 we need fewer assumptions: most importantly, we do not make the assumption that p(z, x) g(z,(x), since this assumption does not hold for the family of diffusion models we consider. Before or during optimization, p(z, |x) might be very far from g(zz|x). After optimization, p(zz x) might be close to q(z,) [x), but we still can't t assume they'r re equal. In addition, we're mostly interested in the properties of the loss function during optimization, since that's when we're using our loss for optimization. We for this reason, our
Theorem 2 is a lot more relevant for optimization.
H
Implementation details
Instead of uniformly sampling f, we applied the low-discrepency sampler of time that was proposed by Kingma et al. [2021], which has been shown effectively reduce the variance of diffusion loss estimator and lead to faster optimization. The model is optimized by Adam [Kingma and Ba, 201 2014] with the default hyperparameter settings. We clipped the learning gradient with a global norm of 1.
For the adaptive noise schedules, we divided the range of [Amin Amax] into 100 evenly spaced bins. During training, maintaicdn exponential moving average of w (A)]le (zxz A)ll with decay rate 0.999, and a constant initialization value of for each bin.
Below we elaborate the implementation details specific for each task.
ImageNet 64x64. For coadscconiiional generation on ImageNet 64x64, we applied the ADM U-Net architecture from Dhariwal and Nichol [2022], with dropout rate 0.1. We didn't ' use any data augmentation. The model was trained with learning rate 1e 4, exponential moving average of 50 million images and learning rate warmup of 10 million images, which mainly follows the configuration of Karras al. [2022]. We employed 128 TPU-v4 chips with batch size of 4096 (32 per chip). We trained the model for 700k iterations and reported the performance of the checkpoint giving the best FID score (checkpoints were saved and evaluated on every 20k iterations). It took around 3 days for single training run. For training noise schedule and sampling noise schedule of DDPM sampler, we set Amin 20 and Amax 20. We fixed the noise schedule used in sampling to the cosine schedule for the DDPM sampler, and the EDM (sampling) schedule for the EDM sampler (see Table for the formulations). We adopted the same hyperparameters of EDM sampler from Karras et al. [2022] with changes (i.c., Table 5 in their work, column 'ImageNet-One Our model'), ) Both DDPM and EDM samplers took 256 sampling steps.
ImageNet 128x128. For class-conditional generation ImageNet 128x128, we heavily followed the setting of simple diffusion [Hoogeboom et al., 2023]. Specifically, we used their 'U-ViT, L' architecture, and followed their learning rate and EMA schedules. The data was augmented with random horizontal flip. The model is trained using 128 TPU-v4 chips with a batch size of 2048 (16 per chip). We trained the model for 700 iterations and evaluated the FID and inception scores every 100k iterations. The results were reported with the checkpoint giving the best FID score. It took around 7 days for single We set Amio 15 s and Xmax 15 + 8, where the shift of the weighting function, with 64 being the base resolution and dl being the model resolution The DDPM sampler used for evaluation used "shifted-cosine noise schedule (Table 4) and took 512 sampling steps.
log(6e (64/d) (d)d) is
(d 128 for this task).
I Relationship with low-bit training
Various earlier work, such as [Kingma and Dhariwal, 2018], found that maximum likelihood training on 5-bit data can lead to perceptually higher visual quality than training on 8-bit data (at the cost of a decrease in color fidelity). This leads to improved visual quality, probably because it allows the model to spend more capacity modeling the bits that are most relevant for human perception.
In [Kingma and Dhariwal, 2018], training on 5-bit data was performed by adding uniform noise the data, before feeding it to the model. It was found that adding Gaussian noise had similar effect as uniform noise. As we have seen, in the case of diffusion models, adding Gaussian noise equivalent to using weighted objective with a monotonic weighting function.
25

---

Therefore, training on 5-bit data is similar to training using a monotonic weighting function in case of diffusion models. We can wonder: which weighting function emulates training on 5-bit data? Here, we'll attempt to answer this question.
I.1
The shape of (();x) for low-hit data
Note that the results of Appendix A.1 can also be written as:
1 = [||e εo(2x:
((:x)
(144)
2
This allows to rewrite the weighted loss as simply:
Amii
C (x)
-
dx ((xxx) w(A) dA
(145)
Amas
To understand the effect of X, we can plot the a. x)) as a function of X.
We'd like to plot A [(); x) dof different choices of bit precision. This will tell us where the different bits 'live' as a function of A. Since training different diffusion models on different bit precisions is very expensive, we instead use an approximation.
particular, we assume that the data x is univariate, with uniform distribution xxx over the 2n possible pixel values, where n is the bit precision. The data is x is, as usual, normalized to [-1, 1]. We then let the model p(2x), for each choice of X, be the optimal model: p( 22))= == f ((x))qZxx)))))))) (x )dz, which is univariate mixture-of- mixture-of-Gaassians where each mixture component is a Gaussian centered one of the 2n possible pixel values. In this case, x(A;x) := DKL(o(ZA)) [((x)).
We start by plotting the function (,(x0) (A(); x)], for differences choices of = 1,.., 8, as shown in Figure 5a. Next, we plot [, ((): x)], ). for each n, in Figure 5b. The more bits included in the data, the more information contained by (x) [ [)) (measured by the area under each curve). We can further visualize the contribution of each additional added bit to the loss, by subtracting the curve for 1 bits from the curve for n bits, as displayed in Figure Sc. The area under each of these curves exactly bit. Interestingly, as more bits being added, the contribution to the loss gradually shifts from low-SNR regions to high-SNR regions.
Now the connection between low-bit training (e.g., 5-bit) and the weighted training objective becomes clear: If we want to "shut down" the last three curves in Figure 5c, a way is to have a sigmoidal 7.5 and X = 10. In fact, the 5-bit unweighted loss curve is very similar to the 8-bit loss curve, when using the following weighting function:
weighting function that goes down sharply between
w (A) = Px ( 2(A - 8.4)))
(146)
where Fyr is the CDF of a standard Gaussian. See Figure 5e for the comparison of the 5-bit unweighted loss and the 8-bit weighted loss, and Figure 5d for the visualization of the above weighting function ('5-bit-like weighting'). Interestingly, this weighting function gives much more weight to low noise levels than the weighting functions used in this paper (i.e. the other curves shown in Figure 5d).
J Fourier analysis of Gaussian noise perturbation
Recall Theorem 1 showed that assuming monotonic weightings, the weighted training objective is equivalent the ELBO with data augmentation, namely Gaussian noise perturbation. One may wonder, why Gaussian noise perturbation can help improve perceptual quality. In this section, we try answer this question by lifting Gaussian noise perturbation the frequency domain.
Given a clean natural image x and a Gaussian white noise image e, the Gaussian-noise -oit----oi----ee-eeeddeeee data is given by zA axx + axe. Denote the Fourier transform operation as F(-). Since F(-) is linear operation, we have
F(Z,) - axF(xx + oxF(c).
(147)
To understand how noise affects the data, we can compare the power spectra of F(x) and ((c).
which is defined as the square amplitude of the Fourier transform: S(-) = [F))) More specifically, given the power spectrum S of an image, we can take the circular average of S for the Fourier
26

---

Equ(c(A.xx) (Expected
EaviEdoa ClA. xi]
1.bil -1
0.
n_bits u:2
nubits -
0LbIS
1LBIts
rubItS
1_bit
0.2
n_bits
n_bits
nubits -8
10
-5
10
15
10
(a) (xc) (A(;; x) a function of A.
(b) (x0) lax [(); x)] a function of A.
Contribution each bit dqui xCA,0)
Weighting furcttions
EL80
tath to Eadidin
0.30
0.2
Skitike weightng
0.6
0.15
0.4
0.00
-10
-5
1.0 15
-7.5
-3.0
-2.5
0.0
2.5
5.0
7.5
10.0
(llg SNR)
(c) Subtraction of the curve for n bits from the (d) Weighting functions, with weighting resembling curve bits in Fig.5b.
5-bit training.
Unweighted 5-bit weighted B-bit
Ea. <(l,) will n. bis
(e) Unweighted 5-bit loss versus weighted 8-bit loss.
Figure Low-bit training analysis. See Section I.1 for description.
components of the same spatial frequency, such that it becomes function of the spatial frequency. Figure 6 displays the average power spectra of axF((x and oxF((e as functions of the spatial frequency, for different X. As À goes from high low, i.e., gradually perturbing the data with more noise, the average power spectrum of F(F) (orange curve) first exceeds the one of axF(x) (blue curve) in the high frequency region, and then followed by the low frequency region. That is, the high frequency components of the data are destroyed more quickly than the low frequency components. When the data perturbed with small amounts of noise, only high frequency details are significantly removed, and most low frequency content remains. Therefore, given that the model capacity is limited, learning data with Gaussian noise perturbation help the model focus components of the data, which often correspond to high level content and global structure that are more crucial human perception. This analysis can also be applied to understand the forward process of diffusion models.
low frequency
27

---

perturbation
Am.5.0
A=.2.5
A.0
A.
A=5.0
Aueroge pooeee spatnum
107
307
107
107
107
107
0,712)
10
(5/14)
10
0,714)
10
10
0,714)
101
10
10
10
10
9y71X
101
10
10
10
101
Spatial frequency
frequency
frequency
frequency
frequency
Figure 6: The circular average power spectrum function of spatial frequency, for the scaled clean image and scaled Gaussian noise image. From the rightmost to the leftmost subfigure, more noise is gradually added.
To make the claim even clearer, we can compute the log signal-to-noise -ogos--tonooss--ratoo (log SNR) À between the power spectra of (x) and (c)), as function of X. Figure shows the average value of that log ratio around different frequency regions. Given a fixed X, is higher for low frequency regions, so that the data are less destroyed by the noise.
fraq11111
(log SNR around frea)
freq=3.91
15
freq-13.62
freq=47.45
10
freq-165.37
-10
-15
(log SNR)
Figure 7: Log-SNR ()) of the power spectra of the scaled clean image and scaled Gaussian noise image, as a function of log-SNR (A) the data space. Given a fixed x, À is higher around low frequency regions.
K Limitations
It is important to emphasize that empirical results, like other deep learning approaches, depend on the choice of prperrparameters. A change in, for example, the dataset or spatial resolution will generally require re-tuning of optimization hyperparameters, architectural choices and/or weighting functions. Such re-tuning can be time consuming and costly.
L Broader impact
While our work primarily focuses on theoretical developments the understanding and optimization of diffusion models, the advancements could have broader implications. some of which could potentially negative. The development of more efficient and effective generative models could, on one hand, propel numerous beneficial applications, such as art and entertainment. However, it is also worth acknowledging the potential misuse of these technologies.
One notable concern is the generation of synthetic media content, for example to mislead. These fraudulent yet realistic-looking images and videos could be used for spreading disinformation for other malicious purposes, such as identity theft or blackmail.
Regarding faimess considerations, generative models typically trained on large datasets and could therefore inherit and reproduce any biases present in the training data. This could potentially result
28

---

in unfair outcomes or perpetuate harmful stereotypes if these models are used in Decision-making processes or content generation.
Mitigation strategies to address these concerns could include the gated release of models, where access the model its outputs is regulated prevent misuse. Additionally, the provision of defenses, such as methods to detect AI generated media, could be included alongside the development of the generative models. Monitoring mechanisms could also be implemented to observe how system is being used and to ensure that it learns from feedback over time in an ethical manner.
These issues should be considered when applying the techniques propose. The community should strive for an open and ongoing discussion about the ethics of Al and the development of strategies to mitigate potential misuse of these powerful technologies.
M Samples from our model trained on 512 x 512 ImageNet
Below we provide random samples from our highest-resolution (512x512) model trained on ImageNet. We did not cherry-pick, except that we removed depictions of humans due ethical guidelines. Samples in Figures and are generated without guidance, while samples in Figures 10 and 11 are generated with guidance strength 4.
29

---

Eurewoh CM T. Rutra
03
OG0
08
DI
Figure 8: Random samples from our 512x512 ImageNet model, without guidance.
30

---

Figure 9: More random samples from our 512x512 ImageNet model, without guidance.
31

---

I
2
PPPE
Figure 10: Random samples from our 512x512 ImageNet model, with guidance strength 4.
32

---

BRUTINLA
N8
P.C
Figure 11: More random samples from our 512x512 ImageNet model, with guidance strength 4.
33


Reviews:
[{"review_id": "cUBTb8pcZ2", "reviewer": "Reviewer_kL4L", "summary": "The paper tackles the loss function of diffusion models from a theoretical point of view. The paper studies the relationship between the weighted loss and the ELBO objective. The paper shows that the diffusion model objectives can be formulated as a special case of weighted loss. In addition to the theoretical contribution, the authors propose monotonic weightings and provide an empirical study of monotonic and non-monotonic weighting.", "questions": "NA", "limitations": "NA", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem studied is important\n- Formulating the different diffusion model loss as weighted loss is interesting\n- The paper presented a detailed theoretical analysis\n- Experimental results support the proposed method", "weaknesses": "- Adding more experimental results could have been helpful but this is not a major point.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback. Please note our top-level comment with additional experimental and theoretical results. Below we address specific questions.\n\n**Q: Adding more experimental results could have been helpful.**\n\nA: Please see our top-level comment and the PDF attached to it. We have added high-resolution ImageNet generation results (i.e. 256x256 and 512x512) to further demonstrate the effectiveness our approach. With our 'EDM-monotonic' function and adaptive noise schedule, we achieved SOTA FID and Inception scores on high resolution ImageNet (i.e., all three resolutions of 128, 256, 512)."}, {"review_id": "biMgnbMSF5", "reviewer": "Reviewer_DESf", "summary": "This is a nice paper that provides a comprensive framework for diffusion objectives, giving the ELBO the driver's seat.\n\nThe paper shows that all of the axes along which diffusion objectives differ (elbo vs Ho-perceptual weighting, noise schedules, different parameterizations like x vs epsilon vs v) can all be seen under a type of weighted and importance sampled ELBO. Moreover, the paper show that the weightings that don't exactly look like the original ELBO can be made equal to it under the right data augmentation. Finally, the authors also present a few techniques for adaptive/learnable importance sampling distributions to help assemble all the aforementioned ingredients in a computationally efficient manner. Finally, the authors show conveying results for image quality for ELBO trained models, whereas people usually result to non-ELBO trained models to do this.", "questions": "Just a brief suggestion:\n\n- around eq 3 it would be good to mention that the f(x,t)-g^2(t) s_\\theta(x,t) form is a backward time SDE since some other authors use the g^2(T-t)s_\\theta(x,T-t) -f(x, T-t) forward time SDE.\n\n- Not necessary to address but I'm quite curious which of the invariance and change of variables bits apply for multivariate (augmented) diffusions like CLD (https://arxiv.org/abs/2112.07068), MDMs (https://arxiv.org/abs/2302.07261), PSLD (https://arxiv.org/abs/2303.01748), and FP-Diffusion (https://arxiv.org/abs/2206.10365). You can't quite take the g^2/sigma^2 weighting out of the squared error like in the scalar diffusion case, but I'm sure there should be some relevant version of your results.", "limitations": "Yes.", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I'm a big fan of papers like this since their primary aim to build solid understanding for people that already work with diffusions but that may have some doubts about how the common moving parts relate to each other (noise schedules, elbo weighting, parameterizations, etc). \n\nOn top of the above, they also present great results for ELBO trained models, whereas there is some common beliefs that non-elbo \"perceptual weighting\" (what the authors call \"reweighted losses\") are the only ones that can be used for good image quality.", "weaknesses": "No comments/weaknesses.\n\nThe adaptive importance sampling also comes up in the following work:\n\nImproved Techniques for Maximum Likelihood Estimation for Diffusion ODEs https://arxiv.org/abs/2305.03935\n\nThis is completely fine since:\n- I can see how both works would arrive at this algorithmic detail independently, given that both dive deeply into related but distinct variance issues in diffusion elbos.\n- both works present a lot of additional insights and techinques beyond this\n\nI do find it appropriate to cite and discuss relationship in either case though, including possible similarities and differences, since it will really help readers of both papers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and kind words to our work. Please note our top-level comment with additional experimental and theoretical results. Below we address specific questions.\n\n**Q: Cite and discuss relationship with “Maximum Likelihood Estimation for Diffusion ODEs” (https://arxiv.org/abs/2305.03935).**\n\nA: We appreciate your pointer to this related work, which also proposes noise schedules with the purpose of reducing the variance of diffusion objectives. They proposed two approaches: 1) design $p(\\lambda)$ such that the coefficient of the mean squared error of normalized v-parametrization is a time-invariant constant; 2) learning a noise schedule network $\\lambda(t)$ by minimizing the variance of the training objective, similar to [Kingma et al, 2021]. While 2) in theory can seek the optimal importance sampling, it introduces extra computational cost with additional gradient operations. We therefore choose a less expensive approach: we maintain the exponential moving average of $w(\\lambda)||\\epsilon - \\hat{\\epsilon}_\\theta||^2_2$, and adaptively update the noise schedule $p(\\lambda)$ to be proportional to that at each iteration. Similar approach was proposed by [Nichol et al, 2021]. \n\nWe will cite this paper and discuss the relationship in our final version, as you suggested.\n\n**Q: Not necessary to address but I'm quite curious which of the invariance and change of variables bits apply for multivariate (augmented) diffusions.**\n\nA: That's a very good point! A more general formulation would be interesting, which we leave open for future work.\n\nThanks for your suggestion on clarifying Eq 3, which we have incorporated to the updated draft. "}, {"review_id": "cnTQRcozFc", "reviewer": "Reviewer_eoU3", "summary": "This paper presents new theoretical insights on denoising diffusion objective. Namely, training a diffusion model with a non-uniform weighting can be expressed as a weighted integral of ELBOs, and monotonic weighting results in training ELBO with additive Gaussian noise augmentation. The authors also propose an adaptive noise schedule to ensure all losses are trained uniformly over time and show the adaptive schedule with a monotonic weighting attains SOTA FID on ImageNet.", "questions": "-", "limitations": "Yes, the authors addressed limitations, potential negative societal impact, and mitigation.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Although not shown in this work, the interpretation that denoising diffusion with monotonic weighting trains ELBO with data augmentation means that the same data augmentation can be applied to other likelihood-based models for a fair comparison to diffusion models. And, we would have a clearer understanding of sample quality and likelihood of various model classes.", "weaknesses": "-", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments. We strongly agree with you that understanding weighted diffusion training objectives as ELBO with data augmentation gives us a way to fairly compare diffusion models with other likelihood-based models, and indicates that ELBO objective is still compatible with high sample quality. Please note our top-level comment with additional experimental and theoretical results."}, {"review_id": "zO2TwneT4E", "reviewer": "Reviewer_bjDS", "summary": "The paper takes a systematic approach to address two important aspects of the implementation of diffusion-based generative models: the objective function and the noise schedule, both of which have been conventionally set in a relatively empirical fashion ($i.e.$, based on sample quality).\n\nThe paper first frames the noise schedule in terms of the log SNR $\\lambda_t$ as a function of time $t$, then writes the diffusion model objective in the form of a weighted denoising score matching objective, Eq. (6). By doing so, the paper unifies the existing works as special cases under a generic form of the objective function with its corresponding noise schedule.\n\nThe paper continues with a theoretical analysis of the diffusion model objective, by first revealing the ELBO objective as one particular realization of the weighted objective, then provides the key theoretical result via Theorem 1, which states that **under the assumption of a monotonically increasing weighting function $w(\\lambda_t)$, the generic diffusion model objective Eq. (6) is equivalent to a negative ELBO objective of noise-perturbed data (a form of data augmentation, controlled by the noise schedule $\\lambda_t$)**. This result has recast the ELBO as the class of objective functions for training diffusion models that are capable of producing high-quality samples.\n\nWith the theoretical results, the paper proposes several new monotonic weighting functions, and demonstrates through experiments that monotonic weighting enables the trained diffusion models to generate high quality samples (competitive with the best diffusion models) for image tasks in terms of FID and IS, including a new SOTA FID score of $1.75$ on class-conditional $128\\times 128$ ImageNet dataset.", "questions": "* Could the authors provide some explanation for the following sentences:\n    * Line 108-109, “During sampling, $p(\\lambda)$ provides the relative amount of time the sampler spends at different noise levels.” It’s not immediately clear what it means, as it seems that by sampling $t\\sim U(0,1)$, each noise level has an equal chance of being sampled.\n    * Line 179-180, “The probability distribution $p_w(t)$ has Dirac delta peak of typically very small mass $w(0)$ at $t = 0$.” Does this mean that the almost noiseless sample contributes very little to the objective computation?\n* Line 262-265: have the authors run experiments for a $v$-prediction model combined with a non-monotonic weighting function?\n* In Figure 1 Right, only the “EDM sampling” noise schedule has $p(\\lambda)$ monotonically decreasing as a function of $\\lambda$. Is there any interpretation of that, $e.g.$, advantages for the EDM sampling scheme to monotonically decrease?\n* It’s a bit confusing on why the noise schedule is casted as a probability distribution $p(\\lambda_t)$ (besides the importance sampling interpretation): $e.g.$, the cosine schedule has the variance at each timestep as a deterministic function of time. Why does $\\lambda$ need to be sampled – is it because timesteps need to be sampled during training?", "limitations": "The authors have discussed both limitations and potential negative social impact via Appendix H and I from Supplementary Material, which speak for the community well.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "* Among all degrees of freedom for the implementation of a diffusion-based generative model, setting the weights at each noise level for the objective function, and noise scheduling, are two crucial aspects for its success that have traditionally been mainly based on empirical choices. By addressing these two aspects with a systematic and rigorous approach, this work fills an important gap for the diffusion model research community, thus is a well-motivated and potentially very beneficial work.\n* By recasting the diffusion model objective under the restriction of the weighting function as an ELBO objective, the work can bring back this important class of objective functions along with the existing literature on training with an ELBO objective.\n* The paper proposes several new monotonic weighting functions for the weighted diffusion objective.\n* The paper is overall well-organized and clearly written: it can be viewed as a good reference that introduces the diffusion-based generative modeling framework collectively; the color-coding for different components of the objective functions is a good design for presenting derivations.", "weaknesses": "* The part of the paper that derives the main conclusion of the paper (the diffusion objective is equivalent to the ELBO with data augmentation), specifically Line 181-185, went a little bit too fast-paced and might lack of some details to facilitate understanding; some descriptions don’t exactly match the mathematical expressions:\n    * The KL divergence $D_{KL}(q(z_{t,\\dots,1}\\vert x)\\vert\\vert p(z_{t,\\dots,1}))$ being “the negative ELBO of noise-perturbed data, plus a constant” is not immediately clear. It perhaps would be nice for the authors to provide the derivations, to establish $D_{KL}(q(z_{t,\\\\dots,1}| x)|\\vert p(z_{t,\\\\dots,1})) =-\\\\mathbb{E}_{q(z_t|x)}[\\\\text{ELBO of }\\\\log p(z_t)]-\\\\mathcal{H}(q(z_t|x))$, where\n\n   \t $\\text{ELBO of }\\\\log p(z_t)=\\\\mathbb{E}_{q(y | z_t, x)}[\\\\log p(z_t, y)-\\\\log q(y | z_t, x)],$ in which $y=z\\_{t+\\\\triangle t,\\dots,1},$ and $\\\\triangle t$ is the time increment.\n\n    \tSimilarly, the inequality in Eq. (9), $D_{KL}(q(z_{t,\\\\dots,1}| x)|\\vert p(z_{t,\\\\dots,1})) = D\\_{KL}(q(z_t | x)|\\vert p(z_t)) + \\\\mathbb{E}_{q(z_t | x)}\\big[D\\_{KL} (q(y | z_t,x)||p(y | z_t))\\big]\\\\ \\geq D\\_{KL}(q(z_t | x)|\\vert p(z_t))$.\n\n    \tIt might be helpful to clarify these terms, since Line 181-182 seems to be the very part that constructs the main claim of the paper $s.t.$ the diffusion objective can be viewed as the (negative) ELBO of noise-perturbed data.\n    * Line 184-185 says that minimizing $\\mathcal{L}\\_w(\\\\theta)$ with a monotonic weighting “is equivalent to maximizing the ELBO of noise-perturbed data”. The expression is a bit strange as it indicates there is only one ELBO term to be optimized, while the objective function is shown in this section to be viewed as a collection of different ELBO terms. Therefore, this simplified description might cause some misinterpretation.\n    * The underbracket text for Eq. (9) and (10) are both “Neg. log-likelihood of noise-perturbed data”, while the expressions are different expectations of $-\\\\log p(z_t)$.\n* The authors mentioned that re-establishing the diffusion objective as an ELBO “allows for a direct apples-to-apples comparison … with other likelihood-based models” (Line 63-64), but it’s unclear on how such comparison can be made: \n    * Does it refer to the training loss when the autoregressive transformer is trained *with the same data augmentation scheme*? If so, it would only make sense if the adaptation of the data augmentation scheme would not worsen the autoregressive transformer’s performance in the first place. \n    * Or does it refer to the computation and comparison of the exact likelihood?\n\n    It would be nice for the authors to elaborate more on how such comparisons can be made, as it occurs noticeably at the end of both the Introduction and the Conclusion section.\n\n> Minor Issues\n* Line 10 in the Abstract says that the weighted loss “equals the ELBO”. Based on Section 3.1, would it be more accurately described as “equivalent to optimizing the ELBO”?\n* Line 39 “This has has led”.\n* Line 42 “training objective” might be “training objectives”.\n* Section 2.5 has both forms for the score network,  “$s_{\\theta}(z;\\lambda)$” and “$s_{\\theta}(z,\\lambda)$”.\n* Line 142 comma between “$x$-prediction” and “$v$-prediction”.\n* Eq. (8) 2nd and 3rd line: the expectation has “$t\\sim U(0,1))$”.\n* Table 5 text box ends with two periods.\n* Line 62 & 295 the expression “The newfound equivalence between monotonic weighting and the ELBO with data augmentation” seems odd; do the authors mean “The newfound equivalence between monotonically weighted diffusion objective and the ELBO with data augmentation”?\n* It’s a bit confusing to first define $\\lambda$ as the log SNR (Line 86), then call it the noise level (Line 91): the sentence in Line 93-94, “At timestep t = 0 the noise level … $\\lambda$ is high, meaning very little noise; at t = 1 the noise level $\\lambda$ is low, meaning a lot of noise”, sounds a bit odd.\n* The term “noise schedule” is overloaded: in the paper, it simultaneously means the log SNR term $\\lambda_t$ itself (Line 149), the function $f_{\\lambda}$ (Line 90), and the distribution $p(\\lambda)$ (Line 510). It’s clear to me that each “noise schedule” would result in a particular value of the variance ${\\sigma_{\\lambda}}^2$ at each timestep, but I’m not sure if such overloading would create any issue in the future, thusI’ll just point it out here.\n* Line 179-180 the weighting function $w$ is simultaneously a function of $\\lambda_t$ and of $t$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and thoughtful comments. Please note our top-level comment with additional experimental and theoretical results. Below we address specific questions.\n\n**Q: The part of the paper that derives the main conclusion of the paper [...], went a little bit too fast-paced [...].**\n\nA: Thank you for your valuable feedback. We have expanded and clarified this section following your suggestion, and added a detailed derivation to the Appendix.\n\n**Q: The authors mentioned that re-establishing the diffusion objective as an ELBO “allows for a direct apples-to-apples comparison … with other likelihood-based models” (Line 63-64), but it’s unclear on how such comparison can be made. [...].**\n\nA: We will clarify this part of our manuscript. As we mentioned in Section 5 (Related work), the weighted loss in Equation 10 is a type of DistAug [Jun et al., 2020], a powerful method of training data distribution augmentation for generative models. In DistAug, the generative model is conditioned on the data augmentation parameter while training, but conditioned on \"no augmentation\" during sampling. It was shown by [Jun et al., 2020] that DistAug can substantially improve sample quality of autoregressive models on images.\n\nIn our case the augmentation is simple additive noise with noise level $\\lambda_t$, and the model $p(z_t)$ as provided by the score network is conditioned on the noise level $\\lambda_t$. After training, the sampling procedure provides samples from $p(z_t)$ at $t=0$, which has virtually no augmentation. So this is directly analogous to DistAug where the model is trained with various types of augmentation (and conditioned on the augmentation indicator), and sampled by setting the augmentation indicator to \"no augmentation\".\n\nSo, we can optimize other types of models similarly using DistAug, where we train with a distribution $t \\sim p_w(t)$, translating to a distribution over levels $\\lambda_t$, where the generative model is conditioned on the augmentation indicator $\\lambda_t$. After training, we sample from the model with the augmentation indicator set to \"no augmentation\", i.e. $t=0$.\n\n**Q: About minor issues.**\n\nA: All points raised here are great. We appreciate your careful reading of our draft, and we will incorporate the suggestions into our updated draft.\n\n**Q: Line 108-109, \"[...], $p(\\lambda)$ provides the relative amount of time the sampler spends at different noise levels\". It’s not immediately clear what it means, [...].**\n\nA: Since the noise level $\\lambda$ is a nonlinear function of time $t$, sampling $t$ uniformly results in a non-uniform distribution of noise levels $\\lambda$. During sampling, $p(\\lambda)$ provides the density of sampling steps at noise level $\\lambda$.\n\n**Q: Line 179-180, “The probability distribution $p_w(t)$ has Dirac delta peak of typically very small mass $w(0)$ at $t=0$.” Does this mean that the almost noiseless sample contributes very little to the objective computation?\"**\n\nA: Yes, that's correct. There's typically still somewhat non-trivial total probability mass over all low noise levels, but the weight over these noise levels is relatively low, since it corresponds to modeling fine pixel details, which are of relatively lower importance for the quality of samples than the more abstract image information, which is modeled at the higher noise levels. Of all the plotted weighting functions, only the ELBO weighting gives high weight to low noise levels, but it also leads to the worst FID scores.\n\n**Q: Line 262-265: have the authors run experiments for a $v$-prediction model combined with a non-monotonic weighting function?**\n\nA: We have not previously. For $v$-parametrized model, we used the $v$-prediction loss $||{\\textbf v} - \\hat{{\\textbf v}}||^2_2$ widely adopted in the literature as the baseline, which corresponds to a monotonic weighting function $\\exp(\\lambda / 2)$. We are running an experiment using $v$-parametrized model and $\\epsilon$-prediction loss ($L_{\\text simple}$), resulting in a non-monotonic weighting function ${\\text sech}(\\lambda /2)$. We will report the result later. \n\n**Q: In Figure 1 Right, only the “EDM sampling” noise schedule $p(\\lambda)$ has monotonically decreasing as a function of $\\lambda$. Is there any interpretation of that, [...]?**\n\nA: We did notice that the \"EDM sampling\" schedule spends a lot more time at the very high and the very high noise levels, and less time at the intermediate noise levels. The reason might be that the optimal design of training and sampling noise schedules can be very different: for training, the optimal schedule seeks to minimize the variance of the training objective, while for evaluation, the optimal design should minimize the \"truncation error\" accumulated along the discretized reverse-ODE or reverse-SDE trajectory, which is also the intuition of EDM sampling schedule. As noted in Appendix D.1 of [Karras et al, 2022], in practice, the truncation error of Euler's method (1st order ODE solver) of the reverse-ODE is monotonically decreasing over $\\lambda$. Therefore, we should allocate more sampling steps in low $\\lambda$ region, to allow for using small step sizes with smaller truncation error. This leads to the monotonically decreasing $p(\\lambda)$.\n\n**Q: It’s a bit confusing on why the noise schedule $p(\\lambda)$ is casted as a probability distribution (besides the importance sampling interpretation); [...]. Why does $p(\\lambda)$ need to be sampled – is it because timesteps need to be sampled during training?**\n\nA: Yes, since time $t$ is sampled randomly during training, the noise level $\\lambda$ has a distribution $p(\\lambda)$ during training. We think the importance sampling interpretation is important, since it allows us to intuitively understand that given a fixed weighting function $w(\\lambda)$, the loss is invariant to the choice of noise schedule $p(\\lambda)$; $p(\\lambda)$ only affects the variance of the Monte Carlo estimator of the loss."}, {"review_id": "E7MRADOef2", "reviewer": "Reviewer_Qo9S", "summary": "The paper studies different weighting schemes of the loss functions of generative diffusion models. In particular, the authors investigate the relationship between a weighted (with respect to noise levels) loss and the proper Evidence Lower Bound (ELBO), discussing how the former can be written as a weighted integral of ELBOs.\n\nIt is further claimed that when the weighting factor is monotonic, there exists an interpretation of the loss as the result of log-likelihood optimization considering data augmentation, in particular noise perturbation.\n\nFinally, the proposed interpretation is leveraged to construct new weighting schemes, corroborating the discussion with empirical validation.\n", "questions": "Given the observations in the \"weaknesses\" section:\n\n* Could you please emphasize the main contribution when compared to prior work such as Kingma2021 (eq. (19) and the corresponding section)?\n\n* Can you clarify the doubts about mathematical details listed in the weaknesses section of the review?\n\n* Can you comment on the generality of the experimental results, and how your conclusion could extend to other datasets?\n\n\n\n==== Post rebuttal message, prior to reviewer discussion ====\n\n\nDear Authors, as a post rebuttal message, I am willing to increase my score to a 6: weak accept. I also have changed the scores to what concerns soundness, presentation and contribution. \n\nI will eventually edit again my official review after the reviewer discussion phase, but for now I would like to thank you for your effort in addressing my concerns.", "limitations": "I couldn't find a limitation nor societal impact section, neither a short paragraph on these points.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* Deepening our understanding of diffusion models is an important endeavour, given the popularity of such methods, and this paper is aligned with research lines on methodological contributions on training such models\n\n* The authors attempted at reproducing the work from Karras et al, and provide an additional insights on the difficulty of obtaining generative performance that is on par with prior work", "weaknesses": "On a technical level, the paper has several problems related to mathematical rigor. Overall, a colloquial exposition is preferred over a clear mathematical language.\n\n* Starting with line 76, the \"series of latent variables\" $z_{0,...,1}$ is introduced without specifying whether the set of time indices is uncountable (i.e. a continuous set between 0 and 1) or whether it is the collection of a finite set of time indices (for example an equally spaced grid between 0 and 1). The term \"conditional distribution\" $q(z_{0,...,1} | x)$ introduced at line 77 is very confusing, and requires a proper discussion, in particular for the case of uncountable time indices.\n\n* This is not just a minor stylistic note, but has concrete impact on the derivation and in the discussions: for example, how is the KL divergence introduced at line 126, $D_{\\text{KL}}\\left(  q(z_{0,..1}),p(z_{0,..1}) \\right)$, defined? Are the authors talking about path measures? In this case, a proper characterization needs the formal tools of stochastic calculus, like for example in Huang2021, where the Girsanov Theorem is used to construct the ELBO.\n\n* Furthermore, in the different derivations I find confusing and counterproductive to switch back and forth between different notations (for example dropping or not the dependency on $t$ of $\\lambda$). Why for example in equation (5) would you use lambda both with and without $t$?\n\n* The paper contains typos, like in the equation under line 138 where the term $x_t$ is never defined.\n\n* On a methodological level, I find the contribution of the paper somehow marginal. Until Section 3.1, the authors present known, or easy-to-derive, results from the literature. The correspondence between the weighted loss and the ELBO is even explicitly stated in Kingma2021 (eq. (19) and the corresponding section)\n\n* The derivation under line 173 is not rigorous for the same motivations mentioned above: the set of time indices is never properly defined.  Even the related proofs in Appendix C does not use a proper mathematical language. \nWhile I understand the change of notation between eqs 12 and 13 ($z_t \\to z_\\lambda$) the presentation needs improvements.\n\n* Finally the experimental validation is limited in several aspects. First, for a paper whose main contribution is on the exploration of different weighting schemes, the authors only experiment on a single dataset (ImageNet 64x64). Is this sufficient to draw any general conclusion about such aspect?\n\n* None of the proposed weighting schemes achieve the performance of the existing baseline (EDM, Karras 2022) which has a FID score of 1.36. In general, the improvements are marginal, since we do not have access to the source code it is not possible to verify whether differences between considered schemes and baselines are due to the weighting schemes or other implementation details. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and valuable comments. Please note our top-level comment with additional experimental and theoretical results. Below we address specific questions.\n\n**Q: I couldn't find a limitation nor societal impact section, neither a short paragraph on these points.**\n\nA: Please find these sections in the PDF under \"Supplementary Material\" link, specifically Sections H and I. At the time of full-paper submission (one week before supplementary submission deadline), we accidentally appended an incomplete Appendix to the main paper. Please ignore that Appendix, and view the complete Appendix, which includes the relevant Sections, under the link of \"Supplementary Material\".\n\n**Q: Could you please emphasize the main contribution when compared to prior work such as [Kingma et al, 2021] (eq. (19)) and the corresponding section?**\n\nA: Our contributions are:\n1. We explicitly derive the weighting functions and noise schedules for a large number of diffusion objectives used in the literature.\n2. We show that the weighted loss is invariant to the noise schedule, except for its endpoints. This was only shown to hold for the unweighted loss (ELBO) by [Kingma et al, 2021].\n3. We show that if the weighting function is monotonic, then the objective is equivalent to maximizing the ELBO combined with a form of DistAug, i.e., a form of data distribution augmentation. This improves our understanding of weighted diffusion training objectives, where the weighting is essential for best sample quality. We'd like to emphasize that this derivation (in Section 3.1 and the Appendix B and C) is non-trivial, hasn't been derived or stated in previous work, and potentially very valuable for our peers.\n4. We experiment with new monotonic weightings, and demonstrate SOTA FID and Inception scores on high-resolution (128x128, 256x256 and 512x512) ImageNet generation. This demonstrates that we do not need to deviate from the ELBO objective to get good qualitative results, as long as we combine it with DistAug. We hope this could inspire future work along this direction.\n\n**Q: Can you clarify the doubts about mathematical details listed in the weaknesses section of the review?**\n\nA: Recall that for a discrete-time model, we have the forward process forming a conditional joint distribution $q(z_0, z_{\\Delta t}, z_{2\\Delta t}, ..., z_{T\\Delta t}|x)$, and a generative model forming a joint distribution $p(z_0, z_{\\Delta t}, z_{2\\Delta t}, ..., z_{T\\Delta t})$, assuming $T\\Delta t= 1$. Like \"Variational Diffusion Models\" [Kingma et al, 2021] and much of the literature on ODEs and SDEs, we treat the continuous-time model as a discrete-time model, but taken in the limit where the number of timesteps $T$ between 0 and 1 goes to infinity. So, $dt = \\lim_{0 \\leftarrow \\Delta t} \\Delta t = \\lim_{T \\to \\infty} 1/T$. This leads to the continuous-time forward process $q(z_{0,...,1}|x)$ and generative model $p(z_{0,...,1})$.\n\nAnalogously, the KL divergence $D_{\\text KL}(q(z_{0,...,1}|x)||p(z_{0,...,1}))$ is defined the same as the KL divergence of the discrete-time model:\n$$E_q\\left[D_{\\text KL}(q(z_1|x)||p(z_1)) + \\sum_{i=1}^{T}D_{\\text KL}(q(z_{(i-1)\\Delta t}|z_{i\\Delta t}, x))||p(z_{(i-1)\\Delta t}|z_{i\\Delta t}))\\right]$$\nbut taken in the limit where the number of timesteps $T$ between 0 and 1 goes to infinity, and time step size $\\Delta t$ goes to $0$. This approach was also taken by [Kingma et al, 2021].\n\nFor notational brevity we didn't include these limits in the notation for the proofs under line 173 and in our Appendix C. However, we can see why this caused unnecessary confusion. We've already improved the notational rigor of our manuscript, and will keep improving for the camera-ready version.\n\n**Q:[...] I find confusing and counterproductive to switch back and forth between different notations (for example dropping or not the dependency on $t$ of $\\lambda$).**\n\nA: In Section 2.2 (line 99), we mentioned *\"Below we sometimes denote $\\lambda$ as $\\lambda_t$ to emphasize that it is a bijective function of $t$.\"* We agree that it is not ideal to switch between these notations. For the final version we will put more effort in finding a way to minimize the need to use different notations.\n\n**Q: Can you comment on the generality of the experimental results, and how your conclusion could extend to other datasets?**\n\nA: As noted in the top-level comment, we expanded our experiments to high-resolution (256x256 and 512x512) ImageNet generation, and found that weighting functions that performed well in earlier lower-resolution experiments also perform well in these high-resolution settings. In our initial draft, we also have results on ImageNet 128x128 generation (Section 6.2), besides ImageNet 64x64. For high-resolution ImageNet generation (128x128, 256x256, 512x512), we achieved SOTA FID and Inception scores among approaches without guidance, and demonstrated significant improvements compared to the baseline approach Simple Diffusion [Hoogeboom et al, 2023] that we built upon. Please check the top-level comment for more details. \n\nGiven these results, it's plausible that the tested weighting functions will also work well on other datasets of images.\n\nIt's hard to say a priori which weighting functions would work well for other modalities in practice, like audio. However, in our newest version of our manuscript, we show that VoiceBox [Le et al., 2023], a state-of-the-art model for audio, uses a training objective that is equivalent to the weighted loss with an exponential weighting $w(-\\lambda/2)$. In experiments we found that this weighting works well for images. Therefore, we now know that this weighting function works well for both images and audio.\n\nOn the other hand, all our theoretical results are agnostic to the types of data used, thus our theoretical conclusions naturally extend to other modalities.\n\n**Q: Typo under line 138.**\n\nA: We have fixed it in our latest draft. Thanks for the careful reading!\n"}]

Author rebuttal:
Dear Reviewers,

**Please see the attached PDF for a one-page PDF with a summary of added experimental results.**

We would like to thank all reviewers for providing constructive feedback that helped us improved the paper. We are encouraged that reviews think our paper:
- "provides a comprensive framework for diffusion objectives" (Reviewer DESf),
- "fills an important gap for the diffusion model research community" (Reviewer bjDS),
- provides "a clearer understanding of sample quality and likelihood of various model classes" (Reviewer eoU3) and was well-supported by both "detailed theoretical analysis" and "experimental results" (Reviewer kL4L).

We have been working diligently on improving the paper on several fronts, addressing your critique. Below, we summarize the changes that we have made in an updated draft.

**1. High-resolution (256x256 and 512x512) ImageNet results**

We provide new experimental results, with state-of-the-art FID and Inception scores for high-resolution ImageNet generation (256x256 and 512x512). Please see the attached PDF for a table of the summary of quantitative evaluations and a figure for generated samples by our approach. 

With the shifted version of ‘EDM-monotonic’ weighting, we achieved state-of-the-art FID and Inception scores on all three resolutions (i.e., 128x128, 256x256, 512x512) of ImageNet generation among all approaches without any guidance. With classifier-free guidance (CFG) [Ho and Salimans, 2022], our method outperforms all diffusion-based approaches on resolutions 128 and 512. On resolution 256, our method only falls a bit behind Gao et al. [2023] and Hang et al. [2023], both of which were build upon the latent space of a pretrained auto-encoder from Stable Diffusion [Rombach et al., 2022a], that was trained on a much larger image dataset than ImageNet, while our model was trained on the ImageNet dataset only. It is worth noting that we achieve significant improvements on the Simple Diffusion baseline model [Hoogeboom et al., 2023] that we built upon, on all resolutions, with and without guidance. In principle it is possible to apply our proposed weighting functions and adaptive noise schedules to other diffusion-based approaches such as Gao et al. [2023], possibly further improve their performance.

We will provide additional samples from our trained models in a new section in the updated Appendix.

**2. More weighting functions**

We expanded Table 1 and the Appendix with more weighting functions used in the literature. Specifically, we added the weighting function corresponding to Flow Matching with Optimal Transport (FM-OT) [Lipman et al., 2022]: $w(\lambda) = \exp(-\lambda/2)$, and the weighting function corresponding to InDI [Delbracio and Milanfar, 2023]: $w(\lambda) = \exp(−\lambda)\text{sech}^2(\lambda/4)$. The detailed derivations will be provided in the Appendix. Interestingly, both weighting functions are monotonic.

**3. More helpful visualizations and analysis**

- We added an analysis on the relationship with low-bit training to the Appendix.
- We added a visualization of integration-by-parts, used in our main proof, to the Appendix.

Please see our reviewer-specific feedback for more information.

Meta-review: After a thorough evaluation by five reviewers and considering the authors' detailed rebuttal, I recommend acceptance of this paper as an oral presentation at NeurIPS 2023.

### Summary
The paper presents a comprehensive theoretical analysis of diffusion models, elucidating the relationships between the weighted loss and the Evidence Lower Bound (ELBO) objective. By investigating monotonic weighting and introducing new monotonic weightings, the authors establish connections between the weighted loss and ELBO with data augmentation. They demonstrate the effectiveness of their approach through experiments on ImageNet, achieving competitive results.

### Strengths
Important Contribution: The work fills a significant gap in the diffusion model research community and provides valuable insights. Reviewer DESf praised it, stating, "I'm a big fan of papers like this since their primary aim [is] to build solid understanding [...] and they also present great results for ELBO trained models."

Theoretical Clarity: The theoretical derivation and insights were appreciated by most reviewers. Reviewer eoU3 noted the contribution as "excellent."

Experimentation: The empirical results were seen positively by several reviewers, and the additional high-resolution ImageNet results presented in the rebuttal further strengthened this aspect.

### Concerns
Mathematical Rigor: Concerns were raised about mathematical rigor and notation, particularly by Reviewer Qo9S. However, the authors provided clarifications in their rebuttal, including insights on their use of continuous-time limits and notation.

Experimental Diversity: Reviewer Qo9S expressed concerns about the lack of diversity in datasets, but the authors provided additional experimental results in their rebuttal and explained the relevance to other modalities.

### Conclusion

Considering the strong theoretical contributions, the efforts to clarify mathematical details, and the effective experimental validation, I believe this paper would make a valuable addition to the conference as an oral presentation. The work builds on existing research in a meaningful way, providing practical insights that could influence the design and understanding of diffusion models. Reviewer DESf's comment encapsulates the sentiment well: "This is a nice paper that provides a comprehensive framework for diffusion objectives, giving the ELBO the driver's seat." The thoughtful rebuttal and additional empirical support further solidify the case for acceptance.
Decision: Accept (oral)