Diffusion Models With Learned Adaptive Noise

Authors: subham sekhar sahoo
aaron  gokaslan
chris  de sa
volodymyr  kuleshov
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:24:45

Abstract: Diffusion models have gained traction as powerful algorithms for synthesizing high-quality images. Central to these algorithms is the diffusion process, a set of equations which maps data to noise 
in a way that can significantly affect performance. 
In this paper, we explore whether the diffusion
process can be learned from data.
Our work is grounded in Bayesian inference and seeks to improve log-likelihood estimation by casting the learned diffusion process as an approximate variational posterior that yields a tighter lower bound (ELBO) on the likelihood.
A widely held assumption is that the ELBO is invariant to the noise process: our work dispels this assumption and proposes multivariate learned adaptive noise (MuLAN), a learned diffusion process that applies noise at different rates across an image. Our method consists of three components: a multivariate noise schedule, adaptive input-conditional diffusion, and auxiliary variables; these components ensure that the ELBO is no longer invariant to the choice of the noise schedule as in previous works.  Empirically, MuLAN sets a new **state-of-the-art** in density estimation on CIFAR-10 and ImageNet while matching the performance of previous state-of-the-art models with **50%** fewer steps.  We provide the code, along with a blog post and video tutorial on the project page: https://s-sahoo.com/MuLAN

Main paper: 
 introduction: 
 Diffusion models, inspired by the physics of heat diffusion, have gained traction as powerful tools for generative modeling, capable of synthesizing realistic, high-quality images [51, 16, 43, 14] . Central to these algorithms is the diffusion process, a gradual mapping of clean images into white noise. The reverse of this mapping defines the data-generating process we seek to learn-hence, its choice can significantly impact performance [22] . The conventional approach involves adopting a diffusion process derived from the laws of thermodynamics, which, albeit simple and principled, may be suboptimal due to its lack of adaptability to the dataset.In this study, we investigate whether the notion of diffusion can be instead learned from data. Our motivating goal is to perform accurate log-likelihood estimation and probabilistic modelling, and we take an approach grounded in Bayesian inference [23] . We view the diffusion process as an approximate variational posterior: learning this process induces a tighter lower bound (ELBO) on the marginal likelihood of the data. Although previous work argued that the ELBO objective of a diffusion model is invariant to the choice of diffusion process [20, 22] , we show that this claim is only true for the simplest types of univariate Gaussian noise: we identify a broader class of noising processes whose optimization yields significant performance gains. Specifically, we propose a new diffusion process, multivariate learned adaptive noise (MuLAN), which augments classical diffusion models [51, 20] with three innovations: a per-pixel polynomial noise schedule, an adaptive input-conditional noising process, and auxiliary latent variables. In practice, this method learns the schedule by which Gaussian noise is applied to different parts of an image, and allows tuning this noise schedule to the each image instance.Our learned diffusion process yields improved log-likelihood estimates on two standard image datasets, CIFAR10 and ImageNet. Remarkably, we achieve state-of-the-art performance with less than half of the training time of previous methods. Our method also does not require any modifications to the underlying UNet architecture, which makes it compatible with most existing diffusion algorithms.Contributions In summary, our paper makes the following contributions:1. We demonstrate that the ELBO of a diffusion model is not invariant to the choice of noise process for many types of noise, thus dispelling a common assumption in the field. 2. We introduce MULAN, a learned noise process that adaptively adds multivariate Gaussian noise at different rates across an image in a way that is conditioned on arbitrary context (including the image itself). 3. We empirically demonstrate that learning the diffusion process speeds up training and matches the previous state-of-the-art models using 2x less compute, and also achieves a new state-of-the-art in density estimation on CIFAR-10 and ImageNet 
background: 
 A diffusion process q transforms an input datapoint denoted by x 0 and sampled from a distribution q(x 0 ) into a sequence of noisy latent variables x t for t ∈ [0, 1] by progressively adding Gaussian noise of increasing magnitude [51, 16, 53] . The marginal distribution of each latent is defined by q(x t |x 0 ) = N (x t ; α t x 0 , σ t I) where the diffusion parameters α t , σ t ∈ R + implicitly define a noise schedule as a function of t, such that ν(t) = α 2 t /σ 2 t is a monotonically decreasing function in t. Given any discretization of time into T timesteps of width 1/T , we define t(i) = i/T and s(i) = (i -1)/T and we use x 0:1 to denote the subset of variables associated with these timesteps; the forward process q can be shown to factorize into a Markov chain q(x 0:1 ) = q(x 0 ) T i=1 q(x t(i) |x s(i) ). The diffusion model p θ is defined by a neural network (with parameters θ) used to denoise the forward process q. Given a discretization of time into T steps, p factorizes as p θ (x 0:1 ) = p θ (x 1 ) T i=1 p θ (x s(i) |x t(i) ). We treat the x t for t > 0 as latent variables and fit p θ by maximizing the evidence lower bound (ELBO) on the marginal log-likelihood given by: log p θ (x 0 ) = ELBO(p θ , q) + D KL [q(x t(1):t(T ) |x 0 )∥p θ (x t(1):t(T ) |x 0 )] ≥ ELBO(p θ , q) (1) In most works, the noise schedule, as defined by ν(t), is either fixed or treated as a hyperparameter [16, 3, 18] . Chen [3] , Hoogeboom et al. [18] show that the noise schedule can have The task of log-likelihood estimation is directly motivated by applied problems such as data compression [31] . 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 [4] . Other applications of log-likelihood estimation include adversarial example detection [52] , semi-supervised learning [5] , and others. Note that our primary focus is density estimation and probabilistic modeling rather than sample quality. The visual appeal of generated images (as measured by e.g., FID) correlates imperfectly with log-likelihood. We focus here on pushing the state-of-the-art in log-likelihood estimation, and while we report FID for completeness, we defer sample quality optimization to future work.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Yes, our model does not get state of the art FID due to it not having a lower frequency bias. See the paper for more details. 
