SlowGuess/NeurIPS_2023_OCR huggingface dataset:
This concludes our proof of Theorem 1.
This result provides us the new insight that any of the objectives with (implied) monotonic weighting, as listed in Table 1, can be understood as equivalent to the ELBO with simple data augmentation, namely additive noise. Specifically, this is a form of Distribution Augmentation (DistAug), where the model is conditioned on the augmentation indicator during training, and conditioned on 'no augmentation' during sampling.
Monotonicity of  $w(\lambda)$  holds for a number of models with state-of-the-art perceptual quality, including VoiceBox for speech generation [Le et al., 2023], and Simple Diffusion for image generation [Hoogeboom et al., 2023].
4.1 Understanding the effect of weighting functions on perceptual quality
In this subsection, we propose possible explanations of the reason that the weighted objective, or equivalently the ELBO with data augmentation assuming monotonic weightings, can effectively improve perceptual quality.
Connection to low-bit training. Earlier work [Kingma and Dhariwal, 2018] found that training on 5-bit data, removing the three least significant bits from the original 8-bit data, can lead to higher fidelity models. A likely reason is the more significant bits are also more important to human perception, and removing the three least significant bits from the data allows the model to spend more capacity on modeling the more significant bits. In [Kingma and Dhariwal, 2018], training on 5-bit data was performed by adding uniform noise (with the appropriate scale) to the data before feeding it to the model, but it was found that adding Gaussian noise had a similar effect. Adding Gaussian noise with a single noise level is a special case of the weighted objective where the weighting function is a step function. The effect of uniform noise can be approximated with a sigmoidal weighting function; in Appendix I we dive into more details.
Fourier analysis. To better understand why additive Gaussian noise can improve perceptual quality, consider the Fourier transform (FT) of a clean natural image perturbed witjh additive Gaussian noise. Since the FT is a linear transformation, the FT of a the sum of a natural image and Gaussian noise, equals the sum of the FT of the clean image and the FT of the Gaussian noise. For natural images, it is known that the power spectrum, i.e. the average magnitude of the FT as a function of spatial frequency, decreases fast as the frequency increases [Burton and Moorhead, 1987], which means that the lowest frequencies have by far the highest magnitudes. On the other hand, the power spectrum of a Gaussian white noise is roughly constant. When adding noise to a clean natural image, the signal-to-noise ratio (i.e. the ratio of the power spectra of the clean and Gaussian noise images) in high frequency regions is lower. As a result, additive Gaussian noise effectively destroys high-frequency information in the data more quickly than the low-frequency information, pushing the model to focus more on the low frequency components of the data, which often correspond to high-level content and global structure that is more crucial to perception. See Appendix J for a more detailed discussion.
Model parameterization	Training noise schedule	Weighting function	DDPM sampler	EDM sampler
Monotonic?	FID ↓	IS ↑	FID ↓	IS ↑
ε-parametrization	Cosine	sech(λ/2) (Baseline)	1.85	54.1 ± 0.79	1.55	59.2 ± 0.78
"	Cosine	sigmoid(-λ + 1)	✓	1.75	55.3 ± 1.23
"	Cosine	sigmoid(-λ + 2)	✓	1.68	56.8 ± 0.85	1.46	60.4 ± 0.86
"	Cosine	sigmoid(-λ + 3)	✓	1.73	56.1 ± 1.36
"	Cosine	sigmoid(-λ + 4)	✓	1.80	55.1 ± 1.65
"	Cosine	sigmoid(-λ + 5)	✓	1.94	53.5 ± 1.12
"	Adaptive	sigmoid(-λ + 2)	✓	1.70	54.8 ± 1.20	1.44	60.6 ± 1.44
"	Adaptive	EDM-monotonic	✓	1.67	56.8 ± 0.90	1.44	61.1 ± 1.80
EDM [Karras et al., 2022]	EDM (training)	EDM (Baseline)	1.36
EDM (our reproduction)	EDM (training)	EDM (Baseline)	1.45	60.7 ± 1.19
"	Adaptive	EDM	1.43	63.2 ± 1.76
"	Adaptive	sigmoid(-λ + 2)	✓	1.55	63.7 ± 1.14
"	Adaptive	EDM-monotonic	✓	1.43	63.7 ± 1.48
v-parametrization	Adaptive	exp(-λ/2) (Baseline)	✓	1.62	58.0 ± 1.56
"	Adaptive	sigmoid(-λ + 2)	✓	1.51	64.4 ± 1.28
"	Adaptive	EDM-monotonic	✓	1.45	64.6 ± 1.35
Table 2: ImageNet 64x64 results. See Section 5.2.
Inspired by the theoretical results in Section 4, in this section we propose several monotonic weighting functions, and present experiments that test the effectiveness of the monotonic weighting functions compared to baseline (non-monotonic) weighting functions. In addition, we test the adaptive noise schedule (Section 3.2). For brevity we had to pick a name for our models. Since we build on earlier results on Variational Diffusion Models (VDM) [Kingma et al., 2021], and our objective is equivalent to the VDM objective combined with data augmentation, we name our models  $VDM++$ .
5.1 Weighting functions
We experimented with several new monotonic weighting functions. First, a sigmoidal weighting of the form  $\mathrm{sigmoid}(-\lambda + k)$ , inspired by the sigmoidal weighting corresponding to low-bit training (see Section 4.1 and Appendix I for details), where  $k$  is a hyper-parameter. In addition, inspired by the EDM weighting function from Karras et al. [2022] (non-monotonic, Table 1), we test a weighting function indicated with 'EDM-monotonic', which is identical to the EDM weighting, except that it is made monotonic by letting  $w(\lambda) = \max_{\lambda} \tilde{w}(\lambda)$  for  $\lambda < \arg \max_{\lambda} \tilde{w}(\lambda)$ , where  $\tilde{w}(\lambda)$  indicates the original EDM weighting function. Put simply, 'EDM-monotonic' equals the EDM weighting, except that it is constant to the left of its peak; see Figure 3.
5.2 ImageNet 64x64
All experiments on ImageNet 64x64 were done with the U-Net diffusion model architecture from [Nichol and Dhariwal, 2021]. We carried out extensive ablation studies over several design choices of diffusion models, namely model parametrization, training noise schedules, weighting functions and samplers. For  $\epsilon$ -parametrization model, we took iDDPM [Nichol and Dhariwal, 2021] as the baseline, which utilized cosine noise schedule and a non-monotonic  $\mathrm{sech}(\lambda /2)$  weighting. For EDM parametrization, we recruited the setting in Karras et al. [2022] as the baseline, with EDM training noise schedule (Figure 2) and non-monotonic EDM weighting (Figure 3). For  $\mathbf{v}$ -parametrization model, we followed Salimans and Ho [2022] taking the  $\mathbf{v}$ -prediction loss as the baseline, which leads to a monotonic  $\exp (-\lambda /2)$  weighting. We always use adaptive noise schedule under this setting. We started by searching the optimal hyperparameter for the sigmoidal weighting under the  $\epsilon$ -parametrization model, and then applied the best sigmoidal weighting and EDM-monotonic weighting under other settings. We evaluated models with two samplers, DDPM and EDM samplers, with their corresponding sampling noise schedules. Table 2 summarizes the FID [Heusel et al., 2017] and Inception scores [Salimans et al., 2016] across different settings. From the table we made several observations as follows.
First, our adaptive noise schedule (Appendix F) works equally well as the hand-tuned fixed noise schedules. We verified this under  $\epsilon$ -parametrization with sigmoid  $(- \lambda + 2)$  weighting and EDM parametrization with EDM weighting. Under both settings, the adaptive noise schedule results in slightly better FID and Inception scores with the EDM sampler, and slightly worse scores with the DDPM sampler for the former setting. Although the adaptive noise schedule did not significantly affect the end result, it allows us to experiment with new weighting functions while avoiding hand-tuning of the noise schedule. In addition, it sometimes leads to faster training; see Appendix F.
Second, it is possible to modify an existing non-monotonic weighting function to a monotonic one with the minimal change, and results in on-par sample quality. Specifically, under EDM parameterization and adaptive noise schedule, we change EDM weighting to EDM-monotonic weighting with the very straightforward approach stated in Section 5.1. It results in nearly the same FID score and slightly better Inception score.
Third, our proposed new weighting functions, sigmoid  $(- \lambda + 2)$  and EDM monotonic weightings, work better than the baseline settings across different model parameterization and samplers in most cases. Larger performance gain has been observed under  $\epsilon$ -parameterization and  $\mathbf{v}$ -parameterization settings. Our-reimplementation of the EDM parameterization baseline setting could not exactly reproduce their reported FID number (1.36), but comes close (1.45). We observed less significant improvement under this setting, possibly because their already exhaustive search of the design space.
Finally, the above observations remain consistent across different model parameterization and samplers, indicating the generalizability of our proposed weighting functions and noise schedules.
5.3 High resolution ImageNet
In our final experiments, we tested whether the weighting functions that resulted in the best scores on ImageNet  $64 \times 64$ , namely sigmoid  $(- \lambda + 2)$  and 'EDM-monotonic', also results in competitive scores on high-resolution generation. As baseline we use the Simple Diffusion model from Hoogeboom et al. [2023], which reported the best FID scores to date on high-resolution ImageNet without sampling modifications (e.g. guidance).
We recruited the large U-ViT model from Simple Diffusion [Hoogeboom et al., 2023], and changed the training noise schedule and weighting function to our proposed ones. Note that for higher-resolution models, Hoogeboom et al. [2023] proposed a shifted version of the cosine noise schedule (Table 4), that leads to a shifted version of the weighting function  $w(\lambda)$ . Similarly, we extended our proposed sigmoidal and 'EDM-monotonic' weightings to their shifted versions (see Appendix D.2.1 for details). For fair comparison, we adopted the same vanilla DDPM sampler as Simple Diffusion, without other advanced sampling techniques such as second-order sampling or rejection sampling. As shown in Table 3, with our adaptive noise schedule for training, the two weighting functions we proposed led to either better or comparable FID and IS scores on ImageNet  $128 \times 128$ , compared to the baseline Simple Diffusion approach.
Next, we test our approach on ImageNet generation of multiple high resolutions (i.e., resolutions 128, 256 and 512), and compare with existing methods in the literature. See Table 3 for the summary of quantitative evaluations and Figure 1 for some generated samples by our approach. With the shifted version of 'EDM-monotonic' weighting, we achieved state-of-the-art FID and IS scores on all three resolutions of ImageNet generation among all approaches without guidance. With classifier-free guidance (CFG) [Ho and Salimans, 2022], our method outperforms all diffusion-based approaches on resolutions 128 and 512. On resolution 256, our method only falls a bit behind Gao et al. [2023] and Hang et al. [2023], both of which were built upon the latent space of a pretrained auto-encoder from latent diffusion models [Rombach et al., 2022a] that was trained on much larger image datasets than ImageNet, while our model was trained on the ImageNet dataset only. It is worth noting that we achieve significant improvements compared to Simple Diffusion which serves as the backbone of our method, on all resolutions, with and without guidance. It is possible to apply our proposed weighting functions and adaptive noise schedules to other diffusion-based approaches such as Gao et al. [2023] to further improve their performance, which we shall leave to the future work.
Method	Without guidance	With guidance
FID ↓	FID ↓
train	eval	IS ↑	train	eval	IS ↑
128 × 128 resolution
ADM [Dhariwal and Nichol, 2022]	5.91	2.97
CDM [Ho et al., 2022]	3.52	3.76	128.8 ± 2.5
RIN [Jabri et al., 2022]	2.75	144.1
Simple Diffusion (U-Net) [Hoogeboom et al., 2023]	2.26	2.88	137.3 ± 2.0
Simple Diffusion (U-ViT, L) [Hoogeboom et al., 2023]	1.91	3.23	171.9 ± 2.5	2.05	3.57	189.9 ± 3.5
VDM++ (Ours), w(λ) = sigmoid(-λ + 2)	1.91	3.41	183.1 ± 2.2
VDM++ (Ours), EDM-monotonic weighting	1.75	2.88	171.1 ± 2.7	1.78	3.16	190.5 ± 2.3
256 × 256 resolution
BigGAN-deep (no truncation) [Brock et al., 2018]	6.9	171.4 ± 2.0
MaskGIT [Chang et al., 2022]	6.18	182.1
ADM [Dhariwal and Nichol, 2022]	10.94	3.94	215.9
CDM [Ho et al., 2022]	4.88	4.63	158.7 ± 2.3
RIN [Jabri et al., 2022]	3.42	182.0
Simple Diffusion (U-Net) [Hoogeboom et al., 2023]	3.76	3.71	171.6 ± 3.1
Simple Diffusion (U-ViT, L) [Hoogeboom et al., 2023]	2.77	3.75	211.8 ± 2.9	2.44	4.08	256.3 ± 5.0
VDM++ (Ours), EDM-monotonic weighting	2.40	3.36	225.3 ± 3.2	2.12	3.69	267.7 ± 4.9
Latent diffusion with pretrained VAE:
DiT-XL/2 [Peebles and Xie, 2022]	9.62	121.5	2.27	278.2
U-ViT [Bao et al., 2023]	3.40
Min-SNR-γ [Hang et al., 2023]	2.06
MDT [Gao et al., 2023]	6.23	143.0	1.79	283.0
512 × 512 resolution
MaskGIT [Chang et al., 2022]	7.32	156.0
ADM [Dhariwal and Nichol, 2022]	23.24	3.85	221.7
RIN [Jabri et al., 2022]	3.95	216.0
Simple Diffusion (U-Net) [Hoogeboom et al., 2023]	4.30	4.28	171.0 ± 3.0
Simple Diffusion (U-ViT, L) [Hoogeboom et al., 2023]	3.54	4.53	205.3 ± 2.7	3.02	4.60	248.7 ± 3.4
VDM++ (Ours), EDM-monotonic weighting	2.99	4.09	232.2 ± 4.2	2.65	4.43	278.1 ± 5.5
Latent diffusion with pretrained VAE:
DiT-XL/2 [Peebles and Xie, 2022]	12.03	105.3	3.04	240.8
LDM-4 [Rombach et al., 2022a]	10.56	103.5 ± 1.2	3.60	247.7 ± 5.6
Table 3: Comparison to approaches in the literature for high-resolution ImageNet generation. 'With guidance' indicates that the method was combined with classifier-free guidance [Ho and Salimans, 2022]. † Models under 'Latent diffusion with pretrained VAE' use the pre-trained VAE from Stable Diffusion [Rombach et al., 2022a], which used a much larger training corpus than the other models in this table.
In summary, we have shown that the weighted diffusion loss, which generalizes diffusion objectives in the literature, has an interpretation as a weighted integral of ELBO objectives, with one ELBO per noise level. If the weighting function is monotonic, then we show that the objective has an interpretation as the ELBO objective with data augmentation, where the augmentation is additive Gaussian noise, with a distribution of noise levels.
Our results open up exciting new directions for future work. The newfound equivalence between monotonically weighted diffusion objectives and the ELBO with data augmentation allows for a direct apples-to-apples comparison of diffusion models with other likelihood-based models. For example, it allows one to optimize other likelihood-based models, such as autoregressive models, towards the same objective as monotonically weighted diffusion models. This would shine light on whether diffusion models are better or worse than other model types, as measured in terms of their held-out objectives as opposed to FID scores. We leave such interesting experiments to future work.
We'd like to thank Alex Alemi and Ben Poole for fruitful discussions and feedback on early drafts. We thank Emiel Hoogeboom for advice and help on the implementation of Simple Diffusion.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22669-22679, 2023.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:1809.11096, 2018.
Geoffrey J Burton and Ian R Moorhead. Color and spatial structure in natural scenes. Applied optics, 26(1):157-170, 1987.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11315-11325, 2022.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan. WaveGrad: Estimating Gradients for Waveform Generation. In ICLR, 2021a.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, Najim Dehak, and William Chan. WaveGrad 2: Iterative Refinement for Text-to-Speech Synthesis. In INTERSPEECH, 2021b.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon. Perception prioritized training of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11472-11481, 2022.
Maurizio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arXiv preprint arXiv:2303.11435, 2023.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. In NeurIPS, 2022.
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffusion for categorical data. arXiv preprint arXiv:2211.15089, 2022.
Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, and Jie Tang. Cogview: Mastering text-to-image generation via transformers, 2021. URL https://arxiv.org/abs/2105.13290.
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is a strong image synthesizer. arXiv preprint arXiv:2303.14389, 2023.
Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy. arXiv preprint arXiv:2303.09556, 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. arXiv preprint arXiv:1706.08500, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.
A Main proof
Here we'll provide a proof of Equation 11.
Note that like in the main text, we use shorthand notation:
$$
\mathcal {L} (t; \mathbf {x}) := D _ {K L} \left(q \left(\mathbf {z} _ {t, \dots , 1} | \mathbf {x}\right) \mid \mid p \left(\mathbf {z} _ {t, \dots , 1}\right)\right) \tag {15}
$$
A.1 Time derivative of  $D_{KL}(q(\mathbf{z}_{t,\dots,1}|\mathbf{x})||p(\mathbf{z}_{t,\dots,1}))$
Let  $dt$  denote an infinitesimal change in time. Note that  $\mathcal{L}(t - dt; \mathbf{x})$  can be decomposed as the sum of a KL divergence and an expected KL divergence:
$$
\mathcal {L} (t - d t; \mathbf {x}) = \mathcal {L} (t; \mathbf {x}) + \mathbb {E} _ {q (\mathbf {z} _ {t} | \mathbf {x})} [ D _ {K L} (q (\mathbf {z} _ {t - d t} | \mathbf {z} _ {t}, \mathbf {x}) | | p (\mathbf {z} _ {t - d t} | \mathbf {z} _ {t}) ] \tag {16}
$$
Due to this identity, the time derivative  $d / dt\mathcal{L}(t;\mathbf{x})$  can be expressed as:
$$
\begin{array}{l} \frac {d}{d t} \mathcal {L} (t; \mathbf {x}) = \frac {1}{d t} (\mathcal {L} (t; \mathbf {x}) - \mathcal {L} (t - d t; \mathbf {x})) (17) \\ = - \frac {1}{d t} \mathbb {E} _ {q (\mathbf {z} _ {t} | \mathbf {x})} [ D _ {K L} (q (\mathbf {z} _ {t - d t} | \mathbf {z} _ {t}, \mathbf {x}) | | p (\mathbf {z} _ {t - d t} | \mathbf {z} _ {t}) ] (18) \\ \end{array}
$$
In Appendix E of [Kingma et al., 2021], it is shown that in our model, this equals:
$$
\begin{array}{l} \frac {d}{d t} \mathcal {L} (t; \mathbf {x}) = - \frac {1}{2} \frac {\operatorname {S N R} (t - d t) - \operatorname {S N R} (t)}{d t} | | \mathbf {x} - \hat {\mathbf {x}} _ {\boldsymbol {\theta}} (\mathbf {z} _ {t}; \lambda_ {t}) | | _ {2} ^ {2} (19) \\ = \frac {1}{2} S N R ^ {\prime} (t) | | \mathbf {x} - \hat {\mathbf {x}} _ {\boldsymbol {\theta}} \left(\mathbf {z} _ {t}; \lambda_ {t}\right) | | _ {2} ^ {2} (20) \\ \end{array}
$$
where  $\mathbf{z}_t = \alpha_\lambda \mathbf{x} + \sigma_\lambda \pmb{\epsilon}$ , and  $SNR(t) \coloneqq \exp(\lambda)$  in our model, and  $SNR'(t) = d / dt$ $SNR(t) = e^{\lambda} d\lambda / dt$ , so in terms of our definition of  $\lambda$ , this is:
$$
\frac {d}{d t} \mathcal {L} (t; \mathbf {x}) = \frac {1}{2} e ^ {\lambda} \frac {d \lambda}{d t} \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} \left[ | | \mathbf {x} - \hat {\mathbf {x}} _ {\boldsymbol {\theta}} (\mathbf {z} _ {t}; \lambda_ {t}) | | _ {2} ^ {2} \right] \tag {21}
$$
In terms of  $\epsilon$ -prediction (see Section E.2), because  $||\epsilon - \hat{\epsilon}_{\theta}||_{2}^{2} = e^{\lambda}||\mathbf{x} - \hat{\mathbf{x}}_{\theta}||_{2}^{2}$  this simplifies to:
$$
\frac {d}{d t} \mathcal {L} (t; \mathbf {x}) = \frac {1}{2} \frac {d \lambda}{d t} \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} \left[ \| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} _ {\boldsymbol {\theta}} (\mathbf {z} _ {t}; \lambda_ {t}) \| _ {2} ^ {2} \right] \tag {22}
$$
where  $\mathbf{z}_{\lambda} = \alpha_{\lambda}\mathbf{x} + \sigma_{\lambda}\boldsymbol{\epsilon}$ . This can be easily translated to other parameterizations; see E.2.
This allows us to rewrite the weighted loss of Equation 5 as:
$$
\begin{array}{l} \mathcal {L} _ {w} (\mathbf {x}) = \frac {1}{2} \mathbb {E} _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (0, \mathbf {I})} \left[ w \left(\lambda_ {t}\right) \cdot - \frac {d \lambda}{d t} \cdot \left| \left| \hat {\epsilon} _ {\boldsymbol {\theta}} \left(\mathbf {z} _ {t}; \lambda\right) - \boldsymbol {\epsilon} \right| \right| _ {2} ^ {2} \right] (23) \\ = \mathbb {E} _ {t \sim \mathcal {U} (0, 1)} \left[ w \left(\lambda_ {t}\right) \cdot - \frac {1}{2} \frac {d \lambda}{d t} \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} \left[ | | \hat {\boldsymbol {\epsilon}} _ {\boldsymbol {\theta}} \left(\mathbf {z} _ {t}; \lambda\right) - \boldsymbol {\epsilon} | | _ {2} ^ {2} \right] \right] (24) \\ = \mathbb {E} _ {t \sim \mathcal {U} (0, 1)} \left[ - \frac {d}{d t} \mathcal {L} (t; \mathbf {x}) w \left(\lambda_ {t}\right) \right] (25) \\ = \int_ {0} ^ {1} - \frac {d}{d t} \mathcal {L} (t; \mathbf {x}) w \left(\lambda_ {t}\right) d t (26) \\ \end{array}
$$
A.2 Integration by parts
Integration by parts is a basic identity, which tells us that:
$$
- \int_ {a} ^ {b} f (t) g ^ {\prime} (t) d t = \int_ {a} ^ {b} f ^ {\prime} (t) g (t) d t + f (a) g (a) - f (b) g (b)
$$
This allows us to further rewrite the expression of the weighted loss in Equation 26 as:
$$
\begin{array}{l} \mathcal {L} _ {w} (\mathbf {x}) = \int_ {0} ^ {1} - \frac {d}{d t} \mathcal {L} (t; \mathbf {x}) w \left(\lambda_ {t}\right) d t (27) \\ = \int_ {0} ^ {1} \frac {d}{d t} w \left(\lambda_ {t}\right) \mathcal {L} (t; \mathbf {x}) d t + w \left(\lambda_ {\max }\right) \mathcal {L} (0; \mathbf {x}) - w \left(\lambda_ {\min }\right) \mathcal {L} (1; \mathbf {x}) (28) \\ \end{array}
$$
$$
- w \left(\lambda_ {\min }\right) \mathcal {L} (1; \mathbf {x}) = - w \left(\lambda_ {\min }\right) D _ {K L} \left(q \left(\mathbf {z} _ {1} \mid \mathbf {x}\right) \| p \left(\mathbf {z} _ {1}\right)\right) \tag {29}
$$
is constant w.r.t.  $\theta$ , since it does not involve the score function, and typically very small, since  $D_{KL}(q(\mathbf{z}_1|\mathbf{x})||p(\mathbf{z}_1))$  is typically small by design.
The term  $w(\lambda_{\mathrm{max}})D_{KL}(q(\mathbf{z}_{0,\dots,1}|\mathbf{x})||p(\mathbf{z}_{0,\dots,1}))$  is typically small, since  $w(\lambda_{\mathrm{max}})$  is typically very small (see Figure 3).
This concludes our proof of Equation 11.
B Visualization
We tried to create a helpful visualization of the result from Section A.2. Note that we can rewrite:
$$
\int_ {0} ^ {1} \frac {d}{d t} w \left(\lambda_ {t}\right) \mathcal {L} (t; \mathbf {x}) d t = \int_ {t = 1} ^ {t = 0} w \left(\lambda_ {t}\right) d \mathcal {L} (t; \mathbf {x}) \tag {30}
$$
The relationship in Equation 28 can be rewritten as: