SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
w \left(\lambda_ {\min }\right) \mathcal {L} (1) + \int_ {t = 1} ^ {t = 0} w \left(\lambda_ {t}\right) d \mathcal {L} (t; \mathbf {x}) = w \left(\lambda_ {\max }\right) \mathcal {L} (0) + \int_ {t = 0} ^ {t = 1} \mathcal {L} (t; \mathbf {x}) d w \left(\lambda_ {t}\right) \tag {31}
$$
The first LHS term equals a weighted prior loss term, and the second LHS term equals the weighted diffusion loss. From a geometric perspective, the two LHS terms together define an area that equals the area given by the right term, as illustrated in the figure below.
On the left, we have a rectangular area that equals a weighted prior loss  $w(\lambda_{\mathrm{min}}) \mathcal{L}(1)$ , plus a curved area equal to the weighted diffusion loss  $\int_{t=1}^{t=0} w(\lambda_t) d\mathcal{L}(t; \mathbf{x})$ . This integral can be intuitively understood as a Riemann sum over many tiny rectangles going from left ( $t = 1$ ) to right ( $t = 0$ ), each with height  $w(\lambda_t)$  and width  $d\mathcal{L}(t; \mathbf{x})$ . On the right, we have the same total area, but divided up into two different subareas: a rectangular area  $w(\lambda_{\mathrm{max}}) \mathcal{L}(0)$  and a curved area that equals the integral  $\int_{t=0}^{t=1} \mathcal{L}(t; \mathbf{x}) dw(\lambda_t)$  going upwards from  $t = 0$  to  $t = 1$ , which can also be intuitively understood as another Riemann sum, with each tiny rectangle having width  $\mathcal{L}(t; \mathbf{x})$  and height  $dw(\lambda_t)$ . The area of each of those tiny rectangles on the right can be understood as the ELBO at each noise level,  $\mathcal{L}(t; \mathbf{x})$ , times the weight of the ELBO at each noise level,  $dw(\lambda_t)$ .
C Relationship between  $D_{KL}(q(\mathbf{z}_{t,\dots,1}|\mathbf{x})||p(\mathbf{z}_{t,\dots,1}))$  and the ELBO
First, note that:
$$
\mathcal {L} (t; \mathbf {x}) = D _ {K L} \left(q \left(\mathbf {z} _ {t, \dots , 1} | \mathbf {x}\right) \| p \left(\mathbf {z} _ {t, \dots , 1}\right)\right) \geq D _ {K L} \left(q \left(\mathbf {z} _ {t} | \mathbf {x}\right) \| p \left(\mathbf {z} _ {t}\right)\right) \tag {32}
$$
More precisely, the joint KL divergence  $D_{KL}(q(\mathbf{z}_{t,\dots,1}|\mathbf{x})||p(\mathbf{z}_{t,\dots,1}))$  is the expected negative ELBO of noise-perturbed data, plus a constant entropy term:
$$
\mathcal {L} (t; \mathbf {x}) = D _ {K L} \left(q \left(\mathbf {z} _ {t, \dots , 1} | \mathbf {x}\right) \| p \left(\mathbf {z} _ {t, \dots , 1}\right)\right) = - \mathbb {E} _ {q \left(\mathbf {z} _ {t} \mid \mathbf {x}\right)} \left[ \operatorname {E L B O} _ {t} \left(\mathbf {z} _ {t}\right) \right] - \underbrace {\mathcal {H} \left(q \left(\mathbf {z} _ {t} \mid \mathbf {x}\right)\right)} _ {\text {c o n s t a n t}} \tag {33}
$$
where the ELBO of noise-perturbed data is:
$$
\operatorname {E L B O} _ {t} \left(\mathbf {z} _ {t}\right) := \mathbb {E} _ {q \left(\tilde {\mathbf {z}} _ {t} \mid \mathbf {z} _ {t}\right)} \left[ \log p \left(\mathbf {z} _ {t}, \tilde {\mathbf {z}} _ {t}\right) - \log q \left(\tilde {\mathbf {z}} _ {t} \mid \mathbf {z} _ {t}\right) \right] \tag {34}
$$
$$
\leq \log p \left(\mathbf {z} _ {t}\right) \tag {35}
$$
where  $\tilde{\mathbf{z}}_t\coloneqq \mathbf{z}_{t + dt,\ldots ,1}$
So,  $\mathcal{L}(t;\mathbf{x})$  is the expected negative ELBO of noise-perturbed data  $\mathbf{z}_t$ :
$$
\mathcal {L} (t; \mathbf {x}) = - \mathbb {E} _ {q \left(\mathbf {z} _ {t} \mid \mathbf {x}\right)} \left[ \operatorname {E L B O} _ {t} \left(\mathbf {z} _ {t}\right) \right] + \text {c o n s t a n t} \geq - \mathbb {E} _ {q \left(\mathbf {z} _ {t} \mid \mathbf {x}\right)} \left[ \log p \left(\mathbf {z} _ {t}\right) \right] + \text {c o n s t a n t} \tag {36}
$$
Therefore, the expression of the weighted loss in Equation 12 can be rewritten as:
$$
\mathcal {L} _ {w} (\mathbf {x}) = \mathbb {E} _ {p _ {w} (t)} [ \mathcal {L} (t; \mathbf {x}) ] + \text {c o n s t a n t} \tag {37}
$$
$$
= - \underbrace {\mathbb {E} _ {p _ {w} (t) , q (\mathbf {z} _ {t} | \mathbf {x})} \left[ \mathrm {E L B O} _ {t} (\mathbf {z} _ {t}) \right]} _ {\text {E L B O o f n o i s e - p e r t u r b e d d a t a}} + \text {c o n s t a n t} \tag {38}
$$
$$
\geq - \underbrace {\mathbb {E} _ {p _ {w} (t) , q \left(\mathbf {z} _ {t} \mid \mathbf {x}\right)} \left[ \log p \left(\mathbf {z} _ {t}\right) \right]} _ {\text {L o g - l i k e l i h o o d o f n o i s e - p e r t u r b e d d a t a}} + \text {c o n s t a n t} \tag {39}
$$
where  $w(\lambda_{\mathrm{min}})$  is constant w.r.t. the diffusion model parameters. Therefore, minimizing  $\mathcal{L}_w(\mathbf{x})$  is equivalent to maximizing this expected ELBO of noise-perturbed data.
D Derivation of weighting functions for previous works
The loss function used in previous works are equivalent to the weighted loss with a certain choice of noise schedule and weighting function  $w(\lambda)$ . In this section, we derive these weighting functions  $w(\lambda)$ .
D.1 'Elucidating Diffusion Models' (EDM) [Karras et al., 2022]
Karras et al. [2022] proposes the following training objective:
$$
\mathcal {L} _ {\mathrm {e d m}} (\mathbf {x}) = \mathbb {E} _ {\tilde {\sigma} \sim p (\tilde {\sigma}), \epsilon \sim \mathcal {N} (0, \mathbf {I})} \left[ \tilde {w} (\tilde {\sigma}) \| \mathbf {x} - \hat {\mathbf {x}} _ {\boldsymbol {\theta}} \left(\mathbf {z} _ {\tilde {\sigma}}; \tilde {\sigma}\right) \| _ {2} ^ {2} \right], \tag {40}
$$
where  $p(\tilde{\sigma})$  and  $\tilde{w} (\tilde{\sigma})$  are defined as:
$$
p (\log \tilde {\sigma}) = \mathcal {N} (\log \tilde {\sigma}; P _ {\text {m e a n}}, P _ {\text {s t d}} ^ {2}), \tag {41}
$$
$$
\tilde {w} (\tilde {\sigma}) = \left(\tilde {\sigma} ^ {2} + \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}\right) / \left(\tilde {\sigma} ^ {2} \cdot \tilde {\sigma} _ {\mathrm {d a t a}} ^ {2}\right), \tag {42}
$$
where  $\tilde{\sigma}$  is equivalent to the standard deviation of noise added to the clean sample  $\mathbf{x}$  in the VP SDE case, so  $\tilde{\sigma}^2 = e^{-\lambda}$ . Karras et al. [2022] used hyperparameters  $P_{\mathrm{mean}} = -1.2$ ,  $P_{\mathrm{std}} = 1.2$  and  $\tilde{\sigma}_{\mathrm{data}} = 0.5$ .
This can be rewritten in terms of  $\lambda = -2\log \tilde{\sigma}$  as:
$$
\mathcal {L} _ {\mathrm {e d m}} (\mathbf {x}) = \mathbb {E} _ {p (\lambda), \boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} \left[ \tilde {w} (\lambda) \| \mathbf {x} - \hat {\mathbf {x}} _ {\boldsymbol {\theta}} (\mathbf {z} _ {\lambda}; \lambda) \| _ {2} ^ {2} \right], \tag {43}
$$
where  $p(\lambda)$  and  $\tilde{w} (\lambda)$  are defined as
$$
p (\lambda) = \mathcal {N} (\lambda ; 2. 4, 2. 4 ^ {2}), \tag {44}
$$
$$
\tilde {w} (\lambda) = \left(e ^ {- \lambda} + 0. 5 ^ {2}\right) / \left(e ^ {- \lambda} \cdot 0. 5 ^ {2}\right), \tag {45}
$$
Comparing Equation 45 with the weighted loss expressed in terms of x-prediction parameterization (Section E.2), we see that (ignoring the constant scaling factor  $1 / 2$ ) the EDM objective is a special case of the weighted loss with weighting function:
$$
w (\lambda) = p (\lambda) e ^ {- \lambda} \cdot \tilde {w} (\lambda) \tag {46}
$$
$$
= p (\lambda) \left(e ^ {- \lambda} + 0. 5 ^ {2}\right) / 0. 5 ^ {2}. \tag {47}
$$
where the division by  $\tilde{\sigma}_{\mathrm{data}}^2$  can be ignored since it's constant. This leads to:
$$
w (\lambda) = \mathcal {N} (\lambda ; 2. 4, 2. 4 ^ {2}) \left(e ^ {- \lambda} + 0. 5 ^ {2}\right) \tag {48}
$$
D.2 The v-prediction loss / 'SNR+1'-weighting [Salimans and Ho, 2022]
Salimans and Ho [2022] introduced the v-parameterization, with a v-prediction model  $\hat{\mathbf{v}}$ , where:
$$
\mathbf {v} := \alpha_ {\lambda} \boldsymbol {\epsilon} - \sigma_ {\lambda} \mathbf {x} \tag {49}
$$
$$
\hat {\mathbf {v}} := \alpha_ {\lambda} \hat {\epsilon} - \sigma_ {\lambda} \hat {\mathbf {x}} \tag {50}
$$
They propose to minimize a v-prediction loss,  $\mathbb{E}[||\mathbf{v} - \hat{\mathbf{v}} ||_2^2 ]$  . Note that  $\hat{\epsilon} = (\mathbf{z}_{\lambda} - \alpha_{\lambda}\hat{\mathbf{x}}) / \sigma_{\lambda}$  . For our general family, this implies:
$$
\left| \left| \mathbf {v} - \hat {\mathbf {v}} \right| \right| _ {2} ^ {2} = \sigma_ {\lambda} ^ {2} \left(e ^ {\lambda} + 1\right) ^ {2} \left| \left| \mathbf {x} - \hat {\mathbf {x}} \right| \right| _ {2} ^ {2} \tag {51}
$$
$$
= \alpha_ {\lambda} ^ {2} \left(e ^ {- \lambda} + 1\right) ^ {2} \left\| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} \right\| _ {2} ^ {2} \tag {52}
$$
In the special case of the variance preserving (VP) SDE, this simplifies to:
$$
\left\| \mathbf {v} - \hat {\mathbf {v}} \right\| _ {2} ^ {2} = (e ^ {\lambda} + 1) \| \mathbf {x} - \hat {\mathbf {x}} \| _ {2} ^ {2} \tag {53}
$$
$$
= \left(e ^ {- \lambda} + 1\right) \| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} \| _ {2} ^ {2}. \tag {54}
$$
Since the  $\epsilon$ -prediction loss corresponds to minimizing the weighted loss with  $w(\lambda) = p(\lambda)$ , the  $\mathbf{v}$ -prediction loss corresponds to minimizing the weighted loss with  $w(\lambda) = (e^{-\lambda} + 1)p(\lambda)$ .
Note that Salimans and Ho [2022] view the loss from the  $\mathbf{x}$ -prediction viewpoint, instead of our ELBO viewpoint; so in their view, minimizing simply  $||\mathbf{x} - \hat{\mathbf{x}} ||_2$  means no weighting. Note that  $e^{\lambda}$  is the signal-to-noise ratio (SNR). Since  $||\pmb {\epsilon} - \hat{\pmb{\epsilon}} ||_2^2 = e^\lambda ||\mathbf{x} - \hat{\mathbf{x}} ||_2^2$ , they call the  $\pmb{\epsilon}$ -prediction loss 'SNR weighting', and since  $||\mathbf{v} - \hat{\mathbf{v}} ||_2^2 = (e^\lambda +1)||\mathbf{x} - \hat{\mathbf{x}} ||_2^2$ , they call this 'SNR+1'-weighting.
Salimans and Ho [2022] propose to use optimize a VP SDE with a cosine schedule  $p(\lambda) = \mathrm{sech}(\lambda / 2) / (2\pi) = 1 / (2\pi \cosh (-\lambda / 2))$  and the v-prediction loss:  $\mathbb{E}[||\mathbf{v} - \hat{\mathbf{v}}||_2^2]$ . This corresponds to minimizing the weighted loss with:
$$
w (\lambda) = \left(e ^ {- \lambda} + 1\right) p (\lambda) \tag {55}
$$
$$
= \left(e ^ {- \lambda} + 1\right) / \left(2 \pi \cosh (- \lambda / 2)\right) \tag {56}
$$
$$
= \pi e ^ {- \lambda / 2} \tag {57}
$$
The factor  $\pi$  can be ignored since it's constant, so we can equivalently use:
$$
w (\lambda) = e ^ {- \lambda / 2} \tag {58}
$$
D.2.1 With shifted cosine schedule
[Hoogeboom et al., 2023] extended the cosine schedule to a shifted version:  $p(\lambda) = \mathrm{sech}(\lambda / 2 - s) / (2\pi)$ , where  $s = \log(64 / d)$ , where 64 is the base resolution and  $d$  is the model resolution (e.g. 128, 256, 512, etc.). In this case the weighting is:
$$
w (\lambda) = \left(e ^ {- \lambda} + 1\right) p (\lambda) \tag {59}
$$
$$
= (2 / \pi) e ^ {- s} e ^ {- \lambda / 2} \tag {60}
$$
Since  $(2 / \pi)e^{-s}$  is constant w.r.t  $\lambda$ , the weighting is equivalent to the weighting for the unshifted cosine schedule.
D.3 Flow Matching with the Optimal Transport flow path (FM-OT) [Lipman et al., 2022]
Flow Matching [Lipman et al., 2022] with the Optimal Transport flow path can be seen as a special case of Gaussian diffusion with the weighted loss.
D.3.1 Noise schedule
Note that in [Lipman et al., 2022], time goes from 1 to 0 as we go forward in time. Here, we'll let time go from 0 to 1 as we go forward in time, consistent with the rest of this paper. We'll also assume
$\sigma_0 = 0$ , for which we can later correct by truncation (see Section E.3.1). In this model, the forward process  $q(\mathbf{z}_t|\mathbf{x})$  is defined by:
$$
\mathbf {z} _ {t} = \alpha_ {t} \mathbf {x} + \sigma_ {t} \epsilon \tag {61}
$$
$$
= (1 - t) \mathbf {x} + t \epsilon \tag {62}
$$
This implies that the log-SNR is given by:
$$
\lambda_ {t} = f _ {\lambda} (t) = \log \left(\alpha_ {t} ^ {2} / \sigma_ {t} ^ {2}\right) \tag {63}
$$
$$
= 2 \log ((1 - t) / t) \tag {64}
$$
Its inverse is given by:
$$
t = f _ {\lambda} ^ {- 1} (\lambda) = 1 / \left(1 + e ^ {\lambda / 2}\right) \tag {65}
$$
$$
= \operatorname {s i g m o i d} (- \lambda / 2) \tag {66}
$$
The derivative, as a function of  $t$ , is:
$$
\frac {d \lambda}{d t} = \frac {d}{d t} f _ {\lambda} (t) = 2 / (- t + t ^ {2}) \tag {67}
$$
This derivative of its inverse, as a function of  $\lambda$ , is:
$$
\frac {d t}{d \lambda} = \frac {d}{d \lambda} f _ {\lambda} ^ {- 1} (\lambda) = \frac {d}{d \lambda} \operatorname {s i g m o i d} (- \lambda / 2) = - \operatorname {s e c h} ^ {2} (\lambda / 4) / 8 \tag {68}
$$
The corresponding density is
$$
p (\lambda) = - \frac {d}{d \lambda} f _ {\lambda} ^ {- 1} (\lambda) = \operatorname {s e c h} ^ {2} (\lambda / 4) / 8 \tag {69}
$$
which is a Logistic distribution; see also Table 4.
D.3.2 Score function parameterization and loss function
Lipman et al. [2022] then propose the following generative model ODE:
$$
d \mathbf {z} = - \hat {\mathbf {o}} \left(\mathbf {z} _ {t}, t\right) d t \tag {70}
$$
The model is then optimized with the Conditional flow matching (CFM) loss:
$$
\mathcal {L} _ {\mathrm {C F M}} (\mathbf {x}) = \mathbb {E} _ {t \sim \mathcal {U} (0, 1), \boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} [ \| \mathbf {o} - \hat {\mathbf {o}} \| _ {2} ^ {2} ] \tag {71}
$$
where they use the parameterization:
$$
\mathbf {o} := \mathbf {x} - \epsilon \tag {72}
$$
D.3.3 Weighting function
What is the weighting function  $w(\lambda)$  corresponding to this loss? Note that this parameterization means that:
$$
\mathbf {z} _ {t} = (1 - t) \mathbf {x} + t \epsilon \tag {73}
$$
$$
= (1 - t) \mathbf {o} + \epsilon \tag {74}
$$
$$
\mathbf {o} = \left(\mathbf {z} _ {t} - \epsilon\right) / (1 - t) \tag {75}
$$
Since  $t = 1 / (1 + e^{\lambda /2})$ , we have that  $1 / (1 - t) = 1 + e^{-\lambda /2}$ , so parameterized as a function of  $\lambda$ , we have:
$$
\mathbf {o} = \left(\mathbf {z} _ {\lambda} - \boldsymbol {\epsilon}\right) \left(1 + e ^ {- \lambda / 2}\right) \tag {76}
$$
Likewise, we can parameterize o-prediction in terms of  $\epsilon$ -prediction:
$$
\hat {\mathbf {o}} (\mathbf {z} _ {\lambda}, \lambda) = (\mathbf {z} _ {\lambda} - \hat {\epsilon} (\mathbf {z} _ {\lambda}, \lambda)) (1 + e ^ {- \lambda / 2}) \tag {77}
$$
We can translate the o-prediction loss to a  $\epsilon$ -prediction loss:
$$
\left| \left| \mathbf {o} - \hat {\mathbf {o}} \left(\mathbf {z} _ {\lambda}, \lambda\right) \right| \right| _ {2} ^ {2} = \left(1 + e ^ {- \lambda / 2}\right) ^ {2} \left| \left| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} \left(\mathbf {z} _ {\lambda}, \lambda\right) \right| \right| _ {2} ^ {2} \tag {78}
$$
Therefore, combining the derivations above, the CFM loss, formulated in terms of the  $\lambda$  parameterization instead of  $t$ , and in terms of the  $\epsilon$ -prediction parameterization instead of the o-prediction parameterization, is:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {C F M}} (\mathbf {x}) = \mathbb {E} _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (0, \mathbf {I})} [ \| \mathbf {o} - \hat {\mathbf {o}} (\mathbf {z} _ {t}, t) \| _ {2} ^ {2} ] (79) \\ = \int_ {0} ^ {1} \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} [ \| \mathbf {o} - \hat {\mathbf {o}} (\mathbf {z} _ {t}, t) \| _ {2} ^ {2} ] d t (80) \\ = \int_ {\lambda_ {\min }} ^ {\lambda_ {\max }} - \frac {d t}{d \lambda} \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} [ \| \mathbf {o} - \hat {\mathbf {o}} (\mathbf {z} _ {\lambda}, \lambda) \| _ {2} ^ {2} ] d \lambda (81) \\ = \int_ {\lambda_ {\min }} ^ {\lambda_ {\max }} \left(\operatorname {s e c h} ^ {2} (\lambda / 4) / 8\right) \mathbb {E} _ {\epsilon \sim \mathcal {N} (0, \mathbf {I})} \left[ \left\| \mathbf {o} - \hat {\mathbf {o}} \left(\mathbf {z} _ {\lambda}, \lambda\right) \right\| _ {2} ^ {2} \right] d \lambda (82) \\ = \int_ {\lambda_ {\min }} ^ {\lambda_ {\max }} \left(\operatorname {s e c h} ^ {2} (\lambda / 4) / 8\right) \left(1 + e ^ {- \lambda / 2}\right) ^ {2} \mathbb {E} _ {\epsilon \sim \mathcal {N} (0, \mathbf {I})} \left[ \left| \left| \epsilon - \hat {\epsilon} \left(\mathbf {z} _ {\lambda}, \lambda\right) \right| \right| _ {2} ^ {2} \right] d \lambda (83) \\ = \frac {1}{2} \int_ {\lambda_ {\min }} ^ {\lambda_ {\max }} w (\lambda) \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} \left[ \left\| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} \left(\mathbf {z} _ {\lambda}, \lambda\right) \right\| _ {2} ^ {2} \right] d \lambda (84) \\ = \frac {1}{2} \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I}), \lambda \sim \tilde {p} (\lambda)} \left[ \frac {w (\lambda)}{\tilde {p} (\lambda)} | | \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} (\mathbf {z} _ {\lambda}, \lambda) | | _ {2} ^ {2} \right] (85) \\ \end{array}
$$
where  $\tilde{p}$  is any distribution with full support on  $[\lambda_{\min},\lambda_{\max}]$ , and where:
$$
\begin{array}{l} w (\lambda) = 2 \left(\operatorname {s e c h} ^ {2} (\lambda / 4) / 8\right) \left(1 + e ^ {- \lambda / 2}\right) ^ {2} (86) \\ = e ^ {- \lambda / 2} (87) \\ \end{array}
$$
Therefore, this weighting is equivalent to the weighting for the v-prediction loss with cosine schedule (Section D.2): the CFM loss is equivalent to the v-prediction loss with cosine schedule.
D.4 Inversion by Direct Iteration (InDI) [Delbracio and Milanfar, 2023]
Delbracio and Milanfar [2023] propose Inversion by Direct Iteration (InDI). Their forward process is identical to the forward process of FM-OT [Lipman et al., 2022] introduced in Section D.3:
$$
\mathbf {z} _ {t} = (1 - t) \mathbf {x} + t \boldsymbol {\epsilon} \tag {88}
$$
As derived in Section D.3 above, this means that the distribution over log-SNR  $\lambda$  is the Logistic distribution:  $p(\lambda) = \mathrm{sech}^2 (\lambda /4) / 8$ . The proposed loss function is the x-prediction loss:
$$
\mathcal {L} _ {\text {I n D I}} (\mathbf {x}) = \mathbb {E} _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (0, \mathbf {I})} [ \| \mathbf {x} - \hat {\mathbf {x}} (\mathbf {z} _ {t}, t) \| _ {2} ^ {2} ] \tag {89}
$$
Since  $||\mathbf{x} - \hat{\mathbf{x}} ||_2^2 = e^{-\lambda}||\boldsymbol {\epsilon} - \hat{\boldsymbol{\epsilon}} ||_2^2$ , and the  $\pmb{\epsilon}$ -prediction loss corresponds to minimizing the weighted loss with  $w(\lambda) = p(\lambda)$ , the  $\mathbf{x}$ -prediction loss above corresponds to minimizing the weighted loss with:
$$
\begin{array}{l} w (\lambda) = e ^ {- \lambda} p (\lambda) (90) \\ = e ^ {- \lambda} \operatorname {s e c h} ^ {2} (\lambda / 4) / 8 (91) \\ \end{array}
$$
Which is a slightly different weighting then the FM-OT weighting, giving a bit more weighting to lower noise levels.
D.5 Perception prioritized weighting (P2 weighting) [Choi et al., 2022]
Choi et al. [2022] proposed a new weighting function:
$$
w (\lambda) = \frac {- d t / d \lambda}{(k + e ^ {\lambda}) ^ {\gamma}} = \frac {p (\lambda)}{(k + e ^ {\lambda}) ^ {\gamma}}, \tag {92}
$$
where empirically they set  $k = 1$  and  $\gamma$  as either 0.5 or 1. Compared to the  $\epsilon$ -prediction objective, where  $w(\lambda) = dt / d\lambda = p(\lambda)$ , this objective put more emphasis on the middle regime of the whole noise schedule, which Choi et al. [2022] hypothesized to be the most important regime for creating content that is sensitive to visual perception. When combined with the most commonly used cosine noise schedule [Nichol and Dhariwal, 2021], the weighting function becomes  $w(\lambda) = \mathrm{sech}(\lambda /2) / (1 + e^{\lambda})^{\gamma}$ .
D.6 Min-SNR-  $\gamma$  weighting [Hang et al., 2023]
Hang et al. [2023] proposed the following training objective:
$$
\begin{array}{l} \mathcal {L} _ {\text {M i n S N R}} (\mathbf {x}) = \mathbb {E} _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (0, \mathbf {I})} \left[ \min  \left\{e ^ {\lambda}, \gamma \right\} \| \mathbf {x} - \hat {\mathbf {x}} \left(\mathbf {z} _ {t}; \lambda\right) \| _ {2} ^ {2} \right] (93) \\ = E _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (0, \mathbf {I})} \left[ \min  \{1, \gamma e ^ {- \lambda} \} \| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} (\mathbf {z} _ {t}; \lambda) \| _ {2} ^ {2} \right] (94) \\ = E _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (0, \mathbf {I})} \left[ \min  \{1, \gamma e ^ {- \lambda} \} \cdot - \frac {d t}{d \lambda} \cdot - \frac {d \lambda}{d t} \| \epsilon - \hat {\epsilon} (\mathbf {z} _ {t}; \lambda) \| _ {2} ^ {2} \right] (95) \\ = E _ {t \sim \mathcal {U} (0, 1), \boldsymbol {\epsilon} \sim \mathcal {N} (0, \mathbf {I})} \left[ \min  \{1, \gamma e ^ {- \lambda} \} p (\lambda) \cdot - \frac {d \lambda}{d t} \| \boldsymbol {\epsilon} - \hat {\boldsymbol {\epsilon}} (\mathbf {z} _ {t}; \lambda) \| _ {2} ^ {2} \right]. (96) \\ \end{array}
$$
Therefore, it corresponds to  $w(\lambda) = \min \{1, \gamma e^{-\lambda}\} p(\lambda)$ . The motivation of the work is to avoid the model focusing too much on small noise levels, since it shares similar hypothesis to [Choi et al., 2022] that small noise levels are responsible for cleaning up details that may not be perceptible. A cosine noise schedule is then combined with the proposed weighting function, leading to  $w(\lambda) = \operatorname{sech}(\lambda / 2) \cdot \min \{1, \gamma e^{-\lambda}\}$ .  $\gamma$  is set as 5 empirically.
E Useful Equations
E.1 SDEs
The forward process is a Gaussian diffusion process, whose time evolution is described by a stochastic differential equation (SDE):
$$
d \mathbf {z} = \underbrace {\mathbf {f} (\mathbf {z} , t)} _ {\text {d r i f t}} d t + \underbrace {g (t)} _ {\text {d i f f u s i o n}} d \mathbf {w} \tag {97}
$$
For derivations of diffusion SDEs, see Appendix B of [Song et al., 2021b]. Their  $\beta(t)$  equals  $\frac{d}{dt}\log(1 + e^{-\lambda_t})$  in our formulation, and their  $\int_0^t\beta(s)ds$  equals  $\log(1 + e^{-\lambda_t})$ , where they assume that  $\lambda \to \infty$  at  $t = 0$ .
E.1.1 Variance-preserving (VP) SDE
A common choice is the variance-preserving (VP) SDE, which generalizes denoising diffusion models [Ho et al., 2020] to continuous time [Song et al., 2021b, Kingma et al., 2021]. In the VP case:
$$
\mathbf {f} (\mathbf {z}, t) = - \frac {1}{2} \left(\frac {d}{d t} \log \left(1 + e ^ {- \lambda_ {t}}\right)\right) \mathbf {z} \tag {98}
$$
$$
g (t) ^ {2} = \frac {d}{d t} \log \left(1 + e ^ {- \lambda_ {t}}\right) \tag {99}
$$
$$
\alpha_ {\lambda} ^ {2} = \operatorname {s i g m o i d} (\lambda) \tag {100}
$$
$$
\sigma_ {\lambda} ^ {2} = \operatorname {s i g m o i d} (- \lambda) \tag {101}
$$
$$
p \left(\mathbf {z} _ {1}\right) = \mathcal {N} (0, \mathbf {I}) \tag {102}
$$
E.1.2 Variance-exploding (VE) SDE
Another common choice of the variance-exploding (VE) SDE. In the VE case:
$$
\mathbf {f} (\mathbf {z}, t) = 0 \tag {103}
$$
$$
g (t) ^ {2} = \frac {d}{d t} \log \left(1 + e ^ {- \lambda_ {t}}\right) \tag {104}
$$
$$
\alpha_ {\lambda} ^ {2} = 1 \tag {105}
$$
$$
\sigma_ {\lambda} ^ {2} = e ^ {- \lambda} \tag {106}
$$
$$
p \left(\mathbf {z} _ {1}\right) = \mathcal {N} \left(0, e ^ {- \lambda_ {\min }} \mathbf {I}\right) \tag {107}
$$
E.2 Possible parameterizations of the score network
There are various ways of parameterizing the score network:
$$
\begin{array}{l} \mathbf {s} _ {\boldsymbol {\theta}} (\mathbf {z}; \lambda) = - \nabla_ {\mathbf {z}} E _ {\boldsymbol {\theta}} (\mathbf {z}, \lambda) \quad \text {(W i t h} \\ = - \hat {\epsilon} _ {\theta} (\mathbf {z}; \lambda) / \sigma_ {\lambda} \quad \text {(W i t h a n o i s e p r e d i c t i o n m o d e l)} (109) \\ = - \sigma_ {\lambda} ^ {- 2} (\mathbf {z} - \alpha_ {\lambda} \hat {\mathbf {x}} _ {\theta} (\mathbf {z}; \lambda)) \quad \text {(W i t h a d a t a p r e d i c t i o n m o d e l)} (110) \\ \end{array}
$$
We can let a neural network output any of  $\mathbf{s}_{\theta}(\mathbf{z}_{\lambda};\lambda),\hat{\epsilon}_{\theta}$  or  $\hat{\mathbf{x}}_{\theta}$ , and we can convert the variables to each other using the equalities above.
The chosen relationship between  $\mathbf{z}_{\lambda}$ ,  $\hat{\mathbf{x}}$ ,  $\hat{\epsilon}$  and  $\mathbf{s}_{\theta}(\mathbf{z}_{\lambda}; \lambda)$  above, are due to the following relationships between  $\mathbf{z}_{\lambda}$ ,  $\mathbf{x}$  and  $\epsilon$ :
$$
\mathbf {z} _ {\lambda} = \alpha_ {\lambda} \mathbf {x} + \sigma_ {\lambda} \epsilon \tag {111}
$$