a forward diffusion process with multivariate adaptive noise: 
 Next, our plan is to define a family of approximate posteriors Q, as well as a family suitably matching reverse processes P, such that the optimization problem max p∈P,q∈Q ELBO(p, q) is tractable and does not suffer from the aforementioned invariance to the choice of q. This subsection focuses on defining Q; the next sections will show how to parameterize and train a reverse model p ∈ P.Notation. Given two vectors a and b, we use the notation ab to represent the Hadamard product (element-wise multiplication). Additionally, we denote element-wise division of a by b as a / b. We denote the mapping diag(.) that takes a vector as input and produces a diagonal matrix as output.Intuitively, a multivariate noise schedule injects noise at different rates for each pixel of an input image. This enables adapting the diffusion process to spatial variations within the image. We will also see that this change is sufficient to make the ELBO no longer invariant in q.Formally, we define a forward diffusion process with a multivariate noise schedule q via the marginal for each latent noise variable x t for t ∈ [0, 1], where the marginal is given by: q(x t |x 0 ) = N (x t ; α t x 0 , diag(σ 2 t )), (2) where x t , x 0 ∈ R d , α t , σ t ∈ R d + and d is the dimensionality of the input data. The α t , σ t denote varying amounts of signal associated with each component (i.e., each pixel) of x 0 as a function of time t(i). We define the multivariate signal-to-noise ratio as ν(t) = α 2 t /σ 2 t and choose α t , σ t so that ν(t) decreases monotonically in t along all dimensions and is differentiable in t ∈ [0, 1]. Let α t|s = α t /α s and σ 2 t|s = σ 2 t -α 2 t|s /σ 2 s with all operations applied elementwise. These marginals induce transition kernels between steps s < t given by (Suppl. 19):EQUATIONIn Sec. 3.5, we argue that this class of diffusion process Q induces an ELBO that is not invariant to q ∈ Q. The ELBO consists of a line integral along the diffusion trajectory specified by ν(t). A line integrand is almost always path-dependent, unless its integral corresponds to a conservative force field, which is rarely the case for a diffusion process [55] . See Sec. 3.5 for details.Next, we extend the diffusion process to support context-adaptive noise. This enables injecting noise in a way that is dependent on the features of an image. Formally, suppose we have access to a context variable c ∈ R m which encapsulates high-level information regarding x 0 . Examples of c could be a class label, a vector of attributes (e.g., features characterizing a human face), or even the input x 0 itself. We define the marginal of the latent x t in the forward process as q(x t |x 0 , c) = N (x t ; α t (c)x 0 , σ 2 t (c)); the reverse process can be similarly derived (Suppl. 19) as:EQUATIONwhere the diffusion parameters α t , σ t are now conditioned on c via a neural network.Specifically, we parameterize the diffusion parameters α t (c), σ t (c), ν(t, c) as α 2 t (c) = sigmoid(-γ ϕ (c, t)), σ 2 t (c) = sigmoid(γ ϕ (c, t)), and ν(c, t) = exp (-γ ϕ (c, t)). Here, γ ϕ (c, t) : R m × [0, 1] → [γ min , γ max ] d is a neural network with the property that γ ϕ (c, t) is monotonic in t. Following Kingma et al. [20] , Zheng et al. [65] , we set γ min = -13.30, γ max = 5.0.We explore various parameterizations for γ ϕ (c, t). These schedules are designed in a manner that guarantees γ ϕ (c, 0) = γ min 1 d and γ ϕ (c, 1) = γ max 1 d , Below, we list these parameterizations. The polynomial parameterization is novel to our work and yields significant performance gains. [20] . We use the monotonic neural network γ vdm (t), proposed in VDM to express γ as a function of t such that γ vdm (t) : [0, 1] → [γ min , γ max ] d . Then we use FiLM conditioning [38] in the intermediate layers of this network via a neural network that maps z. The activations of the FiLM layer are constrained to be positive.We express γ ϕ (c, t) as a monotonic degree 5 polynomial in t . Details about the exact functional form of this polynomial and its implementation can be found in Suppl. E.2.In principle, we can fit a normal diffusion model in conjunction with our proposed forward diffusion process. However, variational inference suggests that the variational and the true posterior ought to have the same dependency structure: that is the only way for the KL divergence between these two distributions to be zero. Thus, we introduce a class of approximate reverse processes P that match the structure of Q and that are naturally suitable for joint optimization max p∈P,q∈Q ELBO(p, q). Formally, we define a diffusion model where the reverse diffusion process is conditioned on the context c. Specifically, given any discretization of t ∈ [0, 1] into T time steps as in Sec. 2, we introduce a context-conditional diffusion model p θ (x 0:1 |c) that factorizes as the Markov chainEQUATIONGiven that the true reverse process is a Gaussian as specified in Eq. 4, the ideal p θ matches this parameterization (the proof mirrors that of regular diffusion models; Suppl. D), which yieldsEQUATIONwhere x θ (x t , t), is a neural network that approximates x 0 . Instead of parameterizing x θ (x t , t) directly using a neural network, we consider two other parameterizations. One is the noise parameterization [16] where ϵ θ (x t , c, t) is the denoising model which is parameterized as ϵ θ (x t , t) = (x t -α t (c)x θ (x t , t, c))/σ t (c); see Suppl. E.1.1 and the other is v-parameterization [45] wherev θ (x t , c, t) is a neural network that models v θ (x t , c, t) = (α t (c)x t -x θ (x t , c, t))/σ t (c); see Suppl. E.1.2.Note that the model p θ (x 0:1 |c) implicitly assumes the availability of c at generation time. Sometimes, this context may be available, such as when we condition on a label. We may then fit a conditional diffusion process with a standard diffusion objective E x0,c [ELBO(x 0 , p θ (x 0:1 |c), q ϕ (x 0:1 |c)], in which both the forward and the backward processes are conditioned on c (see Sec. 3.4).When c is not known at generation time, we may fit a model p θ that does not condition on c. Unfortunately, this also forces us to define p θ (x s |x t ) = N (µ p (x t , t), Σ p (x t , t)) where µ p (x t , t), Σ p (x t , t) is parameterized directly by a neural network. We can no longer use a noise parameterization ϵ θ (x t , t) = (x t -α t (c)x θ (x t , t, c))/σ t (c) because it requires us to compute α t (c) and σ t (c), which we do not know. Since noise parameterization plays a key role in the sample quality of diffusion models [16] , this approach limits performance.We propose an alternative strategy for learning conditional forward and reverse processes p, q that feature the same structure and hence support efficient noise parameterization. Our approach is based on the introduction of auxiliary variables [60] , which lift the distribution p θ into an augmented latent space. Experiments (Suppl. D.3) and theory (Suppl. D) confirm that this approach performs better than parameterizing c using a neural network, c θ (x t , t).Specifically, we introduce an auxiliary latent variable z ∈ R m and define a lifted p θ (x, z) = p θ (x|z)p θ (z), where p θ (x|z) is the conditional diffusion model from Eq. 5 (with context c set to z) and p θ (z) is a simple prior (e.g., unit Gaussian or fully factored Bernoulli). The latents z can be interpreted as a high-level semantic representation of x that conditions both the forward and the reverse processes. Unlike x 0:1 , the z are not constrained to have a particular dimension and can be a low-dimensional vector of latent factors of variation. They can be continuous or discrete. The learning objective for the lifted p θ is given by:EQUATIONEQUATIONwhere ELBO(p θ (x 0:1 |z), q ϕ (x 0:1 |z)) denotes the variational lower bound (VLB) of a diffusion model (defined in Eq. 1) with a forward process q ϕ (x 0:1 |z) (defined in Eq. 4 and Sec. 3.2.2) and and an approximate reverse process p θ (x 0:1 |z) (defined in Eq. 5), both conditioned on z. The distribution q ϕ (z|x 0 ) is an approximate posterior for z parameterized by a neural network with parameters ϕ.Crucially, note that in the learning objective (Eq. 8), the context, which in this case is z, is available at training time in both the forward and reverse processes. At generation time, we can still obtain a valid context vector by sampling an auxiliary latent from p θ (z). Thus, this approach addresses the aforementioned challenges and enables us to use the noise parameterization in Eq. 6.Although we apply Jensen's inequality twice to get (8) , this also enables us to learn the noise process, which significantly offsets any potential increase in ELBO gap reduction and improves ELBO(p θ (x 0:1 |z), q ϕ (x 0:1 |z)) by optimizing over a more expressive class of posteriors. This claim is empirically validated in Table 2 .Next, we derive a precise formula for the learning objective (8) of the auxiliary-variable diffusion model. Using the objective of a diffusion model in (1) we can write (8) as the sum of four terms:EQUATIONThe reconstruction loss, L recons , can be (stochastically and differentiably) estimated using standard techniques; see [23] , L prior = -D KL [q ϕ (x 1 |x 0 , z)∥p θ (x 1 )] is the diffusion prior term, L latent = -D KL [q ϕ (z|x 0 )∥p θ (z)] is the latent prior term, and L diffusion is the diffusion loss term, which we examine below. The complete derivation is given in Suppl. E.3.Discrete-Time Diffusion. We start by defining p θ in discrete time, and as in Sec. 2, we let T > 0 be the number of total time steps and define t(i) = i/T and s(i) = (i -1)/T as indexing variables over the time steps. We also use x 0:1 to denote the subset of variables associated with these timesteps.Starting with the expression in Eq. 1 and following the steps in Suppl. E, we can write L diffusion as:EQUATIONContinuous-Time Diffusion. We can also consider the limit of the above objective as we take an infinitesimally small partition of t ∈ [0, 1], which corresponds to the limit when T → ∞. In Suppl. E we show that taking this limit of Eq. 10 yields the continuous-time diffusion loss:EQUATIONwhere ∇ t γ(z, t) ∈ R d denotes the Jacobian of γ(z, t) with respect to the scalar t. We observe that the limit of T → ∞ yields improved performance, matching the existing theoretical argument by Kingma et al. [20] .We try two different kinds of priors for p θ (z): discrete (z ∈ {0, 1} m ) and continuous (z ∈ R m ).Continuous Auxiliary Latents. In the case where z is continuous, we select p θ (z) as N (0, I m ). This leads to the following KL loss term:D KL (q ϕ (z|x 0 )∥p θ (z)) = 1 2 (µ ⊤ (x 0 )µ(x 0 )) + tr(Σ 2 (x 0 ) -I m ) -log |Σ 2 (x 0 )|).Discrete Auxiliary Latents. In the case where z is discrete, we select p θ (z) as a uniform distribution. Let z ∈ {0, 1} m be a k-hot vector sampled from a discrete Exponential Family distribution p θ (z; θ) with logits θ. Niepert et al. [34] show that z ∼ p θ (z; θ) is equivalent to z = arg max y∈Y ⟨θ + ϵ g , y⟩ where ϵ g denotes the sum of gamma distribution Suppl. F, Y denotes the set of all k-hot vectors of some fixed length m. For k > 1, To differentiate through the arg max we use a relaxed estimator, Identity, as proposed by Sahoo et al. [44] . This leads to the following KL loss term: D KL (q ϕ (z|x 0 )∥p θ (z)) = -m i=1 q ϕ (z|x 0 ) i (log q ϕ (z|x 0 ) i + log m).Having defined our loss, we now return to the question of whether it is invariant to the choice of diffusion process. Notice that we may rewrite Eq. 11 in the following vectorized form:EQUATIONwhere the square is applied elementwise. We seek to rewrite (12) as a line integral b a f (r(t))• d dt r(t)dt for some vector field f and trajectory r(t). Recall that ν(z, t) is monotonically decreasing in each coordinate as a function of t; hence, it is invertible on its image, and we can write t = ν -1 z (ν(z, t)) for some ν -1 z . Let xθ (x ν(z,t) , z, ν(z, t)) ≡ x θ (x ν -1 z (ν(z,t)) , z, ν -1 z (ν(z, t))) and note that for all t, we can write x t as x ν(z,t) ; see Eq. 30, and have xθ (x ν(z,t) , z, ν(z, t)) ≡ x θ (x t , z, t). We can then write the integral in (12) as1 0 (x 0 -xθ (x ν(z,t) , z, ν(z, t))) 2 • d dt ν(z, t)⟩dt, which is a line integral with f (r(t)) ≡ (x 0 -xθ (x ν(z,t) , z, ν(z, t))) 2 and r(t) ≡ ν(z, t).Intuitive explanation. Imagine piloting a plane across a region with cyclones and strong winds, as shown in Fig. 5 . Plotting a direct, straight-line course through these adverse weather conditions requires more fuel and effort due to increased resistance. By navigating around the cyclones and winds, however, the plane reaches its destination with less energy, even if the route is longer.This intuition translates into mathematical and physical terms. The plane's trajectory is denoted by r(t) ∈ R n + , while the forces acting on it are represented by f (r(t)) ∈ R n . The work required to navigate is given by1 0 f (r(t)) • d dt r(t), dt.Here, the work depends on the trajectory because f (r(t)) is not a conservative field. This concept also applies to the diffusion NELBO. From Eq. 12, it's clear that the trajectory r(t) is parameterized by the noise schedule ν(z, t), which is influenced by complex forces, f (analogous to weather patterns), represented by the dimension-wise reconstruction error of the denoising model, (x 0 -x θ (x t , z, t)) 2 . Thus, the diffusion loss, L diffusion , can be interpreted as the work done along the trajectory ν(z, t) in the presence of these vector field forces f . By learning the noise schedule, we can avoid "high-resistance" paths (those where the loss accumulates rapidly), thereby minimizing the overall "energy" expended, as measured by the NELBO. Since the diffusion process corresponds to non-conservative force fields, as noted in Spinney & Ford [55] , different noise schedules should yield different NELBOs-a result supported by our empirical findings. In Suppl. E.5, we show that variational diffusion models are limited to linear trajectories ν(t), rendering their objective invariant to the noise schedule. In contrast, our approach learns a multivariate ν, enabling paths that achieve a better ELBO.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments: 
 This section reports experiments on the CIFAR-10 [25] and ImageNet-32 [58] datasets. We don't employ data augmentation and we use the same architecture and settings as in the VDM model [20] . The encoder, q ϕ (z|x), is modeled using a sequence of 4 ResNet blocks which is much smaller than the denoising network that uses 32 such blocks (i.e., we increase parameter count by only about 10%); the noise schedule γ ϕ is modeled using a two-layer MLP. In all our experiments, we use discrete auxiliary latents with m = 50 and k = 15. A detailed description can be found in Suppl. G.In these experiments, we replace VDM's noise process with MULAN. On CIFAR-10, MULAN attains VDM's likelihood score of 2.65 in just 2M steps, compared to VDM's 10M steps 1). When trained on 4 V100 GPUs, VDM achieves a training rate of 2.6 steps/second, while MULAN trains slightly slower at 2.24 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. On ImageNet-32, VDM integrated with MULAN reaches a likelihood of 3.71 in half the time, achieving this score in 1M steps versus the 2M steps required by VDM.In Table 2 , we also compare MULAN with other recent methods on CIFAR-10 and ImageNet-32. MULAN was trained using v-parameterization for 8M steps on CIFAR-10 and 2M steps on Imagenet-32. During inference, we extract the underlying probability flow ODE and use it to estimate the log-likelihood; see Suppl. I.2. Our algorithm establishes a new state-of-the-art in density estimation on both ImageNet-32 and CIFAR-10. In Table 8 , we also compute variational lower bounds (VLBs) of ≤2.59 and ≤3.71 on CIFAR-10 and ImageNet, respectively. Each bound improves over published results (Table 2 ); our true NLLs (via flow ODEs) are even lower. Concurrent work that seeks to improve log-likelihood estimation by learning the forward diffusion process includes Neural Diffusion Models (NDMs) [1] and DiffEnc [33] . In NDMs, the noise schedule is fixed, but the mean of each marginal q(x t |x 0 ) is learned, while DiffEnc adds a correction term to q. Diffusion normalizing flows (DNFs) represent an earlier effort where q is a normalizing flow trained by backpropagating through sampling. In Table 3 , we compare against NDMs, DiffEnc, and DNFs on the CIFAR-10 dataset, using the authors' published results; note that their published Ima-geNet numbers are either not available or are reported on a different dataset version that is not comparable. Our approach to learned diffusion outperforms previous and concurrent work.Due to the expensive cost of training, we only performed ablation studies on CIFAR-10 with a reduced batch size of 64 and trained the model for 2.5M training steps. In Fig. 2a we ablate each component of MULAN: when we remove the conditioning on an auxiliary latent space from MULAN so that we have a multivariate noise schedule that is solely conditioned on time t, our performance becomes comparable to that of VDM, on which our model is based. Changing to a scalar noise schedule based on latent variable z initially underperforms compared to VDM. This drop aligns with our likelihood formula (Eq. 6) which includes D KL (q ϕ (z|x 0 )|p θ (z)), an extra term not in VDM. The input-conditioned scalar schedule doesn't offer any advantage over the scalar schedule used in VDM. This is due to the reasons outlined in Sec. Perceptual Quality While perceptual quality is not the focus of this work, we report FID numbers for the VDM model and MuLAN (Table 1 ). We use RK45 ODE solver to generate samples by solving the reverse time Flow ODE (Eq. 76). We observe that MuLAN does not degrade FIDs, while improving log-likelihood estimation. 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.Loss curves for different noise schedules. We investigate different parameterizations of the noise schedule in Fig. 2b . Among polynomial, linear, and monotonic neural network, we find that the polynomial parameterization yields the best performance. The polynomial noise schedule is a novel component introduced in our work. The reason why a polynomial function works better than a linear or a monotonic neural network as proposed by VDM is rooted in Occam's razor. In Suppl. E.2, 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. More expressive models (e.g., monotonic 3-layer MLPs) are more difficult to optimize. Examining the noise schedule. Since the noise schedule, γ ϕ (z, t) is multivariate, we expect to learn different noise schedules for different input dimensions and different inputs z ∼ p θ (z).In Fig. 3 , we take our best trained model on CIFAR-10 and visualize the variance of the noise schedule at each point in time for different pixels, where the variance is taken on 128 samples z ∼ p θ (z).We note an increased variation in the early portions of the noise schedule. However, on an absolute scale, the variance of this noise is smaller than we expected. We also tried to visualize noise schedules across different dataset images and across different areas of the same image; refer to Fig. 13 . We also generated synthetic datasets in which each datapoint contained only high frequencies or only low frequencies, and with random masking applied to parts of the data points; see Suppl. H. Surprisingly, none of these experiments revealed human-interpretable patterns in the learned schedule, although we did observe clear differences in likelihood estimation. We hypothesize that other architectures and other forms of conditioning may reveal interpretable patterns of variation; however, we leave this exploration to future work.Replacing the noise schedules in a trained denoising model. We also confirm experimentally our claim that the learning objective is not invariant to the multivariate noise schedule. We replace the noise schedule in the trained denoising model with two alternatives: MULAN with scalar noise schedule, and a linear noise schedule: γ ϕ (z, t) = γ min + t(γ max -γ min )1 d ; see Kingma et al. [20] .For both the noise schedules the likelihood reduces to the same value as that of the VDM: 2.65. 
