Diffusion Models With Learned Adaptive Noise Processes

We want to thank the reviewer for their constructive feedback. We address each concern below.

Concern 1: The strength of the empirical results. Can the authors provide additional empirical evidence?

We have improved our experimental results, and we now achieve a new state-of-the-art in density estimation on CIFAR10 and ImageNet. The table below summarizes our results.

* Partial result obtained on 70% of the data due to time constraints.

Our updated method adds two architectural improvements: velocity reparameterization (as in Zheng et al., 2023) and truncated normal dequantization (Dinh et al., 2017; Salimans et al., 2017; Ho et al., 2019; Zheng et al., 2023). The recent iDODE paper contains additional innovations (e.g., importance sampled training) that could further improve the performance of MuLAN. Lastly, for both iDODE and MuLAN with velocity parameterization, we report an importance-sampled estimate of test NLL with K=20 samples, as in the i-DODE paper (sec A.1., eqn. 32).

Concern 2: Explaining the relationship between Infodiffusion (Wang et al., 2023) and MuLAN.

In brief, auxiliary-variable diffusion is a technique introduced by Wang et al., 2023. MuLAN makes use of this technique in conjunction with several other components. The resulting method is substantially different from InfoDiffusion and solves a different problem: density estimation versus representation learning.

Specifically, the key novel component of MuLAN is a learned adaptive noise process. A widely held assumption is that the ELBO objective of a diffusion model is invariant to the noise process (Kingma et al., 2021). We dispel this assumption: we show that when input-conditioned noise is combined with (a) multivariate noise, (b) a novel polynomial parameterization, and (c) auxiliary variables, a learned noise process yields an improved variational posterior and a tighter ELBO. This approach sets a new state-of-the-art in density estimation.

The table below summarizes the relationship between MuLAN and other methods.

Concern 3: MuLAN without auxiliary variables performs worse than VDM.

This may be a misunderstanding. MuLAN without auxiliary variables performs the same as VDM: in Fig 1a, the blue and gray lines overlap almost perfectly. This is expected: when the noise is not adaptive (on the data via auxiliary variables), there is no improvement.

Perhaps the reviewer meant to ask why MuLAN without multivariate noise is worse than VDM? First, recall that when noise is univariate, the ELBO is invariant to the choice of noise process (Kingma et al., 2021). However, auxiliary variables require applying the ELBO twice: this yields a training objective that bounds the marginal log-likelihood less tightly than in VDM. By itself, this lowers performance; however this performance drop is more than compensated by learning the noise process.

Concern 4: Understanding the time complexity of the method.

When trained on 8 A100 GPUs, VDM achieves a training rate of 11.6 steps/second, while Mulan trains slightly slower at 10.1 steps/second due to the inclusion of an additional encoder network. However, despite this slower training pace, VDM requires 30 days to reach a BPD of 2.65, whereas Mulan achieves the same BPD within a significantly shorter timeframe of 10 days.

Question 1: How was the BPD calculated? Is it VLB or ODE-based? What is the variance of the BPD?

We used both VLB and ODE-based methods to compute BPD.

In the VLB-based approach, we employ Eqn. 9. We use T = 128 in Eqn. 10, discretizing the timesteps [0, 1] into 128 bins. For the ODE-based approach, we follow the exact evaluation procedure in i-DODE: we extract the underlying ODE for the diffusion process and calculate the likelihood using ODE-based exact methods. We also use their importance-weighted bound with K importance samples (note that K=1 corresponds to the ELBO, and K>1 provides a tighter bound on the likelihood).

Below, we report BPD values (mean and 95% Confidence Interval) for MuLAN on CIFAR10 (8M training steps) and ImageNet (2M training steps) using both the VLB-based approach, and the ODE-based approach with K=1 and K-20 importance samples.

* Partial result obtained on 70% of the test data due to time constraints.

Question 2: Did the authors analyze the auxiliary latent space? Was there any interpretable meaning?

