DiffEnc: Variational Diffusion with a Learned Encoder
Authors: Beatrix Miranda Ginn Nielsen
Anders Christensen
Andrea Dittadi
Ole Winther
Venue: ICLR 2024 poster (2024)
Keywords: DDPM
diffusion
image generation
encoder
Primary area: generative models
TL;DR: Adding a learned time dependent encoder to a diffusion model can improve the likelihood bound on image generation

Abstract: Diffusion models may be viewed as hierarchical variational autoencoders (VAEs) with two improvements: parameter sharing for the conditionals in the generative process and efficient computation of the loss as independent terms over the hierarchy. We consider two changes to the diffusion model that retain these advantages while adding flexibility to the model. Firstly, we introduce a data and depth-dependent mean function in the diffusion process, which leads to a modified diffusion loss. Our proposed framework, DiffEnc, achieves a statistically significant improvement in likelihood on CIFAR-10. Secondly, we let the ratio of the noise variance of the reverse encoder process and the generative process be a free weight parameter rather than being fixed to one. This leads to theoretical insights: For a finite depth hierarchy, the evidence lower bound (ELBO) can be used as an objective for a weighted diffusion loss approach and for optimizing the noise schedule specifically for inference. For the infinite-depth hierarchy, on the other hand, the weight parameter has to be one to have a well-defined ELBO.

Published as a conference paper at ICLR 2024
DIFFENC: VARIATIONAL DIFFUSION WITH A LEARNED ENCODER
Beatrix M. G. Nielsen," Anders Christensen, 1.2 Andrea Dittadi, Ole Winther
TTehnical University of Denmark, "Helmholtz AI, Munich, 'University of Copenhagen,
A MMax Planck Institute for Intelligent Systems, CCopenhagen University Hospital
ABSTRACT
Diffusion models may be viewed as hierarchical variational autoencoders (VAEs) with two improvements: parameter sharing for the conditional distributions in the generative process and efficient computation over the hierarchy. We consider two changes to the diffusion model that retain these advantages while adding flexibility to the model. Firstly, we introduce a data- and depth-dependent mean function in the diffusion process, which leads modified diffusion loss. Our proposed framework, DiffEnc, achieves a statistically significant improvement in likelihood CIFAR-10. Secondly, we let the ratio of the noise variance of the reverse encoder process and the generative process be free weight parameter rather than being fixed to This leads theoretical insights: For finite depth hierarchy, the evidence lower bound (ELBO) can be used as an objective for weighted diffusion loss approach and for optimizing the noise schedule specifically for inference. For the infinite-depe depth hierarchy, on the other hand, the weight parameter has be to have a well-defined ELBO..
the loss as independent terms
INTRODUCTION
Diffusion models (Sohl-Dickstein et al., 2015 Song & Ermon 2019 Ho al. 2020) ver- satile generative models that have risen to promi- nence in recent years thanks to their state-of- the-art performance in the generation of images
Standard, delieeieed
Diffasion Process
(Dhariwal & Nichol| 2021 Karras et al. (022).
video (Ho et al. 2022; Höppe et al. [0022] Har- vey et al. 2022), speech (Kong et al. 2020 Jeong 2021 Chen et al. 22200) and music (Huang
P
DiffEnc, learned
et al. al. 2023 Schneider al. (203).. particular, in image generation, diffusion models are state of art both terms of visual quality (Kar-
Trender
broka
Exade
ras al. 2022; Kim et al. 2022a Zheng al.
2022
Hoogeboom et al.] [203 Kingma & Gua
2023
Lou & Ermon 2023) and density estima-
Figure 1: Overview of DiffEnc compared to standard diffusion models. The effect of the encoding has been amplified 5x for the sake of illustration.
tion
Kingma et al. 2021 Nichol & Dhariwal
2021
Song et al. 2021).
Diffusion models can be seen as a time-indexed hierarchy over latent variables generated sequen- tially, conditioning only on the latent vector from the previous step. As such, diffusion models can be understood as hierarchical variational autoencoders (VAEs) (Kingma & Welling) 2013) Rezende al. 2014, [Sonderby al.) 2016) with three restrictions: (1) the forward diffusion process the inference model in variational inference- is fixed and remarkably simple; (2) the generative model Markovian each (time-indexed) layer latent variables generated conditioning only on the previous layer; (3) parameter sharing-ar all steps of the generative model share the same parameters.
"oorespondence to: <bmgi@ddtu.... dtu.dk>.
I Equal advising.

---

