Title: Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Diffusion Model

Evidence Lower Bound

Maximum Likelihood
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

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.

1 Introduction

Diffusion-based generative models, or diffusion models in short, were first introduced by SohlDickstein 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 et 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).

The diffusion model framework has since been successfully applied to text-to-image generation [Rombach et al., 2022b, Nichol et al., 2021, Ramesh et al., 2022, Saharia et al., 2022b, Yu et al., 2022, Nichol and Dhariwal, 2021, Ho et al., 2022, Dhariwal and Nichol, 2022, Ding et al., 2021], image-to-image 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 V AEs, the original diffusion models [Sohl-Dickstein et al., 2015] were 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 Diffusion 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-art likelihoods on image density estimation benchmarks.

However, the best results in terms of sample quality metrics such as FID scores were achieved with other objectives, for example a denoising score matching objective [Song and Ermon, 2019] or a simple noise-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-the-art 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 a 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 as a 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 a 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 IS 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 et al., 2021], including Variational Diffusion Models [Kingma et al., 2021], showed how to optimize continous-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 is 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' at inference time. The type of data augmentation under consideration in this paper, namely additive Gaussian noise, is also a 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.

0.0

0.2

0.4

0.6

0.8

1.0

t

5

0

5

10

t

(log SNR)

Noise schedules

Cosine

EDM training

EDM sampling

5

0

5

10

(log SNR)

0.000

0.025

0.050

0.075

0.100

0.125

0.150

p

(

)

Noise schedules

Cosine

EDM training

EDM sampling

2 Model

Suppose we have a dataset of datapoints drawn from q ( x ) . We wish to learn a generative model p θ ( x ) that approximates q ( x ) . We'll use shorthand notation p := p θ .

The observed variable x might be the output of a pre-trained encoder, as in latent diffusion models [Vahdat et al., 2021, Rombach et al., 2022a], on which the popular Stable Diffusion model is based. Our theoretical analysis also applies to this type of model.

In addition to the observed variable x , we have a series of latent variables z t for timesteps t ∈ [0 , 1] : z 0 ,..., 1 := z 0 , ..., z 1 . The model consists of two parts: a forward process forming a conditional joint distribution q ( z 0 ,..., 1 | x ) , and a generative model forming a joint distribution p ( z 0 ,..., 1 ) .

2.1 Forward process and noise schedule

The forward process is a Gaussian diffusion process, giving rise to a conditional distribution q ( z 0 ,..., 1 | x ) ; see Appendix E.1 for details. For every t ∈ [0 , 1] , the marginal distribution q ( z t | x ) is given by:

In case of the often-used variance preserving (VP) forward process, α 2 λ = sigmoid ( λ t ) and σ 2 λ = sigmoid ( -λ t ) , but other choices are possible; our results are agnostic to this choice. The log signal-to-noise ratio (log-SNR) for timestep t is given by λ = log( α 2 λ /σ 2 λ ) .

