Diffusion Models With Learned Adaptive Noise
Authors: Subham Sekhar Sahoo
Aaron Gokaslan
Christopher De Sa
Volodymyr Kuleshov
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Generative Models
Diffusion Models
Image Diffusion
Primary area: generative_models
TL;DR: A learned noise schedule can improve the log-likelihood of diffusion models.

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

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. its choice Diffusion Models With Learned Adaptive Noise
Subham Sekhar Sahoo Cornell Tech, NYC, USA. ssahooâcs. cornell corpeell edu
Aaron Gokaslan Comell Tech, NYC, USA. akg870cs. wornell.eeeeeeee
Chris De Sa
Comell University, Ithaca, USA.
cdesa@cs.cornell.c cdesaccs. wornell.lleeeeoeeeeoo edu
Volodymyr Kuleshov Cornell Tech, NYC, USA. kuleshoo@@coneell.o edu
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 equa- tions which maps data to noise in way that can significantly affect performance. In this paper, we explore whether the diffusion process can be learned from data. Our work grounded Bayesian inference and seeks to improve log likelihood estimation by casting the learned diffusion process as approximate variational posterior that yields 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 acToss image. Specifically, our method relics on a multivariate noise schedule that is a function of the data to ensure that the ELBO is longer invariant to the choice of the noise schedule in previous works. Empirically, MULAN sets a new state-of-the- -hre----tthhertttt in density estima- tion CIFAR-10 10 and ImageNet and reduces the number of training steps by 50%. We provide the code along with blog post and video tutorial on the project page:
https://--shhoo.co//Muuu..
Introduction
The reverse this mapping defines the data-generating process we seek can significantly impact performance [22]. The conventional approach involves adopting 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 take an approach grounded Bayesian inference [23]. We view the diffusion process as an approximate variational posterior: learning this process induces tighter lower bound (ELBO) on the marginal likelihood of the data. Although previous work argued that the ELBO objective of 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.
hhtps//gihhabcco//sahh////aaa
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

MuLAN
Scalar
BPD
CIFAR-10
3.03
Noise Schedule Properties
Multivariate
x
2.90
2.83
2.65
Learned
X
2.55
Adaptive
x
Pixel CNN
Transformer
Image
Score Flow
Diffusion Models
Reflected
VDM
MuLAN (Ours)
Improves ELBO
x
Figure 1: (Left) Comparison of noise schedule properties: Multivariate Learned Adaptive Noise schedule (MuLAN) (ours) versus typical scalar noise schedule. Unlike scalar noise schedules, MULAN's multivariate and input-adaptive properties improve likelihood. (Right) Likelihood in bits-per-dimension (BPD) CIFAR-10 without data augmentation.
Specifically, we propose a new diffusion process, multivariate learned adaptive noise (MuLAN), which augments classical diffusion models [51, 20] with three innovations: 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, CIFAR1O and ImageNet. Remarkably, we achieve kate-of-the-ee 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 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 sfate-of-the-art in density estimation on CIFAR-10 and ImageNet
2 Background
A diffusion process transforms an input datapoint denoted by Xo and sampled from a distribution (Xx)) into sequence of noisy latent variables x, for [o. 1] by progressively adding Gaussian noise of increasing magnitude [51, 16, 53]. The marginal distribution of each latent is defined by q(x, Xx) = (xx XyXo, II where the diffusion parameters at, de E R+ implicitly define a noise schedule function of l, such that v(() af o is monotonically decreasing function in L. Given any discretization of time into T timesteps of width 1/T, we define ((i) = i/T and s(i) = (i - 1)/T and we use X0.1 to denote the subset of variables associated with these timesteps; the forward process q can be shown factorize into Markov chain (xxaa) q(xu) IT ((x,()) (x.()))
The diffusion model pa is defined by neural network (with parameters 0) used to denoise the forward process q. Given discretization of time into T steps, factorizes as po (Xo:1)
px(x)) I pe (xx(s) (x))). We treat the x; for & 0 latent variables and fit pe by maximizing the evidence lower bound (ELBO) on the marginal log-likelihood given by:
log ELBO(pe, 4) - D (xx) [xo) (Xu1)aa()) xx)] 2 ELBO(po. 4)
In most works, the noise schedule, as defined by v(t), either fixed or treated as a hyper- parameter [16, 3, 18]. Chen [3], Hoogeboom et al. [18] show that the noise schedule can have

---

significant impact on sample quality. Kingma et al. [20] consider learning v(t), but argue that the KL divergence terms in the ELBO are invariant to the choice of function V,, except for the initial values (0); w(1), ), and they set these values to hand-specified constants in their experiments. They only consider learning y for the purpose of minimizing the variance of the gradient of the ELBO. In this work, we show that the ELBO is not invariant to more complex forward processes.
3 Diffusion Models With Multivariate Learned Adaptive Noise
Here, introduce a new diffusion process, multivariate learned adaptive noise (MuLAN), which introduces three innovations: a per per-pixel polynomial noise schedule, a conditional noising process, and auxiliary-variable reverse diffusion. We describe these below.
3.1 Why Learned Diffusion?
Our goal is to perform accurate density estimation and probabilistic modelling, and we take an approach grounded in Bayesian inference [23]. Notice that the gap between the evidence lower bound (L q) and the marginal log-likelihood (MLL) in Eq. 1 is precisely the KL divergence D (xo) (lpoo((()):)c))) (X (x)) between the diffusion process over the latents x, and the true posterior of the diffusion model. The diffusion process plays the role of variational posterior q in EL q): optimizing thus tightens the gap (MLL- ELBO).
This observation suggests that the ELBO can be made tighter by choosing diffusion processes q that is closer to the true posterior (1) (xp)) ) In fact, the key idea of variational inference to optimize max,eo ELBO q) over a family approximate posteriors Q to induce a tighter ELBO [23]. Most diffusion algorithms, however optimize p ELBO(p, q) within some family P with a fixed q. Our work seeks to jointly optimize maxpeP.9e (L q); we will show in our experiments that this improves the likelihood estimation.
The task of log-likelihood estimation is directly motivated by applied problems such as data compres- sion [31]. In that domain, arithmetic coding techniques can take generative model and produce compression algorithm that provably achieves a compression rate (in bits per dimension) that equals the model's ' log-likellnn ligllioel [4]. Other applications of liglliilloood estimation include adversarial example detection [52], semi- superiupee 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 ligeliiid We focus here on pushing the state-of of tatte-of-the-att in log-likeldhood estimation, and while we report FID for completeness, we defer sample quality optimization to future work.
3.2 A Forward Diffusion Process With Multivariate Adaptive Noise
Next, our plan is to define family of approximate posteriors Q, as well as family suitably matching reverse processes P, such that the optimization problem maxpeP.4e . ELBO q) is tractable and does not suffer from the aforementioned invariance to the choice q. This subsection focuses on defining Q; the next sections will show how to parameterize and train reverse model p P.
Notation. Given two vectors and b, (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 input and produces a diagonal matrix output.
the notation ab to represent the Hadamard product
3.2.1 Multivariate Gaussian Noise Schedule
Intuitively, a multivariate noise schedule injects noise 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, for E [0, 1), where the marginal is given by:
qx, xx) . (x;;; Xr, diag(a?))
where x,, Xo E Rdd Ct oi E Rdd and d is the dimensionality of the input data, The Oi,, a, denote varying amounts of signal associated with each component (i.c., each pixel) of xo function of
3

---