We analyzed the effects of the auxiliary variables on the noise process. While we did observe variation in the noise schedule for different auxiliary variables (Figure 2), we did not find any human-interpretable patterns. We defer the human-interpretable analysis to future work. We did not attempt to analyze the ability of the auxiliary variables to perform representation learning; however, that problem has been extensively studied by Wang et al., 2023.

We want to thank the reviewer for their constructive feedback. We address each concern below.

Concern 1: The experimental results are not impressive

We report new experimental results, and we now achieve a new state-of-the-art in density estimation on CIFAR10 and ImageNet. The table below summarizes our results.

* Partial result obtained on 70% of the data due to time constraints

Our updated method adds two architectural improvements: velocity reparameterization (as in Zheng et al., 2023) and truncated normal dequantization (Dinh et al., 2017; Salimans et al., 2017; Ho et al., 2019; Zheng et al., 2023). The recent iDODE paper contains additional innovations (e.g., importance sampled training) that could further improve the performance of MuLAN. Lastly, for both iDODE and MuLAN with velocity parameterization, we report an importance-sampled estimate of test NLL with K=20 samples, as in the i-DODE paper (sec A.1., eqn. 32).

Note also that the magnitude of our improvement over VDM or i-DODE is significant: it is comparable to the progress made by most papers on these benchmarks from 2018-2021.

Concern 2: Understanding the novelty of the paper relative to recent work

MuLAN is the first method to introduce a learned adaptive noise process. A widely held assumption is that the ELBO objective is invariant to the noise process (Kingma et al., 2021). We dispel this assumption: when input-conditioned noise is combined with (a) multivariate noise, (b) a novel polynomial parameterization, and (c) auxiliary variables, a learned noise process yields a tighter ELBO and a new state-of-the-art in density estimation. While (a), (c) were proposed in other contexts, we leverage them as subcomponents of a novel algorithm.

The table below summarizes the relationship between MuLAN and other methods.

Concern 3: The motivation of our work, specifically the intuition behind using adaptive noise from a frequency perspective, is not convincing.

The main motivation of our work comes from variational inference: we view the noising process as a variational posterior, and we learn it to obtain a tighter ELBO (see Section 3.1). This in turn yields state-of-the-art density estimation.

Note that the intuition to which the reviewer refers is not the motivation for our work. We also do not claim that likelihood correlates with image quality. The paragraph in question only provides an example of different datasets that might benefit from different noise characteristics.

Concern 4: The introduction of auxiliary variables does not lower the posterior gap.

Please note that our results demonstrate empirically that introducing both auxiliary variables and learned noise most likely improves the posterior gap: these methods yield significantly better ELBO values on CIFAR10 and ImageNet.

However, if we were to only introduce auxiliary variables without adaptive noise, we would indeed get a worse posterior gap. This is confirmed by our ablation study (orange line in Figure 1a). However, this performance drop is more than compensated by learning the noise process (red line in Figure 1).

We want to thank the reviewer for their constructive feedback. We address each concern below.

Concern 1: Understanding the novelty of the paper relative to recent work.

MuLAN is the first method to introduce a learned adaptive noise process. A widely held assumption is that the ELBO objective is invariant to the noise process (Kingma et al., 2021). We dispel this assumption: when input-conditioned noise is combined with (a) multivariate noise, (b) a novel polynomial parameterization, and (c) auxiliary variables, a learned noise process yields a tighter ELBO and a new state-of-the-art in density estimation. While (a), (c) were proposed in other contexts, we leverage them as subcomponents of a novel algorithm.

The table below summarizes the relationship between MuLAN and other methods.

Concern 2: The inclusion of samples from the model.

We have added the samples to the paper in the appendix G.

Concern 3: Explaining why a polynomial noise schedule performs empirically better