related work: 
 Diffusion models have emerged in recent years as powerful tools for modeling complex distributions [51, 16] , extending flow-based methods [53, 24, 48, 49] The noise schedule, which determines the amount and type of noise added at each step, plays a critical role in diffusion models. Chen [3] empirically demonstrate that different noise schedules can significantly impact the generated image quality using various handcrafted noise schedules. Kingma et al. [20] showed that the likelihood of a diffusion model remains invariant to the noise schedule with a scalar noise schedule. In this work we show that the ELBO is no longer invariant to multivariate noise schedules.Recent works explored multivariate noise schedules (including blurring, masking, etc.) [17, 42, 36, 12 ], yet none have delved into learning the noise schedule conditioned on the input data itself. Likewise, conditional noise processes are typically not learned [26, 39, 62] and their conditioner (e.g., a prompt) is always available. Auxiliary variable models [63, 60] add semantic latents in p, but not in q, and they don't condition or learn q. In contrast, we learn multivariate noise conditioned on latent context.Diffusion normalizing flows (DNFs) [64] learn a q parameterized by a normalizing flow; however, such q do not admit tractable marginals and require sampling full data-to-noise trajectories from q, which is expensive. Concurrent work on neural diffusion models (NDMs) and DiffEnc admits tractable marginals q with learned means and univariate schedules; this yields more expressive q than ours but requires computing losses in a modified space that precludes using a noise parameterization and certain sampling strategies. Empirically, MuLAN performs better with fewer parameters (Suppl. A).Optimal transport techniques seek to learn a noise process that minimizes the transport cost from data to noise, which in practice produces smoother diffusion trajectories that facilitate sampling. Schrondinger bridges [47, 6, 59 , 37] learn expressive q do not admit analytical marginals, require computing full data-to-noise trajectories and involve iterative optimization (e.g., sinkhorn), which can be slow. Rectification [27] seeks diffusion paths that are close to linear; this improves sampling, while our method chooses paths that improve log-likelihood. See Suppl. A for more detailed comparisons. 
