Diffusion Models With Learned Adaptive Noise

Response to reviewer QQVN

## **Concern 1:** Intuition behind ELBO being dependent on the diffusion trajectory Imagine you’re piloting a plane across a region where cyclones and strong winds are present. If you plot a straight line course directly through these adverse weather conditions, the journey requires more fuel and effort due to the resistance encountered. Instead, if you navigate around the cyclones and adverse winds, it takes less energy to reach your destination, even though the path might be longer. This intuition maps to mathematical and physical terms. The trajectory of the plane is denoted by $\mathbf{r}(t) \in \mathbb{R}^n_{+}$ and the forces acting on the plane are given by $\mathbf{f}(\mathbf{r}(t)) \in \mathbb{R}^n_+$. The work used to navigate the plane is $\int_{0}^{1} \mathbf{f}(\mathbf{r}(t)) \cdot \frac{d}{dt}\mathbf{r}(t) \, dt$. The work here is dependent of the trajectory because $\mathbf{f}(\mathbf{r}(t)) \in \mathbb{R}^n_+$ is not a conservative field. In Section 3.5, we argue that the same holds for the diffusion ELBO. The trajectory, $\mathbf{r}(t) \in \mathbb{R}^n_{+}$, is parameterized by the noise schedule which is influenced by complex forces, $\mathbf{f}(\mathbf{r}(t)) \in \mathbb{R}^n_+$ (akin to weather patterns). The diffusion loss can be interpreted as the amount of work done, $\int_{0}^{1} \mathbf{f}(\mathbf{r}(t)) \cdot \frac{d}{dt}\mathbf{r}(t) \, dt$, along the diffusion trajectory in the presence of these forces (the forces are modeled by dimension-wise reconstruction error of the denoising model). By carefully selecting (learning the noise schedule) our path—avoiding “high-resistance” areas where the loss accumulates quickly—we can minimize the overall “energy” expended, measured in terms of the NELBO. ## **Concern 2:** Worse FIDs than the existing diffusion models. Note that MuLAN does not incorporate many tricks that improve FID such as exponential moving averages, truncations, specialized learning schedules, etc.; our FID numbers can be improved in future work using these techniques. However, our method yields state-of-the-art log-likelihoods. We leave the extension to tuning schedules for FID to future work. Optimizing for the likelihood is directly motivated by applied problems such as data compression [2]. In that domain, arithmetic coding techniques can take a generative model and produce a compression algorithm that provably achieves a compression rate (in bits per dimension) that equals the model’s log-likelihood [3]. Other applications of log-likelihood estimation include adversarial example detection [4], semi-supervised learning [5], and others. ## **Concern 3:** Inconsistent FID numbers for VDM as reported in the original paper. Note that in their paper, VDM reported FIDs using the DDPM sampler with `T=1000`, whereas in our paper, we generate samples using an ODE solver that solves the reverse-time Probability Flow ODE, as shown in Eqn (76) for VDM and Eqn (81) for MuLAN. We employ the RK45 ODE solver [6] provided by `scipy.integrate.solve_ivp` with `atol=1e-5` and `rtol=1e-5`. These settings align with prior work by Song et al., 2021 [7]. One method to measure the complexity of the learned noise schedule is by comparing the Number of Function Evaluations (NFEs) required to generate samples using an ODE solver to solve the associated Probability Flow ODE. In Table 1, we aimed to study the difficulty of drawing samples from the diffusion model (measured by NFEs) and the quality of the samples (measured by FID). We observe that MuLAN does not degrade FIDs while improving log-likelihood estimation. --- References: [1] Theis, L., Oord, A. V. D., & Bethge, M. (2015). A note on the evaluation of generative models. *arXiv preprint arXiv:1511.01844*. [2] David JC MacKay. *Information theory, inference and learning algorithms*. Cambridge university press, 2003. [3] Thomas M Cover and Joy A Thomas. Data compression. *Elements of Information Theory*, pp. 103–158, 2005. [4] Yang Song, Taesup Kim, Sebastian Nowozin, Stefano Ermon, and Nate Kushman. Pixeldefend: Leveraging generative models to understand and defend against adversarial examples. *arXiv preprint arXiv:1710.10766*, 2017. [5] Zihang Dai, Zhilin Yang, Fan Yang, William W Cohen, and Russ R Salakhutdinov. Good semi-supervised learning that requires a bad gan. *Advances in neural information processing systems*, 30, 2017. [6] J.R. Dormand and P.J. Prince. A family of embedded runge-kutta formulae. *Journal of Computational and Applied Mathematics*, 6(1):19–26, 1980. ISSN 0377-0427. doi: https://doi. org/10.1016/0771-050X(80)90013-3. URL https://www.sciencedirect.com/science/ article/pii/0771050X80900133. [7] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. *Advances in neural information processing systems*, 34: 1415–1428, 2021.