The reason why a polynomial function works better than a sigmoid or a monotonic neural network as proposed by VDM is rooted in Occam’s razor. In Appendix D2, we show that a degree 5 polynomial is the simplest polynomial that satisfies several desirable properties, including monotonicity and having a derivative that equals zero exactly twice. Simpler noise processes (e.g., scalar, exponential) are not sufficiently expressive to achieve these properties. More expressive models (e.g., monotonic 3-layer MLPs) are more difficult to optimize, hence perform worse.

Question 1: Were the models trained from scratch?

Yes, all the models were trained from scratch.

Addendum: New experimental results

We report new experimental results, and we now achieve a new state-of-the-art in density estimation on CIFAR10 and ImageNet. The table below summarizes our results.

* Partial result obtained on 70% of the data due to time constraints.

Our updated method adds two architectural improvements: velocity reparameterization (as in Zheng et al., 2023) and truncated normal dequantization (Dinh et al., 2017; Salimans et al., 2017; Ho et al., 2019; Zheng et al., 2023). For both iDODE and MuLAN with velocity parameterization, we report an importance-sampled estimate of test NLL with K=20 samples, as in the iDODE paper (sec A.1., eqn. 32).

We want to thank the reviewer for their constructive feedback. We address each concern below.

Concern 1: Multivariate noise requires auxiliary variables to be useful. This raises the question of whether the combined learning of noise schedules and the diffusion model is useful.

The combined learning of the noise schedules and the diffusion model is beneficial, as illustrated in Fig. 1. Neither a multivariate noise schedule nor an auxiliary variable scalar noise schedule individually leads to improvement; however, when these are combined, MuLAN yields a significantly better BPD than VDM.

Note also that we could not learn the noise separately from the diffusion model in MuLAN: they are coupled via the auxiliary variables, which are found in both.

Concern 2: The improvement over VDMs through the use of auxiliary variables is unclear. What is the role of these variables in improving log-likelihood?

Auxiliary variables by themselves do not improve the log-likelihood of VDMs: we show this in Figure 1. Auxiliary variables require applying the ELBO twice: this yields a training objective that bounds the marginal log-likelihood less tightly than in VDM. By itself, this lowers performance (Figure 1; orange line). However this performance drop is more than compensated by learning the noise process (Figure 1; red line). Thus, obtaining good log-likelihood requires using the entire set of components introduced in MuLAN (including noise that is multivariate, adaptive, and learned).

Concern 3: The motivation comes from the manual adjustment of noise schedules in diffusion models. Does mulan lead to improved performance when the obj. Isn’t ELBO?

The main motivation of our work comes from variational inference: we view the noising process as a variational posterior, and we learn it to obtain a tighter ELBO (see Section 3.1). This in turn yields state-of-the-art density estimation.

Note that the manual adjustment of noising schedules in high-resolution diffusion models is not the motivation for our work. We also do not claim that likelihood correlates with image quality. The passage in question is only a citation to relevant concurrent work.

Lastly, the goal of this paper is to study how noise schedule affects the likelihood. We leave to future work the question of whether a learnable noise schedule is also beneficial when the objective function isn’t ELBO. Since there is evidence that manual adjustment of the noise schedule helps with improving the perceptual quality of the images (Chen, 2023; Hoogeboom et al., 2023), methods inspired by MuLAN hold promise there.

Concern 4: Lack of a citation to Song et al., 2020.

We thank the reviewer for bringing this to our attention. We have added the reference to our paper.

Addendum: New experimental results

We report new experimental results, and we now achieve a new state-of-the-art in density estimation on CIFAR10 and ImageNet. The table below summarizes our results.

* Partial result obtained on 70% of the data due to time constraints.

Our updated method adds two architectural improvements: velocity reparameterization (as in Zheng et al., 2023) and truncated normal dequantization (Dinh et al., 2017; Salimans et al., 2017; Ho et al., 2019; Zheng et al., 2023). For both iDODE and MuLAN with velocity parameterization, we report an importance-sampled estimate of test NLL with K=20 samples, as in the iDODE paper (sec A.1., eqn. 32).