SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbf {x} = \alpha_ {\lambda} ^ {- 1} \left(\mathbf {z} _ {\lambda} - \sigma_ {\lambda} \boldsymbol {\epsilon}\right) \tag {112}
$$
$$
\boldsymbol {\epsilon} = \sigma_ {\lambda} ^ {- 1} \left(\mathbf {z} _ {\lambda} - \alpha_ {\lambda} \mathbf {x}\right) \tag {113}
$$
$$
\begin{array}{l} \nabla_ {\mathbf {z} _ {\lambda}} \log q (\mathbf {z} _ {\lambda} | \mathbf {x}) = \nabla_ {\mathbf {z} _ {\lambda}} - \| \mathbf {z} _ {\lambda} - \alpha_ {\lambda} \mathbf {x} \| _ {2} ^ {2} / \left(2 \sigma_ {\lambda} ^ {2}\right) (114) \\ = - \sigma_ {\lambda} ^ {- 2} \left(\mathbf {z} _ {\lambda} - \alpha_ {\lambda} \mathbf {x}\right) (115) \\ = - \sigma_ {\lambda} ^ {- 2} \left(\alpha_ {\lambda} \mathbf {x} + \sigma_ {\lambda} \boldsymbol {\epsilon} - \alpha_ {\lambda} \mathbf {x}\right) (116) \\ = - \epsilon / \sigma_ {\lambda} (117) \\ \end{array}
$$
In addition, there's the v-prediction parameterization  $(\mathbf{v} := \alpha_{\lambda}\epsilon - \sigma_{\lambda}\mathbf{x})$  explained in D.2, and the o-prediction parameterization  $(\mathbf{o} := \mathbf{x} - \epsilon)$  explained in D.3.
Karras et al. [2022] proposed a specific  $\mathbf{F}$ -parametrization, with an  $\mathbf{F}$ -prediction model  $\hat{\mathbf{F}}_{\theta}$ . In the special case of variance explosion (VE) SDE, it is formulated as:
$$
\mathbf {x} = \frac {\tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}}{e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}} \mathbf {z} _ {\lambda} + \frac {e ^ {- \lambda / 2} \tilde {\sigma} _ {\mathrm {d a t a}}}{\sqrt {e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}}} \mathbf {F} \tag {118}
$$
where  $\tilde{\sigma}_{\mathrm{data}} = 0.5$ . Generalizing this to our more general family with arbitrary drift, this corresponds to:
$$
\mathbf {x} = \frac {\tilde {\sigma} _ {\mathrm {d a t a}} ^ {2} \alpha_ {\lambda}}{e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}} \mathbf {z} _ {\lambda} + \frac {e ^ {- \lambda / 2} \tilde {\sigma} _ {\mathrm {d a t a}}}{\sqrt {e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}}} \mathbf {F} \tag {119}
$$
So that we have:
$$
\begin{array}{l} \mathbf {F} = \frac {\sqrt {e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}}}{e ^ {- \lambda / 2} \tilde {\sigma} _ {\mathrm {d a t a}}} \mathbf {x} - \frac {\tilde {\sigma} _ {\mathrm {d a t a}} \alpha_ {\lambda}}{e ^ {- \lambda / 2} \sqrt {e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}}} \mathbf {z} _ {\lambda} (120) \\ = - \frac {\sqrt {e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}}}{\tilde {\sigma} _ {\mathrm {d a t a}}} \epsilon + \frac {e ^ {\lambda / 2} \left(e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2} - \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2} \alpha_ {\lambda} ^ {2}\right)}{\sqrt {e ^ {- \lambda} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}} \tilde {\sigma} _ {\mathrm {d a t a}} \alpha_ {\lambda}} \mathbf {z} _ {\lambda} (121) \\ \end{array}
$$
In summary, given these different parameterizations, the  $\epsilon$ -prediction loss can be written in terms of other parameterizations as follows:
$$
\begin{array}{l} \left| \left| \epsilon - \hat {\epsilon} _ {\boldsymbol {\theta}} \right| \right| _ {2} ^ {2} = e ^ {\lambda} \left| \left| \mathbf {x} - \hat {\mathbf {x}} _ {\boldsymbol {\theta}} \right| \right| _ {2} ^ {2} \quad (\epsilon - \text {p r e d i c t i o n a n d x - p r e d i c t i o n e r r o r}) (122) \\ = \sigma_ {\lambda} ^ {2} \| \nabla_ {\mathbf {z} _ {\lambda}} \log q (\mathbf {z} _ {\lambda} | \mathbf {x}) - \mathbf {s} _ {\boldsymbol {\theta}} \| _ {2} ^ {2} \quad \text {(s c o r e p r e d i c t i o n)} (123) \\ = \alpha_ {\lambda} ^ {- 2} \left(e ^ {- \lambda} + 1\right) ^ {- 2} \left\| \mathbf {v} - \hat {\mathbf {v}} _ {\boldsymbol {\theta}} \right\| _ {2} ^ {2} \quad (\mathbf {v} - \text {p r e d i c t i o n , g e n e r a l}) (124) \\ = \left(e ^ {- \lambda} + 1\right) ^ {- 1} \| \mathbf {v} - \hat {\mathbf {v}} _ {\boldsymbol {\theta}} \| _ {2} ^ {2} \quad (\mathbf {v} - \text {p r e d i c t i o n w i t h V P S D E}) (125) \\ = \left(e ^ {- \lambda} / \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2} + 1\right) ^ {- 1} \| \mathbf {F} - \hat {\mathbf {F}} _ {\boldsymbol {\theta}} \| _ {2} ^ {2} \quad (\mathbf {F} - \text {p r e d i c t i o n}) (126) \\ \end{array}
$$
Interestingly, if we set  $\tilde{\sigma}_{\mathrm{data}}^2 = 1$ , the training objectives of  $\mathbf{F}$ -prediction and  $\mathbf{v}$ -prediction are the same.
Noise schedule name	λ = fλ(t) = ...	t = f-1λ(λ) = ...	p(λ) = -d/dλf-1λ(λ) = ...
Cosine	-2 log(tan(πt/2))	(2/π) arctan(e-λ/2)	sech(λ/2)/(2π)
Shifted cosine	-2 log(tan(πt/2)) + 2s	(2/π) arctan(e-λ/2-s)	sech(λ/2-s)/(2π)
EDM (training)	-FN-1(t; 2.4, 2.42)	FN(-λ; 2.4, 2.42)	N(λ; 2.4, 2.42)
EDM (sampling)	-2ρ log(σ1/ρmax+(1-t)(σ1/ρmin-σ1/ρmax))	1 - e-λ/(2ρ) - σ1/ρmaxσ1/ρmin-σ1/ρmax	e-λ/(2ρ)/2ρ(σ1/ρmax-σ1/ρmin)
Flow Matching with OT (see D.3)	2 log((1-t)/t)	1/(1 + eλ/2)	sech2(λ/4)/8
Table 4: Noise schedules used in our experiments: cosine [Nichol and Dhariwal, 2021], shifted cosine [Hoogeboom et al., 2023], and EDM [Karras et al., 2022] training and sampling schedules. Note that these are the noise schedules before truncation (Section E.3.1).
E.3 Noise schedules
During model training, we sample time  $t$  uniformly:  $t \sim \mathcal{U}(0,1)$ , then compute  $\lambda = f_{\lambda}(t)$ . This results in a distribution over noise levels  $p(\lambda)$ , whose cumulative density function (CDF) is given by  $1 - f_{\lambda}^{-1}(\lambda)$ . For  $\lambda \in [\lambda_{\min}, \lambda_{\max}]$  the probability density function (PDF) is the derivative of the CDF, which is  $p(\lambda_t) = -(d / d\lambda)f_{\lambda}^{-1}(\lambda) = -dt / d\lambda = -1 / f_{\lambda}'(t)$ . Outside of  $[\lambda_{\min}, \lambda_{\max}]$  the probability density is 0.
In table 4 we provide some popular noise schedules: cosine [Nichol and Dhariwal, 2021], shifted cosine [Hoogeboom et al., 2023], and EDM [Karras et al., 2022]. We do not list the 'linear' schedule by [Ho et al., 2020],  $f_{\lambda}(t) = -\log (e^{t^2} - 1)$  which has fallen out of use. Note that:
- In the shifted cosine schedule,  $s = \log (64 / d)$ , where 64 is the base resolution and  $d$  is the used resolution (e.g. 128, 256, 512, etc.).
- In the EDM training schedule, the function  $F_{\mathcal{N}}(\lambda; \mu, \sigma^2)$  is the Normal distribution CDF, and  $\mathcal{N}(\lambda; \mu, \sigma^2)$  is its PDF.
- In the EDM sampling schedule,  $\rho = 7$ ,  $\sigma_{\mathrm{min}} = 0.002$ ,  $\sigma_{\mathrm{max}} = 80$ . The density function  $p(\lambda)$  in the table has support  $\lambda \in [-\log \sigma_{\mathrm{max}}^2, -\log \sigma_{\mathrm{min}}^2]$ . Outside this range,  $p(\lambda) = 0$ .
Note that the table gives the noise schedules before truncation, resulting in  $[\widetilde{f}_{\lambda}(0),\widetilde{f}_{\lambda}(1)] = [\infty , - \infty ]$  The truncation procedure is given in E.3.1.
E.3.1 Truncation
The noise schedules above are truncated, resulting in a noise schedule  $\widetilde{f}_{\lambda}(0)$  whose endpoints have desired values  $[\widetilde{f}_{\lambda}(0),\widetilde{f}_{\lambda}(1)] = [\lambda_{\max},\lambda_{\min}]$ :
$$
\widetilde {f} _ {\lambda} (t) := f _ {\lambda} \left(t _ {0} + \left(t _ {1} - t _ {0}\right) t\right) \tag {127}
$$
$$
\text {w h e r e :} t _ {0} := f _ {\lambda} ^ {- 1} \left(\lambda_ {\max }\right) \tag {128}
$$
$$
t _ {1} := f _ {\lambda} ^ {- 1} \left(\lambda_ {\min }\right) \tag {129}
$$
Its inverse is:
$$
\widetilde {f} _ {\lambda} ^ {- 1} (\lambda) = \left(f _ {\lambda} ^ {- 1} (\lambda) - t _ {0}\right) / \left(t _ {1} - t _ {0}\right) \tag {130}
$$
And the corresponding probability density:
$$
\text {i f} \lambda_ {\min } \leq \lambda \leq \lambda_ {\max }: \tilde {p} (\lambda) = - \frac {d}{d \lambda} \tilde {f} _ {\lambda} ^ {- 1} (\lambda) = - \frac {d}{d \lambda} f _ {\lambda} ^ {- 1} (\lambda) / \left(t _ {1} - t _ {0}\right) = p (\lambda) / \left(t _ {1} - t _ {0}\right) \tag {131}
$$
$$
\text {e l s e :} \quad \tilde {p} (\lambda) = 0 \tag {132}
$$
E.4 Sampling
Anderson [1982] showed that if  $\mathbf{s}_{\theta}(\mathbf{z};\lambda) = \nabla_{\mathbf{z}}\log q_t(\mathbf{z})$ , then the forward SDE is exactly reversed by the following SDE:
$$
d \mathbf {z} = [ \mathbf {f} (\mathbf {z}, t) - g (t) ^ {2} \mathbf {s} _ {\boldsymbol {\theta}} (\mathbf {z}; \lambda) ] d t + g (t) d \mathbf {w} \tag {133}
$$
(b)
Figure 4: Our proposed adaptive noise schedule allows us to freely change the weighting function without needing to handtune a corresponding noise schedule. For some models it leads to slightly worse convergence speed (left) compared to a static noise schedule, probably because the noise schedule was already well-tuned for the weighting function, while in other cases it led to faster convergence (right). The adaptive noise schedule did not significantly affect the end result, but it did allow us to more freely experiment with weighting functions.
Recent diffusion models have used increasingly sophisticated samplers. As an alternative to solving the SDE, [Song et al., 2021b] showed that sampling from the model can alternatively be done by solving the following probability flow ODE:
$$
d \mathbf {z} = [ \mathbf {f} (\mathbf {z}, t) - \frac {1}{2} g (t) ^ {2} \mathbf {s} _ {\boldsymbol {\theta}} (\mathbf {z}; \lambda) ] d t \tag {134}
$$
which, under the assumption that  $\mathbf{s}_{\theta}$  is a conservative vector field, will result in the same marginals  $p(\mathbf{z}_t)$  as the SDE of Equation 133 for every  $t \in [0,1]$ , and therefore also the same marginal  $p(\mathbf{x})$ .
Note that due to the continuous-time nature of the model, any sampling method is necessarily approximate, with the discretization error depending on various factors including the choice of noise schedule. For sampling we can therefore typically use a different noise schedule  $f_{\lambda}$  for sampling than for training, and we can change the SDE drift term; as long as we appropriately rescale the input to the score network, this would still result in correct samples.
F Adaptive noise schedule
The invariance shown in Section 3.2 holds for the loss  $\mathcal{L}_w(\mathbf{x})$ , but not for the Monte Carlo estimator of the loss that we use in training, based on random samples  $t \sim \mathcal{U}(0,1)$ ,  $\epsilon \sim \mathcal{N}(0,\mathbf{I})$ . The noise schedule still affects the variance of this Monte Carlo estimator and its gradients; therefore, the noise schedule affects the efficiency of optimization.
In fact, the noise schedule acts as an importance sampling distribution for estimating the loss integral of Equation 6. Specifically, note that  $p(\lambda) = -1 / (d\lambda /dt)$ . We can therefore rewrite the weighted loss as the following, which clarifies the role of  $p(\lambda)$  as an importance sampling distribution:
$$
\mathcal {L} _ {w} (\mathbf {x}) = \frac {1}{2} \mathbb {E} _ {\epsilon \sim \mathcal {N} (0, \mathbf {I}), \lambda \sim p (\lambda)} \left[ \frac {w (\lambda)}{p (\lambda)} | | \hat {\epsilon} _ {\theta} (\mathbf {z} _ {\lambda}; \lambda) - \epsilon | | _ {2} ^ {2} \right] \tag {135}
$$
In order to avoid having to hand-tune the noise schedule for different weighting functions, we implemented an adaptive noise schedule. The noise schedule  $\lambda_{t}$  is updated online, where we let  $p(\lambda)\propto \mathbb{E}_{\mathbf{x}\sim \mathcal{D},\epsilon \sim \mathcal{N}(0,\mathbf{I})}[w(\lambda)||\hat{\epsilon}_{\theta}(\mathbf{z}_{\lambda};\lambda) - \epsilon ||_2^2 ]$  . This noise schedule ensures that the loss is spread evenly over time, i.e. that the magnitude of the loss  $\mathbb{E}_{\mathbf{x}\sim \mathcal{D},\epsilon \sim \mathcal{N}(0,\mathbf{I})}\left[(w(\lambda) / p(\lambda))||\hat{\epsilon}_{\theta}(\mathbf{z}_{\lambda};\lambda) - \epsilon ||_2^2\right]$  is approximately invariant to  $\lambda$  or  $t$  . We find that this often significantly speeds up optimization.
We implemented an adaptive noise schedule  $p(\lambda)$ , where:
$$
p (\lambda) \propto \mathbb {E} _ {\mathbf {x} \sim \mathcal {D}, \epsilon \sim \mathcal {N} (0, \mathbf {I})} [ w (\lambda) | | \epsilon - \hat {\epsilon} _ {\boldsymbol {\theta}} (\mathbf {z} _ {\lambda}; \lambda) | | _ {2} ^ {2} ] \tag {136}
$$
In practice we approximate this by dividing the range  $[\lambda_{\mathrm{min}},\lambda_{\mathrm{max}}]$  into 100 evenly spaced bins, and during training keep an exponential moving average (EMA) of  $w(\lambda)||\epsilon -\hat{\epsilon}_{\theta}(\mathbf{z}_{\lambda};\lambda)||_2^2$  within each bin. From these EMAs we construct a piecewise linear function  $f_{\lambda}(t)$  such that Equation 136 is approximately satisfied. The EMAs and corresponding noise schedule  $p(\lambda)$  are updated at each training iteration. Similar noise schedules have been proposed in Zheng et al. [2023], Nichol and Dhariwal [2021], Dieleman et al. [2022].
In experiments we measure the effect of changing the fixed noise schedule of existing modules with an adaptive schedule. We found that this lead to approximately equal FID scores. In half of the experiments, optimization was approximately as fast as with the original noise schedule, while in the other half the adaptive noise schedule lead to faster optimization (see Figure 4). The end results were not significantly altered.
G Relationship between the KL divergence and Fisher divergence
We'll use the following definition of the Fisher divergence [Lyu, 2012]:
$$
D _ {F} (q (\mathbf {x}) \| p (\mathbf {x})) := \mathbb {E} _ {q} (\mathbf {x}) [ \| \nabla_ {\mathbf {x}} \log q (\mathbf {x}) - \nabla_ {\mathbf {x}} \log p (\mathbf {x}) \| _ {2} ^ {2} ] \tag {137}
$$
Theorem 2. Assume a model in the family specified in Section 2, and assume the score network encodes a conservative vector field:  $\mathbf{s}_{\theta}(\mathbf{z}_t,\lambda_t) = \nabla_{\mathbf{z}_t}\log p(\mathbf{z}_t)$  (not assumed by the other theorems). Then:
$$
\frac {d}{d \lambda} D _ {K L} \left(q \left(\mathbf {z} _ {t, \dots , 1} | \mathbf {x}\right) \| p \left(\mathbf {z} _ {t, \dots , 1}\right)\right) = \frac {1}{2} \sigma_ {\lambda} ^ {2} D _ {F} \left(q \left(\mathbf {z} _ {t} | \mathbf {x}\right) \| p \left(\mathbf {z} _ {t}\right)\right) \tag {138}
$$
Proof of Theorem 2. Note that (see Equation 117):
$$
\nabla_ {\mathbf {z} _ {t}} \log q (\mathbf {z} _ {t} | \mathbf {x}) = - \epsilon / \sigma_ {\lambda} \tag {139}
$$
And assume the score network encodes a conservative vector field:
$$
\nabla_ {\mathbf {z} _ {t}} \log p (\mathbf {z} _ {t}) = \mathbf {s} _ {\boldsymbol {\theta}} \left(\mathbf {z} _ {t}, \lambda_ {t}\right) = - \hat {\epsilon} _ {\boldsymbol {\theta}} \left(\mathbf {z} _ {t}; t\right) / \sigma_ {\lambda} \tag {140}
$$
So the time derivative of Equation 22 can be expressed as:
$$
\frac {d}{d \lambda} D _ {K L} (q (\mathbf {z} _ {t, \dots , 1} | \mathbf {x}) | | p (\mathbf {z} _ {t, \dots , 1})) = \frac {1}{2} \sigma_ {\lambda} ^ {2} \mathbb {E} _ {q (\mathbf {z} _ {t} | \mathbf {x})} \left[ | | \nabla_ {\mathbf {z} _ {t}} \log q (\mathbf {z} _ {t} | \mathbf {x}) - \nabla_ {\mathbf {z} _ {t}} \log p (\mathbf {z} _ {t}) | | _ {2} ^ {2} \right] \tag {141}
$$
Equation 138 follows from the definition of the Fisher divergence.
G.1 Comparison with Theorem 1 by Lyu [2012]
Lyu [2012] prove a similar result in their Theorem 1. We'll translate their result into our notation. In particular, let the forward process be as in our family, such that  $q(\mathbf{z}_t|\mathbf{x}) = \mathcal{N}(\mathbf{z}_t;\alpha_\lambda \mathbf{x},\sigma_t^2\mathbf{I})$ . The marginal (data) distribution is  $q(\mathbf{x})$ , such that  $q(\mathbf{z}_t) = \int q(\mathbf{z}_t|\mathbf{x})q(\mathbf{x})d\mathbf{x}$ . Similarly, let the generative model have a marginal  $p(\mathbf{x})$ , and  $p(\mathbf{z}_t) = \int p(\mathbf{z}_t|\mathbf{x})p(\mathbf{x})d\mathbf{x}$ . So far the assumptions are the same in our family.
They assume that  $q(\mathbf{z}_t|\mathbf{x}) = \mathcal{N}(\mathbf{z}_t,\mathbf{x},t)$ , which corresponds to a variance exploding (VE) diffusion process, with  $t = \sigma_{\lambda}^{2}$ , so  $\lambda = -\log(t)$ . Importantly, they make the assumption that  $p(\mathbf{z}_t|\mathbf{x}) = q(\mathbf{z}_t|\mathbf{x})$ , i.e. that the forward process for  $p$  equals the forward process for  $q$ . Given these assumptions, Lyu [2012] show that:
$$
\frac {d}{d t} D _ {K L} \left(q \left(\mathbf {z} _ {t}\right) \| p \left(\mathbf {z} _ {t}\right)\right) = - \frac {1}{2} D _ {F} \left(q \left(\mathbf {z} _ {t}\right) \| p \left(\mathbf {z} _ {t}\right)\right) \tag {142}
$$
Which, given the noise schedule  $\lambda = -\log (t)$ , can be rewritten as:
$$
\frac {d}{d \lambda} D _ {K L} \left(q \left(\mathbf {z} _ {t}\right) \| p \left(\mathbf {z} _ {t}\right)\right) = \frac {1}{2} \sigma_ {\lambda} ^ {2} D _ {F} \left(q \left(\mathbf {z} _ {t}\right) \| p \left(\mathbf {z} _ {t}\right)\right) \tag {143}
$$
which looks a lot like our Equation 138. One difference are that in Equation 138, the left-hand-side distributions are joint distributions, and  $q$  conditions on  $\mathbf{x}$ , while Equation 143 is about the
unconditional  $q$ . Another key difference is that for Equation 138 we need fewer assumptions: most importantly, we do not make the assumption that  $p(\mathbf{z}_t|\mathbf{x}) = q(\mathbf{z}_t|\mathbf{x})$ , since this assumption does not hold for the family of diffusion models we consider. Before or during optimization,  $p(\mathbf{z}_t|\mathbf{x})$  might be very far from  $q(\mathbf{z}_t|\mathbf{x})$ . After optimization,  $p(\mathbf{z}_t|\mathbf{x})$  might be close to  $q(\mathbf{z}_t|\mathbf{x})$ , but we still can't assume they're equal. In addition, we're mostly interested in the properties of the loss function during optimization, since that's when we're using our loss for optimization. We for this reason, our Theorem 2 is a lot more relevant for optimization.
Instead of uniformly sampling  $t$ , we applied the low-discrepancy sampler of time that was proposed by Kingma et al. [2021], which has been shown to effectively reduce the variance of diffusion loss estimator and lead to faster optimization. The model is optimized by Adam [Kingma and Ba, 2014] with the default hyperparameter settings. We clipped the learning gradient with a global norm of 1.
For the adaptive noise schedules, we divided the range of  $[\lambda_{\mathrm{min}},\lambda_{\mathrm{max}}]$  into 100 evenly spaced bins. During training, we maintained an exponential moving average of  $w(\lambda)||\epsilon -\hat{\epsilon}_{\pmb{\theta}}(\mathbf{z}_{\lambda};\lambda)||_2^2$  with a decay rate 0.999, and a constant initialization value of 1 for each bin.
Below we elaborate the implementation details specific for each task.
ImageNet 64x64. For class-conditional generation on ImageNet 64x64, we applied the ADM U-Net architecture from Dhariwal and Nichol [2022], with dropout rate 0.1. We didn't use any data augmentation. The model was trained with learning rate  $1e - 4$ , exponential moving average of 50 million images and learning rate warmup of 10 million images, which mainly follows the configuration of Karras et al. [2022]. We employed 128 TPU-v4 chips with a batch size of 4096 (32 per chip). We trained the model for 700k iterations and reported the performance of the checkpoint giving the best FID score (checkpoints were saved and evaluated on every 20k iterations). It took around 3 days for a single training run. For training noise schedule and sampling noise schedule of DDPM sampler, we set  $\lambda_{\mathrm{min}} = -20$  and  $\lambda_{\mathrm{max}} = 20$ . We fixed the noise schedule used in sampling to the cosine schedule for the DDPM sampler, and the EDM (sampling) schedule for the EDM sampler (see Table 4 for the formulations). We adopted the same hyperparameters of EDM sampler from Karras et al. [2022] with no changes (i.e., Table 5 in their work, column 'ImageNet-Our model'). Both DDPM and EDM samplers took 256 sampling steps.
ImageNet 128x128. For class-conditional generation on ImageNet 128x128, we heavily followed the setting of simple diffusion [Hoogeboom et al., 2023]. Specifically, we used their 'U-ViT, L' architecture, and followed their learning rate and EMA schedules. The data was augmented with random horizontal flip. The model is trained using 128 TPU-v4 chips with a batch size of 2048 (16 per chip). We trained the model for 700 iterations and evaluated the FID and inception scores every 100k iterations. The results were reported with the checkpoint giving the best FID score. It took around 7 days for a single run. We set  $\lambda_{\mathrm{min}} = -15 + s$  and  $\lambda_{\mathrm{max}} = 15 + s$ , where  $s = \log(64 / d)$  is the shift of the weighting function, with 64 being the base resolution and  $d$  being the model resolution ( $d = 128$  for this task). The DDPM sampler used for evaluation used 'shifted-cosine' noise schedule (Table 4) and took 512 sampling steps.
I Relationship with low-bit training
Various earlier work, such as [Kingma and Dhariwal, 2018], found that maximum likelihood training on 5-bit data can lead to perceptually higher visual quality than training on 8-bit data (at the cost of a decrease in color fidelity). This leads to improved visual quality, probably because it allows the model to spend more capacity on modeling the bits that are most relevant for human perception.
In [Kingma and Dhariwal, 2018], training on 5-bit data was performed by adding uniform noise to the data, before feeding it to the model. It was found that adding Gaussian noise had a similar effect as uniform noise. As we have seen, in the case of diffusion models, adding Gaussian noise is equivalent to using a weighted objective with a monotonic weighting function.
Therefore, training on 5-bit data is similar to training using a monotonic weighting function in case of diffusion models. We can wonder: which weighting function emulates training on 5-bit data? Here, we'll attempt to answer this question.
I.1 The shape of  $\frac{d}{d\lambda}\mathcal{L}(\lambda ;\mathbf{x})$  for low-bit data
Note that the results of Appendix A.1 can also be written as:
$$
\frac {d}{d \lambda} \mathcal {L} (\lambda ; \mathbf {x}) = \frac {1}{2} \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} \left[ \left| \left| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} _ {\boldsymbol {\theta}} (\mathbf {z} _ {\lambda}; \lambda) \right| \right| _ {2} ^ {2} \right] \tag {144}
$$
This allows us to rewrite the weighted loss as simply:
$$
\mathcal {L} _ {w} (\mathbf {x}) = \int_ {\lambda_ {\max }} ^ {\lambda_ {\min }} \frac {d}{d \lambda} \mathcal {L} (\lambda ; \mathbf {x}) w (\lambda) d \lambda \tag {145}
$$
To understand the effect of  $\lambda$ , we can plot the  $\frac{d}{d\lambda}\mathcal{L}(\lambda ;\mathbf{x})$  as a function of  $\lambda$ .
We'd like to plot  $\frac{d}{d\lambda}\mathcal{L}(\lambda ;\mathbf{x})$  dof different choices of bit precision. This will tell us where the different bits 'live' as a function of  $\lambda$ . Since training different diffusion models on different bit precisions is very expensive, we instead use an approximation. In particular, we assume that the data  $\mathbf{x}$  is univariate, with a uniform distribution  $q(\mathbf{x})$  over the  $2^{n}$  possible pixel values, where  $n$  is the bit precision. The data is  $\mathbf{x}$  is, as usual, normalized to  $[-1,1]$ . We then let the model  $p(\mathbf{z}_{\lambda})$ , for each choice of  $\lambda$ , be the optimal model:  $p(\mathbf{z}_{\lambda}) := \int q(\mathbf{x})q(\mathbf{z}_{\lambda}|\mathbf{x})d\mathbf{z}$ , which is a univariate mixture-of-Gaussians, where each mixture component is a Gaussian centered one of the  $2^{n}$  possible pixel values. In this case,  $\mathcal{L}(\lambda ;\mathbf{x}) := D_{KL}(q(\mathbf{z}_{\lambda}|\mathbf{x})||p(\mathbf{z}_{\lambda}))$ .