Response to reviewer 1B1k (1/3)

We want to thank the reviewer for their constructive feedback. We address each concern below. ## Concern 1: Demonstrating the utility of density estimation on downstream tasks Optimizing for the likelihood is directly motivated by applied problems such as data compression [2]. In that domain, arithmetic coding techniques can take a generative model and produce a compression algorithm that provably achieves a compression rate (in bits per dimension) that equals the model’s log-likelihood [3]. Other applications of log-likelihood estimation include adversarial example detection [4], semi-supervised learning [5], and others. In this work, we take the first step of improving the likelihood, and we aim to explore major downstream applications in future work. We note that one immediate simple application of our work is faster training: we attain strong likelihoods and FIDs on ImageNet using half the training steps (Table 1). Note that MuLAN does not incorporate many tricks that improve FID such as exponential moving averages, truncations, specialized learning schedules, etc.; our FID numbers can be improved in future work using these techniques. However, our method yields state-of-the-art log-likelihoods. We leave the extension to tuning schedules for FID to future work. --- References: [1] Theis, L., Oord, A. V. D., & Bethge, M. (2015). A note on the evaluation of generative models. *arXiv preprint arXiv:1511.01844*. [2] David JC MacKay. *Information theory, inference and learning algorithms*. Cambridge university press, 2003. [3] Thomas M Cover and Joy A Thomas. Data compression. *Elements of Information Theory*, pp. 103–158, 2005. [4] Yang Song, Taesup Kim, Sebastian Nowozin, Stefano Ermon, and Nate Kushman. Pixeldefend: Leveraging generative models to understand and defend against adversarial examples. *arXiv preprint arXiv:1710.10766*, 2017. [5] Zihang Dai, Zhilin Yang, Fan Yang, William W Cohen, and Russ R Salakhutdinov. Good semi-supervised learning that requires a bad gan. *Advances in neural information processing systems*, 30, 2017. ## Concern 2: Discussion on auxiliary latent variables Diffusion models augmented with an auxiliary latent space [1, 2] have been used for representation learning. These auxiliary latents, $\mathbf{z} \in \mathbb{R}^m$, have a smaller dimensionality than the input $\mathbf{x}_0 \in \mathbb{R}^d$ and are used to encode a high-level semantic representation of $\mathbf{x}_0$. In this work we make an observation that an input conditioned multivariate noise schedule can indeed improve likelihoods unlike previously thought so. However, simply conditioning on the noise schedule is challenging because of the reasons mentioned in Sec 10.2 in the appendix. In this work we resolve these challenges by conditioning the noise schedule on the input via an auxiliary latent space. The key differences between these methods: | | Learned noise | Multivariate noise | Input Conditioned noise | Auxiliary latents | Noise parameterization | | --- | --- | --- | --- | --- | --- | | InfoDiffusion [1] | No | No | No | In denoising process | Cosine schedule | | MuLAN (ours) | **Yes** | **Yes** | **Yes** | In **noising** & denoising processes | Polynomial (**novel**) | These auxiliary latents, $\mathbf{z} \sim p_\theta(\mathbf{z})$ , differ from the noisy inputs, $\mathbf{x}_t$ in the sense that these latents are kept fixed throughout the reverse generation process; however, the $\mathbf{x}_t$ is progressively denoised to get clean $\mathbf{x}_0$. References: [1] Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maximizing diffusion models. In *International Conference on Machine Learning*, pp. xxxx–xxxx. PMLR, 2023. [2] Ruihan Yang and Stephan Mandt. Lossy image compression with conditional diffusion models, 2023.