conclusion: 
 We introduced MULAN, a context-adaptive noise process that applies Gaussian noise at varying rates across input data. Our theory challenges the prevailing notion that the likelihood of diffusion models is independent of the noise schedule: this independence only holds true for univariate schedules. Our evaluation of MULAN spans multiple image datasets, where it outperforms state-of-the-art generative diffusion models. We hope our work will motivate further research into the design of noise schedules, not only for improving likelihood estimation but also to improve image quality generation [35, 53] .A stronger fit to the data distribution also holds promise for improving downstream applications of generative modeling, e.g., compression or decision-making [32, 9, 8, 41] . 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We provide this in the paper.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification:Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not believe our method will have a high risk of abuse as our models are not perceptually SOTA, they only provide for state of the art logliklihood.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: We are using standard benchmark datasets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: The noise schedule, which determines the amount and type of noise added at each step, plays a critical role in diffusion models. Chen [3] empirically demonstrate that different noise schedules can significantly impact the generated image quality using various handcrafted noise schedules. Kingma et al. [20] showed that the likelihood of a diffusion model remains invariant to the noise schedule with a scalar noise schedule. In this work we show that the ELBO is no longer invariant to multivariate noise schedules.Recent works, including Hoogeboom & Salimans [17] , Rissanen et al. [42] , Pearl et al. [36] , have explored per-pixel noise schedules (including blurring and other types of noising), yet none have delved into learning or conditioning the noise schedule on the input data itself. The shared components among these models are summarized and compared in Table 4 .Yang et al. [62] proposes noise processes that are conditioned on an external context (e.g., text). We also propose context-conditioned noise processes; however, their setting is that of conditional generation, where the context is always available at training and inference time, and the context represents external data. Our paper instead looks at unconditional generation, and we condition the noising process on the image itself that we want to generate (via latent variable) and learn how to apply noise across an image as a function of the image.Lee et al. [26] , Popov et al. [39] proposed using a data-dependent prior: however, they do not learn q and their noise process is not adaptive to the input x 0 . Thus they propose a fairly different set of methods from what we introduce.Yang & Mandt [63] , Wang et al. [60] have explored diffusion models with an auxiliary latent space, where the denoising network is conditioned on a latent distribution. Our paper also incorporate auxiliary latents, but unlike previous works, we add them to both p and q and we also also focus on learning the process q (as opposed to doing representation learning using the auxiliary learned space). Lastly, our algorithm relies on many other components including a custom noise schedule, multivariate noise, etc. The shared components among these models are summarized and compared in Table 4 .Table 4 : MULAN is a noise schedule that can be integrated into any diffusion model such as VDM [20] or InfoDiffusion [60] . The shared components between MULAN and these models are summarized and compared in this EQUATIONwhere the drift term f θ : R d × R → R d is parameterized by a neural network with parameters θ and the diffusion term g(t) ∈ R + is a scalar constant and w is the standard Brownian motion. However, in MuLAN, the forward process is given byEQUATIONwhere z ∈ {0, 1} m is the auxiliary latent vector, f θ : R m × R → R d and g θ : R m × R → R d are parameterized by a neural network. Notice that the drift term in DNF, f θ (x, t), is a non-linear function in x 0 , and the same holds for MuLAN since in the drift term, f θ (z, t) ⊙ x, z and x depend on x 0 . Additionally, the diffusion coefficient, g θ (z, t), is multivariate and conditioned on x 0 via z.The two parameterizations are different: on one hand, DNF admits more general classes of neural networks because it does not require marginals to be tractable. On the other hand MuLAN admits a more flexible noise model g θ (z, t) and admits more efficient training (see the summarized Table 5 below).MuLAN has the advantage that it is simulation free; which means that given a data x 0 , the noisy sample x t can be computed in closed form; however, in Diffusion Normalizing Flow, to compute x t , one needs to simulate the forward SDE which is resource intensive and limits its scalability to larger denoising models. While MuLAN optimizes the ELBO, DNF optimizes an approximation for the ELBO. In particular, the DNF training objective does not involve a term that accounts for the entropy of the encoder. Thus, the objective is closer to that of a normal auto-encoder in that regard. Other concurrent work seeks to improve log-likelihood estimation by learning the forward diffusion process in a simulation-free setting. In neural diffusion models (NDMs), the noise schedule is fixed, but the mean of each marginal q(x t |x 0 ) is learned. This requires denoising x in a transformed space, which prevents using noise parameterization, a design choice that is important for performance. Their denoising family also induces a parameterization that limits the kinds of samplers that can be sued. Lastly, NDMs use a model that is 2x larger than a regular diffusion model, while ours only adds 10% more parameters.The DiffEnc framework adds an extra learned correction term to q to adjust the mean of each marginal q(x t |x 0 ). This noise choice also requires using certain parameterizations for x that are not compatible with noise parameterization; while their approach supports v-parameterization, it also requires training a mean parameterization network. Similarly to NDMs, the noise schedule remains fixed, while the mean of each marginal is adjusted by the network. Our approach towards learning the noise schedule yields better empirical performance and is, in our opinion, simpler; it can also be combined with this prior work on learning the marginals' means.In techniques based on optimal transport, the goal is to learn a noise process that minimizes the transport cost from data to noise, which in practice produces smoother diffusion trajectories that facilitate sampling.Minimizing Trajectory Curvature (MTC) of ODE-based generative models Lee et al. [27] , the primary goal is to design the forward diffusion process that is optimal for fast sampling; however, MuLAN strives to learn a forward process that optimizes for log-likelihood. In the former, the marginals x t in the forward process are given asEQUATIONwhere x t , x 0 , 𭟋 ∈ R d . However for MuLAN the marginals arex t = α ϕ (z, t) ⊙ x 0 + 1 -α 2 ϕ (z, t) ⊙ ϵ ; ϵ ∼ N (0, I d ) ; z ∼ q ϕ (z|x 0 ) where α ϕ (z, t) : R d × R → R d ≥0 , z ∈ {0, 1} m , ϵ ∈ R dNotice that in the MTC formula, the coefficient of x 0 , the time integral of the drift term, is a scalar and linear function of, and is independent of the input x 0 . In MuLAN, that term is a multivariate non-linear function in t, and conditioned on x 0 via the auxiliary latent variable z. This implies that the forward diffusion process in MuLAN is more expressive than MTC. The simplistic forward process in MTC enables faster sampling whereas the richer / more expressive forward process in MuLAN leads to improved likelihood estimates. Table 6 summarizes the key differences. An alternative approach to learning a forward process that performs optimal transport is via the theory of Schrodinger bridges [47, 6, 59, 37] . Similarly to the DNF framework, these methods do not admit analytical marginals and therefore involve computing full trajectories from noisy and clean data. Additionally, they are typically trained using an iterative procedure that generalizes the sinkhorn algorithm and involves iteratively training q and p. As such, these types of methods are typically more expensive to train and competitive results on standard benchmarks (e.g., CIFAR10, ImageNet)are not yet available to our knowledge.We have a Gaussian diffusion process that begins with the data x 0 , and defines a sequence of increasingly noisy versions of x 0 which we call the latent variables x t , where t runs from t = 0 (least noisy) to t = 1 (most noisy). Given, T , we discretize time uniformly into T timesteps each with a width 1/T . We define t(i) = i/T and s(i) = (i -1)/T .EQUATIONEQUATION)Kingma et al. [20] show that the distribution q(x s |x t , x 0 ) is also gaussian,EQUATIONSince during the reverse process, we don't have access to x 0 , we approximate it using a neural network x θ (x t , t) with parameters θ. Thus,EQUATIONThis corruption process q is the following markov-chain as q(x 0:1 ) = q(x 0 ) T i=1 q(x t(i) |x s(i) ) . In the Reverse Process, or the Denoising Process, p θ , a neural network (with parameters θ) is used to denoise the noising process q. The Reverse Process factorizes as: p θ (x 0:1 ) = p θ (x 1 ) T i=1 p θ (x s(i) |x t(i) ). Let x θ (x t , t) be the reconstructed input by a neural network from x t . Similar to Sohl-Dickstein et al. [51] , Kingma et al. [20] we decompose the negative lower bound (VLB) as:-log p θ (x 0 ) ≤ E q ϕ -log p θ (x t(0):t(T ) ) q(x t(1):t(T ) |x 0 )ELBO(p θ (x0:1),q(x0:1))defined in Eq. 1EQUATIONThe prior loss, L prior , and reconstruction loss, L recons , can be (stochastically and differentiably) estimated using standard techniques; see Kingma & Welling [23] . The diffusion loss, L diffusion , varies with the formulation of the noise schedule. We provide an exact formulation for it in the subsequent sections.For brevity, we use the notation s for s(i) and t for t(i). From Eq. 31 and Eq. 32 we get the following expression for q(x s |x t , x 0 ):D KL (q(x s |x t , x 0 )∥p θ (x s |x t )) = 1 2 (µ q -µ p ) ⊤ Σ -1 θ (µ q -µ p ) + tr Σ q Σ -1 p -I n -log |Σ q | |Σ p | = 1 2 (µ q -µ p ) ⊤ Σ -1 θ (µ q -µ p )Substituting µq, Σq, µp, Σp from equation 20 and equation 19; for the exact derivation see Kingma et al.EQUATION)Thus L diffusion is given byEQUATIONIn practice instead of computing the integral is computed by MC sampling.= - 1 2 E ϵ∼N (0,In),t∼U [0,1] ν ′ (t)∥x 0 -x θ (x t ,EQUATIONEQUATIONChange of variables. We can write x t explicitly in terms of the signal-to-noise ratio, ν(t), and input x 0 in the following manner:ν t = α 2 t σ 2 tWe know α 2 t = 1 -σ 2 t for Variance Preserving process; see Sec. 2.EQUATION)EQUATIONThus, we write x t in terms of the signal-to-noise ratio in the following manner:x ν(t) = α t x 0 + σ t ϵ t ; ϵ t ∼ N (0, I n ) = ν(t) 1 + ν(t) x 0 + 1 1 + ν(t) ϵ tUsing Eq. 28The distribution of x t given x s is given by:q(x s |x t , x 0 ) = N µ q = α t|s σ 2 s σ 2 t x t + σ 2 t|s α s σ 2 t x 0 , Σ q = diag σ 2 s σ 2 t|s σ 2 t (31)Let x θ (x t , t) be the neural network approximation for x 0 . Then we get the following reverse process:EQUATION)For brevity we use the notation s for s(i) and t for t(i). From Eq. 31 and Eq. 32 we get the following expression for q(x s |x t , x 0 ):D KL (q(x s |x t , x 0 )∥p θ (x s |x t )) = 1 2 (µ q -µ p ) ⊤ Σ -1 θ (µ q -µ p ) + tr Σ q Σ -1 p -I n -log |Σ q | |Σ p | = 1 2 (µ q -µ p ) ⊤ Σ -1 θ (µ q -µ p )Substituting µq, µp, Σp from equation 32 and equation 31.= 1 2 σ 2 t|s α s σ 2 t x 0 - σ 2 t|s α s σ 2 t x θ (x t , t) ⊤ diag σ 2 s σ 2 t|s σ 2 t -1 σ 2 t|s α s σ 2 t x 0 - σ 2 t|s α s σ 2 t x θ (x t , t) = 1 2 (x 0 -x θ (x t , t)) ⊤ diag σ 2 t|s α s σ 2 t ⊤ diag σ 2 s σ 2 t|s σ 2 t -1 diag σ 2 t|s α s σ 2 t (x 0 -x θ (x t , t)) = 1 2 (x 0 -x θ (x t , t)) ⊤ diag σ 2 t|s α s σ 2 t ⊙ σ 2 t σ 2 s σ 2 t|s ⊙ σ 2 t|s α s σ 2 t (x 0 -x θ (x t , t)) = 1 2 (x 0 -x θ (x t , t)) ⊤ diag σ 2 t|s α 2 s σ 2 t σ 2 s (x 0 -x θ (x t , t))Simplifying the expression using Eq. 25 and Eq. 26 we get,EQUATIONLike Kingma et al. [20] we train the model in the continuous domain withT → ∞. L diffusion = lim T →∞ 1 2 T i=2 E ϵ∼N (0,In) D KL (q(x s(i) |x t(i) , x 0 )∥p θ (x s(i) |x t(i) )) = lim T →∞ 1 2 T i=2 E ϵ∼N (0,In) (x 0 -x θ (x t(i) , t(i))) ⊤ diag ν s(i) -ν t(i) (x 0 -x θ (x t(i) , t)) = 1 2 E ϵ∼N (0,In) lim T →∞ T i=2 (x 0 -x θ (x t(i) , t(i))) ⊤ diag ν s(i) -ν t(i) (x 0 -x θ (x t(i) , t)) = 1 2 E ϵ∼N (0,In) lim T →∞ T i=2 T (x 0 -x θ (x t(i) , t(i))) ⊤ diag ν s(i) -ν t(i) (x 0 -x θ (x t(i) , t)) 1 TThe forward process for such a method is given as:EQUATIONThe distribution of x t given x s is given by (the derivation is similar to Hoogeboom & Salimans [17] ):EQUATIOND.1 context is available during the inference time.Even though c represents the input x 0 , it could be available during during inference. For example c could be class labels [10] or prexisting embeddings from an auto-encoder [40] .Let x θ (x t , c, t) be an approximation for x 0 . Then we get the following reverse process (for brevity we write x θ (x t , c, t) as x θ ):EQUATIONSimilar to the derivation of multi-variate L diffusion in Eq. 33 we can derive L diffusion for this case too:EQUATIONThis approach is very limited where the diffusion process is only conditioned on class labels. Using pre-existing embeddings like Diff-AE [40] is also not possible in general and is only limited to tasks such as attribute manipulation in datasets.If the context, c is an explicit function of the input x 0 things become challenging because x 0 isn't available during the inference stage. For this reason, Eq. 40 can't be used to parameterize µ p , Σ p in p θ (x s |x t ). Let p θ (x s |x t ) = N (µ p (x t , t), Σ p (x t , t)) where µ p , Σ p are parameterized directly by a neural network. Using Eq. 4 we get the following diffusion loss:EQUATIONDue to the challenges associated with parameterizing µ p , Σ p directly using a neural network we parameterize c using a neural network that approximates c in the reverse process. Let x θ (x t , t) be an approximation for x 0 . Then we get the following reverse Rrocess (for brevity we write x θ (x t , t) as x θ , and c θ denotes an approximation to c in the reverse process.):EQUATIONConsider the limiting case where T → ∞. Let's analyze the 2 terms in Eq. 43 separately.Using Eq. 4 and Eq. 6, term 1 in Eq. 43 simplifies in the following manner:EQUATIONSubstituting 1 / T as δEQUATIONEQUATION)Consider the scalar case: substituting δ = 1/T ,EQUATIONNotice that this equation is in indeterminate for when we substitute δ = 0. One can apply L'Hospital rule twice or break it down into 3 terms below. For this reason let's write it as expression 1:EQUATIONApplying L'Hospital rule in expression 1 we get,EQUATIONEQUATIONEQUATIONThus the final result:EQUATIONFor the second term we have the following:EQUATIONLetp i = σi 2 (c,s) 1- ν i (c,t) ν i (c,s) σi 2 (c θ ,s) 1- ν i (c θ ,t) ν i (c θ ,s)The sequence lim T →∞ : For c = "class labels" or c = x 0 the likelihood estimates are worse than VDM. For c = x 0 , we see that the VLB degrades with increasing T, but for VDM and MULAN, it improves with increasing T. This empirical observation is consistent with our mathematical insights earlier. As these models consistently exhibit inferior performance w.r.t VDM, in line with our initial conjectures, we refrain from training them beyond 300k iterations due to the substantial computational cost involved. Since the forward pass is given by x t = α t (z)x 0 + σ t (z)ϵ t , we can write the noise ϵ t in terms of x 0 , x t in the following manner:T 2 d i=1 (p i -1 -log p i ) converges iff lim T →∞ d i=1 (p i -1 -log p i ) = 0. Notice that the function f (x) = x -1 -log x ≥ 0 ∀x ∈ REQUATIONFollowing Dhariwal & Nichol [10] , Kingma et al. [20] , instead of parameterizing x θ (x t , z, t) using a neural network, we use Eq. 59 to parameterize the denoising model in terms of a noise prediction model ϵ θ (x t , z, t),EQUATIONFollowing Salimans & Ho [45] , Zheng et al. [65] , we explore another parameterization of the denoising network which is given byEQUATIONIn practice, v-parameterization leads to a better performance than noise parameterization; as illustrated in Table 7 .Let f (x; ψ) be a scalar-valued polynomial of degree n with coefficients ψ ∈ R n+1 expressed as:f (x; ψ) = ψ n x n + ψ n-1 x n-1 + • • • + ψ 1 x + ψ 0 ,and denote its derivative with respect to x as d dx f (x; ψ), represented by f ′ (x; ψ). Now we'd like to find least n such that f (x; ψ) satisfies the following properties:1. f (x; ψ) is monotonically increasing, i.e. f ′ (x; ψ) ≥ 0 ∀x ∈ R, ψ ∈ R n+1 . 2. f ′ (x 1 ; ψ) = 0, f ′ (x 2 ; ψ) = 0 ∃x 1 , x 2 ∈ C, x 1 ̸ = x 2 , ∀ψ ∈ R n+1 .For the first condition to hold, we can design f ′ (x; ψ) such that it's a perfect square with real / imaginary roots. That way f ′ (x; ψ) ≥ 0 ∀x ∈ R, ψ ∈ R n+1 . This means that f ′ (x; ψ) is an even degree polynomial, i.e. the degree of f ′ (x; ψ) can take the following values: 2, 4, . . . . Also, note that at least half of the roots of f ′ (x; ψ) are repeated since f ′ (x; ψ) can be expressed as a perfect square, i.e., if f ′ (x; ψ) has a degree 2 then it has exactly 1 unique root (real / imaginary), if f ′ (x; ψ) has a degree 4 then it has at most 2 unique roots (real / imaginary), and so on.For the second condition to hold, f ′ (x; ψ) needs to have at least 2 unique roots ∃ψ ∈ R n+1 . For this reason f ′ (x; ψ) is a polynomial of degree 4. Thus, f ′ (x; ψ) can be written asf ′ (x; ψ) = (ψ 3 x 2 + ψ 2 x+ψ 1 ) 2 . This ensures that ∃ψ ∈ R 5 s.t. f ′ (x; ψ) = 0 twice in x ∈ R, and f ′ (x; ψ) ≥ 0 ∀ψ ∈ R 5 .Thus, f (x; ψ) takes the following functional form:EQUATIONFor the above-mentioned reasons we express γ ϕ (c, t) : R m × [0, 1] → R d as a degree 5 polynomial in t. We define neural networks a ϕ (c) : R m → R d , b ϕ (c) : R m → R df ϕ (c, t) = a 2 ϕ (c) 5 t 5 + a ϕ (c)b ϕ (c) 2 t 4 + b 2 ϕ (c) + 2a ϕ (c)d ϕ (c) 3 t 3 + b ϕ (c)d ϕ (c)t 2 + d 2 ϕ (c)twhere the multiplication and division operations are elementwise. The the noise schedule, γ(c, t), is given as follows:EQUATIONNotice that γ ϕ (c, t) has these interesting properties:• Is increasing in t ∈ [0, 1] which is crucial as mentioned in Sec. 3.5.• γ ϕ (c, t) has end points at t = 0 and t = 1 which the user can specify via γ min and γ max . Specificaly, γ ϕ (c, t = 0) = γ min 1 d and γ ϕ (c, t = 1) = γ max 1 d . • Its time-derivative i.e. ∇ t γ ϕ (c, t) can be zero twice in t ∈ [0, 1]. This isn't a necessary condition but it's nice to have a flexible noise schedule whose time-derivative can be 0 at the beginning and the end of the diffusion process.In this section we derive the VLB. For ease of reading we use the notation x t to denote x t(i) and x t-1 to denote x t(i-1) ≡ x s(i) in the following derivation.-log p θ (x0)EQUATIONSwitching back to the notation used throughout the paper, the VLB is given as:EQUATIONTo derive the diffusion loss, L diffusion in Eq. 9, we first derive an expression for D KL (q ϕ (x s |z, x t , x 0 )∥p θ (x s |z, x t )) using Eq. 4 and Eq. 6 in the following manner (details in Suppl. E):EQUATIONLet lim T →∞ T (ν s (z) -ν t (z)) = -∇ t ν(z, t) be the partial derivative of the vector ν(z, t) w.r.t scalar t. Then we derive the diffusion loss, L diffusion , for the continuous case in the following manner (for brevity we use the notation s for s(i) = (i -1)/T and t for t(i) = i/T ):L diffusion = lim T →∞ 1 2 T i=2 E ϵ∼N (0,In) D KL (q(x s |x t , x 0 , z)∥p θ (x s |x t , z))Using Eq. 66 we get,= lim T →∞ 1 2 T i=2 E ϵ∼N (0,In) (x 0 -x θ (x t , t(i))) ⊤ diag (ν(s(i), z) -ν(t(i), z)) (x 0 -x θ (x t , t(i))) = 1 2 E ϵ∼N (0,In) lim T →∞ T i=2 T (x 0 -x θ (x t , t(i))) ⊤ diag (ν(s(i), z) -ν(t(i), z)) (x 0 -x θ (x t , t(i))) 1 TUsing the fact that limT →∞ T (ν(s, z) -ν(z, t)) = -∇tν(t, z) we get, = - 1 2 E t∼{0,...,1} (x 0 -x θ (x t , t)) ⊤ (∇ t ν t (z)) (x 0 -x θ (x t , t)) Substituting x0 = α -1 t (z)(xt -σt(z)ϵt) from Eq. 59 and Substituting x θ (xt, z, t) = α -1 t (z)(xt -σt(z)ϵ θ (xt, t)) from Eq. 60 we get, = - 1 2 E t∼[0,1] (ϵ t -ϵ θ (x t , t)) ⊤ σ 2 t (z) α 2 t (z) × ∇ t ν t (z) (ϵ t -ϵ θ (x t , t))Let ν -1 (z, t) denote the reciprocal of the values in the vector ν(z, t).= -EQUATIONE.5 Recovering VDM from the Vectorized Representation of the diffusion loss Notice that we recover the loss function in VDM when ν(z, t) = ν(t)1 d where ν t ∈ R + and 1 d represents a vector of 1s of size d and the noising schedule isn't conditioned on z. (b) This concept also applies to the diffusion NELBO. From Eq. 12, it's clear that the trajectory r(t) is parameterized by the noise schedule ν(z, t), which is influenced by complex forces, f (analogous to weather patterns), represented by the dimension-wise reconstruction error of the denoising model, (x 0 -x θ (x t , z, t)) 2 . Thus, the diffusion loss, L diffusion , can be interpreted as the work done along the trajectory ν(z, t) in the presence of these vector field forces f . By learning the noise schedule, we can avoid "high-resistance" paths (those where the loss accumulates rapidly), thereby minimizing the overall "energy" expended, as measured by the NELBO.1 0 ⟨f θ (x 0 , ν(z, t)), d dt ν(t)⟩dt = 1 0 ⟨f θ (x 0 , ν(t)), d dt (ν(t)1 n )⟩dt = 1 0 ⟨f θ (x 0 , ν(t)), 1 d ⟩ν ′ (t)dt = 1 0 ν ′ (t)∥f θ (x 0 , ν(t))∥ 1 1 dt = 1 0 ν ′ (t)∥(x 0 -xθ (x ν(t) , ν(t)))∥ 2 2 dt (68) 1 0 d dt ν(t)∥(x 0 -xθ (x ν(t) , ν(t)))∥ 2The CIFAR-10 dataset [25] is a collection of images consisting of 60,000 32 × 32 color images in 10 different classes, with each class representing a distinct object or scene. The dataset is divided into 50,000 training images and 10,000 test images, with each class having an equal representation in both sets. The classes in CIFAR-10 include: Airplane, Automobile, Bird, Cat, Deer, Dog, Frog, Horse, Ship, Truck.Randomly generated samples for the CIFAR-10 datasaet are provided in Fig. 6a for MULAN and Fig. 6b for VDM. We visualize the noise schedule in Fig. 13 . ImageNet-32 is a dataset derived from ImageNet [7] , where the original images have been resized to a resolution of 32 × 32. This dataset comprises 1,281,167 training samples and 50,000 test samples, distributed across 1,000 labels.Randomly generated samples for the ImageNet datasaet are provided in Fig. 7 for MULAN and Fig. 8 for VDM. We visualize the noise schedule in Fig. 13 . To see if MULAN learns different noise schedules for parts of the images with different frequencies, we modify the images in the CIFAR-10 dataset where we modify an image where we randomly remove the low frequency component an image or remove the high frequency 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 . The corresponding noise schedules is shown in Fig. 13 . As compared to CIFAR-10, we notice that the spatial variation in the noise schedule increases (SNRs for all the pixels form a wider band) while the variance of the noise schedule across instances decreases slightly. To see if MULAN learns different noise schedules for images with different frequencies, we modify the images in the CIFAR-10 dataset where we modify an image where we randomly remove the low frequency component an image or remove the high frequency 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 . The corresponding noise schedules is shown in Fig. 13 . As compared to CIFAR-10, we notice that the spatial variation in the noise schedule increases (SNRs for all the pixels form a wider band) and the variance of the noise schedule across instances increases as well. 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. 12a shows the training samples. MULAN was trained for 500K steps. The samples generated by MULAN is shown in Fig. 12b . The corresponding noise schedules is shown in Fig. 13 . As compared to CIFAR-10, we notice that the spatial variation in the noise schedule slightly increases (SNRs for all the pixels form a wider band) while the variance of the noise schedule across instances decreases. 2 (z, t) = d dt σ 2 t (z) -2σ 2 t (z) d dt log α t (z)Recall that α 2 t (z) = sigmoid(-γ ϕ (z, t)), σ 2 t (z) = sigmoid(γ ϕ (z, t)). Substituting these in the above equations, the expressions for f (z, t) and g 2 (z, t) simplify to the following:f (z, t) = - 1 2 sigmoid(γ ϕ (z, t)) d dt γ ϕ (z, t), g 2 (z, t) = sigmoid(γ ϕ (z, t)) d dt γ ϕ (z, t)The corresponding reverse-time SDE is given as:EQUATIONwhere w is a standard Wiener process when time flows backwards from 1 → 0, and dt is an infinitesimal negative timestep. Given, s θ (x t , z), an approximation to the true score function, ∇ xt log q ϕ (x t |x 0 , z), Song et al. [53] show that the marginals of Eq. 80 can be described by the following Ordinary Differential Equation (ODE):EQUATIONZheng et al. [65] show that the score function, s θ (x t , z), for the noise and the v-parameterization is given as follows:EQUATIONApplying the change of variables formula [2] on Eq. 81, log p θ (x 0 |z) can be computed in the following manner:EQUATIONwhereh θ (xt, z, t) ≡ f (z, t) - 1 2 g 2 (z, t)s θ (xt, z)The expression for log-likelihood (Eq. 8) is as follows,EQUATIONComputing tr (∇ xt h θ (x t , t, z)) is expensive and we follow Chen et al. [2] , Zheng et al. [65] , Grathwohl et al. [15] to estimate it with Skilling-Hutchinson trace estimator [50, 19] . In particular, we haveEQUATIONwhere the random variable ϵ satisfies E p(ϵ) [ϵ] = 0 and Cov p(ϵ) [ϵ] = I. Common choices for p(ϵ) include Rademacher or Gaussian distributions. Notably, the term ∇ xt h θ (x t , t, z)ϵ can be computed efficiently using "Jacobian-vector-product" computation in JAX. In our experiments, we follow the exact evaluation procedure for computing likelihood as outlined in Song et al. [53] , Grathwohl et al. [15] . Specifically, for the computation of Eq. 85, we employ a Rademacher distribution for p(ϵ). To calculate the integral in Eq. 84, we utilize the RK45 ODE solver [11] provided by scipy.integrate.solve_ivp with atol=1e-5 and rtol=1e-5.Real-world datasets for images or texts often consist of discrete data. Attempting to learn a continuous density model directly on these discrete data points can lead to degenerate outcomes [56] and fail to provide meaningful density estimations. Dequantization [46, 16, 65] is a common solution in such cases. To elaborate, let x 0 represent 8-bit discrete data scaled to [-1, 1]. Dequantization methods assume that we have trained a continuous model distribution p θ for x 0 , and define the discrete model distribution byP θ (x 0 ) = [-1 256 , 1 256 ) d p θ (x 0 + u)du.To train P θ (x 0 ) by maximum likelihood estimation, variational dequantization [16, 65] introduces a dequantization distribution q(u|x 0 ) and jointly train p model and q(u|x 0 ) by a variational lower bound:EQUATIONTruncated Normal Dequantization. Zheng et al. [65] show that truncated Normal distribution,EQUATIONwith mean 0, covariance I, and bounds -1 256 , 1 256 along each dimension, leads to a better likelihood estimate. Thus, Eq. 86 simplifies to the following (for details please refer to section A. in Zheng et al. [65] ):EQUATION)), and τ = 3.log p θ x 0 + σϵ αϵ ε is evaluated using Eq. 84.Importance Weighted Estimator. Eq. 87 can also be extended to obtain an importance weighted likelihood estimator to get a tighter bound on the likelihood. The variational bound is given by (for details please refer to section A. in Zheng et al. [65] ):log P θ (x 0 ) ≥E ε(1) ,...,ε (K) ∼T N (0,I,-τ,τ )   log   1 K K i=1 p θ x 0 + σϵ αϵ ε(k) q(ε (i) )     + d log σ ϵ (88) with σϵ αϵ = exp(- 1 2 × 13.3), log σϵ = 1 2 (-13.3 + softplus(-13.3)), q(ε) = 1 (2πZ) 2 exp - 1 2∥ε∥ 2 2 , Z = 0.9974613, and τ = 3.Note that for K = 1, Eq. 88 is equivalent to Eq. 87; see Zheng et al. [65] . log p θ x 0 + σϵ αϵ ε is evaluated using Eq. 84. In Table 8 , we report BPD values for MULAN on CIFAR10 (8M training steps, v-parameterization) and ImageNet (2M training steps, noise parameterization) using both the VLB-based approach, and the ODE-based approach with K = 1 and K = 20 importance samples. Justification: See our introduction for a list of claims including getting SOTA results on density estimation.• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Justification: Our paper is just a diffusion model useful for compression.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).https://github.com/s-sahoo/MuLAN 38th Conference on Neural Information Processing Systems (NeurIPS2024).dt denotes the diffusion loss, L diffusion , as used in VDM; seeKingma et al. [20]. 