The noise schedule is a strictly monotonically decreasing function f λ that maps from the time variable t ∈ [0 , 1] to the corresponding log-SNR λ : λ = f λ ( t ) . We sometimes denote the log-SNR as λ t to emphasize that it is a function of t . The endpoints of the noise schedule are given by λ max := f λ (0) and λ min := f λ (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, f λ is invertible: t = f -1 λ ( λ ) . Given this bijection, we can do a change of variables: a function of the value t can be equivalently written as a function of the corresponding value λ , and vice versa, which we'll make use of in this work.

During model training, we sample time t uniformly: t ∼ U (0 , 1) , then compute λ = f λ ( t ) . This results in a distribution over noise levels p ( λ ) = -dt/dλ = -1 /f ′ λ ( t ) (see Section E.3), which we also plot in Figure 2.

Sometimes it is best to use a different noise schedule for sampling from the model than for training. During sampling, the density p ( λ ) gives the relative amount of time the sampler spends at different noise levels.

2.2 Generative model

The data x ∼ D , with density q ( x ) , plus the forward model defines a joint distribution q ( z 0 , ..., z 1 ) = ∫ q ( z 0 , ..., z 1 | x ) q ( x ) d x , with marginals q t ( z ) := q ( z t ) . The generative model defines a corresponding joint distribution over latent variables: p ( z 0 , ..., z 1 ) .

Let s θ ( z ; λ ) denote a score model , which is a neural network that we let approximate ∇ z log q t ( z ) through methods introduced in the next sections. If s θ ( z ; λ ) = ∇ z log q t ( z ) , then the forward process can be exactly reversed; see Appendix E.4.

For large enough λ max, z 0 is almost identical to x , so learning a model p ( z 0 ) is practically equivalent to learning a model p ( x ) . For small enough λ min, z 1 holds almost no information about x , such that there exists a distribution p ( z 1 ) satisfying D KL ( q ( z 1 | x ) || p ( z 1 )) ≈ 0 . Usually we can use p ( z 1 ) = N (0 , I ) .

If D KL ( q ( z 1 ) || p ( z 1 )) ≈ 0 and s θ ( z ; λ ) ≈ ∇ z log q t ( z ) , then we have a good generative model in the sense that D KL ( q ( z 0 ,..., 1 ) || p ( z 0 ,..., 1 )) ≈ 0 , which implies that D KL ( q ( z 0 ) || p ( z 0 )) ≈ 0 which achieves our goal. So, our generative modeling task is reduced to learning a score network s θ ( z ; λ ) that approximates ∇ z log q t ( z ) .

Sampling from the generative model can be performed by sampling z 1 ∼ p ( z 1 ) , then (approximately) solving the reverse SDE using the estimated s θ ( z ; λ ) . 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 et al. [2020] and the stochastic sampler with Heun'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 s θ ( z ; λ t ) that approximates ∇ z log q t ( z ) , for all noise levels λ t . It was shown by [Vincent, 2011, Song and Ermon, 2019] that this can be achieved by minimizing a denoising score matching objective over all noise scales and all datapoints x ∼ D :

where z t = α λ x + σ λ glyph[epsilon1] .

The glyph[epsilon1] -prediction objective. Most contemporary diffusion models are optimized towards a noiseprediction loss introduced by [Ho et al., 2020]. In this case, the score network is typically parameterized through a noise-prediction ( glyph[epsilon1] -prediction) model: s θ ( z ; λ ) = -ˆ glyph[epsilon1] θ ( z ; λ ) /σ λ . (Other options, such as x -prediction, v -prediction, and EDM parameterizations, are explained in Appendix E.2.) The noise-prediction loss is:

Since || s θ ( z t ; λ t ) - ∇ x t log q ( z t | x ) || 2 2 = σ -2 λ || ˆ glyph[epsilon1] θ ( z t ; λ t ) -glyph[epsilon1] || 2 2 , this is simply a version of the denoising score matching objective in Equation 2, but where ˜ w ( t ) = σ 2 t . 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 a 'linear' to a 'cosine' noise schedule λ t (see Figure 2). This noise-prediction loss with the cosine schedule is currently broadly used.

The ELBO objective. It was shown by [Kingma et al., 2021, Song et al., 2021a] that the evidence lower bound (ELBO) of continuous-time diffusion models simplifies to:

where c is constant w.r.t. the score network parameters.

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

(log SNR)

0.0

0.2

0.4

0.6

0.8

1.0

w

(

) (normalized)

Non-monotonic weighting functions

IDDPM

EDM

P2

min-SNR-

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

(log SNR)

0.0

0.2

0.4

0.6

0.8

1.0

w

(

) (normalized)

Monotonic weighting functions

ELBO

v-prediction and FM-OT

EDM-monotonic (Ours)

sigmoid(

+2) (Ours)

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 et al. [2021] 1 with a different choices of weighting function w ( λ t ) :

See Appendix D for a derivation of the implied weighting functions for all popular diffusion losses. Results are compiled in Table 1, and visualized in Figure 3.

The ELBO objective (Equation 4) corresponds to uniform weighting, i.e. w ( λ t ) = 1 .

The popular noise-prediction objective (Equation 3) corresponds to w ( λ t ) = -1 / ( dλ/dt ) . This is more compactly expressed as w ( λ t ) = p ( λ t ) , i.e., the PDF of the implied distribution over noise levels λ during training. Typically, the noise-prediction objective is used with the cosine schedule λ t , which implies w ( λ t ) ∝ sech ( λ t / 2) . See Section E.3 for the expression of p ( λ t ) for various noise schedules.

3.2 Invariance of the weighted loss to the noise schedule λ t

In Kingma et al. [2021], it was shown that the ELBO objective (Equation 4) is invariant to the choice of noise schedule, except for its endpoints λ min and λ max. This result can be extended to the much

more general weighted diffusion loss of Equation 5, since with a change of variables from t to λ , it can be rewritten to:

Note that this integral does not depend on the noise schedule f λ (the mapping between t and λ ), except for its endpoints λ max and λ min. The shape of the function f λ between λ min and λ max does not affect the loss; only the weighting function w ( λ ) does. Given a chosen weighting function w ( λ ) , the loss is invariant to the noise schedule λ t between t = 0 and t = 1 . This is important, since it means that the only real difference between diffusion objectives is their difference in weighting w ( λ ) .

This invariance does not hold for the Monte Carlo estimator of the loss that we use in training, based on random samples t ∼ U (0 , 1) , glyph[epsilon1] ∼ N (0 , 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. More specifically, we'll show that the noise schedule acts as an importance sampling distribution for estimating the loss integral of Equation 6. Note that p ( λ ) = -1 / ( dλ/dt ) . We can therefore rewrite the weighted loss as the following, which clarifies the role of p ( λ ) as an importance sampling distribution:

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 1. If the weighting w ( λ t ) is monotonic, then the weighted diffusion objective of Equation 5 is equivalent to the ELBO with data augmentation (additive noise).

With monotonic w ( λ t ) we mean that w is a monotonically increasing function of t , and therefore a monotonically decreasing function of λ .

We'll use shorthand notation L ( t ; x ) for the KL divergence between the joint distributions of the forward process q ( z t,... 1 | x ) and the reverse model p ( z t,... 1 ) , for the subset of timesteps from t to 1 :

In Appendix A.1, we prove that 2 :

As shown in Appendix A.1, this allows us to rewrite the weighted loss of Equation 5 as simply:

In Appendix A.2, we prove that using integration by parts, the weighted loss can then be rewritten as:

Now, assume that w ( λ t ) is a monotonically increasing function of t ∈ [0 , 1] . Also, without loss of generality, assume that w ( λ t ) is normalized such that w ( λ 1 ) = 1 . We can then further simplify the weighted loss to an expected KL divergence:

where p w ( t ) is a probability distribution determined by the weighting function, namely p w ( t ) := ( d/dt w ( λ t )) , with support on t ∈ [0 , 1] . The probability distribution p w ( t ) has Dirac delta peak of typically very small mass w ( λ max ) at t = 0 .

Note that:

More specifically, L ( t ; x ) equals the expected negative ELBO of noise-perturbed data, plus a constant; see Section C for a detailed derivation.

This concludes our proof of Theorem 1.

glyph[squaresolid]

This result provides us the new insight that any of the objectives with (implied) monotonic weighting, as 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 ( λ ) holds for a number of models with state-of-the-art perceptual quality, including VoiceBox for speech generation [Le et al., 2023], and Simple Diffusion for image generation [Hoogeboom et al., 2023].

4.1 Understanding the effect of weighting functions on 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 are also more important to human perception, and removing the three least significant bits from the data allows the model to spend more capacity on modeling the more significant bits. In [Kingma and Dhariwal, 2018], training on 5-bit data was performed by adding uniform noise (with the appropriate scale) to the data before feeding it to the model, but it was found that adding Gaussian noise had a similar effect. Adding Gaussian noise with a single noise level is a 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 perturbed witjh additive Gaussian noise. Since the FT is a linear transformation, the FT of a the sum of a 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, it is known that the power spectrum, i.e. 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 a result, additive Gaussian noise effectively destroys high-frequency information in the data more quickly than the low-frequency information, pushing the model to 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.

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 et 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 the form sigmoid ( -λ + k ) , inspired by the sigmoidal weighting corresponding to low-bit training (see Section 4.1 and Appendix I for details), where k is a hyper-parameter. In addition, inspired by the EDM weighting function from Karras et al. [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 w ( λ ) = max λ ˜ w ( λ ) for λ < arg max λ ˜ w ( λ ) , where ˜ w ( λ ) 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 glyph[epsilon1] -parametrization model, we took iDDPM [Nichol and Dhariwal, 2021] as the baseline, which utilized cosine noise schedule and a non-monotonic sech ( λ/ 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 non-monotonic EDM weighting (Figure 3). For v -parametrization model, we followed Salimans and Ho [2022] taking the v -prediction loss as the baseline, which leads to a monotonic exp( -λ/ 2) weighting. We always use adaptive noise schedule under this setting. We started by searching the optimal hyperparameter for the sigmoidal weighting under the glyph[epsilon1] -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 2 summarizes the FID [Heusel et al., 2017] and Inception scores [Salimans et al., 2016] across different settings. From the table we made several observations as follows.

First, our adaptive noise schedule (Appendix F) works equally well as the hand-tuned fixed noise schedules. We verified this under glyph[epsilon1] -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 handtuning of the noise schedule. In addition, it sometimes leads to faster training; see Appendix F.

Second, it is possible to modify an existing non-monotonic weighting function to a monotonic one with the minimal change, and results in on-par sample quality. Specifically, under EDM parameterization and adaptive noise schedule, we change EDM weighting to EDM-monotonic 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 glyph[epsilon1] -parameterization and v -parameterization settings. Our-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 samplers, 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 64 × 64, namely sigmoid ( -λ +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 on high-resolution ImageNet without sampling modifications (e.g. 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 higherresolution 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 w ( λ ) . 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 as second-order sampling or rejection sampling. 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 on ImageNet 128 × 128, compared to the baseline Simple Diffusion approach.

Next, we test our approach on ImageNet generation of multiple high resolutions (i.e., resolutions 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 'EDM-monotonic' weighting, we achieved state-of-the-art 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-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 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. It is possible to apply our proposed weighting functions and adaptive noise schedules to other diffusion-based approaches such as Gao et al. [2023] to further improve their performance, which we shall leave to the future work.

6 Conclusion and Discussion

In summary, we have shown that the weighted diffusion loss, which generalizes diffusion objectives in the literature, has an interpretation as a weighted integral of ELBO objectives, with one ELBO per noise level. If the weighting function is monotonic, then we show that the objective has an interpretation as 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 a direct apples-to-apples comparison of diffusion models with other likelihood-based models. For example, it allows one to optimize other likelihood-based models, such as autoregressive models, towards the same objective as 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.

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 , 12(3):313-326, 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 22669-22679, 2023.

Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:1809.11096 , 2018.

Geoffrey J Burton and Ian R 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-11325, 2022.

Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan. WaveGrad: 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 2: Iterative Refinement for Text-to-Speech Synthesis . In INTERSPEECH , 2021b.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.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.15089 , 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://arxiv.org/abs/2105.13290 .

Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is a strong image synthesizer. arXiv preprint 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.09556 , 2023.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. arXiv preprint arXiv:1706.08500 , 2017.

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

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.11239 , 2020.

Jonathan Ho, Chitwan Saharia, William Chan, David J 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.02808 , 2021.

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

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusionbased generative models. arXiv preprint arXiv:2206.00364 , 2022.

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

Diederik P Kingma and 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 1x1 convolutions. arXiv preprint arXiv:1807.03039 , 2018.

Diederik P 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, et al. Voicebox: Text-guided multilingual universal speech generation at scale. arXiv preprint arXiv:2306.15687 , 2023.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 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 autoregressive modeling with distribution smoothing. arXiv preprint arXiv:2103.15089 , 2021.

Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv preprint arXiv:2102.09672 , 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 arXiv:2112.10741 , 2021.

William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748 , 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 TextConditional 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. In International Conference on Machine Learning , pages 1278-1286, 2014.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 2022a.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. HighResolution Image Synthesis with Latent Diffusion Models. In CVPR , 2022b.

Chitwan Saharia, William Chan, Huiwen Chang, Chris A. Lee, Jonathan Ho, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Palette: Image-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 NeurIPS , 2022b.

Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2022c.

Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512 , 2022.

Tim Salimans, Ian Goodfellow, Wojciech 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 on 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 Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling Through Stochastic Differential Equations. In International Conference on Learning Representations , 2021b.

Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. arXiv preprint arXiv:2106.05931 , 2021.

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

Daniel Watson, Ricardo Chan, William Martin-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. In 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 Text-to-Image Generation. In arXiv:2206.10789 , 2022.

Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Improved techniques for maximum likelihood estimation for diffusion odes. arXiv preprint arXiv:2305.03935 , 2023.

A Main proof

Here we'll provide a proof of Equation 11.

Note that like in the main text, we use shorthand notation:

A.1 Time derivative of D KL ( q ( z t,..., 1 | x ) || p ( z t,..., 1 ))

Let dt denote an infinitesimal change in time. Note that L ( t -dt ; x ) can be decomposed as the sum of a KL divergence and an expected KL divergence:

Due to this identity, the time derivative d/dt L ( t ; x ) can be expressed as:

In Appendix E of [Kingma et al., 2021], it is shown that in our model, this equals:

where z t = α λ x + σ λ glyph[epsilon1] , and SNR ( t ) := exp( λ ) in our model, and SNR ′ ( t ) = d/dt SNR ( t ) = e λ dλ/dt , so in terms of our definition of λ , this is:

In terms of glyph[epsilon1] -prediction (see Section E.2), because || glyph[epsilon1] -ˆ glyph[epsilon1] θ || 2 2 = e λ || x -ˆ x θ || 2 2 this simplifies to:

where z λ = α λ x + σ λ glyph[epsilon1] . This can be easily translated to other parameterizations; see E.2. This allows us to rewrite the weighted loss of Equation 5 as:

A.2 Integration by parts

Integration by parts is a basic identity, which tells us that:

This allows us to further rewrite the expression of the weighted loss in Equation 26 as:

where

is constant w.r.t. θ , since it does not involve the score function, and typically very small, since D KL ( q ( z 1 | x ) || p ( z 1 )) is typically small by design.

The term w ( λ max ) D KL ( q ( z 0 ,..., 1 | x ) || p ( z 0 ,..., 1 )) is typically small, since w ( λ max ) is typically very small (see Figure 3).

This concludes our proof of Equation 11.

B Visualization

We tried to create a helpful visualization of the result from Section A.2. Note that we can rewrite:

The relationship in Equation 28 can be rewritten as:

The first LHS term equals a weighted prior loss term, and the second LHS term equals the weighted diffusion loss. From a geometric perspective, the two LHS terms together define an area that equals the area given by the right term, as illustrated in the figure below.

0

0

t

= 1

t

= 0

t

= 1

t

= 0

w

(

λ

t

)

w

(

λ

t

)

w

(

λ

1

)

· L

(1;

x

)

L

(

t

;

x

)

L

(

t

;

x

)

∫

t

=0

t

=1

w

(

λ

t

)

d

L

(

t

;

x

)

∫

t

=1

t

=0

L

(

t

;

x

)

dw

(

λ

t

)

w

(

λ

0

)

L

(0;

x

)

On the left, we have a rectangular area that equals a weighted prior loss w ( λ min ) L (1) , plus a curved area equal to the weighted diffusion loss ∫ t =0 t =1 w ( λ t ) d L ( t ; x ) . This integral can be intuitively understood as a Riemann sum over many tiny rectangles going from left ( t = 1 ) to right ( t = 0 ), each with height w ( λ t ) and width d L ( t ; x ) . On the right, we have the same total area, but divided up into two different subareas: a rectangular area w ( λ max ) L (0) and a curved area that equals the integral ∫ t =1 t =0 L ( t ; x ) dw ( λ t ) going upwards from t = 0 to t = 1 , which can also be intuitively understood as another Riemann sum, with each tiny rectangle having width L ( t ; x ) and height dw ( λ t ) . The area of each of those tiny rectangles on the right can be understood as the ELBO at each noise level, L ( t ; x ) , times the weight of the ELBO at each noise level, dw ( λ t ) .

C Relationship between D KL ( q ( z t,..., 1 | x ) || p ( z t,..., 1 )) and the ELBO

First, note that:

glyph[squaresolid]

<latexi sh

1_b

64="q

K

F

L

WjV

Y

cQfg5JZTIMp

>A

CRn

DN

v0/

8

u

7E

kG

2

S

H

zB

m

X

O

d

o

+

r

P

9

y

U

w

3

<latexi sh

1_b

64="GQuMYZ

mwPAnd

0

TFN

p

g

>

B

H

c

VDLS

EOz

r

8

3

UCy

J

RX

k

/

I

oK

q

v

2

j

7

W

9

5

+

f

<latexi sh

1_b

64="

R/f8L

K2QNT

7Wq

v

pk

F+PEI

>A

C

H

c

VD

M

m

ZJ

o30

z

UY

B

r

O

S

G

w

j

nX

5

y

u

9

g

d

<latexi sh

1_b

64="q/3Gj cL

N

D

k

MI

Q

>A

CPX

VB

T9

F

w

pR

S

g

Yr

y

+Z

KH

d

fW7n

0

J

O

E2

m

8

5

o

U

<latexi sh

1_b

64="EPD

YJ

gf2L

vUnBRM

SVN

>A

Cd3

c

HW

r

K+

j

yw

Fu

p

q

k

G

9

7

Z

z

5

/

T

O

X

m

0

I

8

Q

o

<latexi sh

1_b

64="Yyw

ME5

V

pvZGo

Wnfg

Tzc

>A

C

X

BN

9

F27

JQI

R

U

D

k

K

q

r

H

/

m

O

+

3

jP

S

8

u

0

L

d

<latexi sh

1_b

64="gECIp7AU

nHW

S

3

z2c

>

PX

VDLTu wF

R

9

ZBN

K

Q

Y

qj

o

+

8

/

r

M

0

v

5

y

J

d

f

m

G

<latexi sh

1_b

64="q

K

F

L

WjV

Y

cQfg5JZTIMp

>A

CRn

DN

v0/

8

u

7E

kG

2

S

H

zB

m

X

O

d

o

+

r

P

9

y

U

w

3

<latexi sh

1_b

64="GQuMYZ

mwPAnd

0

TFN

p

g

>

B

H

c

VDLS

EOz

r

8

3

UCy

J

RX

k

/

I

oK

q

v

2

j

7

W

9

5

+

f

<latexi sh

1_b

64="c/pM7

S

0U

jy

L

CQWPBD

n

>A

d3

V

FH

G

wmv

JY

g

E

R

KI

q

O

Z

k

oz

T2

9

5

N

u

f

r

X

+

8

<latexi sh

1_b

64="q/3Gj cL

N

D

k

MI

Q

>A

CPX

VB

T9

F

w

pR

S

g

Yr

y

+Z

KH

d

fW7n

0

J

O

E2

m

8

5

o

U

<latexi sh

1_b

64="

wZ3

nF

OgE

WRq

A

X7o

>

C

c

VDd

NBGJ2

f

0m

H

5

u

I

Q

S

r

z

9

P

8/

j

K

T

k

+

v

Y

L

M

U

y

p

<latexi sh

1_b

64="Yyw

ME5

V

pvZGo

Wnfg

Tzc

>A

C

X

BN

9

F27

JQI

R

U

D

k

K

q

r

H

/

m

O

+

3

jP

S

8

u

0

L

d

<latexi sh

1_b

64="gECIp7AU

nHW

S

3

z2c

>

PX

VDLTu wF

R

9

ZBN

K

Q

Y

qj

o

+

8

/

r

M

0

v

5

y

J

d

f

m

G

More precisely, the joint KL divergence D KL ( q ( z t,..., 1 | x ) || p ( z t,..., 1 )) is the expected negative ELBO of noise-perturbed data, plus a constant entropy term:

where the ELBO of noise-perturbed data is:

where ˜ z t := z t + dt,..., 1 .

So, L ( t ; x ) is the expected negative ELBO of noise-perturbed data z t :

Therefore, the expression of the weighted loss in Equation 12 can be rewritten as:

where w ( λ min ) is constant w.r.t. the diffusion model parameters. Therefore, minimizing L w ( x ) is equivalent to maximizing this expected ELBO of noise-perturbed data.

D Derivation of weighting functions for previous works

The loss function used in previous works are equivalent to the weighted loss with a certain choice of noise schedule and weighting function w ( λ ) . In this section, we derive these weighting functions w ( λ ) .

D.1 'Elucidating Diffusion Models' (EDM) [Karras et al., 2022]

Karras et al. [2022] proposes the following training objective:

where p (˜ σ ) and ˜ w (˜ σ ) are defined as:

where ˜ σ is equivalent to the standard deviation of noise added to the clean sample x in the VP SDE case, so ˜ σ 2 = e -λ . Karras et al. [2022] used hyperparameters P mean = -1 . 2 , P std = 1 . 2 and ˜ σ data = 0 . 5 .

This can be rewritten in terms of λ = -2 log ˜ σ as:

where p ( λ ) and ˜ w ( λ ) are defined as

Comparing Equation 45 with the weighted loss expressed in terms of x -prediction parameterization (Section E.2), we see that (ignoring the constant scaling factor 1 / 2 ) the EDM objective is a special case of the weighted loss with weighting function:

where the divison by ˜ σ 2 data can be ignored since it's constant. This leads to:

D.2 The v -prediction loss / 'SNR+1'-weighting [Salimans and Ho, 2022]

Salimans and Ho [2022] introduced the v -parameterization, with a v -prediction model ˆ v , where:

They propose to minimize a v -prediction loss, E [ || v -ˆ v || 2 2 ] . Note that ˆ glyph[epsilon1] = ( z λ -α λ ˆ x ) /σ λ . For our general family, this implies:

In the special case of the variance preserving (VP) SDE, this simplifies to:

Since the glyph[epsilon1] -prediction loss corresponds to minimizing the weighted loss with w ( λ ) = p ( λ ) , the v -prediction loss corresponds to minimizing the weighted loss with w ( λ ) = ( e -λ +1) p ( λ ) .

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 signal-to-noise ratio (SNR). Since || glyph[epsilon1] -ˆ glyph[epsilon1] || 2 2 = e λ || x -ˆ x || 2 2 , they call the glyph[epsilon1] -prediction loss 'SNR weighting', and since || v -ˆ v || 2 2 = ( e λ +1) || x -ˆ x || 2 2 , they call this 'SNR+1'-weighting.

Salimans and Ho [2022] propose to use optimize a VP SDE with a cosine schedule p ( λ ) = sech ( λ/ 2) / (2 π ) = 1 / (2 π cosh( -λ/ 2)) and the v -prediction loss: E [ || v -ˆ v || 2 2 ] . This corresponds to minimizing the weighted loss with:

The factor π can be ignored since it's constant, so we can equivalently use:

D.2.1 With shifted cosine schedule

[Hoogeboom et al., 2023] extended the cosine schedule to a shifted version: p ( λ ) = sech ( λ/ 2 -s ) / (2 π ) , where s = log(64 /d ) , where 64 is the base resolution and d is the model resolution (e.g. 128, 256, 512, etc.). In this case the weighting is:

Since (2 /π ) e -s is constant w.r.t λ , the weighting is 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 a special case of Gaussian diffusion with the weighted loss.

D.3.1 Noise schedule

Note that in [Lipman et al., 2022], time goes from 1 to 0 as 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 also assume

σ 0 = 0 , for which we can later correct by truncation (see Section E.3.1). In this model, the forward process q ( z t | x ) is defined by:

This implies that the log-SNR is given by:

The derivative, as a function of t , is:

This derivative of its inverse, as a function of λ , is:

The corresponding density is

Its inverse is given by:

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:

The model is then optimized with the Conditional flow matching (CFM) loss:

where they use the parameterization:

D.3.3 Weighting function

What is the weighting function w ( λ ) corresponding to this loss? Note that this parameterization means that:

Since t = 1 / (1 + e λ/ 2 ) , we have that 1 / (1 -t ) = 1 + e -λ/ 2 , so parameterized as a function of λ , we have:

Likewise, we can parameterize o -prediction in terms of glyph[epsilon1] -prediction:

We can translate the o -prediction loss to a glyph[epsilon1] -prediction loss:

Therefore, combining the derivations above, the CFM loss, formulated in terms of the λ parameterization instead of t , and in terms of the glyph[epsilon1] -prediction parameterization instead of the o -prediction parameterization, is:

where ˜ p is any distribution with full support on [ λ min , λ max ] , and where:

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:

As derived in Section D.3 above, this means that the distribution over log-SNR λ is the Logistic distribution: p ( λ ) = sech 2 ( λ/ 4) / 8 . The proposed loss function is the x -prediction loss:

Since || x -ˆ x || 2 2 = e -λ || glyph[epsilon1] -ˆ glyph[epsilon1] || 2 2 , and the glyph[epsilon1] -prediction loss corresponds to minimizing the weighted loss with w ( λ ) = p ( λ ) , the x -prediction loss above corresponds to minimizing the weighted loss with:

Which is a 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 et al. [2022] proposed a new weighting function:

where empirically they set k = 1 and γ as either 0 . 5 or 1 . Compared to the glyph[epsilon1] -prediction objective, where w ( λ ) = dt/dλ = 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 ( λ ) = sech ( λ/ 2) / (1 + e λ ) γ .

D.6 Min-SNRγ weighting [Hang et al., 2023]

Hang et al. [2023] proposed the following training objective:

Therefore, it corresponds to w ( λ ) = min { 1 , γe -λ } p ( λ ) . The motivation of the work is to 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. A cosine noise schedule is then combined with the proposed weighting function, leading to w ( λ ) = sech ( λ/ 2) · min { 1 , γe -λ } . γ 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 a stochastic differential equation (SDE):

For derivations of diffusion SDEs, see Appendix B of [Song et al., 2021b]. Their β ( t ) equals d dt log(1 + e -λ t ) in our formulation, and their ∫ t 0 β ( s ) ds equals log(1 + e -λ t ) , where they assume that λ →∞ at t = 0 .

E.1.1 Variance-preserving (VP) SDE

A common choice is the variance-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:

E.1.2 Variance-exploding (VE) SDE

Another common choice of the variance-exploding (VE) SDE. In the VE case:

E.2 Possible parameterizations of the score network

There are various ways of parameterizing the score network:

We can let a neural network output any of s θ ( z λ ; λ ) , ˆ glyph[epsilon1] θ or ˆ x θ , and we can convert the variables to each other using the equalities above.

The chosen relationship between z λ , ˆ x , ˆ glyph[epsilon1] and s θ ( z λ ; λ ) above, are due to the following relationships between z λ , x and glyph[epsilon1] :

In addition, there's the v -prediction parameterization ( v := α λ glyph[epsilon1] -σ λ x ) explained in D.2, and the o -prediction parameterization ( o := x -glyph[epsilon1] ) explained in D.3.

Karras et al. [2022] proposed a specific F -parametrization, with an F -prediction model ˆ F θ . In the special case of variance explosion (VE) SDE, it is formulated as:

where ˜ σ data = 0 . 5 . Generalizing this to our more general family with arbitrary drift, this corresponds to:

So that we have:

In summary, given these different parameterizations, the glyph[epsilon1] -prediction loss can be written in terms of other parameterizations as follows:

Interestingly, if we set ˜ σ 2 data = 1 , the training objectives of F -prediction and v -prediction are the same.

And:

E.3 Noise schedules

During model training, we sample time t uniformly: t ∼ U (0 , 1) , then compute λ = f λ ( t ) . This results in a distribution over noise levels p ( λ ) , whose cumulative density function (CDF) is given by 1 -f -1 λ ( λ ) . For λ ∈ [ λ min , λ max ] the probability density function (PDF) is the derivative of the CDF, which is p ( λ t ) = -( d/dλ ) f -1 λ ( λ ) = -dt/dλ = -1 /f ′ λ ( t ) . Outside of [ λ min , λ max ] the probability density is 0.

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 et al., 2020], f λ ( t ) = -log( e t 2 -1) which has fallen out of use. Note that:

In the shifted cosine schedule, s = log(64 /d ) , where 64 is the base resolution and d is the used resolution (e.g. 128, 256, 512, etc.).

In the EDM training schedule, the function F N ( λ ; µ, σ 2 ) is the Normal distribution CDF, and N ( λ ; µ, σ 2 ) is its PDF.

In the EDM sampling schedule, ρ = 7 , σ min = 0 . 002 , σ max = 80 . The density function p ( λ ) in the table has support λ ∈ [ -log σ 2 max , -log σ 2 min ] . Outside this range, p ( λ ) = 0 .

Note that the table gives the noise schedules before truncation, resulting in [ ˜ f λ (0) , ˜ f λ (1)] = [ ∞ , -∞ ] . The truncation procedure is given in E.3.1.

E.3.1 Truncation

The noise schedules above are truncated, resulting in a noise schedule ˜ f λ (0) whose endpoints have desired values [ ˜ f λ (0) , ˜ f λ (1)] = [ λ max , λ min ] :

Its inverse is:

And the corresponding probability density:

E.4 Sampling

Anderson [1982] showed that if s θ ( z ; λ ) = ∇ z log q t ( z ) , then the forward SDE is exactly reversed by the following SDE:

⊔]glyph[arrowtp]˜̂(√√(∐̂∐√√]√˜(√̂[˜̂√⌈˜√



   

⋂√∐]{]{˜(]√˜√∐√]}{

/

/

/

/

/

/

⊕}√√

⊔]glyph[arrowtp]˜̂(√√(∐̂∐√√]√˜(√̂[˜̂√⌈˜√

⊔]glyph[arrowtp]˜̂({}]√˜(√̂[˜̂√⌈˜({}√]{˜}

̂∐√√]√˜({}]√˜(√̂[˜̂√⌈˜({⊗√√√}

(a)



   

⋂√∐]{]{˜(]√˜√∐√]}{

/

/

/

/

⊕}√√

⊔]glyph[arrowtp]˜̂({}]√˜(√̂[˜̂√⌈˜({〉〈⊕}

(b)

̂∐√√]√˜({}]√˜(√̂[˜̂√⌈˜({⊗√√√}

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:

which, under the assumption that s θ is a conservative vector field, will result in the same marginals p ( z t ) as the SDE of Equation 133 for every t ∈ [0 , 1] , and therefore also the same marginal p ( x ) .

Note that due to the continuous-time nature of the model, any sampling method is necessarily approximate, with the discretization error depending on various factors including the choice of noise schedule. For sampling we can therefore typically use a different noise schedule f λ 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 w ( x ) , but not for the Monte Carlo estimator of the loss that we use in training, based on random samples t ∼ U (0 , 1) , glyph[epsilon1] ∼ N (0 , 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 / ( dλ/dt ) . We can therefore rewrite the weighted loss as the following, which clarifies the role of p ( λ ) as an importance sampling distribution:

In order to avoid having to hand-tune the noise schedule for different weighting functions, we implemented an adaptive noise schedule. The noise schedule λ t is updated online, where we let p ( λ ) ∝ E x ∼D , glyph[epsilon1] ∼N (0 , I ) [ w ( λ ) || ˆ glyph[epsilon1] θ ( z λ ; λ ) -glyph[epsilon1] || 2 2 ] . This noise schedule ensures that the loss is spread evenly over time, i.e. that the magnitude of the loss E x ∼D , glyph[epsilon1] ∼N (0 , I ) [ ( w ( λ ) /p ( λ )) || ˆ glyph[epsilon1] θ ( z λ ; λ ) -glyph[epsilon1] || 2 2 ] is approximately invariant to λ or t . We find that this often significantly speeds op optimization.

We implemented an adaptive noise schedule p ( λ ) , where:

In practice we approximate this by dividing the range [ λ min , λ max ] into 100 evenly spaced bins, and during training keep an exponential moving average (EMA) of w ( λ ) || glyph[epsilon1] -ˆ glyph[epsilon1] θ ( z λ ; λ ) || 2 2 within each bin. From these EMAs we construct a piecewise linear function f λ ( t ) such that Equation 136 is approximately satisfied. The EMAs and corresponding noise schedule p ( λ ) 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]:

Theorem 2. Assume a model in the family specified in Section 2, and assume the score network encodes a conservative vector field: s θ ( z t , λ t ) = ∇ z t log p ( z t ) (not assumed by the other theorems). Then:

Proof of Theorem 2. Note that (see Equation 117):

And assume the score network encodes a conservative vector field:

So the time derivative of Equation 22 can be expressed as:

Equation 138 follows from the definition of the Fisher divergence.

G.1 Comparison with Theorem 1 by Lyu [2012]

Lyu [2012] prove a similar result in their Theorem 1. We'll translate their result into our notation. In particular, let the forward process be as in our family, such that q ( z t | x ) = N ( z t ; α λ x , σ 2 t I ) . The marginal (data) distribution is q ( x ) , such that q ( z t ) = ∫ q ( z t | x ) q ( x ) d x . Similarly, let the generative model have a marginal p ( x ) , and p ( z t ) = ∫ p ( z t | x ) p ( x ) d x . So far the assumptions are the same in our family.

They assume that q ( z t | x ) = N ( z t , x , t ) , which corresponds to a variance exploding (VE) diffusion process, with t = σ 2 λ , so λ = -log( t ) . Importantly, they make the assumption that p ( z t | x ) = q ( z t | x ) , i.e. that the forward process for p equals the forward process for q . Given these assumptions, Lyu [2012] show that:

Which, given the noise schedule λ = -log( t ) , can be rewritten as:

which looks a lot like our Equation 138. One difference are that in Equation 138, the left-handside distributions are joint distributions, and q conditions on x , while Equation 143 is about the glyph[squaresolid]

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 t | x ) = q ( z t | x ) , since this assumption does not hold for the family of diffusion models we consider. Before or during optimization, p ( z t | x ) might be very far from q ( z t | x ) . After optimization, p ( z t | x ) might be close to q ( z t | x ) , but we still can't assume they'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 t , we applied the low-discrepency sampler of time that was proposed by Kingma et al. [2021], which has been shown to effectively reduce the variance of diffusion loss estimator and lead to faster optimization. The model is optimized by Adam [Kingma and Ba, 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 [ λ min , λ max ] into 100 evenly spaced bins. During training, we maintaiedn an exponential moving average of w ( λ ) || glyph[epsilon1] -ˆ glyph[epsilon1] θ ( z λ ; λ ) || 2 2 with a decay rate 0.999, and a constant initialization value of 1 for each bin.

Below we elaborate the implementation details specific for each task.

ImageNet 64x64. For class-conditional 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 1 e -4 , exponential moving average of 50 million images and learning rate warmup of 10 million images, which mainly follows the configuration of Karras et al. [2022]. We employed 128 TPU-v4 chips with a 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 a single training run. For training noise schedule and sampling noise schedule of DDPMsampler, we set λ min = -20 and λ max = 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 4 for the formulations). We adopted the same hyperparameters of EDM sampler from Karras et al. [2022] with no changes (i.e., Table 5 in their work, column 'ImageNet-Our model'). Both DDPM and EDM samplers took 256 sampling steps.

ImageNet 128x128. For class-conditional generation on 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 a single run. We set λ min = -15 + s and λ max = 15 + s , where s = log(64 /d ) is the shift of the weighting function, with 64 being the base resolution and d being the model resolution ( d = 128 for this task). The DDPM sampler used for evaluation used 'shifted-cosine' noise schedule (Table 4) and took 512 sampling steps.

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 on 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 to the data, before feeding it to the model. It was found that adding Gaussian noise had a similar effect as uniform noise. As we have seen, in the case of diffusion models, adding Gaussian noise is equivalent to using a weighted objective with a monotonic weighting function.

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 d dλ L ( λ ; x ) for low-bit data

Note that the results of Appendix A.1 can also be written as:

This allows us to rewrite the weighted loss as simply:

To understand the effect of λ , we can plot the d dλ L ( λ ; x ) as a function of λ .

We start by plotting the function E q ( x ) [ L ( λ ; x )] , for differences choices of n = 1 , ..., 8 , as shown in Figure 5a. Next, we plot E q ( x ) [ d dλ L ( λ ; x )] , for each n , in Figure 5b. The more bits included in the data, the more information contained by E q ( x ) [ d dλ L ( λ ; 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 n -1 bits from the curve for n bits, as displayed in Figure 5c. The area under each of these curves is exactly 1 bit. Interestingly, as more bits being added, the contribution to the loss gradually shifts from low-SNR regions to high-SNR regions.

We'd like to plot d dλ L ( λ ; x ) dof different choices of bit precision. This will tell us where the different bits 'live' as a function of λ . Since training different diffusion models on different bit precisions is very expensive, we instead use an approximation. In particular, we assume that the data x is univariate, with a uniform distribution q ( x ) over the 2 n 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 ( z λ ) , for each choice of λ , be the optimal model: p ( z λ ) := ∫ q ( x ) q ( z λ | x ) d z , which is a univariate mixture-of-Gaussians, where each mixture component is a Gaussian centered one of the 2 n possible pixel values. In this case, L ( λ ; x ) := D KL ( q ( z λ | x ) || p ( z λ )) .

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 weighting function that goes down sharply between λ = 7 . 5 and λ = 10 . In fact, the 5-bit unweighted loss curve is very similar to the 8-bit loss curve, when using the following weighting function:

where F N 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 to 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 to answer this question by lifting Gaussian noise perturbation to the frequency domain.

Given a clean natural image x and a Gaussian white noise image glyph[epsilon1] , the Gaussian-noise-perturbed data is given by z λ = α λ x + σ λ glyph[epsilon1] . Denote the Fourier transform operation as F ( · ) . Since F ( · ) is a linear operation, we have

To understand how noise affects the data, we can compare the power spectra of α λ F ( x ) and σ λ F ( glyph[epsilon1] ) , which is defined as the square amplitude of the Fourier transform: S ( · ) = |F ( · ) | 2 . More specifically, given the power spectrum S of an image, we can take the circular average of S for the Fourier

10

5

0

5

10

15

0

1

2

3

4

5

Eq

(

x

)

[

(

,

x

)]

Eq

(

x

)

[

(

,

x

)] (Expected nats as a function of

)

n_bits = 1

n_bits = 2

n_bits = 3

n_bits = 4

n_bits = 5

n_bits = 6

n_bits = 7

n_bits = 8

(a)

E

q

(

x

)

[

L

(

λ

;

x

)]

as a function of

λ

.

10

5

0

5

10

15

0.0

0.1

0.2

0.3

0.4

0.5

Eq

(

x

)

[

d

/

d

(

,

x

)]

Eq

(

x

)

[

d

/

d

(

,

x

)]

n_bits = 1

n_bits = 2

n_bits = 3

n_bits = 4

n_bits = 5

n_bits = 6

n_bits = 7

n_bits = 8

(b)

E

q

(

x

)

[

d

dλ

L

(

λ

;

x

)]

as a function of

λ

.

10

5

0

5

10

15

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

Contribution of each bit to

Eq

(

x

)

[

d

/

d

(

,

x

)]

Contribution of each bit to

Eq

(

x

)

[

d

/

d

(

,

x

)]

bit 1

bit 2

bit 3

bit 4

bit 5

bit 6

bit 7

bit 8

(c) Subtraction of the curve for

n

-

1

bits from the

curve for

n

bits in Fig.5b.

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

(log SNR)

0.0

0.2

0.4

0.6

0.8

1.0

w

(

) (normalized)

Weighting functions

ELBO

v-prediction and FM-OT

EDM-monotonic (Ours)

sigmoid(

+2) (Ours)

5-bit-like weighting

(d) Weighting functions, with a weighting resembling

5-bit training.

10

5

0

5

10

15

0.0

0.1

0.2

0.3

0.4

0.5

Unweighted 5-bit versus weighted 8-bit

Eq

(

x

)

[

d

/

d

(

,

x

)], n_bits = 5

Eq

(

x

)

[

d

/

d

(

,

x

)

w

(

)], n_bits = 8

(e) Unweighted 5-bit loss versus weighted 8-bit loss.

components of the same spatial frequency, such that it becomes a function of the spatial frequency. Figure 6 displays the average power spectra of α λ F ( x ) and σ λ F ( glyph[epsilon1] ) as functions of the spatial frequency, for different λ . As λ goes from high to low, i.e., gradually perturbing the data with more noise, the average power spectrum of σ λ F ( glyph[epsilon1] ) (orange curve) first exceeds the one of α λ F ( 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 is 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 more on low frequency components of the data, which often correspond to high level content and global structure that are more crucial to human perception. This analysis can also be applied to understand the forward process of diffusion models.

10 0

10 1

10 2

Spatial frequency

10

1

10 1

10 3

10 5

10 7

10 9

Average power spetrum

=-5.0

(

x

)

(

)

10 0

10 1

10 2

Spatial frequency

10

1

10 1

10 3

10 5

10 7

10 9

=-2.5

(

x

)

(

)

10 0

10 1

10 2

Spatial frequency

10

1

10 1

10 3

10 5

10 7

10 9

=0.0

(

x

)

(

)

10 0

10 1

10 2

Spatial frequency

10

1

10 1

10 3

10 5

10 7

10 9

=2.5

(

x

)

(

)

10 0

10 1

10 2

Spatial frequency

10

1

10 1

10 3

10 5

10 7

10 9

=5.0

(

x

)

(

)

To make the claim even clearer, we can compute the log signal-to-noise-ratio (log SNR) ˜ λ between the power spectra of α λ F ( x ) and σ λ F ( glyph[epsilon1] ) , as a function of λ . Figure 7 shows the average value of that log ratio around different frequency regions. Given a fixed λ , ˜ λ is higher for low frequency regions, so that the data are less destroyed by the noise.

10

5

0

5

10

15

(log SNR)

15

10

5

0

5

10

15

20

(log SNR around freq)

freq=1.12

freq=3.91

freq=13.62

freq=47.45

freq=165.37

K Limitations

It is important to emphasize that our empirical results, like other deep learning approaches, depend on the choice of hyper-parameters. 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 in the understanding and optimization of diffusion models, the advancements could have broader implications, some of which could potentially be 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 or for other malicious purposes, such as identity theft or blackmail.

Regarding fairness considerations, generative models are typically trained on large datasets and could therefore inherit and reproduce any biases present in the training data. This could potentially result

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 to the model or its outputs is regulated to 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 a 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 we propose. The community should strive for an open and ongoing discussion about the ethics of AI and the development of strategies to mitigate potential misuse of these powerful technologies.

M Samples from our model trained on 512 × 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 to ethical guidelines. Samples in Figures 8 and 9 are generated without guidance, while samples in Figures 10 and 11 are generated with guidance strength 4.

mg.kaa

国

067

国

00

Qm

0G

HTUA

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n* 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\n\nWEAKNESSES:\nOn 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\nQUESTIONS:\nGiven 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.\n\nLIMITATIONS:\nI couldn't find a limitation nor societal impact section, neither a short paragraph on these points.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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$.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe authors have discussed both limitations and potential negative social impact via Appendix H and I from Supplementary Material, which speak for the community well.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: 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.

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:
-

QUESTIONS:
-

LIMITATIONS:
Yes, the authors addressed limitations, potential negative societal impact, and mitigation.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nI'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.\n\nWEAKNESSES:\nNo 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.\n\nQUESTIONS:\nJust 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-14"}

comments: 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.

STRENGTHS:
- The problem studied is important
- Formulating the different diffusion model loss as weighted loss is interesting
- The paper presented a detailed theoretical analysis
- Experimental results support the proposed method

WEAKNESSES:
- Adding more experimental results could have been helpful but this is not a major point.

QUESTIONS:
NA

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-08-01