0.1. Changes Published as a conference paper at ICLR 2024
The simplicity of the forward process (1) and the Markov property of the generative model (2) al- low the evidence lower bound (ELBO) to be ex- pressed as an expectation over the layers of ran- dom variables, i.e., an expectation over time from the stochastic process perspective. Thanks to the heavy parameter sharing in the generative model (3), this expectation can be estimated effectively with a single Monte Carlo sample. These prop- erties make diffusion models highly scalable and flexible, despite the constraints discussed above.
original
0.4
0.6
0.8
1.0
In this work, we relax assumption (1) to improve the flexibility of diffusion models while retaining their scalability. Specifically, we shift away from assuming a constant diffusion process, while still maintaining sufficient simplicity to express the ELBO as an expectation over time. We intro- duce fime-dependent encoder that parameter- izes the mean of the diffusion process: instead of the original image x, the learned denoising model is tasked with predicting X5, which is the encoded image time f. Crucially, this encoder is exclu- sively employed during the training phase and not utilized during the sampling process. As result, the proposed class of diffusion models, DiffEnc, more flexible than standard diffusion models without affecting sampling time. To arrive at the negative log likelihood loss for DiffEnc, Eq. [88, we will first show how we can introduce a time- dependent encoder to the diffusion process and how this introduces an extra term in the loss if we the usual expression for the mean in the generative model, Section 3 We then show how can counter this extra term, using a certain parametrization of the encoder, Section
Figure 2: Changes induced by the encoder on the en- coded image at different timesteps: (xt xx)/(t- for t 0.4, 0.6, 0.8, 1.0 and the channels with red and have been summed blue denoting positive and negative changes, respec- tively. For + 1, global properties such as ap- proximate position of objects are encoded, where for smaller t changes more fine-grained and tend to enhance high-contrast within objects and/or between object and background.
We conduct experiments on MNIST, CIFAR-10 and ImageNet32 with two different parameteriza- tions of the encoder and find that, with a trainable encoder, DiffEnc improves total likelihood on CIFAR-10 and improves the latent loss on all datasets without damaging the diffusion loss. We observe that the changes to x, are significantly different for early and late timesteps, demonstrating the non-trivial, time-dependent behavior of the encoder (see Fig. 2
In addition, we investigate the relaxation of a common assumption in diffusion models: That the variance of the generative process, o.p is equal to the variance of the reverse formulation of the forward diffusion process, o². This introduces an additional term in the diffusion loss, which can be interpreted as a weighted loss (with time-dependent weights wt). We then analytically derive the optimal op. While this is relevant when training in discrete time (i.e., with finite number of layers) or when sampling, we prove that the ELBO is maximized in the continuous- limit when the variances are equal (in fact, the ELBO diverges if the variances are not equal).
Our main contributions can be summarized as follows:
We define new, more powerful class of diffusion modellsnamee DiffEnc-by introduc- ing time-dependent encoder the diffusion process. This encoder improves the flexibility of diffusion models but does not affect sampling time, as it is only needed during training.
We analyse the assumption of forward and backward variances being equal, and prove that (I) by relaxing this assumption, the diffusion loss can be interpreted as a weighted loss, and (2) in continuous time, the optimal ELBO is achieved when the variances are equal- in fact, if the variances are not equal in continuous time, the ELBO is not well-defined.
We perform extensive density estimation experiments and show that DiffEnc achieves a statistically significant improvement in likelihood on CIFAR-10.
The paper is organized as follows: In Section 2 we introduce the notation and framework from Varia- tional Diffusion Models (VDM; Kingma et al. 2021 in Section 3 we derive the general formulation
2

---

Published as a conference paper at ICLR 2024
of DiffEnc by introducing dreptdddependentt encoder; in Section we introduce the encoder param- cterizations used in our experiments and modify the generative model to account for the change in the diffusion loss due to the encoder, in Section 5 we present experimental results.
2 PRELIMINARIES ON VARIATIONAL DIFFUSION MODELS
We begin by introducing the VDM formulation (Kingma al. 2021) of diffusion models. We define a hierarchical generative model with T + layers latent variables:
(x z) = p(
II i=1
pe (Z)(1) ((()))
with x E X' a data point, the model parameters, s(1) = T
i-1 ((i) = i and p(1) (O, In
7
the following, we will drop the index and assume S < <t with marginal distribution:
We define diffusion process
(Zz N(axx
(2)
where / [0, 1] is the time index and at, a₁ are positive scalar functions of t. Requiring Eq. 2 to hold for any and I, the conditionals turns out to be:
a(z (z.) Naaq o1,s1,,
where
ar o of a a
Orla
02
Using Bayes' rule, we can reverse the direction of the diffusion process:
(2, 2 x) = (NQQ,
(3)
with
a
o
11
24
(4)
of
of
We can now express the diffusion process in a way that mirrors the generative model in Eq.
T
q(zlx) = q(z x) II
(5)
and we can define one step of the generative process in the same functional form as Eq. (3)
(z(z,|z,) = N(up,ofI)
with
0 X (2.,t),
of
of
where Xe learned model with parameters O. In a diffusion model, the denoising variance of: is usually chosen to be equal to the reverse diffusion process variance: oî o.. While initially we do make this assumption, we will prove this be optimal the continuous-tus -immee limit. Following VDM, we parameterize the noise schedule through the signal-to-noise ratio (SNR):
SNR(() = of of
and its logarithm: A, log SNR(t). We will use the variance-preserving formulation all our experiments: oj = 1 of sigmoid(Ar).
The evidence lower bound (ELBO) of the model defined above is:
(xo(x (x|z))x((22 FLIBO
log (x) 2 Ey(z/x)
(z((x)
3

---

Published as a conference paper at ICLR 2024
The loss C = -ELBO is the sum of reconstruction (Co), diffusion (Cr), and latent (C1) loss:
L Lu + Cr + L1
Co = - (x(())))) [(x/o())))
C = (x((((22(x))ll((z))
where the expressions for Lo and L1 are derived Appendix Thanks to the matching factor- ization of the generative and reverse noise processes-see Eqs. and -nd the availability of
(z,,) x) in closed form because q is Markov and Gaussian, the diffusion loss Cr can be written as sum or expectation over the layers of random variables:
T
CT(x)=
- Σ i=1
Eelenolie b))
[DKT. (1(Z)(0) (2())s x) x)llp((2e00) [())
T (1..)
(g(s(t)))) Ze())= x) Ipe (z.()) [x()))]
(8)
where U11, }} is the uniform distribution over the indices through T. Since all distributions are Gaussian, the KL divergence has a closed-form expression (see Appendix E
x) lpe (2, 2,)) = (w) log we) +
W/
1p - mel2
(9)
20
where the green part is the difference from using of. # og instead of of. = =o. and we have defined the weighting function
2
W/
and the dependency of oz and o²: on is left implicit, since the step size t - }.2 is fixed. The optimal generative variance can be computed in closed-form (see Appendix F):
1 of o8 =0== d B(ix. [llmp
3 DIFFENC
The main component of DiffEne the time-dependent encoder, which we will define as x x& (Ae), where (x()x) some function with parameters & dependent on x and through A,
=
log SNR(t). The generalized version of Eq. (2) is then:
q x) N(axx, of
(10)
Fig. [ visualizes this change to the diffusion process, and diagram is provided in Appendix A Re-
quiring that the process is consistent upon marginalization, i.e., q( 2 x) = q(Zz/za, (2, x dld..
leads to the following conditional distributions (see Appendix B
g(ze x) (auja2 og(x; xx), o
(11)
where an additional mean shift term introduced by the depth-dependent encoder. As Section can derive the reverse process (see Appendix C):
q(z De, x) = =N(HQ, oQ!
(12)
a
orlo of
PQQ
24
of
x, + ax x,)
(13)
with of given by Eq. We show how we parameterize the encoder in Section
Infinite-dppth limit. Kingma al. (022)) derived the continuous-time limit of the diffusion loss, that is, the loss in the limit of T + oc. We extend that result to case. Using MQ from
Eq. 11) and μp from Eq. 66) the KL divergence in the unweighted case, i.c., 20
||p #4 12 can
be rewritten in the following way, as shown in Appendix G
Ax ASNR
up - pQ 1
ASNR
(z t) - xxx(A) - SNR(s)
20
4

---

Published as a conference paper at ICLR 2024
where Ax xx(A)) - (x,) and similarly for the SNR. In Appendix - 0o, the expression for the optimal op tends to oQ and the additional term in the diffusion loss arising from allowing o3 + og tends 0. This result is accordance with prior work on variational approaches to stochastic processes (Archambeau et al., (007). We have shown that, in the continuous limit, the ELBO has to be an unweighted loss (in the sense that wt = 1). In the remainder of the paper, we will use the continuous formulation and thus set wy = 1. It is of interest to consider optimized weighted losses for a finite number of layers, however, we leave this for future research.
also show that, as
The infinite-eppth limit of the diffusion loss, (x) = -iim--+000 (x(x), becomes (Appendix G
dSNR(()) (ko(Z),)) xx(),)
(xe(
(x)(x) = (
(14)
dde
Loo (x) thus is very similar the standard continuous-time diffusion loss from VDM, though with an additional gradient stemming from the mean shift term. In Section [
we will develop modified
generative model to counter this extra term. In Appendix H we derive the stochastic differential equation (SDE) describing the generative model of DiffEnc in the infinite-depth limit.
4 PARAMETERIZATION OF THE ENCODER AND GENERATIVE MODEL
We now turn to the parameterization of the encoder X (Ar). The reconstruction and latent losses impose constraints on how the encoder should behave at the two ends of the hicrarchy latent vari- ables: The likelihood we use is constructed such that the reconstruction loss, derived in Appendix minimized when X
x (Aa) x. Likewise, the latent loss minimized by x (A1) = O. In between, for 0 < t < 1, non-trivial encoder can improve the diffusion loss.
We propose two related parameterizations of the encoder: trainable one, which we will denote by X and simpler, non-trainable one, Xut, where n/ stands for non-trainable. Let (x(x,, Ax) be neural network with parameters Φb, denoted y ob (Ae) for brevity. We define the trainable encoder as
(Xr) = (1 - of)x ozy (Ar) af (A))(1))
(15)
and the non-trainable encoder as
Xnt (Xe) = o,,x.
(16)
More motivation for these parameterizations can be found in Appendix The trainable encoder x is initialized with y (At) = 0, so at the start of training it acts as the non-trainable encoder Xn, (but differently from the VDM, which corresponds to the identity encoder).
To better fit the infinite-depth diffusion loss in Eq. (4). we define a new mean, Mr, of the gener- ative model po po (2) which is a modification of Eq. (6). Concretely, we would like to introduce
dxa(Ae) This a counterterm in μa that, when taking the continuous limit, approximately counters dA.
term should be expressed in terms of R (>t) rather than X For the non-trainable encoder, we have
dx (Nx) =n, 0.00x =xxxx=xxxmxAA)
dd
Therefore, for the non-trainable encoder, can (Ac) approximation of (A)) The trainable encoder more complicated because it also contains the derivative of yo that
dλ,
cannot as straightforwardly express in terms of &g. We therefore choose to approximate dxa dA, the
same way as (x.) We leave it for future work to explore different strategies for approximating dA, this gradient. Since we use the same approximation for both encoders, in the following we will write (>²) for both.
With the chosen counterterm, which in the continuous limit should approximately cancel out the effect of the mean shift term in Eq. (13), the new mean, up. defined as:
ampa0
a.oi= R (Ac) + tash (0078(((1))
μp =
of
2 2, +
of
(17)
5

---

Published as a conference paper at ICLR 2024
Table 1: Comparison of average bits per dimension (BPD) over seeds CIFAR-10 and ImageNer32 with other work. Types of models are Continuous Flow (Flow), Variational Auto Encoders (VAE), AutoRegressive models (AR) and Diffusion models (Diff). We only compare with results achieved without data augmentation. DiffEnc with trainable encoder improves performance of the VDM on CIFAR-10. Results on ImageNet marked with are the (Van Den Oord al. 2016) version of ImageNet which is no longer officially available. Results without are on the (Chrabaszez al. 2017) version of ImageNet, which is from the official ImageNet website. Results from (Zheng et al.] 2023 are without importance sampling, since importance sampling could also be added to our approach.
Model
Type CIFAR-10
ImageNet x2
Flow Matching OT (Lipman al. 2022)
Flow
2.99
3.53
Stochastic Int. [Albergo & Vanden Vanden-Eijnden
2022
Flow
2.99
3.48*
NVAE Vahdat Kautz 2020)
VAE
2.91
3.92*
Image Transtormer [Parmar et al.
2018
AR
2.90
3.77*
VDVAE (Child, 2020)
VAE
2.87
3.80*
ScoreFlow (Song et al. 2021)
Diff
2.83
3.76
Sparse Transtormer (Child al. 2019)
AR
2.80
Reflected Diffusion Models (Lou & Ermon 2023
Diff
2.68
3.74"
VDM (Kingma et al.
2021 (TOM steps)
Diff
2.65
3.72*
ARDM (Hoogeboom al.
2021
AR
2.64
Flow Matching TN (Zheng et al. 2023)
Flow
2.60
3.45
Our experiments (8M and 1.5M steps, seed avg)
VDM with v-parameterization
Diff
2.64
3.46
DiffEnc Trainable (ours)
Diff
2.62
3.46
Similarly to above, we derive the infinite-depth diffusion loss when the encoder is parameterized by Eq. (15) by taking the limit of Lr for T (see Appendix
1 -- (0
de
/ de) + o, xx(A)) - dx 1
Lox (x) =
(18)
dt
dde
where 2, = a, + o,e with N(0,))
In our experiments use the v-prediction parameterization (Salimans & Ho 2022) for our loss, which means that for the trainable encoder we use the loss
v-parameterization.
(x) -- Ma? v. vo t (x0()))-x((h) xx()e) yo( dy (di()))
1
dd
(19)
and for the non-trainable encoder, we use
1 Loo (x) = - EE,ooUU [oyy [[v, - vo + +a (Xo(Ad) xx
(20)
Derivations of Eqs. (19) and 20 are in Appendix K We note that when using the v-parametrization, as tends to 0, the loss becomes the same as for the e-prediction parameterization. On the other hand, when t tends to 1, the loss has a different behavior depending on the encoder: For the trainable encoder, we have that Vo ~ dA)( suggesting that the encoder can in principle guide the diffusion model. See Appendix L for more detailed discussion.
du,
5 EXPERIMENTS
In this section, we present experimental setup and discuss the results.

---

Published as a conference paper at ICLR 2024
Table 2: Comparison of the different components of the loss for DiftEnc- 32-4 and VDMv-32 with fixed noise schedule on CIFAR-I0. All quantities are bits per dimension (BPD) with standard error over seeds, and models are trained for 8M steps.
Model
Total
Latent
Diffusion
Reconstruction
VDMv-32
2.641 + 0.003
0.0012 + 0.0
2.629 + 0.003
0.01 + (4 10
DiffEne-32-4
2.620 + 0.006
0.0007 + (3 x 10-6)
2.609 + 0.006
0.01 + (4 10-6)
DiffEnc-8-2
DiffEnc-32
VDMv-32
Figure 3: Comparison of unconditional samples of models. The small model struggles to make realistic images, while the large models significantly better, as expected. For some images, details differ between the two large models, others they disagree main element of the image. An example where models make two different cars in column An example where DiffEnc-32 32-4 makes a car and VDMv-32 makes a frog in column? 7.
Experimental Setup. We evaluated variants of DiffEnc against a standard VDM baseline on
MÄIST (LeCun et al. 1998), CIFAR-10 (Krizhevsky al.) 2009) and ImageNet32 (Chrabaszez
et al. 2017). The learned prediction function is implemented as a U-Net (Ronneberger et al.) 2015) consisting of convolutional ResNet blocks without any downsampling, following VDM (Kingmal ct al. 2021). The trainable encoder in DiffEnc is implemented with the same overall U-Net archi- tecture, but with downsampling to resolutions 16x16 and 8x8. We will denote the models trained in our experiments by VDMv-n, DiffEnc-n-m,, and DiffEnc-n-nnt, where: VDMv is VDM model with v parameterization, n and m are the number of ResNet blocks in the "downsampling" part of the v-prediction U-Net and of the encoder U-Net respectively, and nf indicates non-trainable en- coder for DiffEnne On MNIST and CIFAR-10, we trained VDMv-8, DiffEnn-8-2, and DiffEEnc&-nt models. On CIFAR-10 also trained DiffEnc-8-4, VDMv-32 and hiffEne-32-4- On ImageNet32, we trained VDMv-32 and DiffEnc-32-8.
We used a linear log SNR noise schedule: X Amaz (Amaz Xmin)) tt For the large models (VDMv-32, DiftEnn-32-4 and DiftEnc-32-8). we fixed the endpoints, Amar and Amins to the ones Kingma et al. 1 found were optimal. For the small models (VDMv-8, DiffEnc-8-2 and DiffEnc- 8-nt), we also experimented with learning the SNR endpoints. We trained all our models with either 3 seeds depending on the computational cost of the experiments. See more details model structure and training in Appendix Q and on datasets in Appendix R
Results. As we see in Table I the DifiEEnc32-44 model achieves lower BPD score than previous non-flow work and the VDMv-32 on CIFAR-10. Since we do not use the encoder when sampling. this result means that the encoder is useful for learning a better generative model-with higher likelihoods-wii sampling time is not adversely affected. We also see that VDMv-32 after 8M steps achieved better likelihood bound, 2.64 BPD, than the result reported by Kingma et al. (202T) for the 6- c-parameterization after 10M steps, 2.65 BPD. Thus, the --parameterization gives an im- proved likelihood compared to c- e-paramterizatton. Table that the difference in the total loss comes mainly from the improvement in diffusion loss for DiftEnn-32-4, which points to the encoder being helpful in the diffusion process. We provide Fig. 2 since it can be difficult to see what the encoder doing directly from the encodings. From the heatmaps, we see that the encoder has learnt to do something different from how it initialised and that it acts differently over t, making finer changes in earlier timesteps and more global changes in later timesteps. See Appendix [W] for more details. We note that the improvement in total loss is significant, since get p-value of 0.03 for
7

---

Published as a conference paper at ICLR 2024
Table 3: Comparison of the different components of the loss for DiffEnc- 8-2, DiffEnc-8-nt and VDMv-8 on CIFAR-10. All quantities are in bits per dimension (BPD), with standard eroor, 5 seeds, 2M steps. Noise schedules are either fixed or with trainable endpoints.
Model
Noise
Total
Latent
Diffusion
Reconstruction
VDMv-8
fixed
2.783 + 0.004
0.0012 + 0.0
2.772 + 0.004
0.010 + (2 x 10
trainable
2.776 + 0.0006
0.0033 + (2 10-5)
2.770 0.006
0.003 + (5 X 10
DiffEnc-822
fixed
2.783 = 0.004
0.0006 + (3 10-5)
2.772 0.004
0.010 + (3 x 10
trainable
2.783 = 0.003
0.0034 (2 10-5)
2.777
0.003 = (5 x 10
DiffEnc-8-nt
fixed
2.789 + 0.004
(1.6 x 100
2.779
0.010 + (1 x 10-5
trainable
2.786 + 0.004
0,0009 + (1 10
2.782
0.003 + [3 x 10
t-test on whether the mean loss over random seeds is lower for DiffEEne-2-44 than for VDMv-32. Some samples from DiffEnc-8-2, DiffEne-32-4, and VDMv-32 are shown in Fig. More sam- ples from DiffEnc-32-4 and VDMv-32 in Appendix [T] See Fig. 4 in the appendix for examples of encoded MNIST images. DiffEEn-32-4 and VDMv-32 have similar FID scores as shown in Table
For all models with a trainable encoder and fixed noise schedule, we see that the diffusion loss the 4 to
same or better than the VDM baseline (see Tables 2 and encoder being able to preserve the most important signal as + 1. This is supported by the results get from the non-trainable encoder, which only removes signal, where the diffusion loss is always worse than the baseline. We also see that, for a fixed noise schedule, the latent loss of the trainable encoder model is always better than the VDM. When using fixed noise schedule, the minimal and maximal SNR is set to ensure small reconstruction and latent losses. is therefore natural that the diffusion loss (the part dependent on how well the model can predict the noisy image), is the part that dominates the total loss. This means that a lower latent loss does not necessarily have a considerable impact on the total loss: For fixed noise schedule, the DiffEnc-8-2 models on MNIST and CIFAR-10 10 and the DiffEnc- 32-8 model on ImageNet32 all have smaller latent loss than their VDMv counterparts, but since the diffusion loss is the same, the total loss does not show significant change. However, Lin et al. 2023) pointed out that a high latent loss might lead to poor generated samples. Therefore, it might relevant to train a model which has a lower latent loss than another model, if it can achieve the same diffusion loss. From Tables and 4 we see that this is possible using a fixed noise schedule and a small trainable encoder. For results on a larger encoder with a small model see Appendix U
7 We interpret this as the trainable
We only saw an improvement in diffusion loss on the large models trained on CIFAR- 10, and not on the small models. Since ImageNet32 is more complex than CIFAR-10, and we did not an improvement in diffusion loss for the models on ImageNet32, a larger model might be needed on this dataset to see an improvement in diffusion loss. This would be interesting to test in future work.
For the trainable noise schedule, the mean total losses of the models are all lower than or equal their fixed-schedule counterparts. Thus, all models can make some use of this additional flexibility. For the fixed noise schedule, the reconstruction loss is the same for all three types of models, due to how encoder parameterized.
6 RELATED WORK
DDPM: Sohl- et al. (2015) defined score-based diffusion models inspired by nonequi- librium thermodynamics. Ho et al. (2020) showed that diffusion models 1) are equivalent to score- based models and 2) can be viewed as hierarchical variational autoencoders with diffusion encoder and parameter sharing in the generative hierarchy. Song et al. (2020b) defined diffusion models us- ing an SDE.
DDPM with encoder: To the best of our knowledge, only few previous papers consider modifications of the diffusion process encoder. Implicit non-linear diffusion models (Kim et al. 2022b) use an invertible non-linear dime-dependent map, h, to bring the data into latent space where they do linear diffusion. h can be compared to our encoder, however, we do not enforce the encoder to be invertible.
Blurring diffusion models (Hoogeboom & Salimans 2022
Rissanen al. 2022) combines the
8

---

Published as a conference paper at ICLR 2024
added noise with blurring of the image dependent on the timestep. This blurring can be seen as a Gaussian encoder with a mean which is linear in the data, but with a not necessarily iid noise. The encoder parameters are set by the heat dissipation basis (the discrete cosine transform) and time. Our encoder is a learned non-linear function of the data and time and therefore more general than blurring. [Daras et al. (2022) propose introducing more general linear corruption process, where both blurring and masking for example can be added before the noise. Latent diffusion (Rombach et al., 2022) uses a learned depth-independent encoder/decoder to map deterministically between the data and a learned latent space and perform the diffusion in the latent space. (Abstreiter jt al.) (2021) and Preechakul et al. (2022) use an additional encoder that computes small semantic representation of the image. This representation is then used as conditioning in the diffusion model and is therefore orthogonal to our work. Singhal et al. 2023) propose to learn the noising process:
for 2, = 022x + See, they propose to learn Or, and 15g.
Concurrent work: Bartosh et al.] (2023) also propose to add a time-dependent transformation to the data in the diffusion model. However, there is a difference in the target for the predictive function, since in our case X9 predicts the transformed data, X while in their case X9 predicts data x' such that the transformation of x', fx(x', t), is equal to the transformation of the real data fa t). This might, according to their paper, make the prediction model learn something within the data distribution even for / close to 1.
Learned generative process variance: Both Nichol & Dhariwal (202) and Dhariwal & Nichol
(202T) learn the generative process variance, dr- [hhriwall & Nichol (2021) ip observe that it allows for sampling with fewer steps without a large drop in sample quality and Nichol & Dhariwal (202I) argue that could have a positive effect on the likelihood. Neither of these works are in a continuous- time setting, which is the setting we derived our theoretical results for.
LIMITATIONS AND FUTURE WORK
As shown above, adding a trained time-dependent encoder can improve the likelihood of diffusion model, at the cost of a longer training time. Although approach does not increase sampling time, it must be noted that sampling is still significantly slower than, e.g., for generative adversarial
networks (Goodfellow et al. 20114). Techniques for more efficient sampling in diffusion models
(Watson et al.] [021 [Salimans& & Ho 2022; Song et al.] [020aa] Lu et al.. 2022 Berthelot et al.,
2023 Luhman & Luhman 202 Liu
2022) can be directly applied our method.
Introducing the trainable encoder opens up an interesting new direction for representation learning. should be possible distill the time-dependent transformations to get smaller time-dependent representations of the images. It would be interesting to see what such representations could tell us about the data. It would also be interesting to explore whether adding conditioning to the encoder will lead to different transformations for different classes of images.
As shown in TThis et al. 2015) likelihood and visual quality of samples are not directly linked. Thus it is important to choose the application of the model based on the metric it was trained to optimize. Since we show that our model can achieve good results when optimized to maximize likelihood, and likelihood is important in the context of semi-supervised learning, it would be interesting to use this kind of model for classification in a semi-supervised setting.
8 CONCLUSION
We presented DiffEnnc, a generalization of diffusion models with a time-dependent encoder in the diffusion process DiffEnc increases the flexibility of diffusion models while retaining the same computational requirements for sampling. Moreover, we theoretically derived the optimal variance of the generative process and proved that, in the continuous-time limit, it must be equal to the diffusion variance for the ELBO to be well-defined. We defer the investigation of its application to sampling discrete-time training to future work. Empirically, we showed that DiffEnc can improve likelihood on CI and that the data transformation learned by the encoder is non-trivially dependent on the timestep. Interesting avenues for future research include applying improvements to diffusion models that are orthogonal to our proposed method, such as latent diffusion models, model distillation, classifier-free guidance, and different sampling strategies.

---

Published as a conference paper at ICLR 2024
ETHICS STATEMENT
Since diffusion models have been shown to memorize training examples and since is possible to extract these examples (Carlini el al. 2023). diffusion models pose a privacy and copyright risk especially if trained on data scraped from the internet. To the best of our knowledge our work neither improves nor worsens these security risks. Therefore, work still remains on how to responsibly deploy diffusion models with or without a time-dependent encoder.
REPRODUCIBILITY STATEMENT
The presented results are obtained using the setup described in Section More details on models and training are discussed in Appendix Code can be found on GitHui The Readme includes description of setting up the environment with correct versioning. Scripts are supplied for recreating all results present in the paper. The main equations behind these results are Eqs. (19) and (00), which are the diffusion losses used when including our trainable and non-trainable encoder, respectively.
ACKNOWLEDGMENTS
This work was supported by the Danish Pioneer Centre for AI, DNRF grant number PI, and by the Ministry of Education, Youth and Sports of the Czech Republic through the e-INFRA CZ (ID:90254). OW's work was funded in part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF200CC0062666) AC thanks the ELLIS PhD program for support.
REFERENCES
Korbinian Abstreiter, Sarthak Mittal, Stefan Bauer, Bernhard Schölkopf, and Arash Mehrjou. Diffusion-based representation learning. arXiv preprint arXiv:2105. 14255 2021.
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic inter- polants. arXiv preprint arXiv:2209.1552 2022.
Cedric Archambeau, Dan Cornford, Manfred Opper, and John Shawe-Taylor Gaussian process approximations of stochastic differential equations. In Gaussian Processes in Practice, pp. 1-16. PMLR, 2007.
Grigory Bartosh, Dmitry Vetrov, and Christian A. Nacsseth. Neural diffusion models, 2023.
David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbot, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.04248, 2023.
Nicolas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito, and Eric Wallace. Extracting training data from diffusion models. In 32nd USENIX Security Symposium (USENIX Security 23), pp. 5253-5270, 2023.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron Weiss, Mohammad Norouzi, and William Chan. Wave- grad: Estimating gradients for waveform generation. arXiv preprint wwXiv:2009.007/1 2020.
Rewon Child. Very deep generalize autoregressive models and can outperform them images. arXiv preprini arXiv:2011.10633 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. A downsampled variant of imagenet as an alternative to the cifar datasets. arXiv preprint /70.00777 2017.
Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alexandros G Dimakis, and Peyman Milanfar. Soft diffusion: Score matching for general corruptions. arXiv preprint 024442200, 2022.
www. ni /Di ffFnc
10

---

Published as a conference paper at ICLR 2024
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs image synthesis. Ad- vances neural information processing systems, 34:8780-877 8794, 2021.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Fand Farley Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances in Neural Information Processing Systems, 35: 27953-27965, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Asso-
ciates, Inc., 2020. URL https://prroocee.ng.....e www.igs..euiis..... wttprr/is..c/pppp//////EEL.. 4o5ccfec558lat00667........
79ca4b Paper.pdl
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruigi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210002303 2022.
Emiel Hoogeboom and Tim Salimans. Blurring diffusion models. arXiv preprint arXiv:2209.05555, 2022.
Emiel Hoogeboom, Alexey Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:21 2021.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans, simple diffusion: End-to-end diffusion for high resolution images. arXiv preprint arXiv:2301.1199. 2023.
Tobias Höppe, Arash Mehrjou, Stefan Bauer, Didrik Nielsen, and Andrea Dittadi. Diffusion models for video prediction and infilling. Transactions on Machine Learning Research, 2022.
Qingqing Huang, Daniel S Park, Tao Wang, Timo Denk, Andy Ly. Nanxin Chen, Zhengdong Zhang, Zhishuai Zhang, Jiahui Yu, Christian Frank, et al. Moise2musie: Text-conditioned music generation with diffusion models. arXiv preprint arXiv:2302.033: 2023.
Myeonghun Jeong. Hyeongju Kim, Sung Jun Cheon, Byoung Jin Choi, and Nam Soo Kim. Diff-tts: A denoising diffusion model for text-to-speech. arXiv preprint arXiv:2104.01409 2021.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 35:26565-26577, 2022.
Dongjun Kim, Yeongmin Kim, Wanmo Kang, and II-Chul Moon. Refining generative process with discriminator guidance in score-based diffusion models, arXiv preprint arXiv:2211. 2022a.
Dongjun Kim, Byeonghu Na, Se Jung Kwon, Dongsoo Lee, Wanmo Kang, and II-chul Moon. Max- imum likelihood training of implicit nonlinear diffusion model. Advances in Neural Information Processing Systems, 35:32270-32284, 2022b.
Diederik Kingma and Ruiqi Gua. Vdm++: Variational diffusion models for high-quality synthesis. haap///bs/220
arXiv preprint arXiv:2303.00 008488 2023.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems, volume 34, pp. sociates, Inc., 2021. URL https: https:://ppccceedgggg.. neurips cc cctppaprr/22////11ee pdf
Curran As-
Diederik Kingma and Max Welling. 131iv:1312.61 2013.
Auto-encoding variational bayes.
arXiv preprint
11

---

Published as a conference paper at ICLR 2024
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09770 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2) 2324,, 1998.
Shanchuan Lin, Bingchen Liu, Jiashi Li, and Xiao Yang. Common diffusion noise schedules and sample steps are flawed. arXiv preprint arXiv:2305. 08891, 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03000 2022.
Aaron Lou and Stefano Ermon. Reflected diffusion models. arXiv preprint arXiv:2304.04740, 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint vrXiv:2101.02388 2021.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171, PMIR, 2021.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran, Image transformer. International conference on machine learning, pp. 4055- 4064. PMLR, 2018.
Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasom Suwajanakorn. Dif- fusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10619-10629, 2022.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286. PMLR, 2014.
Severi Rissanen, Markus Heinonen, and Amo Solin. Generative modelling with inverse heat dissi- pation. arXiv preprint UrXiv:2206.13397, 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings the IEEF/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical Image Computing and Computer-Assisted Intervention- MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceed- ings, Part III 18, pp. 234-241. Springer, 2015.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint 220.::22000000///// 2022.
Flavio Schneider, Zhijing Jin, and Berhard Schölkopf. Mo), usai: Text-to-music generation with long-context latent diffusion. arXiv preprint arX11:2301.1175,, 2023.
Raghav Singhal, Mark Goldstein, and Rajesh Ranganath. Where to diffuse, how to diffuse. and how to get back: Automated learning for multivariate diffusions. arXiv preprint arXiv:2302.07261, 2023.
12

---

Published as a conference paper at ICLR 2024
Samarth Sinha and Adji Bousso Dieng. Consistency regularization for variational auto-encoders. Advances in Neural Information Processing Systems, 34: 12944-12954, 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing, pp. 2256-2265, PMLR, 2015.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. Advances neural information processing systems, 29, 2016.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- fional Conference on Learning Representations, 2020a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
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 arXiv:201 1. 13456, 2020b.
Yang Song. Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. Advances in Neural Information Processing Systems, 34:1415-1428, 2021.
Lucas Theis, Aäron van den Oord, and Matthias Bethge. A note on the evaluation of generative models. arXiv preprint
2015.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems, 33:1 33:19667-1967, 2020.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:1 1287-1187300 1302, 2021.
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.
Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan. Learning to efficiently sam- ple from diffusion probabilistic models. arXiv preprint arXiv:2/06.03802 2021.
Guangcong Zheng, Shengming Li, Hui Wang, Taiping Yao, Yang Chen, Shouhong Ding, and Xi Li. Entropy-driven driven sampling and training scheme for conditional diffusion generation. In European Conference on Computer Vision, pp. 754-769. Springer, 2022.
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Improved techniques for maximum likelihood estimation for diffusion odes, arXiv preprint arXiv:2305. 03935, 2023.
13

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
Overview of diffusion model with and without encoder
15
B Proof that 2; given x has the correct form
15
Proof that the reverse process has the correct form
16
D The latent and reconstruction loss
17
D.I Latent Loss
17
Reconstruction Loss
E
Diffusion loss
18
E.1 Assuming non-equal variances in the diflusion and generative processes
18
Assuming equal variances the diffusion and generative processes
F
Optimal variance for the generative model
19
G Diflusion loss in continuous time without counterterm
19
DiffEnc as SDE
22
Motivation for choice of parameterization for the encoder
23
Continuous -imm limit of the diffusion loss with an encoder
23
J.1 Rewriting the loss using SNR
23
J.2 Taking the limit
24
Using the Parameterization the DiffEne Loss
25
K.
Rewriting the parameterization
25
K.2
parameterization in continuous diffusion loss
26
K.3 v parameterization of continuous diffusion loss with encoder
26
L
Considering loss early and late timesteps
27
M Detailed Loss Comparison for DiffEne and VDMv MNIST
28
Detailed Loss Comparison for S0HEne--22- and VDMy-32 on CIFAR-10
28
0 Detailed Loss Comparison for DiffEnn and VDMy Iea32Net32
29
P Further Future Work
29
Model Structure and Training
29
R Datasets
30
Encoder examples on MNIST
30
T Samples from models
30
0 Using a Larger Encoder for Small Diffusion Model
31
V FID Scores
32
W Sum heatmap of all timesteps
32
14

---

Published as a conference paper at ICLR 2024
A OVERVIEW OF DIFFUSION MODEL WITH AND WITHOUT ENCODER
The typical diffusion approach can be illustrated with the following diagram:
( (R,)
Z
Ze
Zo
X
p(a, (xs)
g(z,) x)
g(z; x)
where 0 < 8 < S 1. We introduce an encoder fo : x [0, 1] + y' with parameters Φ. that maps and time E [0, 1] latent space V In this work, ) has the same dimensions the original image space. For brevity, we denote the encoded data as X, = (x, t). The following diagram illustrates the process including the encoder:
g(2; m.)
2
2
20
g(z((x))
p(a, l.))
{(x,)x))
(zg (xc))
X,
Xa
(x.U)
(x,)))))
(x(),t)
B PROOF THAT Z. GIVEN HAS THE CORRECT FORM
Proof that we can write
g(z,lx) Marx If for any & when using the definition g(zo|xo) = g( (x) and Eq.
(21)
Proof. By induction:
The definition of q(zu) (o) g(z0) (x) gives us our base case.
To take step, assume (xx, q( (z. x) can be written
q( 2. x) Naaxx, 0
(22)
and take a t
Then sample from g(2, (x) can written as
z, = a,x, OgE, where e, is from standard normal distribution, (O, I) and sample from q(2 /z, Xt, Xt, q(z,) 2g, x) can be written as
(23)
2, = 0: + ox x,) + dtf
(24)
where Egjs is from a standard normal distribution, NN Using the definition of 2,, we get
Z = aa1s(a,x, +0,E,) + (x - x,) + E1l
= axx, + t.p + axx - a:x, + E2j
(25)
= oxXi + O. + o1
15

---

Published as a conference paper at ICLR 2024
Since C3 and Et|s describe two normal distributions, sample from the sum be written as
Vai.00 + ooi
(26)
where is from standard normal distribution, // MN, I). So can write sample 24
Z = aix; + o + t
= a + a 03 + of - a 0 0.E
(27)
= a + JfEt
Thus we get
g(z/|x) Na of
(28)
for any 0 sts 1. We have defined going from z, going through f.
C PROOF THAT THE REVERSE PROCESS HAS THE CORRECT FORM
To that
(22,
= NNQq I
(29)
with
09 0/
o =
(30)
of
2
and
atl. of
0,5
2
t/s
pQQ
21 +
of
X, + (x. - xx)
(31)
is the right form for the reverse process, we take sample from g(z, |z, x.) and sample z from q(zelz). These have the forms:
2 aix; + 0160
(32)
a.o
2, of +
ax(x, - x1) xx) + togco
(33)
of
We show that given 2 we get 2. from ((2, (x) as in Eq. 110
a
2. of
(a,x +,0+)
x1 ax - x1) + ozEa
(34)
of
o2 -o
of
a44 0²
0s
X, + ax(x - x1) + Arli0
2
X1
of
of
orE; + oōEQ
(35)
(36)
Since a, ag we have
aga
02
ag o
o2
aspo
24
o7
9
X1
of
+x -xi)
of
OiE; + 2
oeEq
(37)
0202
2
02
a 2 o²
a,0
X1
o
X1 +
0,
2
x1 + ax,x - -x,) +
oi
O2E; + 0ōE0
(38)
0g
0,X1 (X, x1)
o.0 + oQEQ
(39)
of
!
=,,x, +
b.0 dgco
(40)
of
(41)
16

---

Published as a conference paper at ICLR 2024
- and the sum rule of variances, 2 ox + +o + 2C)V(X,)), We now use that og of where the covariance zero since E, and EQ are independent.
2s xgx,
2 016, + oQE0
(42)
=,,,, +
E1 +
EQ
(43)
-,
oi Es
(44)
0i
(o)
a² o 0
aaxx
1/00
(45)
0
o
of00
a,X, +
(46)
of
-,,,, + 0gEx
(47)
Where 6, is from standard Gaussian distribution.
D THE LATENT AND RECONSTRUCTION LOSS
D.I LATENT LOSS
Since q( xx = (o1x1, I) and p(21) (V I), the latent loss, DKI( (q(21 (x<()x)))p(()) the KL divergence of two normal distributions. For normal distributions No, with means Moo pi and variances Eo, E1. the KL divergence between them is given by
trr (E) d + (m ma)² ΣT (m1 140 log
det ΣΙ
dct Σg ))
DKL. (NolLM))
(48)
where d the dimension. Therefore we have:
ir (oft) - -d 110 - 0xxx112 12xxi11+100 log
(49)
2
det
1
2 (||axx) 112 + (oî llg of -1))
(50)
1
Σ (a?x), + o log of - -1) 11
(51)
2
The last line is used in our implementation.
D.2 RECONSTRUCTION LOss
The reconstruction loss is given by
Co E (2000)) [- log log p(x[2o)]
(52)
We make the simplifying assumption that p(x (z)) factorizes over the elements of x. Let x; be the value of the ith dimension (i.e., pixel) of x and zo,; the corresponding pixel value of Zo:
pxx
II p(z:
(53)
In our case of images, we assume the pixel values are independent given and only dependent on the matching latent component. We construct p(x;/220 20.)) from the variational distribution noting that
q(x(2a) g x()())
(54)
(Zo)
17

---

Published as a conference paper at ICLR 2024
and for high enough SNR at t 0, q(zu|>) will be very peaked around Zo OgX. So we choose
p(z)(20,)) & (z,0, (zo, (20,,,0)0
(55)
where we normalize over all possible values of x4. That is, let v E {0, 255} be the possible pixel values of F;, then for each calculate the density N (agu, oz) at zo,; and then normalise over to get a categorical distribution ((z) that sums to 1.
E DIFFUSION LOSS
The diffusion loss
T
CJ Σ
E Botzroo x) DKL (KL((( Z()) x) ()loo((
(z()))
(56)
where s(t), are the values of S s <t corresponding to the ith timestep.
E. ASSUMING NON-EQUAL VARIANCES IN THE DIFFUSION AND GENERATIVE PROCESSES In this section we let
po (z, 2.) N(ur,
(57)
( 21: = (NHQ,
(58)
where we might have Op + σQ. We then get for the KL divergence, where d is the dimension,
DKI,( Kr, x)
|pa (2. (2.)) = DRI, pQ, of 1) pp, u
Ilpp uQ) log det 07.1 det 0211/
tr
Q
I
d (pP- m22²² oz
p
1
(op) dd
d+
Mp
mQll2 + log
(o
P
d
0
of
1
1 +
log
+
up m2ll2
(59)
20 p
we define
W.
(60)
Then we can write the KL divergence as
DKL (g(gggg(/.))))) x)
(2 x)
=
d
(w: 1 log w/) +
W0'
pp
me)
(61)
2
20
Using our definition of μx
o 2
a.o.
up
o?
zt
+
of
x (2e) - -x)(07300(2))
(62)
we can rewrite the second term of the loss as:
wi
||pp
(63)
1 u (xo(1) - X (Xi) a, ((A, (1)o(8((1) (xx(1.) x (Ae)))
Wt
20
of
Where we have dropped the dependence on À from our notation of &x (X,) and x (Ar), to make the equation fit the page.
18

---

Published as a conference paper at ICLR 2024
E.2 ASSUMING EQUAL VARIANCES IN THE DIFFUSION AND GENERATIVE PROCESSES
If we let
pa(2.|z)) = (Zz,/Ze, = ((H
(64)
(65)
get for the KL divergence
(g 2, i2, x) x)llpo (2,(2z)) De
1p 1Ql12
20
1
a.o (Xo(t) xx()))) + as (() - (1) (xx(1s) xx(),)))
2
(66)
202
F OPTIMAL VARIANCE FOR THE GENERATIVE MODEL
In this section, we compute the optimal variance of. of the generative model in closed-form.
Consider the expectation over the data distribution of the KL divergence in the diffusion loss (Ap- pendix D}:
d 0/6 Eg(x,.)) (ell,,,,)]lin (x)looo.)))l((
1 + log ?
+
Be(xxe)) [llup - pe [1]]
2
0p
207
(67)
and differentiate it w.r.t. o7::
dDRL.
d
0
1
1
207 Ry(xx..))
[1| -
(68)
da p 2
p
0p
-
1
do (g(x.22) [l|eep
me m2ll)2)
(69)
20
The derivative is zero when:
1
of = og + d (aixa [||pe 4all
(70)
Since the second derivative of the KL at this value of oz is positive, this is minimum of the KL divergence.
G DIFFUSION LOSS IN CONTINUOUS TIME WITHOUT COUNTERTERM
In this section. we consider the DiffEnc diffusion process with mean shift term, coupled with the original VDM generative process (see Section B. We show that in the continuous time limit the optimal variance of tends o and the resulting diffusion loss simplifies the standard VDM diffusion loss. We finally derive the diffusion loss in the continuous-time limit.
We start by rewriting the diffusion loss as expectation, using constant step size 7 = 1/T and denot- ing f4 = i/T:
Lr(x) T (E Eytz, (x) [Dk x) (2),
(71)
T
(2) [DrI. (g(e (2, x) (Zc-r 2:))]
(72)
where we dropped indices and directly sample the discrete I.
The KI. divergence can be calculated in closed form (Appendix Gaussian:
because all distributions are
W/ Hp - mel13
DKL(g(Ze [z, x)) ppo (z,-, (2:) = (we - log wi) + 20
(73)
19

---

Published as a conference paper at ICLR 2024
where we have defined the weighting function
Wr
(74)
Ptt
Insert this in the diffusion loss:
d
Wt
Cr(x) EU-U(F
2T (wi 1- log we) + (xx(xx)))
up
Hell
(75)
20
Given the optimal value for the noise variance in the generative model derived above (Appendix
oF = o + (x(22) [llep 2
get the optimal wy
B((x))) [ - [2l13]
ogd
Using the following definitions:
at1.0
0gTifs f/s
pQQ
2 +
x; + txs - xx)
of
of
ago t/s
pp
21
X (2tsl)
of
0
o
2
1/s
o
2
2
1/s
0:
and these intermediate results:
1
a
0
11
SNR(s) SNR(1)
Q 0
0
SNR(s) SNR())
1/
0f
0?0?
SNR(s)
we can write:
pp - mell3
0,0 1/8
+ ax, xxi) 0 r/s
X0 ZZe
202
20
0
0
+
(x. -xi) - (2(,
208 of
SNR(s)
= ASNR
X1
Ax 50 1)
2
ASNR
1
where we used the shorthand ASNR SNR(t) SNR(s) and Ax Xx X
The optimal tends to 1.
The optimal can be rewritten as follows:
w = 1
Bx(xp() [[] #Q [2]
(76)
ogd
ASNR
SNR(s)
(x)x))
X,
Ax
*o (zz,
(77)
ASNR
As T + Do, or equivalently s + and 7 the unweighted case (forward and backward variance equal).
t - O, the optimal w, tends to 1, corresponding to
20

---

Published as a conference paper at ICLR 2024
The first term of diffusion loss tends to zero.
Define:
ASNR
SNR(s)
Ax (z(Ze, tt)
ASNR
such that the optimal wy given by
Then are interested in the term
U 1 log wt
- + log(1 - v)
As + O, we have
ASNR dSNR(t) 0(7))
dt
Ax= (x + +0(72)
dt
V dSNR(1)
Xt (Xe
oo xo 1) 27
(x(().)))
0
dlogg SNR
Since + we can write series expansion around O:
W1 1 log Wi = "²² + +
27 2 *o (2 t)
= T dSNR(!) Bx(xxx))
(Xx( d log SNR
)
+0(7²)
x
d dt
2
= 0
The first term of the weighted diffusion loss Lr is then 0, since as 7 + 0 we get:
d
E.. [w log wi] (())
(78)
2T
Note that, had simply used wy O(T), we would only be able prove that this term in the loss is finite, but not whether it is zero. Here, we showed that the additional term in the loss actually tends zero as T sc.
In fact, we can also observe that, if op * αq and therefore -1- - log O, the first term in the diffusion loss diverges in the continuous-timi limit, so the ELBO is not well-defined.
Continuous- time limit of the diffusion loss. the diffusion loss Lr tends zero. The limit of Ln then becomes
We saw that, as 7 - 0, wt + 1 and the first term in
(x) = lim Cr(x)
(79)
T-x
1
lim E (T.2 Bg((
2ell
(80)
2TOy
1
dSNR(t)
SNR(t)
dx)
lim S 11((
+
de
0x (Zr, t)
T T00
2r
dt
dSNR(C)) dr
(81)
dSNR(t)) dl
(xx(
X1.
+
(82)
2
SNRR SNR
21

---

Published as a conference paper at ICLR 2024
H DIFFENC AS AN SDE
A diffusion model may be seen as a discretization of an SDE. The same is true for the depth depen- dent encoder model. The forward process Eq. (11) can be written
2: = <itl22. 2. + ar (1) (s)) + 02p6
(83)
Let < At such that s + At. If we consider the first term after the equality sign we see that
04 0g
0g + 04 a2 02
(84)
1 + 01 Os At
(85)
ag At
So wc get that
Z1
2,
0t
a2
2, At + (x) - -x)(s) + oz/SE
(86)
a,
Considering the second term, we get
(t( xx At
- (x(xo())xxo(s) = =a,-
(87)
Al
So if we define
0t ag
(t) (t) (s())
X
far(Zs.8)
2, + ai
(88)
A,
Al
we can write
2: -2 fAil(z: s) + Oz
(89)
We will now consider oRe to be able to rewrite Clsc
a
o1i 0,
2
(90)
a
2
o
(91)
0f
0
At
at
(92)
0t
At
Thus, if we define
o²
(sr(s) =
ai
2
af
7
a
At
(93)
can write
2, 2, = ffalza, 8) At + ggr(s) Ate
(94)
We can now take the limit At + using the definition = s + At:
1 da. + axs
dxx(s) dloog a'
n, + as dxx(s)
far(2., 8)
(95)
a, ds
ds
ds
ds
do? doo3/a?
(s)
a?
(96)
ds
So if we use these limits to define the functions
1 day n, at
dx (t)
f(z,, x)
(97)
a, dt
dt
g(() = an /o
(98)
dt
we can write the forward stochastic process when using a time dependent encoder as dz -f(2( t, x)dt + g(t)dw
(99)
22

---

Published as a conference paper at ICLR 2024
where dw is the increment of Wiener process over time At. The diffusion process of DiffEnc in the continuous-ting limit is therefore similar to the usual SDE for diffusion models (Sogg et al. 2020b), with an additional contribution to the drift term.
Given the drift and the diffusion coefficient we can write the generative model as a reverse-time SDE (Song et al. 2020b):
dz = [ x) - V log p(z)) dt + g(t)dw,
(100)
where dw is a reverse-time Wiener process.
1 MOTIVATION FOR CHOICE OF PARAMETERIZATION FOR THE ENCODER
As mentioned in Section 4 we would like our encoding to be helpful for the reconstruction loss at t = and for the latent loss at t = Multiplying the data with a will give us these properties, since it will be close to the identity at t = and send everything to 0 at t 1. Instead of just multiplying with Org, we choose to use o/ since we still get the desirable properties for t = 0 and t = 1, but it makes some of the mathematical expressions nicer (for example the derivative). Thus, we arrive at the non-trainable parameterization
Xne =
(101)
At = 1, we see that all the values of Xnt (At) are very close to 0, which should be easy to ap- proximate from 21 afx + djE, since it will just be the mean of the values in 21. Note that this parameterization gives us a lower latent loss, since the values of ař are closer to zero than the values of a,x. However, this not the same as just using a smaller minimum X in the original formulation, since in the original formulation the diffusion model would still be predicting x and not aix =
1. There is still problem with this formulation, since if we look at what happens between / = 0 and / = we see that at some point, we will be attempting to approximate it = 0₂E 0xXnt (x, Aτ) from very noisy z, while the values of Xnc de) are still very small. In other words, since Z, is a noisy version of Xat (z, Ae) and Xat (z, Ac) has very small values there will not be much signal, but as we move away from t = 1, will also become a worse and worse approximation.
This is why we introduce the trainable encoder
(x(): x -,,x o de)
(102)
=o,x + ofy At)
(103)
Here we allow the inner encoder (x, A2) to add signal dependent on the image at the same pace as we are removing signal via the -o² term. This should give us a better diffusion loss between
t 0 and = but still has (>e) very close to at t 0.
J CONTINUOUS-TIMA LIMIT OF THE DIFFUSION LOSS WITH AN ENCODER
J.1 REWRITING THE LOSS USING SNR
We can express the KL divergence in terms of the SNR:
SNR(t)- of
(104)
a. outside, expand o. and use the definition of the SNR to get:
We pull
a 1
(SNR(s) SNR(t))
(105)
20 o
2
Q
We also see that
11s
o
a o
0,00
SNR(1)
SNR(s) SNR(!)
(106)
t
0300²
SNR(s)
SNR(s)
23

---

Published as a conference paper at ICLR 2024
Inserting this back into Eq. (66), we get:
up MQ l
(107)
2
1
(SNR(s) SNR(
(108)
2
X (Ar) (x) + SNR(s) (() - Ac)o (Ar) - (xx(1e) - X (Arl))
SNR(s) SNR(()
The KL divergence is then:
DxL(o(2,/z (),lpo(2,lz))) 21))
(109)
1 (SNR(s) SNR()))
SNR(s) ((s - (A)) - (xx(1.) - X (Aee))
2e (Nr) - xx) +
SNR(s) SNR(()
with 1-1 and +-
J.2 TAKING THE LIMIT
we rewrite everything in the loss from Eq. (109 to be with respect to A we get
Cx(x) = 2 -B,,-U((.TT [(e - -²
(110)
xo(Ai) xx() εd ((dg - {e)o - (xx(2s) -xx (Ae))) cdo - e²
Where = -i- and i/T. We now want to take the continuous limit. Outside the norm we get the derivative with respect to f, inside the norm, we want the derivative w.r.t. Xz. First we consider
(111)
For T - oc: and see that
de
-
X
(112)
dt
Where X, is the derivative of d w.r.t t. Inside the norm get for - t that
A ), (xx( xx())))
t -1 dxx(
c, - (e>r - -2o,) Ae A
(113)
dete dhi
dd:
de dx
(114)
cde dde
dxp(Xe)
(115)
dh
and
ex
c² -(x X12)
ed xt (2g 2r)0(8((2))
(e et,) 0/ x2))((1))
(116)
(20 oxx0(>:)
(117)
So we get the loss
(x 27
(x) =
E. (e.u Xee (0()))- - (xx( + 02
(118)
2
dd.
24

---

Published as a conference paper at ICLR 2024
K USING THE V PARAMETERIZATION IN THE DIFFENC LOSS
In the following subsections we describe the v-prediction parameterization (Salimans & Ho 2022)
and derive the v-prediction loss for the proposed model, DiffEnc. We start by defining:
Vr- = 0g6 (At) (o(A,) = 0,68 (7x)
(119)
(120)
which give us (see Appendix K.1
(X,) = ai2i - OIVE
(121)
xx() 0474 oivo((2))
(122)
where we learn the v-prediction function ()r) = vo(zx.: 2r). In Appendix using this parameterization in Eq. (18) the loss becomes:
K.2 we show that,
(x) E (00001
Xa vo(A)) vo(A)) ( ot dd
(x)(Ar)
27
(123)
As shown in Appendix K.3
the diffusion loss for the trainable encoder from Eq.
15 becomes:
1 _ (x) - -
Xa vi vo o (Ro(dr) - xx(),) + (d(
dy (Ac)
ddv
(124)
and for the non-trainable encoder:
2 Ee,aoo [xa [|v vo + + (8o())) - xx (A))
L (x)
(125)
Eqs.
124) and 125 are the losses we use in our experiments.
K.1
REWRITING THE PARAMETERIZATION
In the parameterization of the loss from Salimans & Ho 2022 is defined as
0 C2C 0,X
(126)
We use the generalization
11 age 0ix0 (I, A1)
(127)
Note that since
(1 A1) (2 o.E)/al
(128)
and
ai o% =1
(129)
we get
x = (++ orE)/a4
(130)
-(a oi)
(131)
01 + o
+ of
2
0iay
e
(132)
at
Le
a,4 +
2t
010
(133)
at
0t
2
Tt
02
σ
3,0
2t
(134)
0n
at
Tt
a1
(3t
O(E)
(135)
a:
a2
1 (a,E
ox Ac))
(136)
1140
(137)
(138)
25

---

Published as a conference paper at ICLR 2024
So
(2 Xe) = 042t - 010
(139)
Therefore we define
(o(Ai) = a.c0 0xx
(140)
which in the same way gives us
x0(211 A1) 0+ 0,.0
(141)
where we learn Vg-
K.2 V PARAMETERIZATION IN CONTINUOUS DIFFUSION LOSS
For the parameterization have
(A())) = 0gE xx (Ax)
(142)
where c is from standard normal distribution, (O,, I), and
xx(Ai) = a - ovv (Nc)
(143)
So we will set
(A(A = 042t - o.vo((t))
(144)
and
vo(A.) an (dx)
(145)
where we learn (Ac). If we rewrite the second term within the square brackets of Eq. the parameterization, we get:
118 using
X(I)ea+ 20(2)) - (Xe) x2)x (xt)
2
(146)
dd,
1
dx (Me)
2
=(() (vv (A.) +0(x0(11)
(147)
d)
1
dx (Xx)
= x vo(A) 1 (o + 0x00
(148)
0t ddr
So we get the loss
1
((()aa (vo(ke) - (oxoo)) 1 (d)
27
Loo (x)
-2 1e
(149)
ot dd
K.3 V PARAMETERIZATION CONTINUOUS DIFFUSION LOSS WITH ENCODER
We recall our two parameterizations of the encoder
x (N²) oix + o, de)
(150)
at x + o, At)
(151)
and
Xut (de) aix
(152)
We see that
dxs(>i) = of + of
dyo(Ne) of
(153)
dde
dAr
and
dxat
ai 07
(154)
dh
26

---

Published as a conference paper at ICLR 2024
mentioned before. We first consider the loss for trainable encoder. Focusing the part of Eq. (123) inside the norm, and dropping the dependencies on x for brevity, we get
1 (xe
VΦ vo + 0:
(155)
6t d)
( 2 07,0xx + oR dy) - a;o7yo
=v-- v.
-i0 + 0,
(156)
0t
di,
a, - ot dy (x) + ajoNy
=v - 0i0 + 0,80
(157)
dit
dy + tai8s
-o - 1o4 to 91
X9 aix
(158)
dh
dy (x)
Ve - 8/ + oz
X0 a x + (1 - oi)yo (1-o
(159)
dd.
dyo(de)
Ve - 80 to
X9 aix - o + +Y
(160)
dd
dy)(de)
= Vo vi to
Xo - x + ya ddg
(161)
So for the trainable encoder, get the loss
(x) = - Be Bu,00100111
(162)
2
dy (Ac)
27
X(t)at v; - -vo + or (xo(1)) xx(),) + (a()))
dA.
For the non-trainable encoder, if again focus the part of Eq. dropping the dependencies on X for brevity, we get
123 inside the norm, and
1 (Xe)
VΦ vo + 0,800
(163)
0+ dd,
=V Vo + 0.80
(afo}x)
(164)
VΦ vo + 0 80
at (afx)
(165)
= V vo + σe (x x²x)
(166)
V v + ar (Xo X.t)
(167)
So for the non-trainable encoder, get the loss
Loo (x) - [(()) [|,, - vo + oz (xo(2e) - Xm Xne (A:))
(168)
L CONSIDERING LOSS FOR EARLY AND LATE TIMESTEPS
Let us consider what happens the expression inside the norm from our loss Eq. 19 for t close to zero. We see that since a₁ + and or - 0 for t + and Va = αrEo و.80 (de), we get for the trainable encoder
dy dde) X (Ai) xx (yo(A))
|
1
V: vo +
(169)
dd:
+ ||v, vol12 = ||c 2èll?
(170)
and for the non-trainable encoder
|| vo + ae (Xc) - xxx (A:))
(171)
|v voll - è
(172)
27

---

Published as a conference paper at ICLR 2024
Table 4: Comparison of the different components of the loss for DiffEnc- 8-2, DiffEnc-8-ni and VDMv-8 on MNIST. quantities are in bits per dimension (BPD), with standard enoor, seeds, 2M steps. Noise schedules are either fixed or with trainable endpoints.
Model
Noise
Total
Latent
Diffusion
Reconstruction
VDMv-8
fixed
0.370 + 0.002
0.0045 + 0.0
0.360 + 0.002
0.006 + (3 x 10
trainable
0.366 + 0.001
0.0042 + (5 10-3)
0.361 + 0.003
0.001 + (2 x 10
DiffEno-8-2
0.367 0.001
0.0009 (3 10-6
0.360 0.001
0.006 (3 10
trainable
0.363 + 0.002
0.0064 + (8 x 10-5)
0.355 + 0.002
0.001 + (2 10
DiffEnc-8-nt
fixed
0.378 = 0.002
1.6 x 10 0.0
0.371 + 0.002
0.006 + (3 10-5
trainable
0.373 0.001
0.0021 (3 10-5)
0.369 0.001
0.002 (5 10-5)
So we get the same objective as for the epsilon parameterization used in (Kingma et al. both cases. On the other hand, since o₁ + as - 1, we get for the trainable encoder:
2021 in
dy (Ar)
V. vo + (x)()i) xx (ol
(173)
ddc
(de) - (xx(de) vo(Ar) dy
vo 2
(174)
dde
Assuming x (Ne) = Xo (Ae), this loss is small at = if:
vo ~ (xx( + (a)) (de)
(175)
dd
dxe()1)
- 0 (xx, yo(())
(176)
dde
dy dy
~ - - vo (o(
(177)
dAe
dy (X.)
(178)
ddi
So we are saying that at = 1, vo ~ dy Thus the encoder should be able to guide the dhe
diffusion model. For the non-trainable encoder, we get
vv - vo + aa (Ro(Ae) x (Ax))
(179)
xxxx (>,) + (>() x (Ae)11
(180)
= (Ax) (Ax))1
(181)
So in this case, we just saying that &o (Ac) should be close to xxt).). However, note that since X (At) = aix, we have that X (Ar) x X (Ax) = 0 for t = 1. So this is only saying that it should be easy to guess x (xe) As 0 for 1 1, but it will not help the diffusion model guessing the signal, since there is no signal left in this case.
M DETAILED LOss COMPARISON FOR DIFFENC AND VDMV ON MNIST
Table 4 shows the average losses of the models trained on MNIST. We see the same pattern as for the small models trained CIFAR-10: All models with a trainable encoder achieve the same or better diffuuson loss than the VDMv model. For the fixed noise schedules the latent loss is always better for the DiffEnc models than for the VDMv, however for the trainable noise schedule, it seems the DiffEnc with learned encoder sacrifices some latent loss to gain a better diffusion loss.
N DETAILED Loss COMPARISON FOR DIFFENC- 32-2 AND VDMv-30 ON CIFAR-10
To explore the significance the encoder size, we trained a DiffEne-32-2, that is, a large diffusion model with smaller encoder, Table 5 We that after 2M steps the diffusion loss for the
28

---

Published as a conference paper at ICLR 2024
Table 5: Comparison of the different components of the loss for DiftEnc- 32-22 and VDMv-32 with fixed noise schedule on CIFAR-10. All quantities are in bits per dimension (BPD) with standard error over seeds, com- parison at 2M steps.
Model
Total
Latent
Diffusion
Reconstruction
VDMv-32
2.666 + 0.002
0.0012 + 0.0
2.654 + 0.003
0.01 + (4 10
DiffEne-32-2
2.660 + 0.006
0.0007 + (3 x 10-6) 2.649 + 0.006
0.01 + (2 10-6)
Table 6: Comparison of the different components of the loss for DiffEnc-32-8 and VDMv-32 with fixed noise schedule ImageNet32. All quantities are in bits per dimension (BPD) with standard error over 3 sccds, and models trained for 1.5M steps.
Model
Total
Latent
Diffusion
Reconstruction
VDMv-32
3.461 + 0.002
0.0014 = 0.0
3.449 + 0.002
0.01 + (1 10-9)
DiffEnc-32-8
3.461 + 0.002
0.0007 = (9 x 10-7) 3.450 = 0.002
0.01 = (1 x 10 10-5)
DiffEnc model is smaller than for the VDMv, however, not significantly so. When inspecting plot of the losses of the models, the losses seem to be diverging, but one would have to train the DiffEne-32-2 model for longer to be certain. We did not continue this experiment because of the large compute cost.
O DETAILED LOSS COMPARISON FOR DIFFENC AND VDMV ON IMAGENE332
On imagenet32, we see the same pattern in our experiments as for the small models on CIFAR- 10 and MNIST, see Table 6 The diffusion loss is the same for the two models, but the latent loss is better for DiffEnc. Since ImageNet is more complex than CIFAR-10, we might need an even larger base diffusion model achieve difference in diffusion loss.
P FURTHER FUTURE WORK
Our approach could be combined with various existing methods, c.g., latent diffusion (Vahdat et al. 2 Rombach et al. 2022) discriminator guidance (Kim et al. (022a)) If one were to succeed in making the smaller representations from the encoder, one might also combine with consistency regularization (Sinha & Dieng 2021) improve the learned representations.
Q MODEL STRUCTURE AND TRAINING
Code can be found on GitHut
All diffusion models the same overall structure with ResNet blocks, then middle block of 1 ResNet, 1 self attention and ResNet block, and in the end n more ResNet blocks. We train diffusion models with = 8 on MNIST and CIFAR-10 and models with and Im3geNet32. All ResNet blocks in the diffusion models preserve the dimensions the original images (28x28 for MNIST, 32x32 for CIFAR- 10, 32x32 for ImageNet32) and have 128 channels for models on MNIST and CIFAR-10 and 256 out channels for models on ImageNet32 following (Kigmma al. 202 We use both fixed noise schedule with Amaz 13.3 and Amin -5 and trainable noise schedule where we learn Amez and Amin-
32 on CIFAR-10
For our encoder, we use very similar overall structure for the diffusion model. Here we have m ResNet blocks, then middle block of 1 ResNet, self attention and ResNet block, and in the end
httpp:://giiei com wwm.bemmmmcm//ttt. ni /idi FfFnc
29

---

Published as a conference paper at ICLR 2024
Table 7: Comparison of the different components of the loss for DiffEnc-8 and VDMv-8 CIFAR-10 with fixed noise schedule after 1.3M steps. quantities are in bits per dimension (BPD), with standard error, seeds for DiffEnc-8-4, 5 seeds for VDMv-8.
Model
Total
Latent
Diffusion
Reconstruction
VDMv-8
2.794 + 0.004
0.0012 + 0.0
2.782 + 0.004
0.010 + (1 10-3
DiffEnc-8-4
2.789 + 0.002
0.0006 + (2 x 10-9)
2.778 + 0.002
0.010 + (1 10-5)
m more ResNet blocks. However, for the encoder with m = 2, we use maxpooling after each of the first m ResNet blocks and transposed convolution after the last m ResNet blocks, for encoders with 4, we use maxpooling after every other of the first ResNet blocks and transposed convolution after every other the last m ResNet blocks and for encoders with m = 8, we use maxpooling after every fourth of the first m ResNet blocks and transposed convolution after every fourth of the last m ResNet blocks. Thus, for the encoder we downscale to and upscale from resolutions 14x14 and 7x7 MNIST and 16x16 and 8x8 on CIFAR-10 and ImageNet32.
We do experiments with n = 8, m = on MNIST and CIFAR-10, n = 8, m = 2 and n = 32, 4 on CIFAR-10 and 32,
8 on ImageNet32.
We trained 5 seeds for the small models (n = 8), except for the diffusion model size 8 encoder size 4 CIFAR-10 where trained 3 seeds. We trained 3 seeds for the large models (n 32).
For models on MNIST and CIFAR-10 we used a batch size of 128 and no gradient clipping. For models ImageNet32 we used batch size of 256 and gradient clipping.
R DATASETS
We considered three datasets:
MNIST: The MNIST dataset (LeCun al.] 1998) as fetched by the tensorflow datasets
package 60,000 images were used for training and 10,000 images for test. License: Unknown.
CIFAR-10: The CIFAR-10 dataset fetched from the tensorflow datasets package Orig- inally collected by Krizhevsky et al. (2009). 50,000 images were used for training and 10,000 images for test. License: Unknown.
Image.Net 32x32: The official downsampled version of ImageNet (Chrabaszcz et al. from the ImageNet website: https:://mmae: net. org/ hwwp///wwooodddooo php
2017
ENCODER EXAMPLES ON MNIST
Fig. 4 provides example of the encodings we get from MNIST when using DiffEnc with learned encoder.
T SAMPLES FROM MODELS
Examples of samples from our large trained models, DiffEnc-335 -2-4 and VDMy-32, can seen in
Fig. 5
30

---

Published as a conference paper at ICLR 2024
0.00
0.70
0.80
0.90
0.92
0.94
0.96
0.98
1.00
0
0
0
0
0
0
0
0
3
3
3
3
3
3
3
3
3
6
6
6
6
6
6
6
6
6
9
9
9
9
9
9
9
Figure 4: Encoded MNIST images from DiffEnc-8-2. Encoded images are close to the identity up up to t = 0.7.
From & 0.8 to 0.9 the encoder slightly blurs the numbers, and from f 0.9 it makes the hackground lighter, but keeps the high contrast in the middle of the image. Intuitively, the encoder improves the latent loss by bringing the average pixel value close
Figure 5: 100 unconditional samples from DiffEne-32-4 (above) and VDMy-32 (below) after 8 million train- ing steps.
U USING A LARGER ENCODER FOR A SMALL DIFFUSION MODEL
As we can observe in Table 7 when the encoder's size is increased, the average diffusion loss is slightly smaller than that of VDM, albeit not significantly. We propose the following two potential explanations for this phenomenon: (1) Longer training may be needed achieve significant dif- ference. For DiffEnc-322 and VDMv-32, we saw different trends in the loss after about 2 million steps, where the loss of the DiffEnc models decreased more per step. However, it took more training with this trend to achieve a substantial divergence in diffusion loss. (2) a larger diffusion model may be required to fully exploit the encoder.
httas://ww.tenorrfoooor//////aas htte://www...oorrroroooo//aeeees farl
31

---

Published as a conference paper at ICLR 2024
Table 8: Comparison of the FID scores with standard emror for DiffEnc-32-4 and VDMv-32 CIFAR-I with fixed noise schedule 8M steps. seeds. We provide both FID scores on 10K and 50K samples and with respect to both train and test set.
Model
FID 10K train
FID 10K test
FID 50K train
FID 50K test
VDMv-32
14.8 = 0.2
18.9 = 0.2
11.2 0.2
14.9 + 0.2
DiffEnc-32-4
14.6 = 0.8
18.5 = 0.7
11.1 0.8
15.0 + 0.7
V FID SCORES
Although we did not optimize our model for the visual quality of samples, we provide FIDD scores of DiftEnn-32-4 and VDMv-32 on CIFAR-10 in Table We see from these, that the FID scores for the two models are similar and that it makes a big difference to the score whether we use the train or test set to calculate it and how many samples using more samples from the model and (as can expected) better when calculating with respect to the train set that with respect to the test set.
use from the model. The scores are better when
W SUM HEATMAP OF ALL TIMESTEPS
A heatmap the changes x, for all timesteps t and all ten CIFAR-10 classes be found
in Fig. 6 Recall in the following that all pixel values are scaled to the range (-1, 1) before they given as input the model. The DiffEnc model with trainable encoder is initialized with Yo (Ac) = O, that is, with no contribution from the trainable part. This means that if we had made this heatmap at initialization, the images would be blue where values the channels more than 0, red where values are less than 0 and white where values are zero. However, we see that after training. the encoder has a different behaviour around edges for t 0.8. For example, there is a white line in the middle of the cat in the second row which is not subtracted from, probably to preserve this edge in the image, and there is extra "outline" around the whole cat. We also see that for / > 0.8, the encoder gets much more general behaviour. In the fourth row, we see that the encoder adds to the entire middle of the image including the white line on the horse, which would have been subtracted from, if it had had the same behaviour at initialisation. Thus, see that the encoder learns to do something different from how was initialised, and what it learns is different for different timesteps.
32

---

Published as a conference paper at ICLR 2024
original
0.1
0.4
0.5
0.7
0.8
1.0
Figure 6: Change encoded image over range of depths: (x: x,)/(t s) for 0.1, ..., 1.0 and = 0.1. Changes have been summed over the channels with red and blue denoting positive and negative changes, respectively. For t closer to 0 the changes are finer and seem to be enhancing high-contrast edges, but for t + 1 they become more global.
33


Reviews:
[{"review_id": "dtr8Yhrl1x", "reviewer": "Reviewer_g9hi", "summary": "The paper proposes to model the integration of a learnable data and time-dependent mean function for diffusion models. The paper investigates the parametrisation of the diffusion losses and shows optimal variance schedules for the forward and reverse processes. The paper shows density estimation results on CIFAR 10 and MNIST and achieves the best performance compared to relevant baselines.", "questions": "- In section 3, I first missed the added subscript in $x_t$ in eq 13. It might make it easier to read to change the notation from $x_t$ to e.g. $e_t$ or something similar to be clear there.\n- in section 6, I was not fully following the parametrisation of $x(\\lambda_t)$ - how does this related to eq (13)? E.g. does eq. (29) indicate that $q(z_t|x) = N(a_t^3 x; ...)$?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and easy to follow. It adds a simple yet interesting addition to diffusion by introducing a mean shift to the forward diffusion while not being required in the sampling process and therefore ensuring its scalability. The theoretical analysis of the different noise variances adds an interesting flavour too. The results seem to indicate that the added encoder improves the performance in terms of bits per dimension.", "weaknesses": "The paper has very limited evaluation and doesn't compare to some of relevant baselines that are even mentioned in the paper, like latent diffusion and only compares on Cifar-10 and MNIST. Furthermore, it mentions that some methods only show improvement after longer training, hinting at potential inconsistencies in the results in case of slightly different training setups due to not training till convergence. It is hard to judge whether the proposed changes are a significant improvement due to this.\nThe ablation studies of the different methods are interesting but it's unclear what insight is gained from this, e.g. what does it mean that the loss is dominated by the diffusion loss?\nLastly, the theoretical analysis of the variance ratio is interesting but unclear how it fits into the bigger story of the paper.", "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": "7KoOj5f6MN", "reviewer": "Reviewer_neQL", "summary": "This work delves into the realm of diffusion models through the lens of hierarchical Variational Autoencoders (VAEs). The authors put forth a novel iteration of diffusion models that preserves the advantages of VAEs. To be more specific, they introduce a time-dependent DiffEnc, resulting in encoder variations across different time steps. Furthermore, the authors conduct an in-depth exploration of the assumption that forward and backward variances are equal, underpinned by thorough theoretical analysis. In addition, empirical findings are presented to substantiate the efficacy of the proposed DiffEnc.", "questions": "I don't have specific questions regarding the main text, but it's clear that there's a need for improvements in the organization of the content.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "I believe that implementing a trainable encoder within the context of the diffusion model represents a promising avenue for enhancing diffusion models, particularly in terms of ELBO optimization. This paper offers comprehensive insights into the derivation and analysis, rendering it accessible and straightforward to grasp. The experimental findings are not only persuasive but also harmonize effectively with the theoretical framework. For instance, in Figure 1, we observe a logical outcome indicating that during the initial stages, the encoder primarily attends to local patterns, such as edges. Simultaneously, as the value of $t$ increases, it places greater emphasis on capturing global structures.", "weaknesses": "1. The organization of sections is perplexing. It's challenging for me to discern whether Section 2 serves as an introductory section or is meant to highlight one of your contributions.\n2. The absence of a central theorem throughout the paper poses a difficulty for readers in anticipating the direction of the derivations and what to expect.", "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": "mJpHvvJOjJ", "reviewer": "Reviewer_Hg56", "summary": "The paper extends variational diffusion models (VDM) with a time dependent encoder. Instead directly noising the input data points (here images), they are first encoded and then noised. There are two encoding schemes introduced. In one, the encoder has no trainable parameters and is simply a rescaling of the input. In the other, the encoded images are a reweighting of the noise-free image and the output of a time-dependent image-to-image neural network. The authors derive the corresponding loss functions and analyze the model in its infinite depth limit. The results show that the diffusion model equipped with a new encoder achieves state-of-the-art BPD on CIFAR-10.", "questions": "Do you have a sense for what the weaknesses of your gradient approximation (Section 6) are? It seems that the approximation will be very bad for t near 1. Does this effect the learning of the encoder in a serious way? It might be helpful to explore this a bit more if it is one of the more impactful approximations required to your method practical.\n\nWhat is the intended take-away from Figure 1? Is the significance simply that the learned encoder captures the spatial frequencies that we might expect at each time in the diffusion process? Or do you also conclude that there is an implicit form of segmentation happening? The later feels potentially quite profound but needing more justification while the former feels potentially unsurprising, especially if the gradient signal is just incredibly noisy. It could just be learning very little for t near 1.0.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The high-level idea for the paper is quite natural and something that somebody was bound to try because of its potential impact. Overall, I found the writing fairly clear, with some exceptions that I will mention in the next section. The analysis of the method is fairly extensive and supported by lot of details in the appendix, though these are primarily mathematical proofs and not necessarily an exploration of design decisions that have a high level of practical significance.", "weaknesses": "The presentation of the method, namely section 3 and 6 could be improved significantly. There are a lot of variable names, and I had to read through the section many time in order to understand what was happening, even though the final procedure is not that complex. Moving the figure provided in Appendix A to the main text might be helpful in this regard. Or you could include an algorithm, or simply a link to your code, as these would all be easier to parse as someone familiar with common diffusion code and parlance. The presentation is also strong emulating the style of VDM even though the method has interesting consequences and potential applications on diffusion models more broadly, and simplifying the description of the algorithm to rely less on graphical models and equations and simply describe where noising and denoising are happening might make the approach more accessible to a broader audience.\n\nBeyond the presentation, the results are pretty limited. Given that the proposed method is an incremental change to VDM, the performance improvement is not overwhelming. It feels like an improvement on par with what might be possible through better hyperparameter tuning or by increasing the model capacity (which is implicitly part of what this method is doing, at least in the case where the encoder has learnable parameters). The results being limited to CIFAR-10 compound my intuition that the results aren't very notable, because it's not clear if they generalize. Though I think method appears sound and significant and general results seem possible, the current results feel like a negligible improvement on a relatively small dataset. It would be great to see results paralleling the original VDM model. If the improvements spanned all the datasets considered in the original paper, I would be more convinced that the proposed encoder is doing something significant.", "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": "lEO4LXJQQN", "reviewer": "Reviewer_jJ4J", "summary": "The authors propose a version of diffusion models with a simple trainable encoder parameterized as a simple convex combination between  the data $\\mathbf{x}$ and a neural network $f(\\log SNR(t))$. parameterized by the signal to noise ratio ($SNR(t)$), a function of time.", "questions": "It is still not entirely clear to me what is learned by the encoder model. Can the authors directly visualize the output of the learned encoder model at different times $t$?\n\nHow does Section 7 (v-prediction parameterization) relate to the rest of the paper? Is it used for experiments in Section 8?\n\nIn which experiments in Section 8 were a fixed noise schedule used, and in which were a learned noise schedule used?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The paper is generally well-written and easy to read.\n- The formulation for learnable encoders is simple and clear, and involves minimal added components from the original diffusion model.\n- All derivations are well-documented in the appendix.", "weaknesses": "- The benefit of the encoder is unclear. There is limited exploration of the learned encoded space outside of the visualization provided in Figure 1. What is gained by this formulation?\n- In spite of the added theoretical overhead incurred by the learned encoder, there is minimal empirical improvement over existing baselines (BPD 2.65 -> 2.63 on the CIFAR10 dataset).\n- Empirical validation is slim. Model performance is only evaluated in terms of BPD on two simple datasets: MNIST and CIFAR10. MNIST is generally no longer used as a density estimation benchmark. As a result, I wonder is the performance generalizable? Is there improvement in terms of any other metrics? There is a lot of research on how likelihood can be a poor proxy for generative model performance. For example, a lot of improvements in BPD can be attributed to modeling semantically unimportant high frequency noise components [1, 2].\n- Parameterization of the learned encoder is very simplistic. Why is the encoder conditioned only on the time $\\lambda_t$ and not the data $\\mathbf{x}$ itself? While I do appreciate elegant derivations, the justification for this parameterization (Appendix I) -- that this makes the math nice -- does not feel entirely sufficient.\n\n[1] A note on the evaluation of generative models. https://arxiv.org/abs/1511.01844\n\n[2] Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation. https://arxiv.org/abs/2106.05527", "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": {}}]

Author rebuttal:

Decision: Accept (poster)