Decision: The paper presents a variant of diffusion models that excels at likelihood estimation. The method is based on generalizing the scaling and noise variance to a multi-variate (per-pixel) version, conditioning them on the data or a representation of data, and learn the noise schedule along with training. The authors remarked that this introduces additional flexibility in diffusion models for optimizing ELBO, as now the noise schedule also affects the loss in general cases. The method achieves a new state-of-the-art of likelihood estimation on a few image datasets. The reviewers acknowledged the innovation and the empirical achievement, while also raised a few imperfections, e.g., still not resolving the long-standing conflict between generation quality and density estimation, results on higher-resolution images, and more realistic downstream utilities. Nevertheless, the contribution and quality of finish are compelling enough for an accept.

Reviews: review: paper_summary: paper_summary: The authors propose MuLAN, a method to learn a multivariate (pixel-wise for images) noise injection schedule for diffusion models, leading to improved likelihood estimates compared to prior work. They provide extensive experimental results and an ablation study to demonstrate the efficiency of their method.
strength_weakness: strength_weakness: - According to the experimental results provided, the only benefit of MuLAN is improved likelihood estimation. The experiments are conducted exclusively in the image domain, but MuLAN does not yield better FID scores compared to prior works. As discussed in Section 3.1, a model with better likelihood estimation may be useful for tasks like compression or adversarial example detection. However, the authors only provide likelihood estimation results and do not demonstrate improved practical outcomes in downstream tasks. It's not completely clear where MuLAN should be preferred over other models. Including results from downstream tasks, especially in non-image domains where uniform Gaussian noise may not be as effective, would enhance the paper.
- The derivation of the multivariate noise schedule closely follows the derivations from VDM, substituting the original scalar $\gamma$ function with a multivariate (and conditional) one. This approach limits the novelty of the proposed method.
- The significant part of the contribution seems to come from the introduction of the auxiliary latent variable $z$, as indicated by the ablation study in Section 4.4. I would appreciate a more detailed discussion of this technique and its efficacy. This step appears to add an extra hierarchical layer to the variational autoencoder. Since the vanilla diffusion model can be seen as a hierarchical VAE itself, it's not entirely clear what this addition offers. A more detailed discussion of this idea and its connection to other works, such as diffusion in the latent space of VAEs [1] or diffusion where the noising process does not necessarily match the prior at t=1 [2], would be beneficial.
- The authors note that the learned noise injection schedule is challenging to interpret. However, including some visualizations of the learned schedule at different time steps would provide at least some intuition about its characteristics.