time t(2). We define the multivariate signal-to-noise ratio as v(t) a and choose d, so that v() decreases monotonically in along all dimensions and differentiable in [0, 1]. Let G₁ e a./a, and d/ = of / / with all operations applied elementwise.
These marginals induce transition kernels between steps
< t given by (Suppl. 19):
c 2
[(xx|xe, xn) = N
(2x(xxx)) Xsilly
a of
a of Xo, Eg diag
1/s
o
In Sec. 3.5, we argue that this class of diffusion process induces ELBO that is not invariant to E Q. The ELBO consists of line integral along the diffusion trajectory specified by vt)). 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.
3.2.2 Adaptive Noise Schedule Conditioned On Context
Next, extend the diffusion process 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 context variable e E IRm which encapsulates high-level information regarding Xo- Examples of с could be a class label, a vector of attributes (e.g., features characterizing a human face), the input Xo itself. We define the marginal of the latent x, in the forward process as q(x, No. c) = ((x: (c)))x0,
oc(c); the reverse process can be similarly derived (Suppl. 19) as:
qxx,) X No. c)= N -N(m== (c)(()) oc(c) X + oi(c) Xo. Eg diag
()((c)as(c)
of(c)
(4)
where the diffusion parameters ax, og are now conditioned on e via a neural network.
Specifically, we parameterize the diffusion parameters (c), o (c), as (f(c) = sigmoid( (c, ()), of(c) sigmoid(Ta ()), and v(c, () exp (-Yo (c, ()). Here, 7o(c, () RM [0, 1] [imin Tmax id is a neural network with the property that Y t) monotonic in t. Following Kingma et al. [20], Zheng et al. [65], we set Tmin 13.30, 7max 5.0.
We explore various parameterizations for (c, t). These schedules are designed in manner that guarantees (c, 0) = Ymin la and (c,
7max Below, we list these parameterizations. The work and yields significant performance gains.
polynomial parameterization is novel
Monotonic Neural Network [20]. We use the monotonic neural network Tsam (t), proposed in VDM to express y as a function of such that Yedm (t) [0, 1] + [Tmmin ?mmax" ]". 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 be positive.
Polynomial. (Ours) 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.
3.3 Variiiary-Variable
Variable Reverse Diffusson Processes
In principle, we can fit 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 maxpeP. ELBO(p, q)-
Formally, we define a diffusion model where the reverse diffusion process is conditioned on the context Specifically, given any discretization of i E [0, into T' time steps as in Sec. 2, introduce a context-conditions diffusion model po (Xo: |c) that factorizes as the Markov chain
T
Po Px(x po (x (c)
I
pa xx) [((e), c).
(5)

---

Given that the true reverse process is Gaussian as specified in Eq. 4, the ideal po matches this parameterization (the proof mirrors that of regular diffusion models; Suppl. D), which yields
(c)o()) po(xx/c,xx) N Mp (c(c X2
o ((((c)aa(()) xx(x,,)), En diag ( (o:1c)o(,((1//1(c)) of(c)
(6)
where X (x, t), is a neural network that approximates Xu. Instead of parameterizing Xo(X,, t) directly using a neural network, we consider two other parameterizations. One is the noise param- eterization [16] where (x(x,, c) is the denoising model which is parameterized as En (X,, t) = (x, a+ (c)))) Xt, 1, c))/o,())) see Suppl. 1.1.1 and the other is v-parameterization [45] where vo c, 1) is neural network that models vo(x,, c, (ax(c)x; xo t))//a(())) Suppl. 1.1.2
(c); see
3.3.1 Challenges in Conditioning on Context
Note that the model (x implicitly assumes the availability of с at generation time. Sometimes, this context may be available, such as when we condition on label. We may then fit conditional diffusion process with a standard diffusion objective Ex (El.cEEEOO(xp (x-- Pe (Xo:1 [c), (Xo.1 [oo]/c)], in which both the forward and the backward processes are conditioned on (see Sec. 3.4).
When is not known at generation time, we may fit a model that does not condition on c. Unfortu- nately, this also forces us define po (x (x,]x))==> = N (H4, t), Ex(x,, t)) where (x(xx, (x,, t) is parameterized directly by a neural network. We can no longer use noise parameterization cx(X,, t) (x; ag (x) f, c))/o which we do not know. Since noise parameterization plays key role in the sample quality of diffusion models [16], this approach limits performance.
(c)) because it requires us to compute (c and (c).).
3.3.2 Conditioning Noise on an Auxiliary Latent Variable
We propose an alternative strategy for learning conditional forward and reverse processes p, g 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 pe into an augmented latent space. Experiments (Suppl. D.3) and theory (Suppl. D) confirm that this approach performs better than parameterizing using a neural network, Co (x,x,,t). t).
Specifically, we introduce an auxiliary latent variable z RT and define lifted pa(x, z) pe(x|p ( (x/z))x(Z)))))) (z), where pa (x|z)) ) is the conditional diffusion model from Eq. 5 (with context set to z) and po (z) is a simple prior (e.g.. unit Gaussian or fully factored Bernoulli). The latents z can be interpreted as high-level semantic representation of x that conditions both the forward and the reverse processes. Unlike N0:1, the z are not constrained have a particular dimension and can be a low-dimensional vector of latent factors of variation. They can be continuous discrete. The learning objective for the lifted po is given by:
Po (xo)
2 E (M|xa) ) [log po xx 22] - DKL (9s (z)x))))))
(7)
2 E (x|xa) ELBO (o Xx:1 [z), gu (Xoa (z)) (x (z)xo)llpo(z),
(8)
where ELBO (x [z). go (Xo:1 22)) denotes the variational lower bound (VLB) of a diffusion model (defined in Eq. with a forward process qa (x.11 (z.) (defined in Eq. and Sec. 3.2.2) and and an approximate reverse process po [x.a| (z) (defined in Eq. 5), both conditioned on z. The distribution (z((x)) is an approximate posterior for z parameterized by 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, can still obtain valid context vector by sampling an auxiliary latent from Pz(z). Thus, this approach addresses the aforementioned challenges and enables 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 El. EL.BO (xo:1 (2), go Xx:1 1o:1/ (2)) by optimizing over more expressive class of posteriors. This claim is empirically validated in Table 2.
5

---

3.4 Variational Lower Bound
Next, we derive a precise formula for the learning objective (8) of the auxiliary-variaab diffusion model. Using the objective of diffusion model in (1) we can write (8) as the sum of four terms:
log (x 2 Bug [eecom + Ldillusion
Cerior + Liatent):
(9)
The reconstruction loss, Lresons, can be (stochastically and differentiably) estimated using standard techniques; see [23], Lprior
D (x1 [xo, (x( is the diffusion prior term, Listent -DKulqdl (z) (xo) (ppe [z)] is the latent prior term, and Laiffusion is the diffusion loss term, which we
examine below. The complete derivation is given in Suppl. E.3.
3.4.1 Diffusion Loss
Discrete- Time Diffusion. We start by defining Pe 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) 1)/T indexing variables over the time steps. We also use Xo:1 to denote the subset of variables associated with these timesteps. Starting with the expression Eq. and following the steps in Suppl. E, we can write Laiffasion as:
T
Ldiffusion
DKL (xx. Xo, ()ll (x(()) 2)]
1=2
[(e: [o z, Σ
z(i))) (x,
diag (7(z, s(i) - 72, ((z))(())) (E: 2,,() (10)
Continuous-Time Diffusion. We can also consider the limit of the above objective we take an infinitesimally small partition of e [0, 1]. which corresponds to the limit when T oo. In Suppl. we show that taking this limit of Eq. 10 yields the continuous-tin diffusion loss:
Edillusion
1 - c0(x, z, t)) Tdia diag (Ver(z,t) (E: - -x 2))]
(11)
where (ey(2, () Rdd denotes the Jacobian of 7(2, I) with respect the scalar We observe that the limit of T + oo yields improved performance, matching the existing theoretical argument by Kingma et al. [20].
3.4.2 Auxiliary latent loss
We try two different kinds of priors for pe (z): discrete (z E {0, 1) and continuous (z E IR"
Continuous Auxiliary Latents. In the case where z is continuous, we select po(z) as N(, I.).
This leads the following KI. loss term:
(gs (2) x/) (z)) 2 (1 (xx))(xx)) + (x(=²(x))) I... log E² (xo))).
Discrete Auxiliary Latents. In the case where is discrete, we select po as uniform distribution. Let z E {0, t be k-hot vector sampled from discrete Exponential Family distribution (o( ()) with logits 0. Niepert et al. [34] show that ~ po(2; 0) is equivalent to arg maxyey (0 (g. v) where 6g denotes the sum of gamma distribution Suppl. F, Y denotes the set of all h-toot vectors of some fixed length m. For k > 1, To differentiate through the arg max we use relaxed estimator, Identity, as proposed by Sahoo et al. [44]. This leads to the following KI. loss term: DKL (go (z Xo) (x)oooopo (z)) = - go (z] qo xo(x)) + log m.)
3.5 The Variational Lower Bound as a Line Integral Over The Noise Schedule
Having defined our loss, now return to the question of whether it is invariant to the choice of diffusion process. Notice that we may rewrite Eq. 11 the following vectorized form:
Ldiffusion - / (xo x,( Viu(z, t)dt
(12)

---

where the square is applied elementwise. We seek to rewrite (12) as a line integral f" f(r(()) de =r())d for some vector field f and trajectory r(t). Recall that v(z, t) is monotonically decreasing in each coordinate as a function of t; hence, is invertible on its image, and we can write t = Vz (v(2, t)) for some V==" Let X (x.((( 2, v f)) = (() 2, V (v(z, t))) and note that for all t, can write X₂ as (z(z,)) see Eq. 30, and have (z),)( 2, ((2 ()) (x,, 2,1). We then write the integral in (12) as So (xo- (x(xx(e 2, (zz 1)) de
d (z, d/) which is line integral
with f(r( = (xo- X (x)( z, (2, t))}² 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 rtt) E RT. while the forces acting it are represented by f(r(t)) IR.. The work required to navigate is given by Sg f(r(t)) ae &r(t), dt. Here, the work depends on the trajectory because f(r(t)) 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 vz2, t), which is influenced by complex forces, (analogous to weather patterns), represented by the dimension- wise reconstruction error of the denoising model, (X, X (x,, z, ))? Thus, the diffusion loss, Laiffussion can be interpreted as the work done along the trajectory vzz 1) in the presence of these vector field forces f. By learning the noise schedule, we 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 multivariate enabling paths that achieve a better ELBO.
Experiments
This section reports experiments on the CIFAR-10 [25] and ImageNet-32 [58] datasets. We don't t employ data augmentation and we use the same architecture and settings as in the VDM model [20]. The encoder, 9 (z)x), is modeled using a sequence of ResNet blocks which much smaller than the denoising network that uses 32 such blocks (i.e., we increase parameter count by only about 10%); the noise schedule Y 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.
4.1 Training Speed
In these experiments, we replace VDM's noise process with MULAN. On CIFAR-IO. MULAN attains VDM's likelihood score of 2.65 in just 2M steps, compared to VDM's 10M steps 1). When trained on V100 GPUs, VDM achieves a training rate of 2.6 steps/second, while MULAN trains slightly slower 2.24 steps/second due to the inclusion of additional encoder network. However, despite this slower training pace, VDM requires 30 days to reach a BPD 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 IM steps versus the 2M steps required by VDM.
4.2 Likelihood Estimation
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-likelinn lieelikh;ho;d; see Suppl. 1.2. Our algorithm establishes a new state-of-the--rt in density estimation on both ImageNet-32 and CIFAR-10. 10. In Table 8, we also compute variational lower
7

---

Table 1: Likelihood in bits per dimension (BPD) based on the Variational Lower Bound (VLB)
estimate (Suppl. I.1), sample quality (FID scores) and number of function evaluations (NFE) on
CIFAR-IO, for vanilla VDM and VDM when endowed with MULAN, FID and NFE were computed for 10k samples generated using an adaptive-step ODE solver. Both methods use noise parameteriza-
tion (Suppl. F.
CIFAR-10
ImageNet
Model
Steps
VLB (4)
FID (L)
NFE ())
Steps
VLB (4)
FID (L)
NFE (L)
VDM [20]
10M
2.65
23.91
56
2M
3.72
14.26
56
+ MULAN
2M
2.65
18.54
55
1M
3.72
15.00
62
+ MULAN
10M
2.60
17.62
50
2M
3.71
13.19
62
Table 2: Likelihood in bits per dimension (BPD) on the test set of CIFAR-10 and ImageNet. Results with "f" means they are not reported in the original papers. Model types are autoregressive (AR), normalizing flows (Flow), diffusion models (Diff). We only compare with results achieved without data augmentation.
Model
Type CIFAR-10 (L)
ImageNet (4)
PixeICNN [57]
AR
3.03
3.83
Image Transformer [35]
AR
2.90
3.77
DDPM [16]
Diff
$.669
/
ScoreFlow [54]
Diff
2.83
3.76
VDM [20]
Diff
2.65
S 3.72
Flow Matching [28]
Flow
2.99
/
Reflected Diffusion Models [30]
Diff
2.68
3.74
MULAN (Ours)
Diff
2.55 +10-3
3.67
+10
bounds (VLBs) of <2.59 and <3.71 CIFAR-10 and ImageNet, respectively, Each bound improves over published results (Table 2); our true NLLs (via flow ODEs) are even lower.
4.3 Alternative Learned Diffusion Methods
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 fixed, but the mean of each marginal q ((xx(=0) is learned, while DiffEnc adds correction term to q. Diffusion normalizing flows normalizing flow
Table 3: Likelihood in bits per dimension (bpd) on CIFAR-10 for learned diffusion methods.
Model
NLL (L)
(DNFs) represent an earlier effort where trained by backpropagating through sampling. In Table 3, we com- pare 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 are reported on different dataset version that is not comparable. Our approach to learned diffusion outperforms previous and concurrent work.
DNF [64]
3.04
NDM [1]
2.70
DiffEnc [33]
< 2.62
MULAN
2.55
4.4 Ablation Analysis And Additional Experiments
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 ablate each component of MULAN: when we remove the conditioning on an auxiliary latent space from MULAN so that we have multivariate noise schedule that is solely conditioned on time our performance becomes comparable to that of VDM, on which our model is based. Changing to scalar noise schedule based on latent variable z initially underperforms compared to VDM. This drop aligns with our likelihood formula (Eq. 6) which includes DrI. (za(2)xo) (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 the reasons outlined in Sec. 3.5.

---

2.90
2.90
MULAN
Polynomial Monotone NN Linear
Test loss (bits /dim)
2.85
MuLAN w/o aux. latent
Test loss (bits /dim)
2.85
MuLAN w/o multivariate
2.80
MuLAN adaptivity
2.80
2.75
2.75
2.70
2.70
2.65
2.65
0.5M
1M 1.5M
2M
2.5M
0
0.5M
1M
1.5M
2M
2.5M
Iterations
Iterations
(a) In MULAN w/o aux. latent, the noise isn't condi- tioned latent. MULAN w/o multivariate uses scalar noise schedule. MuLAN w/o adaptivity has a linear schedule and no auxiliary latents.
(b) MULAN with different noise schedule parameter- izations: polynomial, monotonic neural network, and linear. Our proposed polynomial parameterization performs the best.
Figure 2: Ablating components of MULAN on CIFAR-10 over 2.5M steps with batch size of 64.
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 liglliicldoood 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 novel component introduced in our work. The reason why a polynomial function works better than a linear monotonic neural network proposed by VDM is rooted in Occam's razor. In Suppl. E.2, show that a degree 5 polynomial 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, Vo (z, () is multivariate, we expect to
CIFAR-10: Each corresponds pixel
500
learn different noise schedules for different in-
put dimensions and different inputs z pe (z). In Fig.
400
take our best trained model on CIFAR-10 and visualize the variance of the noise schedule at each point in time for different pix- els, where the variance is taken on 128 samples z Po (z).
300
100
We note an increased variation in the early portions of the noise schedule. However, on absolute scale, the variance of this noise is smaller than we expected. We also tried visualize noise sched- ules 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 ap- plied to parts of the data points; see Suppl. H. Surprisingly, none of these experiments revealed human-interpretalie 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.
0.0
Figure 3: Noise schedule visualizations for MU- LAN on CIFAR-10. this figure, we plot the variance of n oz tz across different z pa (z) where each curve represents the SNR correspond- ing to input dimension.

---

Replacing the noise schedules in a trained denoising model. We also confirm experimentally our claim that the learning objective is not invariant 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: Yo (2, t) = " + (mmaxx Ymmn) ld; see Kingma et al. [20]. For both the noise schedules the likelihood reduces to the same value as that of the VDM: 2.65.
5 Related Work
Diffusion models have emerged in recent years 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 each step, plays 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 al. [20] showed that the likelihood of diffusion model remains invariant to the noise schedule with a scalar noise schedule. In this work we show that the ELBO 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 the input data itself. Likewise, conditional noise processes are typically not learned [26, 39, 62] and their conditioner (e.g.. prompt) is always available. Auxiliary variable models [63, 60] add semantic latents in p. but not in 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 parameterized by normalizing flow; however, such g 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 than ours but requires computing losses in 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 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-iikeihhod. See Suppl. A for more detailed comparisons.
Conclusion
We introduced MULAN, 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 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-ar 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]. 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, 81].
Acknowledgments and Disclosure of Funding
This work was partially funded by Volodymyr Kuleshov's 's the National Science Foundation under awards 192251,,5,,, CAREER awards 2046760 and 2145577, and the National Institute of Health under award MIRA R35GM151243, and by Christopher De Sa's NSF RI-CAREER award 2046760.
References
[1] Grigory Bartosh, Dmitry Vetrov, and Christian A Naesseth. Neural diffusion models. arXiv preprint arXiv:2310.0833/ 2023.
10

---

[2] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[3] Ting Chen. On the importance of noise scheduling for diffusion models. arXiv preprint 1000::2301..0992,, 2023.
[4] Thomas M Cover and Joy A Thomas. Data compression. Elements of Information Theory, pp- 103-158, 2005.
[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] Valentin De Bortoli, James Thomton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score- based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17700
2021.
[7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009. doi: 10. /O.////PR.00.9.52222.
[8] Shachi Deshpande and Volodymyr Kuleshov. Calibrated uncertainty estimation improves bayesian optimization, 2023.
[9] Shachi Deshpande, Kaiwen Wang, Dhruv Sreenivas, Zheng Li, and Volodymyr Kuleshov. Deep multi-modal structural equations for causal effect estimation with unstructured proxies. Advances in Neurul Information Processing Systems, 35: 10931-10944, 2022.
[10] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780- 8794, 2021.
[11] J.R. Dormand and PJJ. 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. 00000000066077711000000000000000 arricleepps//011100880993333 URL httpt:tt/////ee /tttpp::/www.... sciencedirect. com/
[12] Weitao Du, He Zhang, Tao Yang, and Yuanqi Du A flexible diffusion model. In International Conference on Machine Learning, pp. 8678 8696. PMLR, 2023.
[13] Pavlos S. Efraimidis and Paul G. Spirakis. Weighted random sampling with a reservoir. In- formation Processing Letters, 97(5):181-18,, 2006. ISSN 0020-0190. doi: https://doi.org! 10. 10.1016/f.il..2005.. 1.000.. URL https:://// /www http://.ccc...o..c...c. www.seediencccrcc..ccmmmm com/sciiene//rtiico.oooocoo p11//500201905000Bxxx
[14] Aaron Gokaslan, A Feder Cooper, Jasmine Collins, Landan Seguin, Austin Jacobson, Mihir Patel, Jonathan Frankle, Cory Stephenson, and Volodymyr Kuleshov. Commoncanvas: Open diffusion models trained on creative-commons images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8250-8260, 2024.
[15] Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:
2018.
[16] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[17] Emiel Hoogeboom and Tim Salimans. arXiv:2209.05557. 2022.
Blurring diffusion models.
arXiv preprint
[18] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint arXiv:2301.11093 2023.
[19] Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics Autistics-Simmlatioon and Computation, 18(3):1059 1076, 1989.
11

---

[20] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances neural information processing systems, 34:21696-21707 2021.
[21] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[22] Diederik P Kingma and Ruiqi Gao. Understanding the diffusion objective as a weighted integral of elbos. arXiv preprint arXiv:2303.00848, 2023.
[23] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 2013.
[24] Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions. Advances in neural information processing systems, 31, 2018.
[25] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[26] Sang-gil Lee, Heeseung Kim, Chaehun Shin, Xu Tan, Chang Liu, Qi Meng. Tao Qin, Wei Chen, Sungroh Yoon, and Tie- Yan Liu. Priorgrad: Improving conditional denoising diffusion models with data-dependent adaptive prior. arXiv preprint arXiv:2 (od::2/06.0ee 106.06406, 2021.
[27] Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing trajectory curvature of ode-based generative models. In International Conference on Machine Learning, pp. 18957-18973, PMLR, 2023.
[28] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[29] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100,,2000 2017.
[30] Aaron Lou and Stefano Ermon. Reflected diffusion models. In International Conference on Machine Learning, pp. 22675-22701. PMLR, 2023.
[31] David JC MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
[32] Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 16569-16594. PMLR, 2022. httLt:ts////ooceeiiiii mlr. press//11/2/guuee2222 hhmml.....m........mmmmmm.mmm...
[33] Beatrix MG Nielsen, Anders Christensen, Andrea Dittadi, and Ole Winther. Diffene: Variational diffusion with a learned encoder. arXiv preprint arXiv:2310.1977 2023.
[34] Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit MLE: backpropa- gating through discrete exponential family distributions. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Anmual Conference on Neu- ral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 14567-14579, 2021. URL https: //tepp::///eee http://epoeeeeneeeeeeeeeeeee..ee neurips. cccccc/ppp// httpp/pppp/////////e/// Abstract.. html.
[35] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International conference machine learning, pp- 4055 PMLR, 2018.
[36] Naama Pearl, Yaron Brodsky, Dana Berman, Assaf Zomet, Alex Rav Acha, Daniel Cohen-Or, and Dani Lischinski. Svnr: Spatially-variant noise removal with denoising diffusion. arXiv preprint arXiv:2306.16052 2023.
12

---

[37] Stefano Peluchetti. Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023.
[38] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[39]
Vadim Popov, Ivan Vovk. Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Grad- uts: A diffusion probabilistic model for texl-to-speech. In International Conference on Machine Learning, pp. 8599-8608. PMLR, 2021.
[40] Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasor Suwajanakom. Diffusion autoencoders: Toward a meaningful and decodable representation. In IEEE Confer- ence on Computer Vision and Puttern Recognition (CVPR), 2022.
[41] Richa Rastogi, Yair Schiff, Alon Hacohen, Zhaozhi Li, Ian Lee, Yuntian Deng, Mert R. Sabuneu, and Volodymyr Kuleshov. Semi-parametrie inducing point networks and neural processes. In The Eleventh International Conference on Learning Representations, 2023. URL https: fDrWd5.
/ //opnreviiee wet/foormmiiiie
[42] Severi Rissanen, Markus Heinonen, and Arno Solin. Generative modelling with inverse heat dissipation. arXiv preprint arXiv:2206.. 13397, 2022.
[43] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolutin image synthesis with latent diffusion models, 2021.
[44] Subham Sekhar Sahoo, Anselm Paulus, Marin Vlastelica, Vít Musil, Volodymyr Kuleshov, and Georg Martius. Backpropagation through combinatorial algorithms: Identity with projection works. In The Eleventh International Conference on Learning Representations, 2023. URL btttss//openreviit net/ wtt
[45] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint urXiv:2202.005 2022.
[46] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications. arXiv preprint (Tfi::/700.055555 2017.
[47] Yuyang Shi. Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
[48] Phillip Si, Allan Bishop, and Volodymyr Kuleshov. Autoregressive quantile flows for predictive uncertainty estimation. In International Conference on Learning Representations, 2022.
[49] Phillip Si, Zeyi Chen, Subham Sekhar Sahoo, Yair Schiff, and Volodymyr Kuleshov. Semi- autoregressive energy flows: exploring likelihood-free training of normalizing flows. In Interna- tional Conference on Machine Learning, pp. 31732-31753, PMLR, 2023.
[50] John Skilling. The eigenvalues of mega-dimensional matrices. 1989. URL https://api. semanticscholar arg/CorpusII::1e 117844915.
[51] Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies, 2015.
[52] Yang Song, Taesup Kim, Sebastian Nowozin, Stefano Ermon, and Nate Kushman. Pixeldefend: Leveraging generative models to understand and defend against adversarial examples. arXiv preprint 17Xiv:1710. 10766, 2017.
[53] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 2.Xiv:20111.3456 2020.
[54] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances neural information processing systems, 34: 1415-1428, 2021.
13

---

[55] Richard E. Spinney and Ian J. Ford. Fluctuation relations: pedagogical overview, 2012.
[56] Benigno Uria, lain Murray, and Hugo Larochelle. Rnade: The real-valued neural autoregressive density-estimator. Advances in Neural Information Processing Systems, 26, 2013.
[57] Aaron Van den Oord, Nal Kalchbrenner, Lasse Espcholt, Oriol Vinyals, Alex Graves, al. Con- ditional image generation with pixelenn decoders. Advances in neural information processing systems, 29, 2016.
[58] Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pp. 1747-1756. PMLR, 2016.
[59] Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International conference on machine learning, pp. 10794-10804. PMLR, 2021.
[60] Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang. Christopher Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maximizing diffusion models. International Conference on Machine Learning, pp. XXXX-XXXX PMLR, 2023.
[61] Sang Michael Xie and Stefano Ermon. Reparameterizable subset sampling via continuous relaxations. arXiv preprint 1rXiv:1901.105/1 2019.
[62] Ling Yang, Zhilong Zhang, Zhaochen Yu, Jingwei Liu, Minkai Xu, Stefano Ermon, and Bin Cui. Cross-modal contextualized diffusion models for text-guided visual generation and editing. arXiv preprint arXiv:2402.16622 16627, 2024.
[63] Ruihan Yang and Stephan Mandt. Lossy image compression with conditional diffusion models, 2023.
[64] Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Informa- fion Processing Systems, 34:16280-16291 2021.
[65] Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Improved techniques for maximum likelihood estimation for diffusion odes. arXiv preprint arXiv:2305.033: 033555 2023.
14

---

Contents
Introduction
2 Background
3 Diffusion Models With Multivariate Learned Adaptive Noise
3.1 Why Learned Diffusion?
3.2 A Forward Diffusion Process With Multivariate Adaptive Noise
3.3 Auxiliary- Variable Reverse Diffusion Processes
3,4
Variational Lower Bound
3.5
The Variational Lower Bound as a Line Integral Over The Noise Schedule
4
Experiments
4.1
Training Speed
4.2
Likelihood Estimation
4.3 Alternative Learned Diffusion Methods
8
4.4
Ablation Analysis And Additional Experiments
5
Related Work
10
6
Conclusion
10
Appendices
16
Appendix A Comparing to Previous Work
16
A.1 Diffusion Models with Custom Noise
17
A.2 Advanced Diffusion Models
17
A.3 Learned Diffusion
18
A.4 Optimal Transport..
19
Appendix B Standard Diffusion Models
20
B.1 Forward Process
20
B.2
Reverse Process
20
B.3
Variational Lower Bound
21
B.4
Diffusion Loss
Appendix C
Multivariate noise schedule
22
C.I Forward Process...
22
C.2 Reverse Process
23
C.3 Diffusion Loss
23
C,4 Vectorized Representation of the diffusion loss
24
C.5 Log likelihood and Noise Schedules: A Thermodynamics perspective
24
15

---

Appendix D
Multivariate noise schedule conditioned on context
24
D.1 context is available during the inference time.
25
D.2 context isn't available during the inference time.
25
D.3 Experimental results
28
Appendix E McLAN: MUltivariate Latent Auxiliary variable Noise Schedule
29
E,1
Parameterization in the reverse process
29
E,2
Polynomial Noise Schedule
30
E.3
Variational Lower Bound
30
E,4
Diffusion Loss
E.5
Recovering VDM from the Vectorized Representation of the diffusion loss
32
Appendix
Subset Sampling
33
Appendix G Experiment Details
34
G.1 Model Architecture
34
G.2 Hardware.
G.3 Hyperparameters
34
Appendix H Datasets and Visualizations
34
H.1
CIFAR-10
35
................................
H.2 ImageNet-32
35
H.3 Frequency
37
H.4 Frequency-2...
38
H.5
CIFAR-10: Intensity
38
H.6
Mask
39
Appendix I
Likelihood Estimation
42
L1
VLB Estimate
42
I.2 Exact likelihood computation using Probability Flow ODE
42
Appendices
Appendix A Comparing to Previous Work
MULAN is the first method to introduce learned adaptive noise process. A widely held assumption that the ELBO objective of diffusion model is invariant the noise process [20]. We dispel this assumption: we show that when input-conditioned noise is combined with (a) multivariate noise, (b) a novel polynomial parameterization, and (c) auxiliary variables, learned noise process yields an improved variational posterior and a tighter BLBO. This approach sets new state-of-the-an in density estimation. While (a), (c) were proposed other contexts, we leverage them as subcomponents of novel algorithm. We elaborate further this below.
16

---

A.1 Diffusion Models with Custom Noise
The noise schedule, which determines the amount and type of noise added each step, plays a critical role in diffusion models. Chen [3] empirically demonstrate that different noise schedules can sigmificantly impact the generated image quality using various handcrafted noise schedules. Kingma al. [20] showed that the likelihood of diffusion model remains invariant to the noise schedule with a scalar noise schedule. this work we show that the ELBO is no longer invariant to multivariate noise schedules.
Recent works, including Hoogeboom Salimans [17]. Rissanen 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.
A.2 Advanced Diffusion Models
Yang et al. [62] proposes noise processes that are conditioned on an external context (e.g.. 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 unconditional generation, and condition the noising process on the image itself that we want to generate (via latent variable) and learn how to apply noise across image
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 the input xo. Thus they propose 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 and and we also also focus on learning the process (as opposed to doing representation learning using the auxiliary learned space). Lastly, our algorithm relies many other components including a custom noise schedule, multivariate noise, etc. The shared components among these models are summarized and compared in Table 4.
17

---

Table 4: MuLAN is 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 table.
learned
multivariate input conditioned
auxiliary latents
noise schedule
Method
noise
noise
noise
Monotonic
VDM [20]
Yes
No
No
No
neural network
Blurring
Frequency scaling
Diffusion
No
Yes
No
No
Model [17]
Inverse Heat
No
Yes
No
No
Exponential
Dissipation [42]
SVNR [36]
No
Yes
No
No
Linear
In denoising
InfoDiffusion [60]
No
No
No
Cosine
process
Lossy Compression [63]
denoising process
Linear / Cosine
No
No
No
In noising and denoising process
MULAN (Ours)
Yes
Yes
Yes
Polynomial
A.3 Learned Diffusson
Diffusion Normalizing Flow (DNF) uses the following forward process:
dx, = fx t)dt + gt)dw,
(13)
where the drift term fo Rd x R +
Rd is parameterized by neural network with parameters 0 and the diffusion term (t) IR+ + is scalar constant and is the standard Brownian motion. However, in MuLAN, the forward process is given by
dx, = fz 1) @ x,di + go go 1) O dw; z ~ (o(2
(14)
where z 10, 3m is the auxiliary latent vector, R" X + Rd and ge R parameterized by neural network. Notice that the drift term in DNF, fo(x, (), is non-linear function in Xo, and the same holds for MuLAN since in the drift term, fo fz 1) @ x, z and x depend on Xo. Additionally, the diffusion coefficient, (2( l), is multivariate and conditioned on Xo via 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 more flexible noise model gz(z, t) and admits more efficient training (see the summarized Table below).
Bd
MuLAN has the advantage that it is simulation free; which means that given data X the noisy sample can be computed in closed form; however, in Diffusion Normalizing Flow, to compute x one needs to simulate the forward SDE which is resource intensive and limits scalability to larger denoising models. While MuLAN optimizes the ELBO, DNF optimizes approximation for the ELBO. particular, the DNF training objective does not involve term that accounts for the entropy of the encoder. Thus, the objective is closer to that of normal auto-encoder in that regard.
18

---

Table The key differences between MULAN and DNF is listed below.
Aspect
Property
Diffusion Normalizing Flow
MULAN
Multivariate
Yes
Yes
Drift Term
Adaptive
Yes
Yes
Learnable
Yes
Yes
Multivariate
No
Yes
Diffusion Term
Adaptive
No
Yes
Learnable
No
Yes
Simulation Free Training
No
Yes
Exact ELBO Optimization
No
Yes
NLL (4)
CIFAR-10
3.04
2.55
Other concurrent work seeks to improve log-likelihood estimation by learning the forward diffusion process in a simulation-fres setting. neural diffusion models (NDMs), the noise schedule is fixed, but the mean of each marginal q(4, Xx) is learned. This requires denoising in transformed space, which prevents using noise parameterization, design choice that important for performance. Their denoising family also induces 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 g adjust the mean of each marginal (xx()x)). This noise choice also requires using certain parameterizations for x that are not compatible with noise parameterization; while their approach supports v-parameterization, also requires training 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.
A.4 Optimal Transport
In techniques based on optimal transport, the goal is learn 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- 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 forward process that optimizes for log- likeliood In the former, the marginals x, in the forward process given
X2 = (1 ()-t)xu tz; (z(x))
(15)
(z,(z,t) O Xe [2) (x)2((xx) where ag (z, Rd X Rd E (0, 1]mh E Tdd Notice that in the MTC formula, the coefficient of xg the time integral of the drift term, is scalar and linear function of, and is independent of the input X,. In MuLAN, that term is a multivariate non-linear function I, and conditioned on X₂ via the auxiliary latent variable z. This implies that the forward diffusion process in Mul.NN is more expressive than MTC. The simplistic forward process MTC enables faster sampling whereas the richer/ / more expressive forward process in MuLAN leads to improved likelihood estimates. Table 6 summarizes the key differences.
where X2, Xo, Rd. However for MuLAN the marginals
/ -1=a3(z,))ccee
19

---

Table Comparison between Minimizing Trajectory Curvature and MuLAN methods.
Aspect
Property
Minimizing Trajectory Curvature
MuLAN
Goal
Design faster sampler
Improve likelihood
log-
Learnable
No
Yes
Linearity
Linear in time t, linear in X
Non-linear in time t, Non-linear (and hence X²)
Drift Term
Dimensionality
Scalar
Multivariate
Adaptive
No
Yes
Lincarity
Linear in time t
Non-linear in time
Dimensionality
Multivariate
Multivariate
Diffusion Term
Learnable
Yes
Yes
Adaptive
Yes
Yes
An alternative approach learning 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 and p. As such, these types of methods are typically more expensive to train and competitive results on standard benchmarks (e.g., CIFARIO, ImageNet) not yet available our knowledge.
Appendix B Standard Diffusion Models
We have a Gaussian diffusion process that begins with the data Xg, and defines sequence of increasingly noisy versions of which we call the latent variables where runs from = 0 (least noisy) to 1 (most noisy). Given, T, we discretize time uniformly into T timesteps each with width 1/T. We define t(i) i/T and s(i) = (i 1)/T.
B.1 Forward Process
(xx(xx,) = NN(u X: Tn)
(16)
where
0/
(17)
ax
a 1|s
0
0,
(18)
B.2 Reverse Process
Kingma et al. [20] show that the distribution g(X, |x:, Xo) is also gaussian,
02.0s +
0.00 E,
Orlo0 g(x,/xr, Xr. xo) = N M4
of
Xx
of xo:
of
I
(19)
Since during the reverse process, we don't have access to Xo. we approximate using a neural network (x,,,, with parameters 0. Thus.
orlo0
po (xx(xx) = N [4p
X1
+
x
t),
Ep
In
(20)
of
of
oi
20

---

B.3 Variational Lower Bound
This corruption process is the following markov-chain as ((x0:1) = (xx)) (IT ((x(()) (xx()))
In the Reverse Process, or the Denoising Process, Po. neural network (with parameters 0) is used to denoise the noising process 9. The Reverse Process factorizes as: pa (xx:1)
(o(x1) IT Pe (Xa(i) X((e)) Let Xy xg(X,, be the reconstructed input by a neural network from Xt. Similar to Sohl-Dickstein et al. [51], Kingma et al. [20] we decompose the negative lower bound (VLB) as:
Po (x((O))))))
- log (xoo <
Eg
log
(()((I) xx)
EL H.nO(po (x0:))poXox))) deffned in Eq.
(x (xo) log po (x ((())]
+ E (xxc) (20 DKL. (500)000200000000000001000 x0)]
2=
+ DxI. (Po (x) (x1x(x)
Ex (xx (xo) - log po [(())
T B. DEL lpo (xx()) (x(()) (x))) xx)]
Carmioo
+ DxL Ppo (x) (xx(x))))())
(21)
Epiie
The prior loss, Epriors and reconstruction loss, Crecons can be (stochastically and differentiably) estimated using standard techniques; see Kingma & Welling [23]. The diffusion loss, Laiffusion varies with the formulation of the noise schedule. We provide exact formulation for in the subsequent sections.
B.4 Diffusion Loss
For brevity, we use the notation for sii) and for t(i). From Eq. 31 and Eq. 32 we get the following expression for qxx, /x,, xp):
DKL x,, (x,
1
Σ
(Mg - -pp) TEp (My - pp) tr (E, IR) - log
q
2
ΣS
= (wg (1p)) TEx (eU - Mp) -5
Substituting p4gs Egg Map En from equation 20 and equation 19; for the exact derivation Kingma al. [20]
(v(s) v()) [xxx xx(xe, 1))ll
(22)
2
Thus Laittssion is given by
Laittusion
lim T- 2
Dx..uo (00l Xo)]
=lim
Σ
BeNN(O v(()) ||x xx(xx,
T-400 1m2
T
Σ
(p(s) v(()) ||xo - (x(x
2
T-XX
t=2
21

---

T
1
lim T-
(v(ss) - -(t)) l|xo - (x f)12 T
2
==
d di u(t =v((t); see Kingma al. [20]
Substituting
lim T-000
-
So
(((]]xx xo (x, e)1l dt
(23)
2
In practice instead of computing the integral is computed by MC sampling.
- [t)( xx [ 2 Appendix C
(24)
Multivariate noise schedule
For multivariate noise schedule we have at, o, E R4 where E (0 1]. at, o, are vectors. The timesteps 8, f satisfy S < t S 1. Furthermore, we use the following notations where arithmetic division represents element wise division between vectors:
01 02s
(25)
0 0 d
(26)
C.1 Forward Process
((x(xx,) = N (atj,x,
(27)
Change of variables. We can write x, explicitly in terms of the signal-to-noise ratio, v(t), and input X in the following manner:
V:
We know ai of for Variance Preserving process; see Sec,
1 o
=
0,
o²
and a 2 Vt
(28)
1
1 1t
of
of
We know a = 1 of for Variance Preserving process; see Sec. 2.
1 1- o7
=V
of
1
14 and 0F
of 1 14
(29)
1/
Thus, write x, in terms of the signal-to-noise, ratio in the following manner:
Xx(i) 0,80 + oiEs E, ~ IO
vu(l) + V1+ v(t) Et
1
Using Eq. 28
(30)
V1+
22

---

C.2 Reverse Process
The distribution of x, given x, is given by:
ai X,
02:0
(1x xo) = N (1y
Xu. Eg diag
t/s
(31)
ii
o
i
Let (x(x,,)) t) be the neural network approximation for Xo.
Then we get the following reverse process:
of X, + o a
o
))
po(x, = N Pp
X (x,, t), E, = diag
(32)
of
o
C.3
Diffusion Loss
For brevity we use the notation for s(i) and t for t(i). From Eq. 31 and Eq. 32 we get the following expression for (xx,, [x,, xo):
(xx(g(x, 1o) xx))
log Eg
(MQ Hppp Eg (Ma Hp) ir EE I.)
2 (Pu - -1p)) Egg (Mq Hp)
Substituting M4g Up, Ep from equation 32 and equation 31.
_ 1
axs
Xg
o as
x
(x(,,))
t)
T
diag
tle
a2
Xo
agjaa a
xx(x),1)
o
o²
(xo xx 1)) diag
o as
diag
o 2
f|s
diag
0 oi.
ass
(xo
xo(xi, 1))
o
o
0
1/s as
(xo xx t)) diag
(xo
x6 (x(, t))
i oi
=x(x0 (x,, t)} diag
xo t)) (xo
Simplifying the expression using Eq. 25 and Eq. 26 we get,
af
Xo (X, t)) diag
(xo t)),
o?
Using the relation v1)) = a² o² we get,
(x xx ) diag (M(s) vtt)) (xo x))
(33)
Like Kingma et al. [20] we train the model in the continuous domain with T + ao.
Laillusion
T
-imm
E--NOo. (x Xo)PpxxX
T-+00 2
2=
lim
R BE- - (0(x((0))(i))) diag (Ve(0) v2(9) (Xo - xo(x1()) 1))
T
2
2=2
1
EE-N(0LL))
lim
(xn xx(xx(), diag (Va() 2(()) (xo xx(x((())))
2
T-
lim T-
T(xo xx( (()))))))))))) diag (Ma(0) - Me()) (xo Xa xx(x))))
23

---

Thus Laitrisionn can be interpreted as the amount of work done along the trajectory v(0) + (11 in the presence of a vector field fg (xo, mzzzt), t)). From the perspective of thermodynamics, this is precisely equal to the amount of heat lost into the environment during the process of transition between Let T(va()) - (u()) = df (() denote scalar derivative of vector v(t) wart T4000
(..N(O. [,, (xo xx t)) diag
(
d dt (() (xo xo(x,, t))dt
(34)
In practice instead of computing the integral is computed by MC sampling.
(x t)) diag |i ((1) )
(xo -0o
(xo -xxxxxxx100 0))
(35)
C.4 Vectorized Representation of the diffusion loss
Let (t) be the vectorized representation of the diagonal entries of the matrix vt). We can rewrite the integral in 34 in the following vectorized form where @ denotes element wise multiplication and ) denotes dot product between vectors.
Ldittusion
d { (xo xo(xi, diag
dt v(t)) (xo xo(x,, t)
2 20
1 //'' ( ( (xx--xx xxxxx t)) @ (xo x xx t)),
d
v(t) dt
20
dt
Using change of variables as mentioned in See. 3.2 we have
= - ( (xx0- xo xo(Xx(e) v())) (xo (x( v (t))) d
di ult) dt
Let fx(xx, v11)) = (xo Xo(XX(() (t))) @ (xo - Xx(Xx()), (t)))
1/' (6(x0)x(1) dt u(() dt
d
(36)
equilibria via the noise schedule specified by vt)).
C,5 Log likelihood and Noise Schedules: A Thermodynamics perspective
A diffusion model characterizes a quasi-static process that occurs between two equilibrium distri- butions: g(xo) > q(x)) ), via a stochastic trajectory [51]. According to Spinney & Ford [55],
demonstrated that the diffusion schedule or the noising process plays pivotal role in determining the
"measure of irreversibility" for this stochastic trajectory which is expressed as log Pa(xi:)) (xox))) (x.x)))
represents the probability of observing the forward path X6:1 and (x(:):)) represents the probability Pr(xa.) of observing the reverse path X1:0- It's worth noting that log P=(x10) corresponds precisely to the ELBO Eq. that we optimize when training diffusion model. Consequently, thermodynamics asserts that the noise schedule indeed has an impact on the log-likelihood of the diffusion model which contradicts Kingma al. [20].
Appendix D Multivariate noise schedule conditioned on context
Let's say we have context variable c E R that captures high level information about Xg- (c(c), (c) E Ry are vectors. The timesteps s, t satisfy 0 s < s Furthermore, we use the following notations:
(c)
(cj(c) = (c(c)
(37)
a
[/,(() = of(c)
(38)
(c
24

---

The forward process for such method is given as:
(x( c) .V (au,(c)x,, oi.(c))
(39)
The distribution of x given x, given by (the derivation is similar Hoogeboom & Salimans [17]):
(x Xxx xo. c)
(1)a(C) (c(c)
()((c)a,(c)
= =N May
X₁ +
of(c) x0, E. = diag
(40)
o7(c)
D.1 context is available during the inference time.
Even though represents the input Xo, it could be available during during inference. For example could be class labels [10] or prexisting embeddings from an auto-encoder [40].
D.1.1 Reverse Process:
Approximate
Let (x,,,, c, t) be an approximation for X6- write xo(X,, C, as X6):
Then we get the following reverse process (for brevity
az (c)
o ())(c)a(())
px(xx/x,,o x, c) = N Hp oi(c)
X₁ +
xo. Σ, diag
cc(c)
oi(c)
(41)
D.1.2 Diffusion Loss
Similar to the derivation of multi-variate Lailussonn in Eq. 33 we can derive Luitlusion for this case too:
di v(()) (xo - xo(x,, c, ())
d
Liiffissonn
Ee ((--) (xo cx, diag 2
(42)
D.1.3 Limitations of this method
This approach 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.
D.2 context isn't available during the inference time.
If the context, is an explicit function of the input things become challenging because X isn available during the inference stage. For this reason, Eq. 40 can't be used parameterize μp, E, in (x(x, (x)). Let Po xx,) = A (x,, t), (x,, t)) where pl, Dp are parameterized directly by neural network. Using Eq. 4 we get the following diffusion loss:
Ldiffusion = T Ei-Ut,, DKL (g(xx()) (x()) xo) (X,(0) (x(()))
= Egg
1 ( mp) " (Mg up)
T
T
tr ir(E, -
In.
log Egl Σp
(43)
2
temm
term 2
D.2.1
Reverse Process: Approximate
Due to the challenges associated with parameterizing Mp, Ep directly using neural network parameterize using a neural network that approximates c the reverse process. Let X. (x;, () be an approximation for Xg- Then get the following reverse Rrocess (for brevity we write xo xx(x;, as Xg. and cy denotes an approximation to c in the reverse process.):
(x(x)(xx)
25

---

-N(1e [1 (co,cco)0((o) of(co)
of (,,co))a,co)) X1 (ct)
( (G(())
o (c))) o((co)
xp, En diag
(44)
Consider the limiting case where T + co. 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:
T pu - Up) " (M up))
lim T-000
2
T ((Hg); - (1ep)s)²
lim T-+00
(45)
(Eo);
i=1
Substituting 1/ as
lim
6-0+
80 2 (xo,t - (1- (xo..)
(x
(x, 1 - 8) 1x(x, 2h + (x((xx (0)0
(1__vixx,1) vi(x, 34
vi(xxt -1-2,(x,-- 8)
(x(x,t) vi(x,t - 8)
ax -6) (x 22 ox:(xo; -5) (1
1x(
(200)
(46)
Dz (x0, (x,x)t
8)
2x(x0, t - 8)
Consider the scalar case: substituting 5 = 1/T,
1
lim 8-+
x
d0 8xo 0) (1 (1- v(xe.t)
v(x0)1-1)
a(xxt -8) 8) t(x,t)
a(x,t) v(x,t - 8) 2 + t-x) (1 -xxtt
v(x,t)
v(x,t - 8) x
2
a(xa,t- - 8) v(xo())
t) t) v(xe,t - 8) 21 + axxx, 8) (1- v(xo,t)
(47)
v(xe, - 8)
Notice that this equation indeterminate for when we substitute One can apply I.' L'Hospital rule twice or break it down into 3 terms below. For this reason let's write it as
1 axxxtt= -0) v(x,t)
21 + a(x,t- (1
v(xx,) (1)-__(())) v(x,1 -6) x
expression 1: lim x
8-+0 8-
a(x,t) v(x,t - 8)
a(xo,t - 0) v(xo,t) t)
=x a(x0, (- 0) (1- u(xo,t) t)
x0
(48)
((xo,1) 1) v(xo,1 8)
u(xo; 1. - 8)
1
[a(x,1- -8) 8) v(x,1) ()
+ alxx 5) (1
v(xx1)
expression 2: lim 8-40
v(x.,.)
a(x,t) v(x,t 5)
v(x,t
8))
[(x).1-0)
2
a(xo, -0) 6) v(x0,() ()
24 alxo, 8) (1- u(xo;1) 1)
x0
(49)
a(xo,t) v(xo,t 8)
v(xo, t 8)
Applying L' L'Hospital rule in expression I we get, d
d8 ( (a(x,t- 8) v(x,t)
t)x,t) v(x, -xx,t-8) 8)) )
v(x t) t)a t) + a(x, t), n/((x,t)
a(x,t)
6=
26

---

(x(xxtt) 1/ (x,,))
(50)
a(x,1) 1)
vxxx1) 1)
a(x, 5) (1- u(x,t)
(1-_((xit) v(x.t 0))
=aa(x,1 t) (x),))
(51)
2(xx
6=0
((
-x((xxt) (x(,t)
a(xx1) 1(x,1) + a' txx,
1/ (xx,t) t)
Ze
(52)
a(xo,1)
v(xo, 1)
-a(x, 1) 1/ (x, a(xo. 1) v' 11(x0)t) x0 v(x,t) t)
u(x,1)
v(xo.t)
x
(xx,)) t)
(53)
Thus the final result:
ai'(xxx tx,
() + vx(xx.t)
+ a (xo,t)
vi((xo.) vi(xo, 1)
btt
1x(x,t)
(xo,x.
oxi(x 1) Mr (xx
Mi 1x 1)
Mi 1)
tx
oxixxo,))
22 1) x0
us (x,(
pix,
A diag
v(xx()) 1)
where
[( [[=a (xx.))
v/ (xx
a (xo.
w (x
x) us arx 1) (x 1
(x (xa. ex(xy. 6)
(54)
((x,()
m(x.))
a(x..t)
((xx
For the second term we have the following:
T
Ezg
lim T-x
tr
-I,) log
Epp
T
tr diag (c,cs)) (1- (1- p(c,t)
v(c, ((c,t))
diag ('(( s) (1 p(co,t)
)) In
_-lim
lim
v(Co,s) s)
diag (c,(c)))(1 w(e.))
log
diag (o²(co,))( 8)(1
vics.!)
v(cn.))
0, (c, 8) (1-
eile,t)
(c,
vi(c.))
T Σ
vi(c,s)
vi(c,s)
lim
- 1 - log
(55)
T
o (co, 8)
vi(ea.t)
σ (co, 8)
vi(ev.t)
i=1
(i(Co,8)
Ve(Ca.8)
(56)
8, (c,0))
k((0.1)
Let p:
Pi(e.o]
a. (c(()))))
1-
()(9g14)
The sequence 1i d d
i= (D; 11
log px) converges iff limy Ed (p: - log p:)
Notice that the function
f [x)
= r 1 1- -log r 2 0 Var E IR and the equality holds for x = 1. Thus,
the condition lim
d
(p: 1 1-log pc) holds iff lim pi Vi E (1,...,d) d). Thus,
lim pi T-000
27

---

o. (c, s)
(c(c)))
Mi (c,8)
lim T+000
=1
o, (ce, s)
(i(())))
volco,
Substituting I/T as 6,
σ. (c,
vilc,t)
vile,t-6)
lim 6-00+
a;²(co,
v.(eo.0)
1i(0011-6)
vi(c,t)
(c, lim o. (co, t) 6-00
v.(e,1-6)
(i(oo.t)
vilcy.t-83
Applying L. L'Hospital rule,
-e.(e.))
σ. tc,
pilc,t)
0. (ca,t) t)
vi(co.))
a. (c,
[)((c,t)),(co:))
(57)
o. (co,t)
(.(c,t)u/(coott))
In the vector form the above equation can be written
o(cc),((o)Vvvvccc.. 1a
(58)
Eq. 58 holds if:
Zz 2g i.c. the unet can perfectly map x, to xo e [0, 1] which is unrealistic.
Clever parameterizations for or ar that ensure Eq. 58 holds.
Because of aforementioned challenges we evaluate this method with finite T = 1000. We demonstrate the performance of the model empirically in Fig. 4.
D.2.2 Recovering VDM
If we substitute (e(c),,((c)) with u(t) (since the SNR isn't conditioned the context c), (r(Ca), (c) with a, and ar(co), with ar, Eq. 45 reduces to the intermediate loss in VDM ie. (xo xo) T Vau(t) (Xe- - -x)) and Eq. 55 reduces to 0.
D.3 Experimental results
In Fig. we demonstrate that the multivariate diffusion processes where с "class labels" or c Xo perform worse than VDM. Since continuous time formulation i.e. T - co for the case when c = X isn't possible (unlike MuLAN or VDM) we evaluate these models the discrete time setting where we use T = 1000. Furthermore we also ablate T = 10k, 100k for Xy to show that the VLB degrades with increasing whereas for VDM and MULAN improves for increasing T; see Kingma et al. [20]. This empirical observation is consistent with our mathematical insights earlier. As these models consistently exhibit inferior performance w.r.t VDM, line with our initial conjectures, we refrain from training them beyond 300k iterations due to the substantial computational cost involved.
28

---

class labels, T=1k
Xo. T=100k
Test loss (bits dim)
c Xo. T=10k
c Xa, T=1k
VDM
50k
100k
150k Iterations
200k
250k
Figure 4: For с = "class labels" or = X the likelihood estimates are worse than VDM. For Xo- we 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. VDM, in line with our initial conjectures, refrain from training them beyond 300k iterations due to the substantial computational cost involved.
Table 7: Likelihood in bits per dimension (BPD) (mean and 95% confidence interval), on the test set of CIFAR- 10 computed using VLB estimate.
parameterization
Num training steps
CIFAR- 10 (L)
Noise parameterization v-paraaeeeriizttion
10M
2.60 + 10
8M
2.59 10-3
Appendix E MULAN: MUltivariate Latent Auxiliary variable Noise Schedule
E.1
Parameterization in the reverse process
E.1.1 Noise parameterization
Since the forward pass is given by (z(2))0 + (z(2),, we can write the noise c in terms of Xp, X; in the following manner:
X ar (Z)x)
E,
(59)
(z)
Following Dhariwal & Nichol [10], Kingma et al. [20], instead of parameterizing Xy x,(x,, z, t) using a neural network, Eq. 59 parameterize the denoising model in terms of noise prediction model 6o (x4, z, t),
(xx) (2)xx 2,
(60)
(z)
E.1.2 Velocity parameterization
Following Salimans Ho [45], Zheng et [65], we explore another parameterization of the denoising network which is given by
vo(x,, 2,1) (z(z)x; Xx 2,
(61)
(z(z)
practice, v-parameterization leads to better performance than noise parameterization; as illustrated in Table 7.
29

---

E.2 Polynomial Noise Schedule
Let (x; be a scalar-valued polynomial of degree with coefficients & R'+1 expressed as:
f(z; (b) vh,"" Vn- Vn-12"
+ 2 4/b.
and denote its derivative with respect to d (a; v) represented by f"(x; 1b). Now we' like to find least n such that [(x; v) satisfies the following properties:
1. (ax; v/ is monotonically increasing, i.e. fx(;; v)) 2 0 Vr E R. E
2. (x11 1p) O, (22; vb) = 0 12,, 22 E C, 21 # 22. Vel E ER++!
For the first condition to hold, can design (x; t) such that it's perfect square with real imaginary roots. That way (=; vb) 2 0 Va E R, da E R This means that f'(z; 4) an even degree polynomial, i.e. the degree of (2 rh) take the following values: 2, ..... Also, note that at least half the roots of f((x; u/) are repeated since fz; vp can be expressed as perfect square, i.e., if ["(x;; +) has degree then it has exactly unique root (real / imaginary), if ["(2; vb) has degree then it has at most 2 unique roots (real imaginary), and so on.
For the second condition to hold, f'(x;; th) needs to have at least unique roots 3ee E IR++1 For this reason fz(:; v/) a polynomial of degree 4. Thus, fr; 1) can be written as f'(;; v) (dyaa
d22 + )2 This ensures that 3y E IR s.t. f'(;; x = 0 twice in x E R, and f'(x; v)) 2 0 Vi R
Thus, f(z; (5 takes the following functional form:
[x; = f (Waa² + 1220 + ddx
2
0/30/22
4/2 + 24/381
5 x
+
x² +
x + o + v/ + constant.
(62)
2
3
For the above-mentionee reasons we express Yolc, t) Rm 10 + Rd as degree 5 polynomial We define neural networks a (c) Rm + Rd, (c(c): Rm - IRd and dæ(c) R Rd with parameters Φ. Let fø Rm x [0, TRd be defined as:
folc,t) ac 5 as (cibbb/l bz(c) + 2as(c)do(c) & by(c)ds(c)E + dc(c)t
2
where the multiplication and division operations are elementwise. The the noise schedule, 7 t), is given follows:
(c, , 1) = Ymin + (Tiax 7min) folc,t)
(63)
Solc,(=1)
Notice that 1) has these interesting properties:
Is increasing in t E [0, 1] which is crucial as mentioned in Sec. 3.5.
(c (c, t) has end points at = and = which the user can specify via Ymin and "max- Specificaly, (c, 7 = 0) Nmin la and (c, 1) Ymax la.
Its time-derivative i.e. Veyo (c. f) can be zero twice in t E [0, 1]. This isn't 't a necessary condition but it's nice have flexible noise schedule whose time-derivative can be 0 the beginning and the end of the diffusion process.
E.3
Variational Lower Bound
In this section we derive the VLBB For ease of reading we use the notation x; to denote Xi(O) and X;-1 to denote Xx(i-1) = Xx(e) in the following derivation.
loog -oxo)
Po (z, Xo:T)
< Egg
log
(2, X1.r (xv)
log po (xxoo 1 (z, XT)
Egg
log (xxr) log (z)
(z,2, X1:1
30

---

Po PoxoTT-- -1, xr)
1
Eyg
log
log (x(2)xu) log po (xcr)
log pa (x)
4o 4x(x1:7 14, Xu)
Po (x0:7 |z, xT)
po (2)
log
llog po(p) - -log Qu(2i (xu)
Qo (X1.T 22, (2Xu)
Po Po 1xxxx-12,, Xx)
po (z)
E
log
log log (= Xo)
q (x, |X-- 1, Xn)
log Po (Xa 2, X1
log Po (o(x(-11) 1 Xe)
log pxer) - log (2)
Egg
go (X1 Xo,, 2))
(x X 1,2, Xo)
no (2) Xa)
log (xo |., xx1
log Po 10 7, 16) (x,, (.,
log pa (xer) - log Do
En
(X1 [oo, z)
(X: Xr, 2, XD Xp X
go x/xx))
log (x x1)
log Po ( xx)
9o ( xo)
(o
Eus
log
loog ps (x()
loog
(Xi [Xo,
(x - Xx)
9o (x: (2, Xu)
Qol (xo)
2=
log (x (, x1)
log Po (xx -1 x:)
(xx1 |z, xo)
log po (z)
Eug
log
log po(xx)
go (xi (X0, 2.)
9o (xi-1 |X,, 2,
9o (xT Xo)
go (z xo)
Pol X X1)
log po
Egg
log Pa Xo 12, x1)
log
log
log po px(T))
2o (x- 1 Xt, 2, Xo)
Qoi (XT z Xo)
(z (x)
Pe 1
po (xr)
log (z)
Egg
log po (xo 1z, xI)
log
log
4n (2x X1, Xo)
4: (xT is, xo)
4id (24 (x))
Eng
log ps (Xo 22,
DeL [pa(ppa 12, xx) 1xo
+ Egs
(xe (xxxxxx))llxxxxx///) xo)]
DEL (z)) ((()xoo)
(64)
Guten
Switching back to the notation used throughout the paper, the VLB is given as:
.log pa (Xo)
T
< Egg
log (x +
DKL (x)() x11))) (x)))))))))) z, xu)]
Lamsoo
E
DRL DKalpo(X1)) (x 12, xo)]
DRL. (po (lao(()))))
(65)
Eyii
Puaen
E.4
Diffusson Loss
To
derive the diffusion loss, LaiiTusion in Eq. 9, we first derive an expression for DKL (x. Xx, xo) 2, x:)) using Eq. 4 and Eq. 6 in the following manner (details in
Suppl. E):
Xe, /z, xt))
1
(Mgo Hp) E (Hg, (p)) + tr Ep In) loog
Σ n9
2
Ep1
1
((xo - -x9) diag(v(Zz, v(z, (x) xo))
(66)
2
31

---

((t))att Let limr-> (f)))) - be the partial derivative of the vector v(z, t) w.r. scalar Then we derive the diffusion loss, Laiffusionn for the continuous case in the following manner (for brevity we use the notation for a(i) = (i 1)/T and I. for 1(i) = i/T):
Ldiftusion
-iim
Dei xo, (x,x)|pp(xx|xxx
T-+00
im2
Using Eq. 66 we get,
lim T+000
BonN xo(x;;((1)) diag (u(s(i), 2) v(((i),2) (xo - xo(xx, (()))
1=
T ΣT(x- xx(xx,((1)) diag (n(s(i), 2) 2(((1),2)) (xo xo(x,, ((i)) 7
1
lim BenN(()
-2
T+000
Using the fact that lim (v((22) v(z,()) V.e((l,)) get.
2 (0.... [xx (x t)) (Ve)a(z)) (xo - xx0-xoox8 (x,,
Substituting ai (x) - (z))z from Eq. 59 and
Substituting 2, ai. (z) - (z) f)) from Eq. 60 we get,
oi(z)
2- Bo (e - ())
x Vevr(2) (e: - cx(xx; 0)) xx
Let v - (z, f) denote the reciprocal of the values in the vector vZz, t).
- ((E: Ex(x;, t)) diag (M (())))V (Er - Ex(xx, t))]
Substituting v(z, f) = exp(- t)) from Sec. E.1.1
Em [(e:
cx t)) T diag (exp (r t) exp (-> wwwpveexpp-p//..//p/ (er (el-<e(xx;))) co(x),t))
2 En-o..1 [(: - cx(,,, t)) diag (cxp (7(z, t)) exp (0)((2)()))))))) (a (
2 Be-1o.11 [(E: - cx(xx, t)) diag (Vey(z,t)) (x
(67)
E.5 Recovering VDM from the Vectorized Representation of the diffusion loss
Notice that we recover the loss function in VDM when vzz, t) v(t)la where is R+ and la represents a vector of of size and the noising schedule isn't 't conditioned on
d f (66(x u(z, t)), dt
/''' (fo, u(t),, d (v(t) dt
/ (Fo(xo, vtt)), 1a/e/
v'(t)]lfo(xo (dt
[ v(())((X0 (x( (()
(68)
Je' के (x()))(( - (x (i denotes the diffusion loss, Laitiusionn as used in VDM; see Kingma et al. [20].
32

---

(a)
d
(b)
1=1
(x xo 2x), 8)) d
Work Done
f(r(t)
r dt)
NELBO
uzz t2,t t)
t=
dt
t=
dt
f(r(t)) R²
r(t)E R²
(x - X0 (Xe, 2, t)) R20
((2,))ERS
Destination
o0ooo
Figure 5: (a) Imagine piloting a plane across a region with cyclones and strong winds, as shown in Fig. 5. Plotting 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 's trajectory is denoted by r(t) E RT, while the forces acting it are represented by f(r(t)) E R". The work required to navigate is given by fo f(r(t)) 4 conservative field.
a rr))), dt. Here, the work depends on the trajectory because f(r(t)) is not a
(b) This concept also applies to the diffusion NELBO. From Eq. 12, it's clear that the trajectory r(t) parameterized by the noise schedule (z,,t,, t), which is influenced by complex forces, f (analogous to weather patterns), represented by the dimension wise reconstruction error of the denoising model, (Xo - X (X,, 2, ())² Thus, the diffusion loss, Laiffussion can be interpreted as the work done along the trajectory vzz, 1) 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.
Appendix F Subset Sampling
Sampling a subset of k items from a collection of collection of ni items, 2t, 22, , .3 belongs to a category of algorithms called reservoir algorithms. In weighted reservoir sampling, every I; associated with a weight w; 2 0. The probability associated with choosing the sequence Swis [ii, .ee be tuple of indices. Then the probability associated with sampling this sequence is
p(Swis w)
Wi,
Wia
Wis
(69)
Z
Z Wi.
Z
Wi;
Efraimidis & Spirakis [13] give an algorithm for weighted reservoir sampling where each item
is assigned a random key ri u; where is drawn from a uniform distribution [0, 1] and w, is the weight of item Vy. Let TopK(r, k) which takes keys r r= [r,, T2, ........ r'] and returns a sequence [i,, .2, ix]. Efraimidis & Spirakis [13] proved that TopK(r, k) is distributed according to p(Swss/ve w).
Let's represent a subset S E {0, 13"h with exactly k non-zero elements that are equal to 1. Then the probability associated with sampling S is given as,
p(S[w)
Σ p(Swi/w)
(70)
SAn
where I(S) denotes all possible permutations of the sequence S. By ignoring the ordering of the elements in Swm we can sample using the same algorithm. Xie & Ermon [61] show that this sampling
algorithm is equivalent to TFpK(F, k) where = [f,, f..,, An] where f; = - -og( log(r:))
log + Gumbel(O, 1). This holds true because the monotonic transformation -log(- Log(-loog( log(#:))
preserves the ordering of the keys and thus TopK(r, k) TopK(f, k).
33

---

Sum of Gamma Distribution. Niepert et al. [34] show that adding SOG noise instead of Gumbel noise leads better performance.
Niepert et al. [34] show that z ~ (o; z;; 0) is equivalent to z = arg maXyEV (0 + Egg y) where Eg is
sample from Sum-of-Gamma distribution given by
SoG(k,
Gamma
log
(71)
where a positive integer and Gamma(a, 3) is the Gamma distribution with (a, 3) as the shape and scale parameters.
And hence, given logits log w, sample k-hot vector using TopK(log + c). We choose categorical prior with uniform distribution across classes. Thus the KL loss term is given by:
Wi
W
log
Y1
(72)
7
7
Appendix G Experiment Details
G.1 Model Architecture
Denoising network. Our model architecture is extremely similar to VDM. The UNet of our pixel- space diffusion has unchanged architecture from Kingma et al. [20].This structure is specifically designed for optimal performance in maximum likelihood training. We employ features from VDM such the elimination of internal downsampling/upsaney features to enhance fine-scale prediction accuracy. In alignment with the configurations suggested by Kingma et al. (2021), our approach varies depending the dataset: For CIFAR- 10, employ U-Net with depth of 32 and I 28 channels; for ImageNet-32. the U-Net also has depth of 32, but the channel count is increased 256. Additionally, all these models incorporate dropout rate of 0. in their intermediate layers.
processes and the integration of Fourier
Encoder network. (x(z(x) is modeled using sequence of 4 Resnet blocks with channel count of 128 for CIFAR-10 and 256 for ImageNet-32 with a drop out of 0.1 in their intermediate layers.
Noise schedule. For polynomial noise schedule, to ag (z), (z), c(z); see Eq. E.2 for details. The MLP has 2 hidden layers size 3072 with swish activation function. The final layer is linear mapping 3072 values corresponding as (z), bz(2)) c(z). Note that as (z), (z), c(z) have the same dimensionality 3072.
MLP that maps the latent vector
G.2 Hardware.
For the ImageNet experiments, we used a single GPU node with 8-A100s. For the cifar- 10 experi- ments, the models were trained on GPUs spanning several GPUs types like V100, A5000s, A6000s, and 3090s with float32 precision.
G.3 Hyperparameters
We follow the same default training settings as Kingma al. [20]. For all our experiments, we use the Adam [21] optimizer with learning rate 2 10 exponential decay rates of 31 = 0.9, 32 = 0.99 and decoupled weight decay [29] coefficient of 0.01. We also maintain an exponential moving average (EMA) of model parameters with an EMA rate of 0.9999 for evaluation. For other hyperparameters, we use fixed start and end times which satisfy Ymin = 13.3, Tmax = 5.0, which used in Kingma et al. [20], Zheng et al. [65].
Appendix H Datasets and Visualizations
In this section we provide a brief description of the datasets used in the paper and visualize the generated samples and the noise schedules.
34

---

H.1 CIFAR-10
The CIFAR-10 dataset [25] is a collection of images consisting of 60,000 32 x 32 color images in 10 different classes, with each class representing distinct object scene. The dataset is divided into 50,000 training images and 10,000 test images, with each class having am 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 datasact provided in Fig. 6a for MULAN and Fig. 6b for VDM. We visualize the noise schedule in Fig. 13.
(a) MULAN with velocity reparameterization after
(b) VDM after 10M training iterations.
8M training iterations.
Figure CIFAR-10 samples generated by different methods.
H.2 ImageNet-32
ImageNet-32 is dataset derived from ImageNet [7], where the original images have been resized to a resolution of 32 x 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 Fig. for MULAN and Fig. 8 for VDM. We visualize the noise schedule in Fig. 13.
35

---

Figure 7: MULAN with noise parameterization after 2M training iterations.
36

---

Figure VDM after 2M training iterations.
H.3 Frequency
To see if MULAN learns different noise schedules for parts of the images with different frequencies, we modify the images in the CIFAR-10 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
CIFAR-10, we notice that the spatial variation in the noise schedule increases (SNRs for all the
pixels form wider band) while the variance of the noise schedule across instances decreases slightly.
37

---

(a) Training samples.
(b) Samples generated by MULAN with noise pa- rameterization after 500K training iterations.
Figure 9: Frequency Split CIFAR- 10 dataset.
H.4 Frequency-2
To see MULAN learns different noise schedules for images with different frequencies, we modify the images in the CIFAR-10 dataset where modify image where randomly remove the low frequency component an image or remove the high frequency with equal probability. Fig. 9a shows the training samples. MULAN 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. 10, notice that the spatial variation in the noise schedule increases (SNRs for all the pixels form wider band) and the variance of the noise schedule across instances increases as well.
PLANA &
AYEEN THE
WIL DE BAN
(a) Training samples.
(b) Samples generated by MULAN with noise pa- rameterization after 500K training iterations.
Figure 10: Frequency Split-2 CIFAR-10 dataset.
H.S CIFAR-10: Intensity
To see MUL.AN learns different noise schedules for images with different intensities, we modify the images in the CIFAR-10 10 dataset where randomly convert an image into a low intensity or
38

---

high intensity image with equal probability. Originally, the CIFARIO images are in the range [0, 255]. To convert an image into low intensity image we multiply all pixel values by 0.5. To convert an image into high intensity image multiply all the pixel values by 0.5 and add 127.5 to them. Fig. 11a shows the training samples. MuLAN was trained for 500K steps. The samples generated by MULAN is shown in Fig. 11b. 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 wider band) while the variance of the noise schedule across instances slightly decreases.
(a) Training samples.
(b) Samples generated by MULAN with noise pa- rameterization after 500K training iterations.
Figure Intensity CIFAR-10 dataset.
H.6 Mask
We modify the CIFAR-10 dataset where we randomly mask (i.e. replace with Os) the top of an image or the bottom half of an image with equal probability. Fig. 12a shows the training samples. MULAN was trained 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 wider band) while the variance of the noise schedule across instances decreases.
39

---

(a) Training samples.
(b) Samples generated by MULAN with noise pa- rameterization after 500K training iterations.
Figure 12: Intensity CIFAR-10 dataset.
40

---

x10
meanz Vz(Z. t)
x102
varz vo(z, t)
4
CIFAR-10
2
0
x10
6
ImageNet
4-
0.5
2
0
0.0
x10
x102
3
Frequency Split
CIFAR-10 10
x10
x102
7.5
Frequency Split
CIFAR-10
4
5.0
2.5
0.0
x10
x101
CIFAR-10 Intensity Split
7.5
5.0
2.5
0
0.0
x10
X101
1.0
CIFAR-10 Mask
0.5
2
0.0
0.2
0.5
1.0
0.2
0.5
r
Figure 13: signal-to-noise ratio for different datasets.
41

---

Appendix I
Likelihood Estimation
We used both Variance Lower Bound (VLB) and ODE-based methods to compute BPD.
I.1
VLB Estimate
In the VLB-based approach, we employ Eq. 9. To compute Laifiuusion we use 7 = 1228 in Eq. 10, discretizing the timesteps, E [0, 1] into 128 bins.
I.2
Exact likelihood computation using Probability Flow ODE
diffusion process whose marginal given by (the same as in Eq. 2),
(x)) (xx;;; 0x diag(o])), Xo (xx);;
(73)
can be modeled as the solution to an Itō Stochastic Differential Equation (SDE):
dx = d( + g(t)dw,, Xa ~ (x(x),;
(74)
where f(t) E RRd g()) E I2d take the following expressions [53]]
d f(t) log az. dt
dt o 2 d
d
g² (t)
log a2
dt
The corresponding reverse process, Eq. 4, can also be modelled by equivalent reverse-time SDE:
dx, = [ft) V log (xx//xo)]dt + g(t) ~ px(x)),
(75)
where W standard Wiener process when time flows backwards from 1 - 0, and dt is an infinitesimal negatíve timestep. Song et al. [53] show that the marginals of Eq. 75 can be described by the following Ordinary Differential Equation (ODE) in the reverse process:
dx, [(()x) 2 (t) log (xxx) dt.
(76)
This ODE, also called the probablity flow ODE, allows to compute the exact likelihood any input data via the instantaneous change of variables formula as proposed in Chen et al. [2]. Note that during the reverse process, the term g(x,) is unknown and is approximated by parameterized by
(xx).. For the probability flow defined in Eq. 76, Chen et al. [2] show that the log- likelihood of (xo) can be computed using the following equation:
fr ir (VV,ha ()) de,
log pa (xo) = log pa (x1)
(77)
1
where ho 1) = f(I)x: 78 (1) (x))
1.2.1 Probability Flow ODE for MULAN.
Similarly for the forward process conditioned on the auxiliary latent variable, z,
(x, (xo z) N/ (xx; (z(z)),, diag( (z))), Xo ~ qu(xx)) z qo (z/xo)
(78)
can extend Eq. 74 in the following manner,
dx, fzz, l)x,dt + g(z, f)dw,, Xo ~ 40 (Xo), z ~ (2(x))))
(79)
to obtain the corresponding SDE formulation. Notice that the random variable z in the above equation
doesn't have subscript i, and hence, z is drawn from gs (z(x)) once and the same z is used as xo
diffuses to x The expressions for f(z, t) R [0, 1] as follows:
+ Rd, g(z,t) t) Rm 10, 1] Id is given
d f(z, t) log (z).
42

---

d
g" (2, t) dl o2 20
log cx:(z)
dl
Recall that (z) = sigmoid(- (z, 1)), [2(2) = sigmoid(7s 1)). Substituting these in the
above equations, the expressions for f(z, 1) and (z, t) simplify to the following:
f(z, t) = - sigmoid (7 t)) d
dt 7o(2,t)
d
g (z, 1) sigmoid(7 1)) dp (o(2,))
The corresponding reverse-time SDE is given as:
(80) where W a standard Wiener process when time flows backwards from 1 - 0, and dt is an infinitesimal negative timestep. Given, se s8(x,, z), approximation the true score function, Xx, log 90 X No. z). Song et al. [53] show that the marginals of Eq. 80 can be described by the following Ordinary Differential Equation (ODE):
dx; = [t() - Vx, log do (x, z)]dt + g(t)dW; x1 ~ po(x1), ~ po(z),
1 f(z.t) 2 28 (z, (x,, dt.
dx;-
(81)
Zheng et al. [65] show that the score function, So (x,, 2), for the noise and the v-parameterization is given follows:
(x,, t)
Noise parameterization; see Sec. E.1.1 (82a)
(z)
(Xx, 2)
X exp (oleett vo(xe,, z.t)
v-parameteriiatioon see Sec. E.1.2
(82b)
Applying the change of variables formula [2] on Eq. 81, log (o (z) can be computed in the following manner:
log Pe (xolz)= = log px())) f tr (r(V (Vx, no(,,, 2, 1)) dt,
(83)
where ho(x,, t, = f(z,t) z f)
The expression for log-likelihood (Eq. 8) is as follows,
log (xo))(X)) 2 (2)() (log (xo [x)] DKI (x)
Using Eq. 83,
= = (2((((
logpo(x1) - tr (r(Vx, ho 2,) de DKu (ys (z/xo)llpo
t=
(84)
Computing tr (Vx, ho(x,, 2)) is expensive and we follow Chen et al. [2], Zheng et al. [65], Grath- wohl et [15] to estimate with Skilling-Hutchinng Hutchinson trace estimator [50, 19]. In particular, we have
V (x,, 2)) Ep(c) [c T ho
(85)
where the random variable satisfies Ep(e]le] and Cov [E] = I. Common choices for p(e) include Rademacher or Gaussian distributions. Notably, the term Vx, (x can be computed efficiently using "Jacobian-veede veettn-veeo---eee vector-product" computation in JAX. In our experiments, we follow the exact evaluation procedure for computing likelihood as outlined Song al. [53], Grathwohl et al. [15]. Specifically, for the computation of Eq. 85, we employ a Rademacher distribution for p(e). To calculate the integral in Eq. 84, we utilize the RK45 ODE solver [11] provided by scipy. int.eraatee...... solve _ivvp with atol=1e-8 and rtol=1e-5.
43

---

1.2.2 Dequantization.
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 degenerate outcomes [56] and fail to provide meaningful density estimations. Dequantization [46, 16, 65] is a common solution in such cases. To elaborate, let To represent 8-bit discrete data scaled to [-1, 1J. Dequantization methods assume that we have trained continuous model distribution po for z0, and define the discrete model distribution by
Px(xx) = -
Px + u)du.
258
To train (x() by maximum likelihood estimation, variational dequantization [16, 65] introduces a dequantization distribution g(u(xo) and jointly train Pmodel and q(((x)) (xo) by variational lower bound:
log (x 2 Bx(((x)) (xo + u) log g(u[x))].
(86)
Truncated Normal Dequantization. Zheng et [65] show that truncated Normal distribution,
g(u(xx) TN 0,
256 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]):
with mean O, covariance I, and bounds
256 256
Xu + 232)
ae )
log Po(xo) 2BE-TN(O1I logpo Po
+ (1 + log(2mo?)) 0.01522 x (87)
with = exp( 13.3).
a. = sqrt(qigmoois - 13.3)). and r = 3.
log X0 + e+ is evaluated using Eq. 84.
Importance Weighted Estimator. Eq. 87 can also extended obtain an importance weighted likelihood estimator to get a tighter bound on the likelihood. The variational bound is given by (for details please refer section A. in Zheng et [65]):
K
Pa
Xg + de 2(t)
))
log (xo) ZEg0 "..,,,.., (.., log
+ d dog
(88)
K
(((()))
1=1
with expl- x 133.3), log =
( 2
(3.3 + softplus(-13:)
((C) (2 2) exp (-211eni) ||21| 0.9974613, and r 3.
Note that for K = 1, Eq. 88 equivalent to Eq. 87; see Zheng et al. [65]. log pe Xo + d= is evaluated using Eq. 84. In Table 8, report BPD values for MULAN CIFAR10 (8M training steps, v-parameterizationn) and ImageNet (2M training steps, noise parameterization) using both the VLB-based approach, and the ODE-based approach with K and K 20 importance samples.
44

---

Table NLL (mean and 95% Confidence Interval for MULAN) on CIFARIO (8M training steps. "-parameterization) and ImageNet (2M training steps, noise parameterization) using both the VLB- based approach, and the ODE-based approach. estimator since Eq. 88 is equivalent to Eq. 87 for this case; see Zheng et al. [65].
means that we do not use importance weighted
Likelihood Estimation type
CIFAR-10 (4)
Imagenet (4)
VLB-based
2.59 =10
3.71 =10
ODE-based (K = 1; Eq. 87)
2.59 3 x 10
3.71 +10
ODE-based (K = 20; Eq. 88)
2.55 = 10
3.67 =10
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: See our introduction for list of claims including getting SOTA results on density estimation.
2. Limitations
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 lower frequency bias. See the paper for more details.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create 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 few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
45

---

Answer:
[Yes]
Justification:
Please see our detailed proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: Not only do show all equations and train standard datasets, we will open source the code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
46

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
We will open source after paper acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: /nips. httpss.ccppbblcccccecc httpppullliuudeeeepppeppppppeeel CoodeSumiis ssionPolccy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we include all hyperparameters in the paper and will open source code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: we report the deviatations for BPD in Table
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
47

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation 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 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 emrr 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.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : http:p:////rppppppppppcc/////e/e
Answer: [Yes]
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. 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 no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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.
48

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. 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, 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 such risks.
Released models that have 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 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 require this, but we encourage authors to take this into account and make 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 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 dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.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, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
49

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsource or research with human subjets
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
50


Reviews:
[{"review_id": "Wm3ET19ubM", "reviewer": "Reviewer_1B1k", "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.", "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.\n- 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?\n- 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?\n- 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?\n- 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": "The authors have adequately addressed the limitations.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow.\n- The authors present a theoretical motivation for making the noise schedule multivariate and learnable from a variational inference perspective, showing that it can impact the ELBO of the model unlike a univariate schedule.\n- The proposed method demonstrates state-of-the-art likelihood estimation results.\n- The authors conduct an extensive ablation study to thoroughly investigate the proposed method.\n\nOverall, I believe MuLAN is a well-motivated and logical extension of diffusion models that could be useful in practical applications.", "weaknesses": "- 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.\n- 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.\n- 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.\n- 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.\n\n[1] Vahdat, A., Kreis, K., & Kautz, J. (2021). Score-based generative modeling in latent space.\n\n[2] Lee, S., Kim, B., & Ye, J. C. (2023, July). Minimizing trajectory curvature of ode-based generative models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A7hUwdagya", "reviewer": "Reviewer_QQVN", "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%.", "questions": "Please refer to the Weakness section.", "limitations": "I am not aware of any potential negative societal impact of this work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The ELBO perspective of diffusion models has drawn much attention since the diffusion models' introduction. This paper made a non-trivial extension to the current framework and showed that the well-known assumption is no longer held that noise schedule does not alter EBLO. This observation gives additional flexibility to increase ELBO, resulting in a better performance in density estimation.", "weaknesses": "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. \n\n2. The FIDs of the proposed model are significantly worse than the existing diffusion models. \n\n3. The FIDs reported for VDM seem inconsistent with the one reported in the original paper. Can authors clarify how the FIDs are obtained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UmLhUvhU8v", "reviewer": "Reviewer_22Fq", "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.", "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": "Yes. \n\nNot 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I believe this is overall a good paper. The research question is relevant, the method is theoretically sound, well motivated, and explained clearly and in detail. The experiments are sufficiently extensive and the results are good, although only on 2 benchmark datasets that are relatively simple by today's standards.", "weaknesses": "- 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.\n- 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).\n- 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.\n- 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.\n- 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.\n- Typo: at line 54, $\\sigma_t$ should be $\\sigma_t^2$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to 22Fq\n\nWe thank the reviewer for their detailed and thorough feedback. We address their concerns below. \n\n## Concern 1: Evaluating on larger scale experiments (ImageNet-64)\n\nUnfortunately, it is not feasible for us to train an ImageNet-64 model within our academic research group. Using the official VDM Jax codebase, it takes 8 days on 8xA100 to train a VDM model on ImageNet-32. We would need to train multiple models on the larger 64x64 dataset, and this is outside the scope of what out group has access to. However, the 32x32 CIFAR10 and ImageNet datasets on which we report our results are standard benchmarks in the field and we report state-of-the-art performance on these benchmarks.\n\nWe anticipate our method to scale to larger datasets with similar computational complexity as existing diffusion models. Mulan requires an encoder that is much smaller than the denoising model, and requires 5-10% of additional memory. \n\n## Concern 2: Comparisons with other learned diffusion algorithms\n\nOur method can be seen as a special case of a diffusion normalizing flow (DNF), which uses the following forward process: $\\text{d}\\mathbf{x} = \\mathbf{f}_\\theta(\\mathbf{x}, t)\\text{d}t +  g(t)\\text{d}\\mathbf{w}$ . However, for such general processes, training requires backpropagating through sampling from the model, which is computationally expensive. In practice, this reduces scalability and produces worse results. However, there may be a class of processes between full normalizing flows and our simpler noise processes that will admit scalable training and improved performance.\n\nOur work shares similarity with the NDM and DiffEnc methods, which add an additive learned term to the noise process instead of an multiplicative one like in our paper. They still admit efficient sampling. These methods can potentially be combined with our work: in theory they are orthogonal, but we anticipate that getting them to work will require non-trivial engineering work around architecture tuning and optimization strategies.\n\nWe do provide a brief discussion in lines 583-589. We’ll add discussions to the updated manuscript. \n\n## Concern 3: Other Concerns\n\nWe thank the reviewer for pointing out the typo. The reviewers comments on the presentation style such as consolidating sec 3.2.1 and 3.2.2 with self-contained definition for the forward and backward processes is also very useful. We’ll factor these comments into the next version of the manuscript. \n\n## Question 1: Clarification regarding confidence Intervals\n\nTo compute the confidence interval, we calculate the variance, $\\sigma$, of the likelihood across all the data points in the validation/test set and report the `95% confidence interval` using the formula: $\\pm 1.96 \\frac{\\sigma}{\\sqrt{n}}$, where $n$ denotes the number of data points. Given the very small confidence interval of `0.001` on the likelihood, we did not further randomize the computation of likelihood using additional seeds.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gaP1V8xWe7", "reviewer": "Reviewer_uSJa", "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%.\n\nI 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?\n\nThe 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.\n\nOverall, 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.", "questions": "see summary", "limitations": "see summary", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "see summary", "weaknesses": "see summary", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to uSJa\n\nWe thank the reviewer for their detailed and thorough feedback. We address their concerns below. In the appendix, we have provided numerous experiments trying to explore how the noise schedule relates to different aspects of an image namely:\n\n1. Frequency distribution. \n2. Intensity.\n\n## Concern 1: Finding Interpretable Spatio-Temporal Variation in the Noise Schedule\n\nWe have provided a thorough analysis trying to relate the learnt 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:\n\n1. CIFAR-10 Frequency Split\n    \n    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.\n    \n    **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.\n    \n2. CIFAR-10 Masking\n    \n    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. \n    \n    **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.\n    \n3. CIFAR-10 Intensity Split\n    \n    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.\n    \n    **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.\n    \n\n**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$. \n\nFurthermore, the periphery of the image did not exhibit more inhomogeneity than the other parts of the image. \n\n## Concern 2: Optimal schedule is spatially inhomogeneous is empirical\n\nWe use principles of physics to motivate spatial inhomogeneity of an optimal noise schedule. In Section 3.5, we show that the diffusion loss is a line integral and 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$,\nalong the diffusion trajectory $\\mathbf{r}(t) \\in \\mathbb{R}^n_{+}$  in the presence of a force field $\\mathbf{f}(\\mathbf{r}(t)) \\in \\mathbb{R}^n_+$. 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 [1].  The path $\\mathbf{r}(t)$ is parameterized by the noise schedule, and thus by learning the noising process, we are effectively learning a path that incurs the least work done. Since this force field can be “arbitrary” as it is parameterized by the denoising neural network, it is highly unlikely that the path of least work done is a straight line joining $\\mathbf{r(0)}$ and $\\mathbf{r(1)}$. In other words, it is highly “unlikely” that the optimal noise schedule is scalar. Note that a straight line trajectory represents a scalar noise schedule as described in the appendix sec 11.5.\n\n----\nReference:\n\n[1] Richard E. Spinney and Ian J. Ford. Fluctuation relations: a pedagogical overview, 2012.\n\n\n## Concern 3: Inhomogeneity in the drift / diffusion term\n\n\nIn this work we consider variance preserving diffusion process which means that we can’t possibly disentangle the inhomogeneity in the drift and the diffusion term. The forward process takes the form \n$\\mathbf{x}_t = \\mathbf{\\alpha}\\_{t|s}(\\mathbf{z}) \\mathbf{x}_s + \\sqrt{1 -  \\mathbf{\\alpha}\\_{t|s}^2(\\mathbf{z})} \\epsilon$\nwith $s < t$. Removing the adaptivity in the drift or the diffusion term would mean that the marginals, $\\mathbf{x}_t \\sim q(\\mathbf{x}_t | \\mathbf{x}_0)$ are intractable to compute and would require simulating the diffusion process to obtain intermediate latents $\\mathbf{x}_t$ such as Diffusion Normalizing Flows [1].\n\n----\nReference:\n\n[1] Zhang, Q. and Chen, Y., 2021. Diffusion normalizing flow. *Advances in neural information processing systems*, *34*, pp.16280-16291.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (spotlight)