Response to reviewer 1B1K (2/3)

## Concern 3: Visualizations of the learned schedule We have provided visualizations for the noise schedule for different pixels in Fig. 11 (left column) in the appendix for various datasets such as CIFAR-10 and ImageNet. Furthermore, we have numerous experiments trying to relate the learned noise schedule to different properties of the input image such as intensity and frequency in sec. 14 in the Appendix. We summarize three relevant experiments hereby: 1. CIFAR-10 Frequency Split To see if MULAN learns different noise schedules for pixels in an image with different frequencies, we modify the images in the CIFAR-10 dataset where we split an image into high frequency and low frequency components using gaussian blur. **Observation** We do notice two separate bands in the noise schedule (see fig. 11) however these bands don’t necessarily correspond to the low / high freq regions in the image. Upon comparing the schedules to that of the original CIFAR10-dataset, we do notice an increase in the inter-spatial variation of the noise schedules. 2. CIFAR-10 Masking We modify the CIFAR-10 dataset where we randomly mask (i.e. replace with 0s) the top of an image or the bottom half of an image with equal probability. Fig. 10a shows the training samples. MULAN was trained for 500K steps. The samples generated by MULAN is shown in Fig. 10b. **Observation** We don’t observe any interpretable patterns in the noise schedule for the masked pixels vs non-masked pixels. Upon comparing the schedules to that of the original CIFAR10-dataset, we do notice an increase in the inter-spatial variation of the noise schedules. 3. CIFAR-10 Intensity Split To see if MULAN learns different noise schedules for images with different intensities, we modify the images in the CIFAR-10 dataset where we randomly convert an image into a low intensity or a high intensity image with equal probability. Fig. 9a shows the training samples. MULAN was trained for 500K steps. The samples generated by MULAN is shown in Fig. 9b. **Observation** We don’t observe any interpretable patterns in the noise schedules for the high intensity images vs low intensity images. Upon comparing the schedules to that of the original CIFAR10-dataset, we do notice an increase in the inter-spatial variation of the noise schedules. **Temporal Variation** In fig. 11 (left) in the appendix, we visualize the noise schedules for all the pixels and we observe that the schedules are spatially identical in the beginning and the end of the diffusion processes i.e. at $t=0$ and $t=1$. ## Concern 4: Comparisons to other works ### **Variational Diffusion Models** We'd like to emphasize the fact that this work **changes the understanding of noise schedules** and how they affect the likelihood of a diffusion model. It discards the existing notion that the likelihood of a diffusion model is invariant to the choice of noise schedules. To achieve this we introduce the concept of adaptive noise schedules where the noise schedule is conditioned on the input via auxiliary latent variables. We include a comparison to VDM hereby that illustrates the key differences between these methods. | | Learned noise | Multivariate noise | Input Conditioned noise | Auxiliary latents | Noise parameterization | | --- | --- | --- | --- | --- | --- | | VDM (Kingma et al., 2021) | Yes | No | No | No | Monotonic neural network | | MuLAN (ours) | Yes | Yes | Yes | In noising & denoising processes | Polynomial, sigmoid | ### **Score-based generative modeling in latent space.** Latent Score-based Generative Model (LSGM) [1] defines a diffusion process in the latent space, rather than directly on the input data. In contrast, MuLAN defines the diffusion process on the input data but conditions the forward and reverse processes on the auxiliary latent space. Furthermore, we’d like to highlight that MuLAN achieves a Negative Log Likelihood (NLL) of 2.65 on CIFAR-10, which is significantly better than LSGM’s 2.87. ### **Minimizing trajectory curvature of ode-based generative models.** In Minimizing Trajectory Curvature (MTC) of ODE-based generative models, the primary goal is to design the forward diffusion process that is optimal for fast sampling; however, MuLAN strives to learn a more expressive forward process that optimizes for log-likelihood. In the appendix, we have provided a detailed explanation in sec 7.4 and Tab. 6 highlights the key differences between the methods. References: [1] Vahdat, A., Kreis, K., & Kautz, J. (2021). Score-based generative modeling in latent space. [2] Lee, S., Kim, B., & Ye, J. C. (2023, July). Minimizing trajectory curvature of ode-based generative models.