[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.
questions: questions: - You claim (e.g., on lines 190-193) that MuLAN leads to a tighter NLL-ELBO gap. Do you report the gap for MuLAN and compare it with prior works? In Table 2, to which you refer, you only provide likelihood estimation results, without an ELBO comparison.
- At the end of the day, the polynomial parameterization of the $\gamma$ function involves just three parameters, which seems significantly less flexible than the parameterization from VDM. Am I correct in understanding that you claim that the main reason the polynomial parameterization performs better is because it is easier to optimize?
- As you discuss, the polynomial parameterization may have two points with zero derivative. Do you restrict the polynomial parameterization to have zero derivative at t=0 and t=1?
- Could you clarify how you calculate the likelihood? In regular diffusion, we may integrate the density starting from the unit Gaussian distribution and following the ODE dynamics. However, in MuLAN, you introduce auxiliary variables z, and it’s unclear how to incorporate them into the likelihood estimation pipeline. Do you simply sample z from its prior and then integrate the ODE starting from the data point x, conditionally on z?
- In Section 4.3, where you discuss connections with other approaches that also learn the forward process, you mention that these papers do not provide experimental results on the same version of the ImageNet dataset. However, from what I see, NDM references the same version of ImageNet. I think it’s worth including this comparison in Table 3. While NDM reports a lower NLL score, as you discussed, your approach is significantly different, so I don’t believe it characterizes your work negatively.
limitations: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_summary: This paper introduces an enhanced framework for variational diffusion models (VDM). Rather than using a uniform noise scheduler for all pixels, the new approach assigns different schedulers to individual pixels, adapting to the data distributions. The authors highlight the novelty of this extension, noting that the Evidence Lower Bound (ELBO) depends on the entire trajectory induced by the scheduler, not just the endpoints, as is the case in traditional VDM. Empirical studies demonstrate that their framework achieves state-of-the-art density estimation on CIFAR10 and ImageNet, reducing the number of training steps by 50%.
strength_weakness: strength_weakness: 1. In section 3.5, the authors discussed why the generalization makes the ELBO rely on the entire trajectory; however, they state the fact without giving any intuitive explanations. The authors should discuss this issue more. In addition, providing some toy examples to clearly show how the extended framework is differentiated from the existing one could significantly improve the paper's presentation. 

2. The FIDs of the proposed model are significantly worse than the existing diffusion models. 

3. The FIDs reported for VDM seem inconsistent with the one reported in the original paper. Can authors clarify how the FIDs are obtained.
questions: questions: Please refer to the Weakness section.
limitations: limitations: I am not aware of any potential negative societal impact of this work.
paper_summary: paper_summary: This paper proposes to extend diffusion models by learning a per-pixel noise schedule for the forward noising process, that can be conditional on a context or on an auxiliary variable. This leads to faster convergence and SOTA results on density estimation on the simple benchmark datasets CIFAR-10 and ImageNet-32.
strength_weakness: strength_weakness: - On tables 2, 7, 8 there are confidence intervals on the means, but it's unclear how many random seeds were used, and what's the impact of randomness (I guess weight initialization). Also, on the plots (I'm looking at those with ablations) there seem to be no CIs.
- Sections 3.2.1 and 3.2.2 are a bit repetitive, since 3.2.2 extends 3.2.1 only in a minor way. Especially Eqs. (3) and (4) are exactly the same (except for the context dependency).
- It would be helpful to have a clear self-contained definition for the forward and backward processes, after introducing everything, to summarize the overall setup.
- In the forward process, the dependency on the context $\mathbf{c}$ or the latent variables $\mathbf{z}$ seems to be only through the element-wise noise schedule. In particular, the mentioned NDM and DiffEnc methods appear to be at least partially orthogonal and might lead to further improvements. As far as I can tell, this is not explicitly discussed.
- Only up to 32x32 images are considered here. I think ImageNet-64 should really be included at this point. Since by now density estimation on these benchmarks works very well, it might be interesting and relevant to scale it up to even higher resolutions, but as a proof of concept I think something like the standard ImageNet-64 (used for example in the VDM paper, too) should be sufficient.
- Typo: at line 54, $\sigma_t$ should be $\sigma_t^2$.
questions: questions: Why are the results in Table 3 reported separately? I think it would be clearer to include them in Table 2 (maybe separated by a horizontal line from older methods).
limitations: limitations: Yes. 

Not necessarily a limitation per se, but I would potentially mention that the forward diffusion process depends on the context or latent variables only through the pixel-wise noise schedule. There are orthogonal methods (which are already mentioned in the paper) that learn the mean of the forward process too.
paper_summary: paper_summary: This paper introduces the Multivariate Learned Adaptive Noise (MULAN) model, a novel diffusion process that adapts the noise schedule to different regions of an image. The authors claim significant improvements in log-likelihood estimation and training efficiency, challenging the conventional assumption that the ELBO is invariant to the noise process, and demonstrate state-of-the-art performance in density estimation on CIFAR-10 and ImageNet, reducing training steps by 50%.

I appreciate the paper because it is based on the presumption, which I share, that a significant improvement in generative diffusion approaches can be achieved by adjusting the forward process to the data. Specifically, the authors suggest adapting both drift and noise to the input data. My understanding is that the adaptivity of the MULAN model is achieved by utilizing more degrees of freedom associated with the forward process within the ELBO-style learning/training. However, the claim that the “optimal” schedule is spatially inhomogeneous is based purely on empirical evidence. I would like to see the authors turn this empirical finding into a more systematic exploration. For instance, can we draw conclusions about the spatio-temporal details of the optimal schedule, such as whether the earlier stages of the forward process are more spatially inhomogeneous than the later ones? Do we see more inhomogeneity on the periphery of the image? How does the inhomogeneity correlate with a spatially coarse-grained (filtered) version of the image? Is it more beneficial to have inhomogeneity in the drift (deterministic) or diffusion (stochastic) part of the forward process?

The paper also discusses using labels and other high-level features of the image, but it is unclear how these affect or interplay with the spatial inhomogeneity.

Overall, while the paper presents interesting ideas, I feel it is not yet ready for prime time. The empirical claims need to be backed by a more systematic and thorough analysis.
strength_weakness: strength_weakness: see summary
questions: questions: see summary
limitations: limitations: see summary

Comments: title: Thank you for the detailed review
Thank you for your feedback
Thank you for you invaluable feedback
Current rating doesn't reflect the updated score of 6
Thank you for the detailed responses.
Response to reviewer 1B1k (3/3)
Response to reviewer 1B1K (2/3)
Response to reviewer 1B1k (1/3)
Response to reviewer QQVN
General Response to reviewers
comment: Dear reviewer,

We are immensely grateful to you for providing such detailed and thorough feedback on our paper. We will incorporate these discussions into the next version of the manuscript.

Thanks,
authors
Dear Reviewer, 

We are immensely grateful to you for providing detailed and thorough feedback on our manuscript. We'll incorporate the discussion on the intuition on ELBO into the final version of the manuscript. 

Thanks,
authors.
Dear Reviewer,

We truly appreciate you providing invaluable feedback on our manuscript.

Thanks,
authors
Dear Reviewer,

Thank you very much for evaluating our manuscript and providing invaluable feedback.

We would like to bring to your attention that the current rating still reflects the previous score of 4 instead of the **updated score of 6**. We wanted to mention this in case there was any confusion.
Good response.  My new score is 6.
I appreciate the authors' thorough responses. All my concerns have been addressed. However, I would appreciate it if the authors could incorporate discussions on the intuition behind ELBO into the final version. I have adjusted the score accordingly. Nice work!
Thank you for addressing my questions and providing a thorough and insightful rebuttal to all reviewers. I am satisfied with your responses and recommend acceptance, though I will maintain my current score for now.
I thank the authors for their comprehensive rebuttal. I will maintain my score.
## 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.
## 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.
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.
## **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.
We thank the reviewers for their helpful feedback. Here is a summary of our clarifications in response to their concerns:

### 1. **Visualization of the learned noise schedule.**
We’d like to point out that the visualizations for the noise schedule has been provided 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. 

### 2. **The intuition behind ELBO is 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.

### 3. **FIDs do not yet match those of 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.

Finally, 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.

---
**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.