Response to reviewer 1B1k (3/3)

## Concern 5: Other concerns > ### Polynomial Parameterization Although the Monotonic Neural Network parameterization proposed in VDM is intended to be more expressive, the design of the neural network restricts its expressivity. Specifically, the weights must always be positive, and it utilizes sigmoid activations to maintain the monotonicity of the noise schedule. We observed that in deeper layers, the activations could easily saturate, thus limiting the expressivity of the function. In contrast, the polynomial parameterization features only 3 free variables that can be set arbitrarily while still ensuring the noise schedule’s monotonicity. Furthermore, as the reviewer noted, the polynomial parameterization may have two zero derivatives between $t=0$ and $t=1$. However, we do not explicitly set the derivatives at the endpoints to zero. > ### Imagenet likelihood for NDM The reason we did not include the likelihood numbers for NDM on ImageNet is that they employed numerous tricks to achieve improved likelihood, which the baselines reported in Table 2 did not use, namely: 1. Data augmentation. NDM pre-processes their data using flipping which is a commonly used trick to improve the likelihood. For ex, on CIFAR-10, VDM w/o data augmentation achieves an NLL of 2.65 but with data augmentation, the NLL improves to 2.49; see [1]. 2. Parameter Count. NDM [2] uses an encoder model that is of the same size as that of the denoising model (see sec 4.1 in [2]) and hence they use almost 2x more parameters than ours. In contrast, MuLAN’s encoder is only ~10% size of the denoising model. References: [1] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. *Advances in neural information processing systems*, 34:21696–21707, 2021. [2] Grigory Bartosh, Dmitry Vetrov, and Christian A Naesseth. Neural diffusion models. *arXiv preprint arXiv:2310.08337*, 2023. > ### Clarification on Likelihood computation To compute the likelihood of a datapoint $\mathbf{x}_0$ we need to evaluate the following equation (eqn. (84) in the paper): $- \mathbb{E}\_{q\_\phi(\mathbf{z}|\mathbf{x}\_0)} \log p\_\theta(\mathbf{x}\_0 | \mathbf{z}) + \mathbb{D}\_{\text{KL}} \left( q\_\phi(\mathbf{z}|\mathbf{x}\_0) \parallel p_\theta(\mathbf{z}) \right)$ where $\mathbf{x}\_0$ denotes the input data, $\mathbf{z}$ denotes the auxiliary latent variable, $q_\phi(\mathbf{z}|\mathbf{x}\_0)$ and $p_\theta(\mathbf{z})$ denote the approximate posterior and the prior distributions for the auxiliary latent variables respectively, $p_\theta(\mathbf{x}_0 | \mathbf{z})$ is the likelihood of the diffusion model conditioned on the auxiliary latent. 1. To compute the term $\mathbb{E}\_{q\_\phi(\mathbf{z}|\mathbf{x}\_0)} \log p_\theta(\mathbf{x}\_0 | \mathbf{z})$, we draw a sample $\mathbf{z} \sim q_\phi(\mathbf{z} | \mathbf{x}_0)$ and then use an ODE solver to compute the ODE dynamics in Eqn (83) in the paper. 2. The term $\mathbb{D}\_{\text{KL}} \left( q_\phi(\mathbf{z}|\mathbf{x}\_0) \parallel p_\theta(\mathbf{z}) \right)$ is computed in closed form; see line 223 in the paper.