Heavy-Tailed Diffusion Models

## Abstract

Abstract Diffusion models achieve state-of-the-art generation quality across many applications, but their ability to capture rare or extreme events in heavy-tailed distributions remains unclear. In this work, we show that traditional diffusion and flow-matching models with standard Gaussian priors fail to capture heavy-tailed behavior. We address this by repurposing the diffusion framework for heavy-tail estimation using multivariate Student-t distributions. We develop a tailored perturbation kernel and derive the denoising posterior based on the conditional Student-t distribution for the backward process. Inspired by γ 𝛾 \gamma italic_γ -divergence for heavy-tailed distributions, we derive a training objective for heavy-tailed denoisers. The resulting framework introduces controllable tail generation using only a single scalar hyperparameter, making it easily tunable for diverse real-world distributions. As specific instantiations of our framework, we introduce t-EDM and t-Flow , extensions of existing diffusion and flow models that employ a Student-t prior. Remarkably, our approach is readily compatible with standard Gaussian diffusion models and requires only minimal code changes. Empirically, we show that our t-EDM and t-Flow outperform standard diffusion models in heavy-tail estimation on high-resolution weather datasets in which generating rare and extreme events is crucial.

## 1 Introduction

In many real-world applications, such as weather forecasting, rare or extreme events—like hurricanes or heatwaves—can have disproportionately larger impacts than more common occurrences. Therefore, building generative models capable of accurately capturing these extreme events is critically important (Gründemann et al., ). However, learning the distribution of such data from finite samples is particularly challenging, as the number of empirically observed tail events is typically small, making accurate estimation difficult.

One promising approach is to use heavy-tailed distributions, which allocate more density to the tails than light-tailed alternatives. In popular generative models like Normalizing Flows (Rezende & Mohamed, ) and Variational Autoencoders (VAEs) (Kingma & Welling, ), recent works address heavy-tail estimation by learning a mapping from a heavy-tailed prior to the target distribution (Jaini et al., ; Kim et al., ).

While these works advocate for heavy-tailed base distributions, their application to real-world, high-dimensional datasets remains limited, with empirical results focused on small-scale or toy datasets. In contrast, diffusion models (Ho et al., ; Song et al., ; Lipman et al., ) have demonstrated excellent synthesis quality in large-scale applications. However, it is unclear whether diffusion models with Gaussian priors can effectively model heavy-tailed distributions without significant modifications.

In this work, we first demonstrate through extensive experiments that traditional diffusion models—even with proper normalization, preconditioning, and noise schedule design (see Section 4)—fail to accurately capture the heavy-tailed behavior in target distributions (see Fig. 1 for a toy example). We hypothesize that, in high-dimensional spaces, the Gaussian distribution in standard diffusion models tends to concentrate on a spherical narrow shell, thereby neglecting the tails. To address this, we adopt the multivariate Student-t distribution as the base noise distribution, with its degrees of freedom providing controllability over tail estimation. Consequently, we reformulate the denoising diffusion framework using multivariate Student-t distributions by designing a tailored perturbation kernel and deriving the corresponding denoiser. Moreover, we draw inspiration from the $\gamma$-power Divergences (Eguchi, ; Kim et al., ) for heavy-tailed distributions to formulate the learning problem for our heavy-tailed denoiser.

We extend widely adopted diffusion models, such as EDM (Karras et al., ) and straight-line flows (Lipman et al., ; Liu et al., ), by introducing their Student-t counterparts: *t-EDM* and *t-Flow*. We derive the corresponding SDEs and ODEs for modeling heavy-tailed distributions. Through extensive experiments on the HRRR dataset (Dowell et al., ), we train both unconditional and conditional versions of these models. The results show that standard EDM struggles to capture tails and extreme events, whereas *t-EDM* performs significantly better in modeling such phenomena. To summarize, we present,

- •
Heavy-tailed Diffusion Models. We repurpose the diffusion model framework for heavy-tail estimation by formulating both the forward and reverse processes using multivariate Student-t distributions. The denoiser is learned by minimizing the $\gamma$-power divergence (Kim et al., ) between the forward and reverse posteriors.
- •
Continuous Counterparts. We derive continuous formulations for heavy-tailed diffusion models and provide a principled approach to constructing ODE and SDE samplers. This enables the instantiation of *t-EDM* and *t-Flow* as heavy-tailed analogues to standard diffusion and flow models.
- •
Empirical Results. Experiments on the HRRR dataset (Dowell et al., ), a high-resolution dataset for weather modeling, show that *t-EDM* significantly outperforms EDM in capturing tail distributions for both unconditional and conditional tasks.
- •
Theoretical Connections. To theoretically justify the effectiveness of our approach, we present several theoretical connections between our framework and existing work in diffusion models and robust statistical estimators (Futami et al., ).

Figure: Figure 1: Toy Illustration. Our proposed diffusion model (*t-Diffusion*) captures heavy-tailed behavior more accurately than standard Gaussian diffusion, as shown in the histogram comparisons (top panel, x-axis). The framework allows for controllable tail estimation using a hyperparameter $\nu$, which can be adjusted for each dimension. Lower $\nu$ values model heavier tails, while higher values approach Gaussian diffusion. (Best viewed when zoomed in; see App. C.3 for details)
Refer to caption: extracted/5961946/figures/toy_fig.png

## 2 Background

As prerequisites underlying our method, we briefly summarize Gaussian diffusion models (as introduced in (Ho et al., ; Sohl-Dickstein et al., )) and multivariate Student-t distributions.

### 2.1 Diffusion Models

Diffusion models define a forward process (usually with an affine drift and no learnable parameters)
to convert data ${\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0}),{\mathbf{x}}_{0}\in\mathbb{R}^{d}$ to noise. A learnable reverse process is then trained to generate data from noise. In the discrete-time setting, the training objective for diffusion models can be specified as,

$$ $\displaystyle\mathbb{E}_{q}\bigg{[}\underbrace{D_{\mathrm{KL}}\!\left(q({ \mathbf{x}}_{T}|{\mathbf{x}}_{0})~{}\|~{}p({\mathbf{x}}_{T})\right)}_{L_{T}}+ \sum_{t>\Delta t}\underbrace{D_{\mathrm{KL}}\!\left(q({\mathbf{x}}_{t-\Delta t }|{\mathbf{x}}_{t},{\mathbf{x}}_{0})~{}\|~{}p_{\theta}({\mathbf{x}}_{t-\Delta t }|{\mathbf{x}}_{t})\right)}_{L_{t-1}}\underbrace{-\log p_{\theta}({\mathbf{x}} _{0}|{\mathbf{x}}_{\Delta t})}_{L_{0}}\bigg{]},$ (1) $$

where $T$ denotes the trajectory length while $\Delta t$ denotes the time increment between two consecutive time points. $D_{\text{KL}}$ denotes the Kullback-Leibler (KL) divergence defined as, $D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)=\int q({\mathbf{x}})\log\frac{q({
\mathbf{x}})}{p({\mathbf{x}})}d{\mathbf{x}}$.
In the objective in Eq. 1, the trajectory length $T$ is chosen to match the generative prior $p({\mathbf{x}}_{T})$ and the forward marginal $q({\mathbf{x}}_{T}|{\mathbf{x}}_{0})$. The second term in Eqn. 1 proposes to minimize the KL divergence between the forward posterior $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$ and the learnable posterior $p_{\theta}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})$ which corresponds to learning the *denoiser* (i.e. predicting a less noisy state from noise). The forward marginals, posterior, and reverse posterior are modeled using Gaussian distributions, which exhibit an analytical form of the KL divergence. The discrete-time diffusion framework can also be extended to the continuous time setting (Song et al., ; ; Karras et al., ). Recently, Lipman et al. (); Albergo et al. () proposed stochastic interpolants (or flows), which allow flexible transport between two arbitrary distributions.

### 2.2 Student-t Distributions

The multivariate Student-t distribution $t_{d}(\mu,{\bm{\Sigma}},\nu)$ with dimensionality d, location $\mu$, scale matrix ${\bm{\Sigma}}$ and degrees of freedom $\nu$ is defined as,

$$ $t_{d}(\mu,{\bm{\Sigma}},\nu)=C_{\nu,d}\Big{[}1+\frac{1}{\nu}({\mathbf{x}}-{\bm {\mu}})^{\top}{\bm{\Sigma}}^{-1}({\mathbf{x}}-{\bm{\mu}})\Big{]}^{-\frac{\nu+d }{2}},$ (2) $$

where $C_{\nu,d}$ is the normalizing factor. Since the multivariate Student-t distribution has polynomially decaying density, it can model heavy-tailed distributions. Interestingly, for $\nu=1$, the Student-t distribution is analogous to the Cauchy distribution. As $\nu\rightarrow\infty$, the Student-t distribution converges to the Gaussian distribution. A Student-t distributed random variable ${\mathbf{x}}$ can be reparameterized as (Andrews & Mallows, ), ${\mathbf{x}}={\bm{\mu}}+{\bm{\Sigma}}^{1/2}{\mathbf{z}}/\sqrt{\kappa}$, with ${\mathbf{z}}\sim{\mathcal{N}}(0,{\bm{I}}_{d}),\kappa\sim\chi^{2}(\nu)/\nu$
where $\chi^{2}$ denotes the Chi-squared distribution.

## 3 Heavy-Tailed Diffusion Models

We now repurpose standard diffusion models using multivariate Student-t distributions. The main idea behind our design is the use of heavy-tailed generative priors (Jaini et al., ; Kim et al., ) for learning a transport map towards a potentially heavy-tailed target distribution. From Eqn. 1 we note three key requirements for training diffusion models: choice of the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$, form of the target denoising posterior $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$ and the parameterization of the learnable reverse posterior $p_{\theta}({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})$. Therefore, we begin our discussion in the context of discrete-time diffusion models and later extend to the continuous regime. This has several advantages in terms of highlighting these three key design choices, which might be obscured by the continuous-time framework of defining a forward and a reverse SDE (Karras et al., ) while at the same time leading to a simpler construction.

### 3.1 Noising Process Design with Student-t Distributions.

Our construction of the noising process involves three key steps.

- 1.
Firstly, given two consecutive noisy states ${\mathbf{x}}_{t}$ and ${\mathbf{x}}_{t-\Delta t}$, we specify a joint distribution $q({\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})$.
- 2.
Secondly, given $q({\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})$, we construct the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=\int q({\mathbf{x}}_{t},{\mathbf{x}}_{t-
\Delta t}|{\mathbf{x}}_{0})d{\mathbf{x}}_{t-\Delta t}$ which can be used as a noising process during training.
- 3.
Lastly, from Steps 1 and 2, we construct the forward denoising posterior $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})=\frac{q({
\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})}{q({\mathbf{x}}_{t
}|{\mathbf{x}}_{0})}$. We will later utilize the form of $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$ to parameterize the reverse posterior.

It is worth noting that our construction of the noising process bypasses the specification of the forward transition kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{t-\Delta t})$. This has the advantage that we can directly specify the form of the perturbation kernel parameters $\mu_{t}$ and $\sigma_{t}$ as in Karras et al. () unlike Song et al. (); Ho et al. (). We next highlight the noising process construction in more detail.

Specifiying the joint distribution $q({\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})$.
We parameterize the joint distribution $q({\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})$ as a multivariate Student-t distribution with the following form,

$$ $q({\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})=t_{2d}({\bm{\mu }},{\bm{\Sigma}},\nu),\quad{\bm{\mu}}=[\mu_{t};\mu_{t-\Delta t}]{\mathbf{x}}_{ 0},\quad{\bm{\Sigma}}=\begin{pmatrix}\sigma_{t}^{2}&\sigma_{12}^{2}(t)\\ \sigma_{21}^{2}(t)&\sigma_{t-\Delta t}^{2}\end{pmatrix}\otimes{\bm{I}}_{d},$ $$

where $\mu_{t},\sigma_{t},\sigma_{12}(t),\sigma_{21}(t)$ are time-dependent scalar design parameters. While the choice of the parameters $\mu_{t}$ and $\sigma_{t}$ determines the perturbation kernel used during training, the choice of $\sigma_{12}(t)$ and $\sigma_{21}(t)$ can affect the ODE/SDE formulation for the denoising process and will be clarified when discussing sampling.

Constructing the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$: Given the joint distribution $q({\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})$ specified as a multivariate Student-t distribution, it follows that the perturbation kernel distribution $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ is also a Student-t distribution (Ding, ) parameterized as, $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^
{2}{\bm{I}}_{d},\nu)$ (proof in App. A.1).
We choose the scalar coefficients $\mu_{t}$ and $\sigma_{t}$ such that the perturbation kernel at time $t=T$ converges to a standard Student-t generative prior $t_{d}(0,{\bm{I}}_{d},\nu)$. We discuss practical choices of $\mu_{t}$ and $\sigma_{t}$ in Section 3.5.

Estimating the reference denoising posterior. Given the joint distribution $q({\mathbf{x}}_{t},{\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{0})$ and the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$, the denoising posterior can be specified as (see Ding ()),

$$ $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})=t_{d}(\bar{{\bm {\mu}}}_{t},\frac{\nu+d_{1}}{\nu+d}\bar{\sigma}_{t}^{2}{\bm{I}}_{d},\nu+d),$ (3) $$

$$ $\bar{{\bm{\mu}}}_{t}=\mu_{t-\Delta t}{\mathbf{x}}_{0}+\frac{\sigma_{21}^{2}(t) }{\sigma_{t}^{2}}({\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}),\qquad\bar{\sigma} _{t}^{2}=\Big{[}\sigma_{t-\Delta t}^{2}-\frac{\sigma_{21}^{2}(t)\sigma_{12}^{2 }(t)}{\sigma_{t}^{2}}\Big{]},$ (4) $$

where $d_{1}=\frac{1}{\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|^{2}$.
Next, we formulate the training objective for heavy-tailed diffusions.

### 3.2 Parameterization of the Reverse Posterior

Following Eqn. 3, we parameterize the reverse (or the denoising) posterior distribution as:

$$ $p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})=t_{d}({\bm{\mu}}_{ \bm{\theta}}({\mathbf{x}}_{t},t),\bar{\sigma}_{t}^{2}{\bm{I}}_{d},\nu+d),$ (5) $$

where the denoiser mean ${\bm{\mu}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ is further parameterized as follows:

$$ $\displaystyle{\bm{\mu}}_{\theta}({\mathbf{x}}_{t},t)=\frac{\sigma_{21}^{2}(t)} {\sigma_{t}^{2}}{\mathbf{x}}_{t}+\Big{[}\mu_{t-\Delta t}-\frac{\sigma_{21}^{2} (t)}{\sigma_{t}^{2}}\mu_{t}\Big{]}{\bm{D}}_{\theta}({\mathbf{x}}_{t},\sigma_{t }).$ (6) $$

While we adopt the "${\mathbf{x}}_{0}$-prediction" parameterization (Karras et al., ), it is also possible to parameterize the posterior mean as an $\epsilon$-prediction objective instead (Ho et al., ) (See App. A.2). Lastly, when parameterizing the reverse posterior scale, we drop the data-dependent coefficient $\frac{\nu+d_{1}}{\nu+d}$
This aligns with prior works in diffusion models (Ho et al., ; Song et al., ; Karras et al., ) where it is common to only parameterize the denoiser mean. However, heteroskedastic modeling of the denoiser is possible in our framework and could be an interesting direction for future work. Next, we reformulate the training objective in Eqn. 1 for heavy-tailed diffusions.

### 3.3 Training with Power Divergences

The optimization objective in Eqn. 1 primarily minimizes the KL-Divergence between a given pair of distributions. However, since we parameterize the distributions in Eqn. 1 using multivariate Student-t distributions, using the KL-Divergence might not be a suitable choice of divergence. This is because computing the KL divergence for Student-t distributions does not exhibit a closed-form expression. An alternative is the $\gamma$-Power Divergence (Eguchi, ; Kim et al., ) defined as,

$$ $D_{\gamma}\!\left(q~{}\|~{}p\right)=\frac{1}{\gamma}\big{[}{\mathcal{C}}_{ \gamma}(q,p)-{\mathcal{H}}_{\gamma}(q)\big{]},\quad\gamma\in(-1,0)\cup(0,\infty)$ $$

$$ ${\mathcal{H}}_{\gamma}(p)=-\left\|p\right\|_{1+\gamma=}-\bigg{(}\int p({ \mathbf{x}})^{1+\gamma}d{\mathbf{x}}\bigg{)}^{\frac{1}{1+\gamma}}\qquad{ \mathcal{C}}_{\gamma}(q,p)=-\int q({\mathbf{x}})\bigg{(}\frac{p({\mathbf{x}})} {\left\|p\right\|_{1+\gamma}}\bigg{)}^{\gamma}d{\mathbf{x}},$ $$

where, like Kim et al. (), we set $\gamma=-\frac{2}{\nu+d}$ for the remainder of our discussion. Moreover, ${\mathcal{H}}_{\gamma}$ and ${\mathcal{C}}_{\gamma}$ represent the $\gamma$-power entropy and cross-entropy, respectively. Interestingly, the $\gamma$-Power divergence between two multivariate Student-t distributions, $q_{\nu}=t_{d}({\bm{\mu}}_{0},{\bm{\Sigma}}_{0},\nu)$ and $p_{\nu}=t_{d}({\bm{\mu}}_{1},{\bm{\Sigma}}_{1},\nu)$, can be tractably computed in closed form and is defined as (see Kim et al. () for a proof),

$$ $\displaystyle\begin{split}\mathcal{D}_{\gamma}[q_{\nu}||p_{\nu}]&=\textstyle- \frac{1}{\gamma}C_{\nu,d}^{\frac{\gamma}{1+\gamma}}\left(1+\frac{d}{\nu-2} \right)^{-\frac{\gamma}{1+\gamma}}\Bigl{[}\,-|{\bm{\Sigma}}_{0}|^{-\frac{ \gamma}{2(1+\gamma)}}\left(1+\frac{d}{\nu-2}\right)\\ &\textstyle\quad+|{\bm{\Sigma}}_{1}|^{-\frac{\gamma}{2(1+\gamma)}}\left(1+ \frac{1}{\nu-2}\text{tr}\left({\bm{\Sigma}}_{1}^{-1}{\bm{\Sigma}}_{0}\right)+ \frac{1}{\nu}({\bm{\mu}}_{0}-{\bm{\mu}}_{1})^{\top}{\bm{\Sigma}}_{1}^{-1}({\bm {\mu}}_{0}-{\bm{\mu}}_{1})\right)\Bigr{]}.\end{split}$ (7) $$

Therefore, analogous to Eqn. 1, we minimize the following optimization objective,

$$ $\displaystyle\mathbb{E}_{q}\bigg{[}D_{\gamma}\!\left(q({\mathbf{x}}_{T}|{ \mathbf{x}}_{0})~{}\|~{}p({\mathbf{x}}_{T})\right)+\sum_{t>1}D_{\gamma}\!\left (q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})~{}\|~{}p_{ \theta}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})\right)-\log p_{\theta}({ \mathbf{x}}_{0}|{\mathbf{x}}_{1})\bigg{]}.$ (8) $$

Here, we note a couple of caveats. Firstly, while replacing the KL-Divergence with the $\gamma$-Power Divergence in the objective in Eqn. 1 might appear to be due to computational convenience, the $\gamma$-power divergence has several connections with robust estimators (Futami et al., ) in statistics and provides a tunable parameter $\gamma$ which can be used to control the model density assigned at the tail (see Section 5). Secondly, while the objective in Eqn. 1 is a valid ELBO, the objective in Eq. 8 is not. However, the following result provides a connection between the two objectives (see proof in App. A.3),

###### Proposition 1 .

For arbitrary distributions $q$ and $p$, in the limit of $\gamma\rightarrow 0$, $D_{\gamma}\!\left(q~{}\|~{}p\right)$ converges to $D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)$. Consequently, for a finite-dimensional dataset with ${\mathbf{x}}_{0}\in{\mathbb{R}}^{d}$ and $\gamma=-\frac{2}{\nu+d}$, under the limit of $\gamma\rightarrow 0$, the objective in Eqn. 8 converges to the objective in Eqn. 1.

Therefore, under the limit $\gamma\rightarrow 0$, the standard diffusion model framework becomes a special case of our proposed framework. Moreover, for $\gamma=-2/(\nu+d)$, this also explains the tail estimation moving towards Gaussian diffusion for an increasing $\nu$ (See Fig. 1 for an illustration.)

Simplifying the Training Objective. Plugging the form of the forward posterior $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$ in Eqn. 3, the reverse posterior $p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})$ in the optimization objective in Eqn. 8, we obtain the following simplified training loss (proof in App. A.4),

$$ ${\mathcal{L}}(\theta)\;\;=\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_ {0})}\mathbb{E}_{t\sim p(t)}\mathbb{E}_{\epsilon\sim{\mathcal{N}}(0,{\bm{I}}_{ d})}\mathbb{E}_{\kappa\sim\frac{1}{\nu}\chi^{2}(\nu)}\Big{\|}{\bm{D}}_{\bm{ \theta}}\big{(}\mu_{t}{\mathbf{x}}_{0}+\sigma_{t}\frac{\epsilon}{\sqrt{\kappa} },\sigma_{t}\big{)}-{\mathbf{x}}_{0}\Big{\|}^{2}_{2}.$ (9) $$

Intuitively, the form of our training objective is similar to existing diffusion models (Ho et al., ; Karras et al., ). However, the only difference lies in sampling the noisy state ${\mathbf{x}}_{t}$ from a Student-t distribution based perturbation kernel instead of a Gaussian distribution. Next, we discuss sampling from our proposed framework under discrete and continuous-time settings.

**Table 1: Comparison between different modeling components for constructing Gaussian vs Heavy-Tailed diffusion models. Under the limit of $\nu\rightarrow\infty$, our proposed t-Diffusion framework converges to Gaussian diffusion models.**
| Component | Gaussian Diffusion | (Ours) t-Diffusion |
| --- | --- | --- |
| Perturbation Kernel ($q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$) | ${\mathcal{N}}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^{2}{\bm{I}}_{d})$ | $t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^{2}{\bm{I}}_{d},\nu)$ |
| Forward Posterior ($q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$) | ${\mathcal{N}}(\bar{{\bm{\mu}}}_{t},\bar{\sigma}_{t}^{2}{\bm{I}}_{d})$ | $t_{d}(\bar{{\bm{\mu}}}_{t},\frac{\nu+d_{1}}{\nu+d}\bar{\sigma}_{t}^{2}{\bm{I}}<br>_{d},\nu+d)$ |
| Reverse Posterior ($p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})$) | ${\mathcal{N}}({\bm{\mu}}_{\bm{\theta}}({\mathbf{x}}_{t},t),\bar{\sigma}_{t}^{2<br>}{\bm{I}}_{d})$ | $t_{d}({\bm{\mu}}_{\bm{\theta}}({\mathbf{x}}_{t},t),\bar{\sigma}_{t}^{2}{\bm{I}<br>}_{d},\nu+d)$ |
| Divergence Measure | $D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)$ | $D_{\gamma}\!\left(q~{}\|~{}p\right)$ |

###### Proposition 1 .

### 3.4 Sampling

Discrete-time Sampling. For discrete-time settings, we can simply perform ancestral sampling from the learned reverse posterior distribution $p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})$. Therefore, following simple re-parameterization, an ancestral sampling update can be specified as,

$$ $\displaystyle{\mathbf{x}}_{t-\Delta t}$ $\displaystyle={\bm{\mu}}_{\bm{\theta}}({\mathbf{x}}_{t},t)+\bar{\sigma}_{t} \frac{{\mathbf{z}}}{\sqrt{\kappa}},\quad{\mathbf{z}}\sim{\mathcal{N}}(0,{\bm{I }}_{d}),\;\;\kappa\sim\chi^{2}(\nu+d)/(\nu+d),$ $\displaystyle=\frac{\sigma_{21}^{2}(t)}{\sigma_{t}^{2}}{\mathbf{x}}_{t}+\Big{[ }\mu_{t-\Delta t}-\frac{\sigma_{21}^{2}(t)}{\sigma_{t}^{2}}\mu_{t}\Big{]}{\bm{ D}}_{\theta}({\mathbf{x}}_{t},\sigma_{t})+\bar{\sigma}_{t}\frac{{\mathbf{z}}}{ \sqrt{\kappa}}.$ $$

Continuous-time Sampling. Due to recent advances in accelerating sampling in continuous-time diffusion processes (Pandey et al., ; Zhang & Chen, ; Lu et al., ; Song et al., ; Xu et al., ), we reformulate discrete-time dynamics in heavy-tailed diffusions to the continuous time regime. More specifically, we present a family of continuous-time processes in the following result (Proof in App. A.5).

###### Proposition 2 .

The posterior parameterization in Eqn. 5 induces the following continuous-time dynamics,

where $f:{\mathbb{R}}^{+}\times{\mathbb{R}}^{+}\rightarrow{\mathbb{R}}$ and $g:{\mathbb{R}}^{+}\times{\mathbb{R}}^{+}\rightarrow{\mathbb{R}}^{+}$ are scalar-valued functions, $\beta_{t}\in{\mathbb{R}}^{+}$ is a scaling coefficient such that the following condition holds,

$$ $\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\beta(t)g(\sigma_{t },\dot{\sigma}_{t})\Delta t\big{)}-1=f(\sigma_{t},\dot{\sigma}_{t})\Delta t,$ $$

where $\dot{\mu}_{t},\dot{\sigma}_{t}$ denote the first-order time-derivatives of the perturbation kernel parameters $\mu_{t}$ and $\sigma_{t}$ respectively and the differential $d{\bm{S}}_{t}\sim t_{d}(0,dt,\nu+d)$.

Based on the result in Proposition 2, it is possible to construct deterministic/stochastic samplers for heavy-tailed diffusions. It is worth noting that the SDE in Eqn. 10 implies adding heavy-tailed stochastic noise during inference (Bollerslev, ). Next, we provide specific instantiations of the generic sampler in Eq. 10.

Sampler Instantiations. We instantiate the continuous-time SDE in Eqn. 10 by setting $g(\sigma_{t},\dot{\sigma}_{t})=0$ and $\sigma_{12}(t)=\sigma_{t}\sigma_{t-\Delta t}$. Consequently, $f(\sigma_{t},\dot{\sigma}_{t})=-\frac{\dot{\sigma}_{t}}{\sigma_{t}}$. In this case, the SDE in Eqn. 10 reduces to an ODE, which can be represented as,

$$ $\frac{d{\mathbf{x}}_{t}}{dt}=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}- \bigg{[}-\frac{\dot{\sigma}_{t}}{\sigma_{t}}+\frac{\dot{\mu}_{t}}{\mu_{t}} \bigg{]}({\mathbf{x}}_{t}-\mu_{t}{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t}, \sigma_{t})).$ (11) $$

Summary. To summarize our theoretical framework, we present an overview of the comparison between Gaussian diffusion models and our proposed heavy-tailed diffusion models in Table 1.

Figure: Algorithm 1 Training (t-EDM)

###### Proposition 2 .

### 3.5 Specific Instantiations: t-EDM

Karras et al. () highlight several design choices during training and sampling which significantly improve sample quality while reducing sampling budget for image datasets like CIFAR-10 (Krizhevsky, ) and ImageNet (Deng et al., ). With a similar motivation, we reformulate the perturbation kernel as $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(s(t){\mathbf{x}}_{0},s(t)^{2}\sigma
(t)^{2}{\bm{I}}_{d},\nu)$ and denote the resulting diffusion model as *t-EDM*.

Training. During training, we set the perturbation kernel, $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$, parameters $s(t)=1$, $\sigma(t)=\sigma\sim\mathrm{LogNormal}(P_{\text{mean}},P_{\text{std}})$. Moreover, we parameterize the denoiser ${\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})$ as

$$ $D_{\theta}({\mathbf{x}},\sigma)=c_{\text{skip}}(\sigma,\nu){\mathbf{x}}+c_{ \text{out}}(\sigma,\nu){\bm{F}}_{\bm{\theta}}(c_{\text{in}}(\sigma,\nu){ \mathbf{x}},c_{\text{noise}}(\sigma)).$ $$

Our denoiser parameterization is similar to Karras et al. () with the difference that coefficients like $c_{\text{out}}$ additionally depend on $\nu$. We include full derivations in Appendix A.6. Consequently, our denoising loss can be specified as follows:

$$ ${\mathcal{L}}(\theta)\;\;\propto\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({ \mathbf{x}}_{0})}\mathbb{E}_{\sigma}\mathbb{E}_{{\mathbf{n}}\sim t_{d}(0, \sigma^{2}{\bm{I}}_{d},\nu)}\big{[}\lambda(\sigma,\nu)\|{\bm{D}}_{\bm{\theta}} ({\mathbf{x}}_{0}+{\mathbf{n}},\sigma)-{\mathbf{x}}_{0}\|^{2}_{2}\big{]},$ (12) $$

where $\lambda(\sigma,\nu)$ is a weighting function set to $\lambda(\sigma,\nu)=1/c_{\text{out}}(\sigma,\nu)^{2}$.

Sampling.
Interestingly, it can be shown that the ODE in Eqn. 11 is equivalent to the deterministic dynamics presented in Karras et al. () (See Appendix A.7 for proof). Consequently, we choose $s(t)=1$ and $\sigma(t)=t$ during sampling, further simplifying the dynamics in Eqn. 11 to

$$ $d{\mathbf{x}}_{t}/dt=({\mathbf{x}}_{t}-{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t} ,t))/t.$ $$

We adopt the timestep discretization schedule and the choice of the numerical ODE solver (Heun’s method (Ascher & Petzold, )) directly from EDM.

Figure 2 illustrates the ease of transitioning from a Gaussian diffusion framework (EDM) to t-EDM. Under standard settings, transitioning to t-EDM requires as few as two lines of code change, making our method readily compatible with existing implementations of Gaussian diffusion models.

Extension to Flows. While our discussion has been restricted to diffusion models, we can also construct flows (Albergo et al., ; Lipman et al., ) using heavy-tailed base distributions. We denote the resulting model as *t-Flow* and discuss its construction in more detail in App. B.

## 4 Experiments

Figure: Figure 3: Sample 1-d histogram comparison between EDM and t-EDM on the test set for the Vertically Integrated Liquid (VIL) channel. t-EDM captures heavy-tailed behavior more accurately than other baselines. INC: Inverse CDF Normalization, PCP: Per-Channel Preconditioning
Refer to caption: x1.png

To assess the effectiveness of the proposed heavy-tailed diffusion and flow models, we demonstrate experiments using real-world weather data for both unconditional and conditional generation tasks. We include full experimental details in App. C.

Datasets. We adopt the High-Resolution Rapid Refresh (HRRR) (Dowell et al., ) dataset, which is an operational archive of the US km-scale forecasting model. Among multiple dynamical variables in the dataset that exhibit heavy-tailed behavior, based on their dynamic range, we only consider the *Vertically Integrated Liquid* (VIL) and *Vertical Wind Velocity* at level 20 (denoted as w20) channels (see App. C.1 for more details). It is worth noting that the VIL and w20 channels have heavier right and left tails, respectively (See Fig. 6)

Tasks and Metrics. We consider both unconditional and conditional generative tasks relevant to weather and climate science. For unconditional modeling, we aim to generate the VIL and w20 physical variables in the HRRR dataset. For conditional modeling, we aim to generatively predict the hourly evolution of the target variable for the next lead-time ($\tau+1$) hour-ahead evolution of VIL and w20 based on information only at the current state time $\tau$; see more details in the appendix and see Pathak et al. () for discussion of why hour-ahead, km-scale atmospheric prediction is a stochastic physical task appropriate for conditional generative models. To quantify the empirical performance of unconditional modeling, we rely on comparing 1-d statistics of generated and train/test set samples. More specifically, for quantitative analysis, we report the *Kurtosis Ratio* (KR), the *Skewness Ratio* (SR), and the *Kolmogorov-Smirnov* (KS)-2 sample statistic (at the tails) between the generated and train/test set samples. For qualitative analysis, we compare 1-d histograms between generated and train/test set samples. For the conditional task, we adopt standard probabilistic prediction score metrics such as the *Continuous Ranked Probability Score* (CRPS), the *Root-Mean Squared Error* (RMSE), and the *skill-spread ratio* (SSR); see, e.g., Mardani et al. (); Srivastava et al. (). A more detailed explanation of our evaluation protocol is provided in App. C.

Methods and Baselines. In addition to standard diffusion (EDM (Karras et al., )) and flow models (Albergo et al., ) based on Gaussian priors, we introduce two additional baselines that are variants of EDM. To account for the high dynamic-range often exhibited by heavy-tailed distributions, we include *Inverse CDF Normalization* (INC) as an alternative data preprocessing step to z-score normalization. Using the former reduces dynamic range significantly and can make the data distribution closer to Gaussian. We denote this preprocessing scheme combined with standard EDM training as *EDM + INC*. Alternatively, we could instead modulate the noise levels used during EDM training as a function of the dynamic range of the input channel while keeping the data preprocessing unchanged. The main intuition is to use more heavy-tailed noise for large values. We denote this modulating scheme as *Per-Channel Preconditioning* (PCP) and denote the resulting baseline as *EDM + PCP*. We elaborate on these baselines in more detail in App. C.1.2

**Table 2: t-EDM outperforms standard diffusion models for unconditional generation on the HRRR dataset. For all metrics, lower is better. Values in bold indicate the best results in a column.**
|  |  |  | VIL (Train) | VIL (Test) |  | w20 (Train) | w20 (Test) |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Method | $\nu$ | KR $\downarrow$ | SR $\downarrow$ | KS $\downarrow$ | KR $\downarrow$ | SR $\downarrow$ | KS $\downarrow$ | $\nu$ | KR $\downarrow$ | SR $\downarrow$ | KS $\downarrow$ | KR $\downarrow$ | SR $\downarrow$ | KS $\downarrow$ |
| Baselines | EDM | $\infty$ | 210.11 | 10.79 | 0.997 | 45.35 | 5.23 | 0.991 | $\infty$ | 12.59 | 0.89 | 0.991 | 5.01 | 0.38 | 0.978 |
| +INC | $\infty$ | 11.33 | 2.29 | 0.987 | 1.70 | 0.74 | 0.95 | $\infty$ | 1.80 | 0.18 | 0.909 | 0.23 | 0.13 | 0.763 |  |
| +PCP | $\infty$ | 2.12 | 0.72 | 0.800 | 0.31 | 0.09 | 0.522 | $\infty$ | 2.17 | 0.70 | 0.838 | 0.40 | 0.24 | 0.648 |  |
| Ours | t-EDM | 3 | 1.06 | 0.43 | 0.431 | 0.54 | 0.23 | 0.114 | 3 | 2.44 | 0.65 | 0.683 | 0.52 | 0.21 | 0.286 |
| t-EDM | 5 | 29.66 | 4.07 | 0.955 | 5.73 | 1.68 | 0.888 | 5 | 8.55 | 1.77 | 0.895 | 3.22 | 1.03 | 0.774 |  |
| t-EDM | 7 | 24.35 | 4.14 | 0.959 | 4.57 | 1.72 | 0.908 | 7 | 7.03 | 1.58 | 0.82 | 2.55 | 0.89 | 0.622 |  |

### 4.1 Unconditional Generation

We assess the effectiveness of different methods on unconditional modeling for the VIL and w20 channels in the HRRR dataset. Fig. 3 qualitatively compares 1-d histograms of sample intensities between different methods for the VIL channel. We make the following key observations. Firstly, though EDM (with additional tricks like noise conditioning) can improve tail coverage, t-EDM covers a broader range of extreme values in the test set. Secondly, in addition to better dynamic range coverage, t-EDM qualitatively performs much better in capturing the density assigned to intermediate intensity levels under the model. We note similar observations from our quantitative results in Table 2, where t-EDM outperforms other baselines on the KS metric, implying our model exhibits better tail estimation over competing baselines for both the VIL and w20 channels. More importantly, unlike traditional Gaussian diffusion models like EDM, t-EDM enables controllable tail estimation by varying $\nu$, which could be useful when modeling a combination of channels with diverse statistical properties. On the contrary, standard diffusion models like EDM do not have such controllability. Lastly, we present similar quantitative results for t-Flow in Table 3. We present additional results for unconditional modeling in App. C.1

**Table 3: t-Flow outperforms standard Gaussian flows for unconditional generation on the HRRR dataset. For all metrics, lower is better. Values in bold indicate the best results in a column.**
|  |  |  | VIL (Train) | VIL (Test) |  | w20 (Train) | w20 (Test) |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Method | $\nu$ | KR $\downarrow$ | SR$\downarrow$ | KS$\downarrow$ | KR $\downarrow$ | SR $\downarrow$ | KS$\downarrow$ | $\nu$ | KR $\downarrow$ | SR $\downarrow$ | KS $\downarrow$ | KR $\downarrow$ | SR $\downarrow$ | KS $\downarrow$ |
| Baselines | Gaussian<br>Flow | $\infty$ | 0.46 | 0.09 | 0.897 | 0.67 | 0.52 | 0.704 | $\infty$ | 2.03 | 0.36 | 0.294 | 0.34 | 0.01 | 0.384 |
| Ours | t-Flow | 3 | 1.39 | 0.37 | 0.711 | 0.47 | 0.27 | 0.275 | 5 | 1.08 | 0.21 | 0.333 | 0.07 | 0.42 | 0.512 |
| t-Flow | 5 | 3.30 | 0.75 | 0.857 | 0.05 | 0.07 | 0.633 | 7 | 3.24 | 0.36 | 0.259 | 0.87 | 0.01 | 0.300 |  |
| t-Flow | 7 | 3.36 | 0.84 | 0.844 | 0.04 | 0.02 | 0.603 | 9 | 5.47 | 0.41 | 0.478 | 1.86 | 0.034 | 0.289 |  |

### 4.2 Conditional Generation

Next, we consider the task of conditional modeling, where we aim to predict the hourly evolution of a target variable for the next lead time ($\tau+1$) based on the current state at time $\tau$. Table 4 illustrates the performance of EDM and t-EDM on this task for the VIL and w20 channels. We make the following key observations. Firstly, for both channels, t-EDM exhibits better CRPS and SSR scores, implying better probabilistic forecast skills and ensemble than EDM. Moreover, while t-EDM exhibits under-dispersion for VIL, while it is well-calibrated for w20, with its SSR close to an ideal score of 1. On the contrary, the baseline EDM model exhibits under-dispersion for both channels, thus implying overconfident predictions. Secondly, in addition to better calibration, t-EDM is better at tail estimation (as measured by the KS statistic) for the underlying conditional distribution. Lastly, we notice that different values of the parameter $\nu$ are optimal for different channels, which suggests a more data-driven approach to learning the optimal $\nu$ directly. We present additional results for conditional modeling in App. C.2.

**Table 4: t-EDM outperforms EDM for conditional next frame prediction for the HRRR dataset. Values in bold indicate the best results in each column. We note that VIL has a higher dynamic range over w20, and thus, the gains for VIL are more apparent (see hist plots in Fig. 6). ($\rightarrow 1$) indicates values near 1 are better.**
|  |  |  | VIL (Test) |  | w20 (Test) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Method | $\nu$ | CRPS $\downarrow$ | RMSE $\downarrow$ | SSR ($\rightarrow 1$) | KS $\downarrow$ | $\nu$ | CRPS $\downarrow$ | RMSE $\downarrow$ | SSR ($\rightarrow 1$) | KS $\downarrow$ |
| Baselines | EDM | $\infty$ | 1.696 | 4.473 | 0.203 | 0.715 | $\infty$ | 0.304 | 0.664 | 0.865 | 0.345 |
| Ours | t-EDM | 3 | 1.649 | 4.526 | 0.255 | 0.419 | 3 | 0.295 | 0.734 | 1.045 | 0.111 |
| t-EDM | 5 | 1.609 | 4.361 | 0.305 | 0.665 | 5 | 0.301 | 0.674 | 0.901 | 0.323 |  |

## 5 Discussion and Theoretical Insights

To conclude, we propose a framework for constructing heavy-tailed diffusion models and demonstrate their effectiveness over traditional diffusion models on unconditional and conditional generation tasks for a high-resolution weather dataset. Here, we highlight some theoretical connections that help gain insights into the effectiveness of our proposed framework while establishing connections with prior work.

Figure: Figure 4: Variation of the mean and standard deviation of the ratio $(\nu+d_{1})/(\nu+d)$ with $\nu$ across diffusion sampling trajectory for the toy dataset. As $\nu$ decreases, the mean ratio and its standard deviation increase, leading to large score multiplier weights.
Refer to caption: extracted/5961946/figures/means.png

Exploring distribution tails during sampling. The ODE in Eq. 11 can re-formulated as,

$$ $\frac{d{\mathbf{x}}_{t}}{dt}=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}+ \sigma_{t}^{2}\Big{(}\frac{\nu+d^{\prime}_{1}}{\nu+d}\Big{)}\bigg{[}\frac{\dot {\mu}_{t}}{\mu_{t}}-\frac{\dot{\sigma}_{t}}{\sigma_{t}}\bigg{]}\nabla_{{ \mathbf{x}}}\log p({\mathbf{x}}_{t},t),$ (13) $$

where $d^{\prime}_{1}=(1/\sigma_{t}^{2})\|{\mathbf{x}}_{t}-\mu_{t}{\bm{D}}_{\bm{
\theta}}({\mathbf{x}}_{t},\sigma_{t})\|_{2}^{2}$. By formulating the ODE in terms of the score function, we can gain some intuition into the effectiveness of our model in modeling heavy-tailed distributions. Figure 4 illustrates the variation of the mean and variance of the multiplier
$(\nu+d^{\prime}_{1})/(\nu+d)$ along the diffusion trajectory across 1M samples generated from our toy models. Interestingly, as the value of $\nu$ decreases, the mean and variance of this multiplier increase significantly, which leads to large score multiplier weights. We hypothesize that this behavior allows our proposed model to explore more diverse regions during sampling (more details in App. A.10).

Enabling efficient tail coverage during training. The optimization objective in Eq. 8 has several connections with robust statistical estimators. More specifically, it can be shown that (proof in App. A.11),

$$ $\nabla_{\theta}D_{\gamma}\!\left(q~{}\|~{}p_{\theta}\right)=-\int q({\mathbf{x }})\bigg{(}\frac{p_{\theta}({\mathbf{x}})}{\left\|p_{\theta}\right\|_{1+\gamma }}\bigg{)}^{\gamma}\Big{(}\nabla_{\theta}\log p_{\theta}({\mathbf{x}})-\mathbb {E}_{\tilde{p}_{\theta}({\mathbf{x}})}[\nabla_{\theta}\log p_{\theta}({\mathbf {x}})]\Big{)}d{\mathbf{x}},$ $$

where $q$ and $p_{\theta}$ denote the forward ($q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$) and reverse diffusion posteriors ($p_{\theta}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})$), respectively. Intuitively, the coefficient $\gamma$ weighs the likelihood gradient, $\nabla_{\theta}\log p_{\theta}({\mathbf{x}})$, and can be set accordingly to ignore or consider outliers when modeling the data distribution. Specifically, when $\gamma>1$, the model would learn to ignore outliers (Futami et al., ; Fujisawa & Eguchi, ; Basu et al., ) since data points on the tails would be assigned low likelihood. On the contrary, a negative value of $\gamma$ (as is the case in this work since we set $\gamma=-2/(\nu+d)$), the model can assign more weights to capture these extreme values.

We discuss some other connections to prior work in heavy-tailed generative modeling and more recent work in diffusion models in App. F.1 and some limitations of our approach in App. F.2.

## Reproducibility Statement

We include proofs for all theoretical results introduced in the main text in Appendix A. We describe our complete experimental setup (including data processing steps, model specification for training and inference, description of evaluation metrics, and extended experimental results) in Appendix C.

## Ethics Statement

We develop a generative framework for modeling heavy-tailed distributions and demonstrate its effectiveness for scientific applications. In this context, we do not think our model poses a risk of misinformation or other ethical biases associated with large-scale image synthesis models. However, we would like to point out that similar to other generative models, our model can sometimes hallucinate predictions for certain channels, which could impact downstream applications like weather forecasting.

## Appendix A Proofs

### A.1 Derivation of the Perturbation Kernel

###### Proof.

By re-parameterization of the distribution $q({\mathbf{x}}|{\mathbf{x}}_{0})$, we have,

$$ ${\mathbf{x}}={\bm{\mu}}+{\bm{\Sigma}}^{1/2}{\mathbf{z}}/\sqrt{\kappa},\;\;{ \mathbf{z}}\sim\mathcal{N}(0,{\bm{I}}_{2d})\;\;\text{and}\;\;\kappa\sim\chi^{2 }(\nu)/\nu$ (14) $$

This implies that the conditional distribution $q({\mathbf{x}}|{\mathbf{x}}_{0},\kappa)=\mathcal{N}(\mu,{\bm{\Sigma}}/\kappa)$. Therefore, following properties of gaussian distributions, $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0},\kappa)=\mathcal{N}(\mu_{t}{\mathbf{x}}_{0
},\sigma_{t}^{2}/\kappa{\bm{I}}_{d})$. Therefore, from reparameterization,

$$ ${\mathbf{x}}_{t}|\kappa=\mu_{t}{\mathbf{x}}_{0}+\sigma_{t}{\mathbf{z}}/\sqrt{\kappa}$ (15) $$

which implies that $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^
{2}{\bm{I}}_{d},\nu)$. This completes the proof.
∎

###### Proof.

### A.2 On the Posterior Parameterization

The perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ for Student-t diffusions is parameterized as,

$$ $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^ {2}{\bm{I}}_{d},\nu)$ (16) $$

Using re-parameterization,

$$ ${\mathbf{x}}_{t}=\mu_{t}{\mathbf{x}}_{0}+\sigma_{t}\frac{\epsilon}{\sqrt{ \kappa}},\quad\epsilon\sim{\mathcal{N}}(0,{\bm{I}}_{d}),\kappa\sim\chi^{2}(\nu )/\nu$ (17) $$

During inference, given a noisy state ${\mathbf{x}}_{t}$, we have the following estimation problem,

$$ ${\mathbf{x}}_{t}=\mu_{t}\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]+\sigma_{ t}\mathbb{E}\big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}$ (18) $$

Therefore, the task of denoising can be posed as either estimating $\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]$ or $\mathbb{E}\big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}$. With this motivation, the posterior $p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})$ can be parameterized appropriately. Recall the form of the forward posterior

$$ $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})=t_{d}(\bar{{\bm {\mu}}}_{t},\frac{\nu+d_{1}}{\nu+d}\bar{\sigma}_{t}^{2}{\bm{I}}_{d},\nu+d)$ (19) $$

$$ $\bar{{\bm{\mu}}}_{t}=\mu_{t-\Delta t}{\mathbf{x}}_{0}+\frac{\sigma_{21}^{2}(t) }{\sigma_{t}^{2}}({\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}),\qquad\bar{\sigma} _{t}^{2}=\Big{[}\sigma_{t-\Delta t}^{2}-\frac{\sigma_{21}^{2}(t)\sigma_{12}^{2 }(t)}{\sigma_{t}^{2}}\Big{]}$ (20) $$

where $d_{1}=\frac{1}{\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|^{2}$. Further simplifying the mean $\bar{{\bm{\mu}}}_{t}$,

$$ $\displaystyle\bar{{\bm{\mu}}}_{t}({\mathbf{x}}_{t},{\mathbf{x}}_{0},t)$ $\displaystyle=\mu_{t-\Delta t}{\mathbf{x}}_{0}+\frac{\sigma_{21}^{2}(t)}{ \sigma_{t}^{2}}({\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0})$ (21) $\displaystyle=\mu_{t-\Delta t}{\mathbf{x}}_{0}+\frac{\sigma_{21}^{2}(t)}{ \sigma_{t}^{2}}{\mathbf{x}}_{t}-\frac{\sigma_{21}^{2}(t)}{\sigma_{t}^{2}}\mu_{ t}{\mathbf{x}}_{0}$ (22) $\displaystyle=\Big{(}\mu_{t-\Delta t}-\mu_{t}\frac{\sigma_{21}^{2}(t)}{\sigma_ {t}^{2}}\Big{)}{\mathbf{x}}_{0}+\frac{\sigma_{21}^{2}(t)}{\sigma_{t}^{2}}{ \mathbf{x}}_{t}$ (23) $$

Therefore, the mean ${\bm{\mu}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ of the reverse posterior $p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})$ can be parameterized as,

$$ $\displaystyle\bar{{\bm{\mu}}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ $\displaystyle=\Big{(}\mu_{t-\Delta t}-\mu_{t}\frac{\sigma_{21}^{2}(t)}{\sigma_ {t}^{2}}\Big{)}\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]+\frac{\sigma_{21} ^{2}(t)}{\sigma_{t}^{2}}{\mathbf{x}}_{t}$ (24) $\displaystyle\approx\Big{(}\mu_{t-\Delta t}-\mu_{t}\frac{\sigma_{21}^{2}(t)}{ \sigma_{t}^{2}}\Big{)}D_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})+\frac{ \sigma_{21}^{2}(t)}{\sigma_{t}^{2}}{\mathbf{x}}_{t}$ (25) $$

where $\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]$ is learned using a parametric estimator $D_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})$. This corresponds to the ${\mathbf{x}}_{0}$-prediction parameterization presented in Eq. 6 in the main text. Alternatively, From Eqn. 18,

$$ $\displaystyle\bar{{\bm{\mu}}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ $\displaystyle=\Big{(}\mu_{t-\Delta t}-\mu_{t}\frac{\sigma_{21}^{2}(t)}{\sigma_ {t}^{2}}\Big{)}\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]+\frac{\sigma_{21} ^{2}(t)}{\sigma_{t}^{2}}{\mathbf{x}}_{t}$ (26) $\displaystyle=\frac{1}{\mu_{t}}\Big{(}\mu_{t-\Delta t}-\mu_{t}\frac{\sigma_{21 }^{2}(t)}{\sigma_{t}^{2}}\Big{)}\Big{(}{\mathbf{x}}_{t}-\sigma_{t}\mathbb{E} \big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}\Big{)}+\frac{ \sigma_{21}^{2}(t)}{\sigma_{t}^{2}}{\mathbf{x}}_{t}$ (27) $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}-\frac{\sigma_{t }}{\mu_{t}}\Big{(}\mu_{t-\Delta t}-\mu_{t}\frac{\sigma_{21}^{2}(t)}{\sigma_{t} ^{2}}\Big{)}\mathbb{E}\big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t} \big{]}$ (28) $\displaystyle\approx\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}-\frac{ \sigma_{t}}{\mu_{t}}\Big{(}\mu_{t-\Delta t}-\mu_{t}\frac{\sigma_{21}^{2}(t)}{ \sigma_{t}^{2}}\Big{)}\epsilon_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})$ (29) $$

where $\mathbb{E}\big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}$ is learned using a parametric estimator $\epsilon_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})$. This corresponds to the $\epsilon$-prediction parameterization (Ho et al., ).

### A.3 Proof of Proposition 1

We restate the proposition here for convenience,

###### Proposition 1 .

For arbitrary distributions $q$ and $p$, in the limit of $\gamma\rightarrow 0$, $D_{\gamma}\!\left(q~{}\|~{}p\right)$ converges to $D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)$. Consequently, the objective in Eqn. 8 converges to the DDPM objective stated in Eqn. 1.

###### Proof.

We present our proof in two parts:

- 1.
Firstly, we establish the following relation between the $\gamma$-Power Divergence and the KL-Divergence between two distributions $q$ and $p$.
$D_{\gamma}\!\left(q~{}\|~{}p\right)=D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)+{
\mathcal{O}}(\gamma)$
(30)
- 2.
Next, for the choice of $\gamma=-\frac{2}{\nu+d}$, we show that in the limit of $\gamma\rightarrow 0$, the optimization objective in Eqn. 8 converges to the optimization objective in Eqn. 1

Relation between $D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)$ and $D_{\gamma}\!\left(q~{}\|~{}p\right)$.
The $\gamma$-Power Divergence as stated in (Kim et al., ) assumes the following form:

$$ $D_{\gamma}\!\left(q~{}\|~{}p\right)=\frac{1}{\gamma}\big{[}{\mathcal{C}}_{ \gamma}(q,p)-{\mathcal{H}}_{\gamma}(q)\big{]}$ (31) $$

where $\gamma\in(-1,0)\cup(0,\infty)$, and,

$$ ${\mathcal{H}}_{\gamma}(p)=-\left\|p\right\|_{1+\gamma}=-\bigg{(}\int p({ \mathbf{x}})^{1+\gamma}d{\mathbf{x}}\bigg{)}^{\frac{1}{1+\gamma}}\qquad{ \mathcal{C}}_{\gamma}(q,p)=-\int q({\mathbf{x}})\bigg{(}\frac{p({\mathbf{x}})} {\left\|p\right\|_{1+\gamma}}\bigg{)}^{\gamma}d{\mathbf{x}}$ (32) $$

For subsequent analysis, we assume $\gamma\rightarrow 0$. Under this assumption, we simplify ${\mathcal{H}}_{\gamma}(q)$ as follows. By definition,

$$ $\displaystyle{\mathcal{H}}_{\gamma}(q)=-\left\|q\right\|_{1+\gamma}$ $\displaystyle=-\bigg{(}\int q({\mathbf{x}})^{1+\gamma}d{\mathbf{x}}\bigg{)}^{ \frac{1}{1+\gamma}}$ (33) $\displaystyle=-\bigg{(}\int q({\mathbf{x}})q({\mathbf{x}})^{\gamma}d{\mathbf{x }}\bigg{)}^{\frac{1}{1+\gamma}}$ (34) $\displaystyle=-\bigg{(}\int q({\mathbf{x}})\exp(\gamma\log q({\mathbf{x}}))d{ \mathbf{x}}\bigg{)}^{\frac{1}{1+\gamma}}$ (35) $\displaystyle=-\bigg{(}\int q({\mathbf{x}})\big{[}1+\gamma\log q({\mathbf{x}}) +{\mathcal{O}}(\gamma^{2})\big{]}d{\mathbf{x}}\bigg{)}^{\frac{1}{1+\gamma}}$ (36) $\displaystyle=-\bigg{(}\int q({\mathbf{x}})d{\mathbf{x}}+\gamma\int q({\mathbf {x}})\log q({\mathbf{x}})d{\mathbf{x}}+{\mathcal{O}}(\gamma^{2})\bigg{)}^{ \frac{1}{1+\gamma}}$ (37) $\displaystyle=-\bigg{(}1+\gamma\int q({\mathbf{x}})\log q({\mathbf{x}})d{ \mathbf{x}}+{\mathcal{O}}(\gamma^{2})\bigg{)}^{\frac{1}{1+\gamma}}$ (38) $$

Using the approximation $(1+\delta x)^{\alpha}\approx 1+\alpha\delta x$ for a small $\delta$ in the above equation, we have,

$$ $\displaystyle{\mathcal{H}}_{\gamma}(q)=-\left\|q\right\|_{1+\gamma}$ $\displaystyle\approx-\bigg{(}1+\frac{\gamma}{1+\gamma}\Big{[}\int q({\mathbf{x }})\log q({\mathbf{x}})d{\mathbf{x}}+{\mathcal{O}}(\gamma)\Big{]}\bigg{)}$ (39) $\displaystyle\approx-\bigg{(}1+\gamma(1-\gamma)\Big{[}\int q({\mathbf{x}})\log q ({\mathbf{x}})d{\mathbf{x}}+{\mathcal{O}}(\gamma)\Big{]}\bigg{)}$ (40) $$

where we have used the approximation $\frac{1}{1+\gamma}\approx 1-\gamma$ in the above equation. This is justified since $\gamma$ is assumed to be small enough. Therefore, we have,

$$ ${\mathcal{H}}_{\gamma}(q)=-\left\|q\right\|_{1+\gamma}\approx-\bigg{(}1+\gamma \int q({\mathbf{x}})\log q({\mathbf{x}})d{\mathbf{x}}+{\mathcal{O}}(\gamma^{2} )\bigg{)}$ (41) $$

Similarly, we now obtain an approximation for the power-cross entropy as follows. By definition,

$$ $\displaystyle{\mathcal{C}}_{\gamma}(q,p)$ $\displaystyle=-\int q({\mathbf{x}})\bigg{(}\frac{p({\mathbf{x}})}{\left\|p \right\|_{1+\gamma}}\bigg{)}^{\gamma}d{\mathbf{x}}$ (42) $\displaystyle=-\int q({\mathbf{x}})\bigg{(}1+\gamma\log\bigg{(}\frac{p({ \mathbf{x}})}{\left\|p\right\|_{1+\gamma}}\bigg{)}+{\mathcal{O}}(\gamma^{2}) \bigg{)}d{\mathbf{x}}$ (43) $\displaystyle=-\bigg{(}\int q({\mathbf{x}})d{\mathbf{x}}+\gamma\int q({\mathbf {x}})\log\bigg{(}\frac{p({\mathbf{x}})}{\left\|p\right\|_{1+\gamma}}\bigg{)}+{ \mathcal{O}}(\gamma^{2})\bigg{)}$ (44) $\displaystyle=-\bigg{(}1+\gamma\int q({\mathbf{x}})\log p({\mathbf{x}})d{ \mathbf{x}}-\gamma\int q({\mathbf{x}})\log\left\|p\right\|_{1+\gamma}d{\mathbf {x}}+{\mathcal{O}}(\gamma^{2})\bigg{)}$ (45) $$

From Eqn. 41, it follows that,

$$ $\left\|p\right\|_{1+\gamma}\approx\bigg{(}1+\gamma\int p({\mathbf{x}})\log p({ \mathbf{x}})d{\mathbf{x}}+{\mathcal{O}}(\gamma^{2})\bigg{)}$ (46) $$

Therefore,

$$ $\displaystyle\log\left\|p\right\|_{1+\gamma}$ $\displaystyle=\log\bigg{(}1+\gamma\int p({\mathbf{x}})\log p({\mathbf{x}})d{ \mathbf{x}}+{\mathcal{O}}(\gamma^{2})\bigg{)}$ (47) $\displaystyle\approx\gamma\int p({\mathbf{x}})\log p({\mathbf{x}})d{\mathbf{x}}$ (48) $$

where the above result follows from the logarithmic series and ignores the terms of order ${\mathcal{O}}(\gamma^{2})$ or higher. Plugging the approximation in Eqn. 48 in Eqn. 45, we have,

$$ $\displaystyle{\mathcal{C}}_{\gamma}(q,p)$ $\displaystyle=-\bigg{(}1+\gamma\int q({\mathbf{x}})\log p({\mathbf{x}})d{ \mathbf{x}}-\gamma\int q({\mathbf{x}})\log\left\|p\right\|_{1+\gamma}d{\mathbf {x}}+{\mathcal{O}}(\gamma^{2})\bigg{)}$ (49) $\displaystyle\approx-\bigg{(}1+\gamma\int q({\mathbf{x}})\log p({\mathbf{x}})d {\mathbf{x}}+{\mathcal{O}}(\gamma^{2})\bigg{)}$ (50) $$

Therefore,

$$ $\displaystyle D_{\gamma}\!\left(q~{}\|~{}p\right)$ $\displaystyle=\frac{1}{\gamma}\big{[}{\mathcal{C}}_{\gamma}(q,p)-{\mathcal{H}} _{\gamma}(q)\big{]}$ (51) $\displaystyle=\frac{1}{\gamma}\big{[}\gamma\bigg{(}\int q({\mathbf{x}})\log q( {\mathbf{x}})d{\mathbf{x}}-\int q({\mathbf{x}})\log p({\mathbf{x}})d{\mathbf{x }}+{\mathcal{O}}(\gamma^{2})\bigg{)}\big{]}$ (52) $\displaystyle=\int q({\mathbf{x}})\log\frac{q({\mathbf{x}})}{p({\mathbf{x}})}d {\mathbf{x}}+{\mathcal{O}}(\gamma)$ (53) $\displaystyle=D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)+{\mathcal{O}}(\gamma)$ (54) $$

This establishes the relationship between the KL and $\gamma$-Power divergence between two distributions. Therefore, for two distributions $q$ and $p$, the difference in the magnitude of $D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)$ and $D_{\gamma}\!\left(q~{}\|~{}p\right)$ is of the order of ${\mathcal{O}}(\gamma)$. In the limit of $\gamma\rightarrow 0$, the $D_{\gamma}\!\left(q~{}\|~{}p\right)\rightarrow D_{\mathrm{KL}}\!\left(q~{}\|~{
}p\right)$. This concludes the first part of our proof.

Equivalence between the objectives under $\gamma\rightarrow 0$. For $\gamma=-\frac{2}{\gamma+d}$ and a finite-dimensional dataset with ${\mathbf{x}}_{0}\in{\mathbb{R}}^{d}$, it follows that $\gamma\rightarrow 0$ implies $\nu\rightarrow\infty$. Moreover, in the limit of $\nu\rightarrow\infty$, the multivariate Student-t distribution converges to a Gaussian distribution. As already shown in the previous part, under this limit, $D_{\gamma}\!\left(q~{}\|~{}p\right)$ converges to $D_{\mathrm{KL}}\!\left(q~{}\|~{}p\right)$. Therefore, under this limit, the optimization objective in Eqn. 8 converges to the standard DDPM objective in Eqn. 1. This completes the proof.
∎

###### Proposition 1 .

###### Proof.

### A.4 Derivation of the simplified denoising loss

Here, we derive the simplified denoising loss presented in Eq. 9 in the main text. We specifically consider the term $D_{\gamma}\!\left(q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0
})~{}\|~{}p_{\theta}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})\right)$ in Eq. 8. The $\gamma$-power divergence between two Student-t distributions is given by,

$$ $\displaystyle\begin{split}\mathcal{D}_{\gamma}[q_{\nu}||p_{\nu}]&=\textstyle- \frac{1}{\gamma}C_{\nu,d}^{\frac{\gamma}{1+\gamma}}\left(1+\frac{d}{\nu-2} \right)^{-\frac{\gamma}{1+\gamma}}\Bigl{[}\,-|{\bm{\Sigma}}_{0}|^{-\frac{ \gamma}{2(1+\gamma)}}\left(1+\frac{d}{\nu-2}\right)\\ &\textstyle\quad+|{\bm{\Sigma}}_{1}|^{-\frac{\gamma}{2(1+\gamma)}}\left(1+ \frac{1}{\nu-2}\text{tr}\left({\bm{\Sigma}}_{1}^{-1}{\bm{\Sigma}}_{0}\right)+ \frac{1}{\nu}({\bm{\mu}}_{0}-{\bm{\mu}}_{1})^{\top}{\bm{\Sigma}}_{1}^{-1}({\bm {\mu}}_{0}-{\bm{\mu}}_{1})\right)\Bigr{]}.\end{split}$ (55) $$

where $\gamma=-\frac{2}{\nu+d}$. Recall, the definitions of the forward denoising posterior,

$$ $q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})=t_{d}(\bar{{\bm {\mu}}}_{t},\frac{\nu+d_{1}}{\nu+d}\bar{\sigma}_{t}^{2}{\bm{I}}_{d},\nu+d)$ (56) $$

$$ $\bar{{\bm{\mu}}}_{t}=\mu_{t-\Delta t}{\mathbf{x}}_{0}+\frac{\sigma_{21}^{2}(t) }{\sigma_{t}^{2}}({\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}),\qquad\bar{\sigma} _{t}^{2}=\Big{[}\sigma_{t-\Delta t}^{2}-\frac{\sigma_{21}^{2}(t)\sigma_{12}^{2 }(t)}{\sigma_{t}^{2}}\Big{]}$ (57) $$

and the reverse denoising posterior,

$$ $p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})=t_{d}({\bm{\mu}}_{ \bm{\theta}}({\mathbf{x}}_{t},t),\bar{\sigma}_{t}^{2}{\bm{I}}_{d},\nu+d)$ (58) $$

where the denoiser mean ${\bm{\mu}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ is further parameterized as follows:

$$ $\displaystyle{\bm{\mu}}_{\theta}({\mathbf{x}}_{t},t)=\frac{\sigma_{21}^{2}(t)} {\sigma_{t}^{2}}{\mathbf{x}}_{t}+\Big{[}\mu_{t-\Delta t}-\frac{\sigma_{21}^{2} (t)}{\sigma_{t}^{2}}\mu_{t}\Big{]}{\bm{D}}_{\theta}({\mathbf{x}}_{t},\sigma_{t})$ (59) $$

Since we only parameterize the mean of the reverse posterior, the majority of the terms in the $\gamma$-power divergence are independent of $\theta$ and can be ignored (or treated as scalar coefficients). Therefore,

$$ $\displaystyle D_{\gamma}\!\left(q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{ \mathbf{x}}_{0})~{}\|~{}p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}} _{t})\right)$ $\displaystyle\propto(\bar{{\bm{\mu}}}_{t}-{\bm{\mu}}_{\theta}({\mathbf{x}}_{t} ,t))^{\top}(\bar{{\bm{\mu}}}_{t}-{\bm{\mu}}_{\theta}({\mathbf{x}}_{t},t))$ (60) $\displaystyle\propto\|\bar{{\bm{\mu}}}_{t}-{\bm{\mu}}_{\theta}({\mathbf{x}}_{t },t)\|_{2}^{2}$ (61) $\displaystyle\propto\Big{[}\mu_{t-\Delta t}-\frac{\sigma_{21}^{2}(t)}{\sigma_{ t}^{2}}\mu_{t}\Big{]}^{2}\|{\mathbf{x}}_{0}-D_{\bm{\theta}}({\mathbf{x}}_{t}, \sigma_{t})\|_{2}^{2}$ (62) $$

For better sample quality, it is common to ignore the scalar multiple in prior works (Ho et al., ; Song et al., ). Therefore, ignoring the time-dependent scalar multiple,

$$ $D_{\gamma}\!\left(q({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t},{\mathbf{x}}_{0 })~{}\|~{}p_{\bm{\theta}}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})\right) \propto\|{\mathbf{x}}_{0}-D_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})\|_{2}^{2}$ (63) $$

Therefore, the final loss function ${\mathcal{L}}_{\theta}$ can be stated as,

$$ ${\mathcal{L}}(\theta)=\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})} \mathbb{E}_{t\sim p(t)}\mathbb{E}_{\epsilon\sim{\mathcal{N}}(0,{\bm{I}}_{d})} \mathbb{E}_{\kappa\sim\frac{1}{\nu}\chi^{2}(\nu)}\Big{\|}{\bm{D}}_{\bm{\theta} }\big{(}\mu_{t}{\mathbf{x}}_{0}+\sigma_{t}\frac{\epsilon}{\sqrt{\kappa}}, \sigma_{t}\big{)}-{\mathbf{x}}_{0}\Big{\|}^{2}_{2}$ (64) $$

### A.5 Proof of Proposition 2

We restate the proposition here for convenience,

###### Proposition 2 .

The posterior parameterization in Eqn. 5 induces the following continuous-time dynamics,

$$ $d{\mathbf{x}}_{t}=\Bigg{[}\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}-\bigg{ [}f(\sigma_{t},\dot{\sigma}_{t})+\frac{\dot{\mu}_{t}}{\mu_{t}}\bigg{]}({ \mathbf{x}}_{t}-\mu_{t}{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},t))\Bigg{]}dt+ \sqrt{\beta(t)g(\sigma_{t},\dot{\sigma}_{t})}d{\bm{S}}_{t}$ (65) $$

where $f:{\mathbb{R}}^{+}\times{\mathbb{R}}^{+}\rightarrow{\mathbb{R}}$ and $g:{\mathbb{R}}^{+}\times{\mathbb{R}}^{+}\rightarrow{\mathbb{R}}^{+}$ are scalar-valued functions, $\beta_{t}\in{\mathbb{R}}^{+}$ is a scaling coefficient such that the following condition holds,

$$ $\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\beta(t)g(\sigma_{t },\dot{\sigma}_{t})\Delta t\big{)}-1=f(\sigma_{t},\dot{\sigma}_{t})\Delta t$ (66) $$

where $\dot{\mu}_{t},\dot{\sigma}_{t}$ denote the first-order time-derivatives of the perturbation kernel parameters $\mu_{t}$ and $\sigma_{t}$ respectively and the differential $d{\bm{S}}_{t}\sim t_{d}(0,dt,\nu+d)$.

###### Proof.

We start by writing a single sampling step from our learned posterior distribution. Recall

$$ $p_{\theta}({\mathbf{x}}_{t-\Delta t}|{\mathbf{x}}_{t})=t_{d}({\bm{\mu}}_{ \theta}({\mathbf{x}}_{t},t),\bar{\sigma}_{t}^{2}{\bm{I}}_{d},\nu+d)$ (67) $$

where (using the $\epsilon$-prediction parameterization in App. A.2),

$$ ${\bm{\mu}}_{\bm{\theta}}({\mathbf{x}}_{t},t)=\frac{\mu_{t-\Delta t}}{\mu_{t}}{ \mathbf{x}}_{t}+\frac{1}{\sigma_{t}}\bigg{[}\sigma_{21}^{2}(t)-\frac{\mu_{t- \Delta t}}{\mu_{t}}\sigma_{t}^{2}\bigg{]}\epsilon_{\theta}({\mathbf{x}}_{t}, \sigma_{t})$ (68) $$

From re-parameterization, we have,

$$ $\displaystyle{\mathbf{x}}_{t-\Delta t}$ $\displaystyle=\mu_{\theta}({\mathbf{x}}_{t},t)+\frac{\bar{\sigma}_{t}}{\sqrt{ \kappa}}{\mathbf{z}}\quad{\mathbf{z}}\sim\mathcal{N}(0,{\bm{I}}_{d}),\;\; \kappa\sim\chi^{2}(\nu+d)/(\nu+d)$ (69) $\displaystyle{\mathbf{x}}_{t-\Delta t}$ $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\frac{1}{\sigma _{t}}\bigg{[}\sigma_{21}^{2}(t)-\frac{\mu_{t-\Delta t}}{\mu_{t}}\sigma_{t}^{2} \bigg{]}\epsilon_{\theta}({\mathbf{x}}_{t},\sigma_{t})+\frac{\bar{\sigma}_{t}} {\sqrt{\kappa}}{\mathbf{z}}$ (70) $$

Moreover, we choose the posterior scale to be the same as the forward posterior $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$ i.e.

$$ $\bar{\sigma}_{t}^{2}=\Big{[}\sigma_{t-\Delta t}^{2}-\frac{\sigma_{21}^{2}(t) \sigma_{12}^{2}(t)}{\sigma_{t}^{2}}\Big{]}{\bm{I}}_{d}$ (71) $$

This implies,

$$ $\sigma_{21}^{2}(t)=\frac{\sigma_{t}^{2}}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t- \Delta t}^{2}-\bar{\sigma}_{t}^{2}\big{)}$ (72) $$

Substituting this form of $\sigma_{21}^{2}(t)$ into Eqn. 70, we have,

$$ $\displaystyle{\mathbf{x}}_{t-\Delta t}$ $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\frac{1}{\sigma _{t}}\bigg{[}\frac{\sigma_{t}^{2}}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t }^{2}-\bar{\sigma}_{t}^{2}\big{)}-\frac{\mu_{t-\Delta t}}{\mu_{t}}\sigma_{t}^{ 2}\bigg{]}\epsilon_{\theta}({\mathbf{x}}_{t},\sigma_{t})+\frac{\bar{\sigma}_{t }}{\sqrt{\kappa}}{\mathbf{z}}$ (73) $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}\bigg {[}\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\bar{\sigma}_{t} ^{2}\big{)}-\frac{\mu_{t-\Delta t}}{\mu_{t}}\bigg{]}\epsilon_{\theta}({\mathbf {x}}_{t},\sigma_{t})+\frac{\bar{\sigma}_{t}}{\sqrt{\kappa}}{\mathbf{z}}$ (74) $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}\bigg {[}\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\bar{\sigma}_{t} ^{2}\big{)}-1+1-\frac{\mu_{t-\Delta t}}{\mu_{t}}\bigg{]}\epsilon_{\theta}({ \mathbf{x}}_{t},\sigma_{t})+\frac{\bar{\sigma}_{t}}{\sqrt{\kappa}}{\mathbf{z}}$ (75) $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}\bigg {[}\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\bar{\sigma}_{t} ^{2}\big{)}-1+\frac{\dot{\mu}_{t}}{\mu_{t}}\Delta t\bigg{]}\epsilon_{\theta}({ \mathbf{x}}_{t},\sigma_{t})+\frac{\bar{\sigma}_{t}}{\sqrt{\kappa}}{\mathbf{z}}$ (76) $$

where in the above equation we use the first-order approximation $\dot{\mu}_{t}=\frac{\mu_{t}-\mu_{t-\Delta t}}{\Delta t}$. Next, we make the following design choices:

- 1.
Firstly, we assume the following form of the reverse posterior variance $\bar{\sigma}_{t}^{2}$:
$\bar{\sigma}_{t}^{2}=\beta(t)g(\sigma_{t},\dot{\sigma}_{t})\Delta t$
(77)
where $g:\mathbb{R}^{+}\times\mathbb{R}^{+}\rightarrow\mathbb{R}^{+}$ and $\beta(t)\in\mathbb{R}^{+}$ represents a time-varying scaling factor chosen empirically which can be used to vary the noise injected at each sampling step. It is worth noting that a positive $\dot{\sigma}_{t}$ (as indicated in the definition of g) is a consequence of a monotonically increasing noise schedule $\sigma_{t}$ in diffusion model design.
- 2.
Secondly, we make the following design choice:
$\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\bar{\sigma}_{t}^{2
}\big{)}-1=f(\sigma_{t},\dot{\sigma}_{t})\Delta t$
(78)
where $f:\mathbb{R}^{+}\times\mathbb{R}^{+}\rightarrow\mathbb{R}$

With these two design choices, Eqn. 76 simplifies as:

$$ $\displaystyle{\mathbf{x}}_{t-\Delta t}$ $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}\bigg {[}\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\bar{\sigma}_{t} ^{2}\big{)}-1+\frac{\dot{\mu}_{t}}{\mu_{t}}\Delta t\bigg{]}\epsilon_{\theta}({ \mathbf{x}}_{t},\sigma_{t})+\frac{\bar{\sigma}_{t}}{\sqrt{\kappa}}{\mathbf{z}}$ (79) $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}\bigg {[}f(\sigma_{t},\dot{\sigma}_{t})\Delta t+\frac{\dot{\mu}_{t}}{\mu_{t}}\Delta t \bigg{]}\epsilon_{\theta}({\mathbf{x}}_{t},\sigma_{t})+\sqrt{\beta(t)g(\sigma_ {t},\dot{\sigma}_{t})}\sqrt{\Delta t}\frac{{\mathbf{z}}}{\sqrt{\kappa}}$ (80) $\displaystyle=\frac{\mu_{t-\Delta t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}\bigg {[}f(\sigma_{t},\dot{\sigma}_{t})+\frac{\dot{\mu}_{t}}{\mu_{t}}\bigg{]} \epsilon_{\theta}({\mathbf{x}}_{t},\sigma_{t})\Delta t+\sqrt{\beta(t)g(\sigma_ {t},\dot{\sigma}_{t})}\sqrt{\Delta t}\frac{{\mathbf{z}}}{\sqrt{\kappa}}$ (81) $\displaystyle{\mathbf{x}}_{t-\Delta t}-{\mathbf{x}}_{t}$ $\displaystyle=\big{(}\frac{\mu_{t-\Delta t}}{\mu_{t}}-1\big{)}{\mathbf{x}}_{t} +\sigma_{t}\bigg{[}f(\sigma_{t},\dot{\sigma}_{t})+\frac{\dot{\mu}_{t}}{\mu_{t} }\bigg{]}\epsilon_{\theta}({\mathbf{x}}_{t},\sigma_{t})\Delta t+\sqrt{\beta(t) g(\sigma_{t},\dot{\sigma}_{t})}\sqrt{\Delta t}\frac{{\mathbf{z}}}{\sqrt{\kappa}}$ (82) $\displaystyle{\mathbf{x}}_{t-\Delta t}-{\mathbf{x}}_{t}$ $\displaystyle=-\Bigg{[}\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}-\sigma_{t }\bigg{[}f(\sigma_{t},\dot{\sigma}_{t})+\frac{\dot{\mu}_{t}}{\mu_{t}}\bigg{]} \epsilon_{\theta}({\mathbf{x}}_{t},\sigma_{t})\Bigg{]}\Delta t+\sqrt{\beta(t)g (\sigma_{t},\dot{\sigma}_{t})}\sqrt{\Delta t}\frac{{\mathbf{z}}}{\sqrt{\kappa}}$ (83) $$

In the limit of $\Delta t\rightarrow 0$:

$$ $d{\mathbf{x}}_{t}=\Bigg{[}\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}-\sigma _{t}\bigg{[}f(\sigma_{t},\dot{\sigma}_{t})+\frac{\dot{\mu}_{t}}{\mu_{t}}\bigg{ ]}\epsilon_{\theta}({\mathbf{x}}_{t},\sigma_{t})\Bigg{]}dt+\sqrt{\beta(t)g( \sigma_{t},\dot{\sigma}_{t})}\frac{dW_{t}}{\sqrt{\kappa}}$ (84) $$

which gives the required SDE formulation for the diffusion posterior sampling. Next, we discuss specific choices of $f(\sigma_{t},\dot{\sigma}_{t})$ and $g(\sigma_{t},\dot{\sigma}_{t})$.

On the choice of $f(\sigma_{t},\dot{\sigma}_{t})$ and $g(\sigma_{t},\dot{\sigma}_{t})$: Recall our design choices:

$$ $\bar{\sigma}_{t}^{2}=\beta(t)g(\sigma_{t},\dot{\sigma}_{t})\Delta t$ (85) $$

$$ $\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\bar{\sigma}_{t}^{2 }\big{)}-1=f(\sigma_{t},\dot{\sigma}_{t})\Delta t$ (86) $$

Substituting the value of $\bar{\sigma}_{t}^{2}$ from the first design choice to the second yields the following condition:

$$ $\frac{1}{\sigma_{12}^{2}(t)}\big{(}\sigma_{t-\Delta t}^{2}-\beta(t)g(\sigma_{t },\dot{\sigma}_{t})\Delta t\big{)}-1=f(\sigma_{t},\dot{\sigma}_{t})\Delta t$ (87) $$

This concludes the proof. It is worth noting that the above equation provides two degrees of freedom, i.e., we can choose two variables among $\sigma_{12}(t),g,f$ and automatically determine the third. However, it is more convenient to choose $\sigma_{12}(t)$ and $g$, since both these quantities should be positive. Different choices of these quantities yield different instantiations of the SDE in Eqn. 84 as illustrated in the main text.
∎

###### Proposition 2 .

###### Proof.

### A.6 Deriving the Denoiser Preconditioner for t-EDM

Recall our denoiser parameterization,

$$ $D_{\theta}({\mathbf{x}},\sigma)=c_{\text{skip}}(\sigma,\nu){\mathbf{x}}+c_{ \text{out}}(\sigma,\nu){\bm{F}}_{\bm{\theta}}(c_{\text{in}}(\sigma,\nu){ \mathbf{x}},c_{\text{noise}}(\sigma))$ (88) $$

Karras et al. () highlight the following design choices, which we adopt directly.

- 1.
Derive $c_{\text{in}}$ based on constraining the input variance to 1
- 2.
Derive $c_{\text{skip}}$ and $c_{\text{out}}$ to constrain output variance to 1 and additionally minimizing $c_{\text{out}}$ to bound scaling errors in ${\bm{F}}_{\theta}({\mathbf{x}},\sigma)$.

The coefficient $c_{\text{in}}$ can be derived by setting the network inputs to have unit variance. Therefore,

$$ $\displaystyle\mathrm{Var}_{{\mathbf{x}}_{0},{\mathbf{n}}}\big{[}c_{\text{in}}( \sigma)({\mathbf{x}}_{0}+{\mathbf{n}})\big{]}$ $\displaystyle=$ $\displaystyle 1$ (89) $\displaystyle c_{\text{in}}(\sigma,\nu)^{2}~{}\mathrm{Var}_{{\mathbf{x}}_{0},{ \mathbf{n}}}\big{[}{\mathbf{x}}_{0}+{\mathbf{n}}\big{]}$ $\displaystyle=$ $\displaystyle 1$ (90) $\displaystyle c_{\text{in}}(\sigma,\nu)^{2}\big{(}\sigma_{\text{data}}^{2}+ \frac{\nu}{\nu-2}\sigma^{2}\big{)}$ $\displaystyle=$ $\displaystyle 1$ (91) $\displaystyle c_{\text{in}}(\sigma,\nu)$ $\displaystyle=$ $\displaystyle 1\big{/}\sqrt{\frac{\nu}{\nu-2}\sigma^{2}+\sigma_{\text{data}}^{ 2}}\text{.}$ (92) $$

The coefficients $c_{\text{skip}}$ and $c_{\text{out}}$ can be derived by setting the training target to have unit variance. Similar to Karras et al. () our training target can be specified as:

$$ ${\bm{F}}_{\text{target}}=\frac{1}{c_{\text{out}}(\sigma,\nu)}\big{(}{\mathbf{x }}_{0}-c_{\text{skip}}(\sigma,\nu)({\mathbf{x}}_{0}{+}{\mathbf{n}})\big{)}$ (93) $$

$$ $\displaystyle\mathrm{Var}_{{\mathbf{x}}_{0},{\mathbf{n}}}\big{[}F_{\text{ target}}({\mathbf{x}}_{0},{\mathbf{n}};\sigma)\big{]}$ $\displaystyle=$ $\displaystyle 1$ (94) $\displaystyle\mathrm{Var}_{{\mathbf{x}}_{0},{\mathbf{n}}}\Big{[}\tfrac{1}{c_{ \text{out}}(\sigma,\nu)}\big{(}{\mathbf{x}}_{0}-c_{\text{skip}}(\sigma,\nu)({ \mathbf{x}}_{0}+{\mathbf{n}})\big{)}\Big{]}$ $\displaystyle=$ $\displaystyle 1$ (95) $\displaystyle\tfrac{1}{c_{\text{out}}(\sigma,\nu)^{2}}\mathrm{Var}_{{\mathbf{x }}_{0},{\mathbf{n}}}\big{[}{\mathbf{x}}_{0}-c_{\text{skip}}(\sigma,\nu)({ \mathbf{x}}_{0}+{\mathbf{n}})\big{]}$ $\displaystyle=$ $\displaystyle 1$ (96) $\displaystyle c_{\text{out}}(\sigma,\nu)^{2}$ $\displaystyle=$ $\displaystyle\mathrm{Var}_{{\mathbf{x}}_{0},{\mathbf{n}}}\big{[}{\mathbf{x}}_{ 0}-c_{\text{skip}}(\sigma,\nu)({\mathbf{x}}_{0}+{\mathbf{n}})\big{]}$ (97) $\displaystyle c_{\text{out}}(\sigma,\nu)^{2}$ $\displaystyle=$ $\displaystyle\mathrm{Var}_{{\mathbf{x}}_{0},{\mathbf{n}}}\Big{[}\big{(}1-c_{ \text{skip}}(\sigma,\nu)\big{)}~{}{\mathbf{x}}_{0}+c_{\text{skip}}(\sigma,\nu) ~{}{\mathbf{n}}\Big{]}$ (98) $\displaystyle c_{\text{out}}(\sigma,\nu)^{2}$ $\displaystyle=$ $\displaystyle\big{(}1-c_{\text{skip}}(\sigma,\nu)\big{)}^{2}~{}\sigma_{\text{ data}}^{2}+c_{\text{skip}}(\sigma,\nu)^{2}\frac{\nu}{\nu-2}\sigma^{2}\text{.}$ (99) $$

Lastly, setting $c_{\text{skip}}(\sigma,\nu)$ to minimize $c_{\text{out}}(\sigma,\nu)$, we obtain,

$$ $c_{\text{skip}}(\sigma,\nu)=\sigma_{\text{data}}^{2}/\bigg{(}\frac{\nu}{\nu-2} \sigma^{2}+\sigma_{\text{data}}^{2}\bigg{)}\text{.}$ (100) $$

Consequently $c_{\text{out}}(\sigma,\nu)$ can be specified as:

$$ $c_{\text{out}}(\sigma,\nu)=\sqrt{\frac{\nu}{\nu-2}}\sigma\cdot\sigma_{\text{ data}}\big{/}\sqrt{\frac{\nu}{\nu-2}\sigma^{2}+\sigma_{\text{data}}^{2}}\text{.}$ (101) $$

### A.7 Equivalence with the EDM ODE

Similar to Karras et al. (), we start by deriving the optimal denoiser for the denoising loss function. Moreover, we reformulate the form of the perturbation kernel as $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^
{2}{\bm{I}}_{d},\nu)=t_{d}(s(t){\mathbf{x}}_{0},s(t)^{2}\sigma(t)^{2}{\bm{I}}_
{d},\nu)$ by setting $\mu_{t}=s(t)$ and $\sigma_{t}=s(t)\sigma(t)$. The denoiser loss can then be specified as follows,

$$ $\displaystyle{\mathcal{L}}(D,\sigma)\;\;$ $\displaystyle=\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\mathbb {E}_{{\mathbf{n}}\sim t_{d}(0,\sigma^{2}{\bm{I}}_{d},\nu)}\big{[}\lambda( \sigma,\nu)\|{\bm{D}}(s(t)[{\mathbf{x}}_{0}+{\mathbf{n}}],\sigma)-{\mathbf{x}} _{0}\|^{2}_{2}\big{]}$ (102) $\displaystyle=\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\mathbb {E}_{{\mathbf{x}}\sim t_{d}(s(t){\mathbf{x}}_{0},s(t)^{2}\sigma^{2}{\bm{I}}_{d },\nu)}\big{[}\lambda(\sigma,\nu)\|{\bm{D}}({\mathbf{x}},\sigma)-{\mathbf{x}}_ {0}\|^{2}_{2}\big{]}$ (103) $\displaystyle=\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\Big{[} \int t_{d}({\mathbf{x}};s(t){\mathbf{x}}_{i},s(t)^{2}\sigma^{2}{\bm{I}}_{d}, \nu)\big{[}\lambda(\sigma,\nu)\|{\bm{D}}({\mathbf{x}},\sigma)-{\mathbf{x}}_{0} \|^{2}_{2}\big{]}d{\mathbf{x}}\Big{]}$ (104) $\displaystyle=\;\;\frac{1}{N}\sum_{i}\int t_{d}({\mathbf{x}};s(t){\mathbf{x}}_ {i},s(t)^{2}\sigma^{2}{\bm{I}}_{d},\nu)\big{[}\lambda(\sigma,\nu)\|{\bm{D}}({ \mathbf{x}},\sigma)-{\mathbf{x}}_{i}\|^{2}_{2}\big{]}d{\mathbf{x}}$ (105) $$

where the last result follows from assuming $p({\mathbf{x}}_{0})$ as the empirical data distribution. Thus, the optimal denoiser can be specified by setting $\nabla_{D}{\mathcal{L}}(D,\sigma)=0$. Therefore,

$$ $\nabla_{D}{\mathcal{L}}(D,\sigma)=0$ (106) $$

Consequently,

$$ $\displaystyle\nabla_{D}\frac{1}{N}\sum_{i}\int t_{d}({\mathbf{x}};s(t){\mathbf {x}}_{i},s(t)^{2}\sigma^{2}{\bm{I}}_{d},\nu)\big{[}\lambda(\sigma,\nu)\|{\bm{D }}({\mathbf{x}},\sigma)-{\mathbf{x}}_{i}\|^{2}_{2}\big{]}d{\mathbf{x}}$ $\displaystyle=0$ (107) $\displaystyle\frac{1}{N}\sum_{i}\int t_{d}({\mathbf{x}};s(t){\mathbf{x}}_{i},s (t)^{2}\sigma^{2}{\bm{I}}_{d},\nu)\big{[}\lambda(\sigma,\nu)({\bm{D}}({\mathbf {x}},\sigma)-{\mathbf{x}}_{i})\big{]}d{\mathbf{x}}$ $\displaystyle=0$ (108) $\displaystyle\int\sum_{i}t_{d}({\mathbf{x}};s(t){\mathbf{x}}_{i},s(t)^{2} \sigma^{2}{\bm{I}}_{d},\nu)({\bm{D}}({\mathbf{x}},\sigma)-{\mathbf{x}}_{i})d{ \mathbf{x}}$ $\displaystyle=0$ (109) $$

The optimal denoiser ${\bm{D}}({\mathbf{x}},\sigma)$, can be obtained from Eq. 109 as,

$$ ${\bm{D}}({\mathbf{x}},\sigma)=\frac{\sum_{i}t_{d}({\mathbf{x}};s(t){\mathbf{x} }_{i},s(t)^{2}\sigma^{2}{\bm{I}}_{d},\nu){\mathbf{x}}_{i}}{\sum_{i}t_{d}({ \mathbf{x}};s(t){\mathbf{x}}_{i},s(t)^{2}\sigma^{2}{\bm{I}}_{d},\nu)}$ (110) $$

We can further simplify the optimal denoiser as,

$$ $\displaystyle{\bm{D}}({\mathbf{x}},\sigma)$ $\displaystyle=\frac{\sum_{i}t_{d}({\mathbf{x}};s(t){\mathbf{x}}_{i},s(t)^{2} \sigma^{2}{\bm{I}}_{d},\nu){\mathbf{x}}_{i}}{\sum_{i}t_{d}({\mathbf{x}};s(t){ \mathbf{x}}_{i},s(t)^{2}\sigma^{2}{\bm{I}}_{d},\nu)}$ (111) $\displaystyle=\frac{\sum_{i}t_{d}(\frac{{\mathbf{x}}}{s(t)};{\mathbf{x}}_{i}, \sigma^{2}{\bm{I}}_{d},\nu){\mathbf{x}}_{i}}{\sum_{i}t_{d}(\frac{{\mathbf{x}}} {s(t)};{\mathbf{x}}_{i},\sigma^{2}{\bm{I}}_{d},\nu)}$ (112) $\displaystyle={\bm{D}}\Big{(}\frac{{\mathbf{x}}}{s(t)},\sigma\Big{)}$ (113) $$

Next, recall the ODE dynamics (Eqn. 11) in our formulation,

$$ $\frac{d{\mathbf{x}}_{t}}{dt}=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}- \bigg{[}-\frac{\dot{\sigma}_{t}}{\sigma_{t}}+\frac{\dot{\mu}_{t}}{\mu_{t}} \bigg{]}({\mathbf{x}}_{t}-\mu_{t}{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t}, \sigma(t)))$ (114) $$

Reparameterizing the above ODE by setting $\mu_{t}=s(t)$ and $\sigma_{t}=s(t)\sigma(t)$,

$$ $\displaystyle\frac{d{\mathbf{x}}_{t}}{dt}$ $\displaystyle=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}-\bigg{[}-\frac{ \dot{\sigma}_{t}}{\sigma_{t}}+\frac{\dot{\mu}_{t}}{\mu_{t}}\bigg{]}({\mathbf{x }}_{t}-s(t){\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},\sigma(t)))$ (115) $\displaystyle=\frac{\dot{s}(t)}{s(t)}{\mathbf{x}}_{t}-\bigg{[}\frac{\dot{s}(t) }{s(t)}-\frac{\dot{\sigma}(t)s(t)+\sigma(t)\dot{s}(t)}{\sigma(t)s(t)}\bigg{]}( {\mathbf{x}}_{t}-s(t){\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},\sigma(t)))$ (116) $\displaystyle=\frac{\dot{s}(t)}{s(t)}{\mathbf{x}}_{t}-\bigg{[}-\frac{\dot{ \sigma}(t)}{\sigma(t)}\bigg{]}({\mathbf{x}}_{t}-s(t){\bm{D}}_{\bm{\theta}}({ \mathbf{x}}_{t},\sigma(t)))$ (117) $\displaystyle=\frac{\dot{s}(t)}{s(t)}{\mathbf{x}}_{t}+\frac{\dot{\sigma}(t)}{ \sigma(t)}({\mathbf{x}}_{t}-s(t){\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},\sigma (t)))$ (118) $$

Lastly, since Karras et al. () propose to train the denoiser ${\bm{D}}_{\bm{\theta}}$ using unscaled noisy state, from Eqn. 113, we can re-write the above ODE as,

$$ $\displaystyle\frac{d{\mathbf{x}}_{t}}{dt}$ $\displaystyle=\Big{[}\frac{\dot{s}(t)}{s(t)}+\frac{\dot{\sigma}(t)}{\sigma(t)} \Big{]}{\mathbf{x}}_{t}-\frac{\dot{\sigma}(t)s(t)}{\sigma(t)}{\bm{D}}_{\bm{ \theta}}\Big{(}\frac{{\mathbf{x}}_{t}}{s(t)},\sigma(t)\Big{)}$ (119) $$

The form of the ODE in Eqn. 119 is equivalent to the ODE presented in Karras et al. () (Algorithm 1 line 4 in their paper). This concludes the proof.

### A.8 Conditional Vector Field for t-Flow

Here, we derive the conditional vector field for the *t-Flow* interpolant. Recall, the interpolant in *t-Flow* is derived as follows,

$$ ${\mathbf{x}}_{t}=t{\mathbf{x}}_{0}+(1-t)\frac{\epsilon}{\sqrt{\kappa}},\quad \epsilon\sim{\mathcal{N}}(0,{\bm{I}}_{d}),\kappa\sim\chi^{2}(\nu)/\nu$ (120) $$

It follows that,

$$ ${\mathbf{x}}_{t}=t\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]+(1-t)\mathbb{E }\big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}$ (121) $$

Moreover, following Eq. 2.10 in Albergo et al. (), the conditional vector field ${\bm{b}}({\mathbf{x}}_{t},t)$ can be defined as,

$$ ${\bm{b}}({\mathbf{x}}_{t},t)=\mathbb{E}[\dot{{\mathbf{x}}}_{t}|{\mathbf{x}}_{t }]=\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]-\mathbb{E}\big{[}\frac{ \epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}$ (122) $$

From Eqns. 121 and 122, the conditional vector field can be simplified as,

$$ ${\bm{b}}({\mathbf{x}}_{t},t)=\frac{{\mathbf{x}}_{t}-\mathbb{E}\big{[}\frac{ \epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}}{t}$ (123) $$

This concludes the proof.

### A.9 Connection to Denoising Score Matching

We start by defining the score for the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$. The pdf for the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^
{2}{\bm{I}}_{d},\nu)$ can be expressed as (ignoring the normalization constant):

$$ $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})\;\;\propto\;\;\bigg{[}1+\frac{1}{\nu \sigma_{t}^{2}}\big{(}{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\big{)}^{\top} \big{(}{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\big{)}\bigg{]}^{\frac{-(\nu+d) }{2}}$ (124) $$

Therefore,

$$ $\displaystyle\nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ $\displaystyle=-\frac{\nu+d}{2}\nabla_{{\mathbf{x}}_{t}}\log\bigg{[}1+\frac{1}{ \nu\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|^{2}_{2}\bigg{]}$ (125) $\displaystyle=-\frac{\nu+d}{2}\frac{1}{1+\frac{1}{\nu\sigma_{t}^{2}}\|{\mathbf {x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|^{2}_{2}}\frac{2}{\nu\sigma_{t}^{2}}\bigg{(} {\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\bigg{)}$ (126) $$

Denoting $d_{1}=\frac{1}{\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|^{2}
_{2}$ for convenience, we have,

$$ $\nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=-\bigg{(} \frac{\nu+d}{\nu+d_{1}}\bigg{)}\frac{1}{\sigma_{t}^{2}}\bigg{(}{\mathbf{x}}_{t }-\mu_{t}{\mathbf{x}}_{0}\bigg{)}$ (127) $$

In Denoising Score Matching (DSM) (Vincent, ), the following objective is minimized,

$$ $L_{\text{DSM}}=\mathbb{E}_{t\sim p(t),{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0}) ,{\mathbf{x}}_{t}\sim q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})}\Bigg{[}\gamma(t)\| \nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})-{\mathbf{s} }_{\theta}({\mathbf{x}}_{t},t)\|_{2}^{2}\Bigg{]}$ (128) $$

for some loss weighting function $\gamma(t)$. Parameterizing the score estimator ${\mathbf{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ as:

$$ ${\mathbf{s}}_{\theta}({\mathbf{x}}_{t},t)=-\bigg{(}\frac{\nu+d}{\nu+d_{1}} \bigg{)}\frac{1}{\sigma_{t}^{2}}\bigg{(}{\mathbf{x}}_{t}-\mu_{t}D_{\bm{\theta} }({\mathbf{x}}_{t},\sigma_{t})\bigg{)}$ (129) $$

With this parameterization of ${\mathbf{s}}_{\theta}({\mathbf{x}}_{t},t)$ and some choice of $\gamma(t)$, the DSM objective can be further simplified as follows,

$$ $\displaystyle L_{\text{DSM}}$ $\displaystyle=\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\mathbb{E}_ {t}\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I),\kappa\sim\chi^{2}(\nu)/\nu}\Big{[ }\Big{(}\frac{\nu+d}{\nu+d_{1}}\Big{)}^{2}\Big{\|}{\mathbf{x}}_{0}-D_{\bm{ \theta}}(\mu_{t}{\mathbf{x}}_{0}+\sigma_{t}\frac{\epsilon}{\sqrt{\kappa}}, \sigma_{t})\Big{\|}^{2}_{2}\Big{]}$ (130) $\displaystyle=\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\mathbb{E}_ {t}\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I),\kappa\sim\chi^{2}(\nu)/\nu}\Big{[ }\lambda({\mathbf{x}}_{t},\nu,t)\Big{\|}{\mathbf{x}}_{0}-D_{\bm{\theta}}(\mu_{ t}{\mathbf{x}}_{0}+\sigma_{t}\frac{\epsilon}{\sqrt{\kappa}},\sigma_{t})\Big{\| }^{2}_{2}\Big{]}$ (131) $$

where $\lambda({\mathbf{x}}_{t},\nu,t)=\Big{(}\frac{\nu+d}{\nu+d_{1}}\Big{)}^{2}$, which is equivalent to a scaled version of the simplified denoising loss Eq. 9. This concludes the proof. As an additional caveat, the score parameterization in Eq. 129 depends on $d_{1}=\frac{1}{\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|^{2}
_{2}$, which can be approximated during inference as, $d_{1}\approx\frac{1}{\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}{\bm{D}}_{\theta
}({\mathbf{x}}_{t},\sigma_{t})\|^{2}_{2}$

### A.10 ODE Reformulation and Connections to Inverse Problems.

Recall the ODE dynamics in terms of the denoiser can be specified as,

$$ $\frac{d{\mathbf{x}}_{t}}{dt}=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}- \bigg{[}-\frac{\dot{\sigma}_{t}}{\sigma_{t}}+\frac{\dot{\mu}_{t}}{\mu_{t}} \bigg{]}({\mathbf{x}}_{t}-\mu_{t}{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t}, \sigma_{t}))$ (132) $$

From Eq. 129, we have,

$$ ${\mathbf{x}}_{t}-\mu_{t}D_{\bm{\theta}}({\mathbf{x}}_{t},\sigma_{t})=-\sigma_{ t}^{2}\bigg{(}\frac{\nu+d_{1}}{\nu+d}\bigg{)}\nabla_{{\mathbf{x}}_{t}}\log p({ \mathbf{x}}_{t},t)$ (133) $$

where $d_{1}=\frac{1}{\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|^{2}
_{2}$. Substituting the above result in the ODE dynamics, we obtain the reformulated ODE in Eq. 13.

$$ $\frac{d{\mathbf{x}}_{t}}{dt}=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}+ \sigma_{t}^{2}\Big{(}\frac{\nu+d_{1}}{\nu+d}\Big{)}\bigg{[}\frac{\dot{\mu}_{t} }{\mu_{t}}-\frac{\dot{\sigma}_{t}}{\sigma_{t}}\bigg{]}\nabla_{{\mathbf{x}}} \log p({\mathbf{x}}_{t},t)$ (134) $$

Since the term $d_{1}$ is data-dependent, we can estimate it during inference as $d_{1}\approx d^{\prime}_{1}=\frac{1}{\sigma_{t}^{2}}\|{\mathbf{x}}_{t}-\mu_{t}
{\bm{D}}_{\theta}({\mathbf{x}}_{t},\sigma_{t})\|^{2}_{2}$. Thus, the above ODE can be reformulated as,

$$ $\frac{d{\mathbf{x}}_{t}}{dt}=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}+ \sigma_{t}^{2}\Big{(}\frac{\nu+d^{\prime}_{1}}{\nu+d}\Big{)}\bigg{[}\frac{\dot {\mu}_{t}}{\mu_{t}}-\frac{\dot{\sigma}_{t}}{\sigma_{t}}\bigg{]}\nabla_{{ \mathbf{x}}}\log p({\mathbf{x}}_{t},t)$ (135) $$

Tweedie’s Estimate and Inverse Problems. Given the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^
{2}{\bm{I}}_{d},\nu)$, we have,

$$ ${\mathbf{x}}_{t}=\mu_{t}{\mathbf{x}}_{0}+\sigma_{t}\frac{\epsilon}{\sqrt{ \kappa}},\quad\epsilon\sim{\mathcal{N}}(0,{\bm{I}}_{d}),\kappa\sim\chi^{2}(\nu )/\nu$ (136) $$

It follows that,

$$ $\displaystyle{\mathbf{x}}_{t}$ $\displaystyle=\mu_{t}\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]+\sigma_{t} \mathbb{E}\Big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\Big{]}$ (137) $\displaystyle\mathbb{E}[{\mathbf{x}}_{0}|{\mathbf{x}}_{t}]$ $\displaystyle=\frac{1}{\mu_{t}}\Bigg{(}{\mathbf{x}}_{t}-\sigma_{t}\mathbb{E} \Big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\Big{]}\Bigg{)}$ (138) $\displaystyle\approx\frac{1}{\mu_{t}}\Bigg{(}{\mathbf{x}}_{t}+\sigma_{t}^{2} \Big{(}\frac{\nu+d^{\prime}_{1}}{\nu+d}\Big{)}\nabla_{{\mathbf{x}}}\log p({ \mathbf{x}}_{t},t)\Bigg{)}$ (139) $$

which gives us an estimate of the predicted ${\mathbf{x}}_{0}$ at any time $t$. Moreover, the framework can also be extended for solving inverse problems using diffusion models. More specifically, given a conditional signal ${\mathbf{y}}$, the goal is to simulate the ODE,

$$ $\displaystyle\frac{d{\mathbf{x}}_{t}}{dt}$ $\displaystyle=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}^{2}\Big {(}\frac{\nu+d^{\prime}_{1}}{\nu+d}\Big{)}\bigg{[}\frac{\dot{\mu}_{t}}{\mu_{t} }-\frac{\dot{\sigma}_{t}}{\sigma_{t}}\bigg{]}\nabla_{{\mathbf{x}}}\log p({ \mathbf{x}}_{t}|{\mathbf{y}})$ (140) $\displaystyle=\frac{\dot{\mu}_{t}}{\mu_{t}}{\mathbf{x}}_{t}+\sigma_{t}^{2}\Big {(}\frac{\nu+d^{\prime}_{1}}{\nu+d}\Big{)}\bigg{[}\frac{\dot{\mu}_{t}}{\mu_{t} }-\frac{\dot{\sigma}_{t}}{\sigma_{t}}\bigg{]}\Big{[}w_{t}\nabla_{{\mathbf{x}}} \log p({\mathbf{y}}|{\mathbf{x}}_{t})+\nabla_{{\mathbf{x}}}\log p({\mathbf{x}} _{t})\Big{]}$ (141) $$

where the above decomposition follows from $p({\mathbf{x}}_{t}|y)\propto p({\mathbf{y}}|{\mathbf{x}}_{t})^{w_{t}}p({
\mathbf{x}}_{t})$ and the weight $w_{t}$ represents the *guidance weight* of the distribution $p({\mathbf{y}}|{\mathbf{x}}_{t})$.
The term $\log p({\mathbf{y}}|{\mathbf{x}}_{t})$ can now be approximated using existing posterior approximation methods in inverse problems (Chung et al., ; Song et al., ; Mardani et al., ; Pandey et al., )

### A.11 Connections to Robust Statistical Estimators

Here, we derive the expression for the gradient of the $\gamma$-power divergence between the forward and the reverse posteriors (denoted by $q$ and $p_{\theta}$, respectively for notational convenience) i.e., $\nabla_{\theta}D_{\gamma}\!\left(q~{}\|~{}p_{\theta}\right)$. By the definition of the $\gamma$-power divergence,

$$ $D_{\gamma}\!\left(q~{}\|~{}p_{\theta}\right)=\frac{1}{\gamma}\big{[}{\mathcal{ C}}_{\gamma}(q,p_{\theta})-{\mathcal{H}}_{\gamma}(q)\big{]},\quad\gamma\in(-1, 0)\cup(0,\infty)$ (142) $$

$$ ${\mathcal{H}}_{\gamma}(p)=-\left\|p\right\|_{1+\gamma}=-\bigg{(}\int p({ \mathbf{x}})^{1+\gamma}d{\mathbf{x}}\bigg{)}^{\frac{1}{1+\gamma}}\qquad{ \mathcal{C}}_{\gamma}(q,p)=-\int q({\mathbf{x}})\bigg{(}\frac{p({\mathbf{x}})} {\left\|p\right\|_{1+\gamma}}\bigg{)}^{\gamma}d{\mathbf{x}}$ (143) $$

Therefore,

$$ $\nabla_{\theta}D_{\gamma}\!\left(q~{}\|~{}p_{\theta}\right)=\frac{1}{\gamma} \nabla_{\theta}{\mathcal{C}}_{\gamma}(q,p_{\theta})$ (144) $$

Consequently, we next derive an expression for $\nabla_{\theta}{\mathcal{C}}_{\gamma}(q,p_{\theta})$.

$$ $\displaystyle\nabla_{\theta}{\mathcal{C}}_{\gamma}(q,p_{\theta})$ $\displaystyle=-\nabla_{\theta}\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({ \mathbf{x}})}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma}d{\mathbf{ x}}$ (145) $\displaystyle=-\gamma\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({\mathbf{x}} )}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma-1}\nabla_{\theta} \Bigg{(}\frac{p_{\theta}({\mathbf{x}})}{\left\|p_{\theta}\right\|_{1+\gamma}} \Bigg{)}d{\mathbf{x}}$ (146) $\displaystyle=-\gamma\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({\mathbf{x}} )}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma-1}\frac{\left\|p_{ \theta}\right\|_{1+\gamma}\nabla_{\theta}p_{\theta}({\mathbf{x}})-p_{\theta}({ \mathbf{x}})\nabla_{\theta}\left\|p_{\theta}\right\|_{1+\gamma}}{\left\|p_{ \theta}\right\|_{1+\gamma}^{2}}d{\mathbf{x}}$ (147) $$

From the definition of $\left\|p_{\theta}\right\|_{1+\gamma}$,

$$ $\displaystyle\left\|p_{\theta}\right\|_{1+\gamma}$ $\displaystyle=\bigg{(}\int p_{\theta}({\mathbf{x}})^{1+\gamma}d{\mathbf{x}} \bigg{)}^{\frac{1}{1+\gamma}}$ (148) $\displaystyle\nabla_{\theta}\left\|p_{\theta}\right\|_{1+\gamma}$ $\displaystyle=\nabla_{\theta}\bigg{(}\int p_{\theta}({\mathbf{x}})^{1+\gamma}d {\mathbf{x}}\bigg{)}^{\frac{1}{1+\gamma}}$ (149) $\displaystyle=\frac{1}{1+\gamma}\bigg{(}\int p_{\theta}({\mathbf{x}})^{1+ \gamma}d{\mathbf{x}}\bigg{)}^{-\frac{\gamma}{1+\gamma}}\int\nabla_{\theta}\big {(}p_{\theta}({\mathbf{x}})^{1+\gamma}\big{)}d{\mathbf{x}}$ (150) $\displaystyle=\frac{1}{1+\gamma}\bigg{(}\int p_{\theta}({\mathbf{x}})^{1+ \gamma}d{\mathbf{x}}\bigg{)}^{-\frac{\gamma}{1+\gamma}}(1+\gamma)\int p_{ \theta}({\mathbf{x}})^{\gamma}\nabla_{\theta}p_{\theta}({\mathbf{x}})d{\mathbf {x}}$ (151) $\displaystyle\nabla_{\theta}\left\|p_{\theta}\right\|_{1+\gamma}$ $\displaystyle=\bigg{(}\int p_{\theta}({\mathbf{x}})^{1+\gamma}d{\mathbf{x}} \bigg{)}^{-\frac{\gamma}{1+\gamma}}\int p_{\theta}({\mathbf{x}})^{\gamma} \nabla_{\theta}p_{\theta}({\mathbf{x}})d{\mathbf{x}}$ (152) $\displaystyle=\frac{\left\|p_{\theta}\right\|_{1+\gamma}}{\bigg{(}\int p_{ \theta}({\mathbf{x}})^{1+\gamma}d{\mathbf{x}}\bigg{)}}\int p_{\theta}({\mathbf {x}})^{\gamma}\nabla_{\theta}p_{\theta}({\mathbf{x}})d{\mathbf{x}}$ (153) $\displaystyle=\frac{\left\|p_{\theta}\right\|_{1+\gamma}}{\bigg{(}\int p_{ \theta}({\mathbf{x}})^{1+\gamma}d{\mathbf{x}}\bigg{)}}\int p_{\theta}({\mathbf {x}})^{1+\gamma}\nabla_{\theta}\log p_{\theta}({\mathbf{x}})d{\mathbf{x}}$ (154) $\displaystyle=\left\|p_{\theta}\right\|_{1+\gamma}\int\underbrace{\frac{p_{ \theta}({\mathbf{x}})^{1+\gamma}}{\bigg{(}\int p_{\theta}({\mathbf{x}})^{1+ \gamma}d{\mathbf{x}}\bigg{)}}}_{=\tilde{p}_{\theta}({\mathbf{x}})}\nabla_{ \theta}\log p_{\theta}({\mathbf{x}})d{\mathbf{x}}$ (155) $\displaystyle=\left\|p_{\theta}\right\|_{1+\gamma}\mathbb{E}_{\tilde{p}_{ \theta}({\mathbf{x}})}[\nabla_{\theta}\log p_{\theta}({\mathbf{x}})]$ (156) $$

where we denote $\tilde{p}_{\theta}({\mathbf{x}})=\frac{p_{\theta}({\mathbf{x}})^{1+\gamma}}{
\bigg{(}\int p_{\theta}({\mathbf{x}})^{1+\gamma}d{\mathbf{x}}\bigg{)}}$ for notational convenience. Substituting the above result in Eq. 147, we have the following simplification,

$$ $\displaystyle\nabla_{\theta}{\mathcal{C}}_{\gamma}(q,p_{\theta})$ $\displaystyle=-\gamma\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({\mathbf{x}} )}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma-1}\frac{\left\|p_{ \theta}\right\|_{1+\gamma}\nabla_{\theta}p_{\theta}({\mathbf{x}})-p_{\theta}({ \mathbf{x}})\nabla_{\theta}\left\|p_{\theta}\right\|_{1+\gamma}}{\left\|p_{ \theta}\right\|_{1+\gamma}^{2}}d{\mathbf{x}}$ (157) $\displaystyle=-\gamma\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({\mathbf{x}} )}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma-1}\frac{\left\|p_{ \theta}\right\|_{1+\gamma}\nabla_{\theta}p_{\theta}({\mathbf{x}})-p_{\theta}({ \mathbf{x}})\left\|p_{\theta}\right\|_{1+\gamma}\mathbb{E}_{\tilde{p}_{\theta} ({\mathbf{x}})}[\nabla_{\theta}\log p_{\theta}({\mathbf{x}})]}{\left\|p_{ \theta}\right\|_{1+\gamma}^{2}}d{\mathbf{x}}$ (158) $\displaystyle=-\gamma\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({\mathbf{x}} )}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma-1}\frac{\nabla_{ \theta}p_{\theta}({\mathbf{x}})-p_{\theta}({\mathbf{x}})\mathbb{E}_{\tilde{p}_ {\theta}({\mathbf{x}})}[\nabla_{\theta}\log p_{\theta}({\mathbf{x}})]}{\left\| p_{\theta}\right\|_{1+\gamma}}d{\mathbf{x}}$ (159) $\displaystyle=-\gamma\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({\mathbf{x}} )}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma-1}\frac{p_{\theta}({ \mathbf{x}})\nabla_{\theta}\log p_{\theta}({\mathbf{x}})-p_{\theta}({\mathbf{x }})\mathbb{E}_{\tilde{p}_{\theta}({\mathbf{x}})}[\nabla_{\theta}\log p_{\theta }({\mathbf{x}})]}{\left\|p_{\theta}\right\|_{1+\gamma}}d{\mathbf{x}}$ (160) $$

$$ $\displaystyle\nabla_{\theta}{\mathcal{C}}_{\gamma}(q,p_{\theta})$ $\displaystyle=-\gamma\int q({\mathbf{x}})\bigg{(}\frac{p_{\theta}({\mathbf{x}} )}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{)}^{\gamma}\Big{(}\nabla_{\theta }\log p_{\theta}({\mathbf{x}})-\mathbb{E}_{\tilde{p}_{\theta}({\mathbf{x}})}[ \nabla_{\theta}\log p_{\theta}({\mathbf{x}})]\Big{)}d{\mathbf{x}}$ (162) $$

Plugging this result in Eq. 144, we have the following result,

$$ $\nabla_{\theta}D_{\gamma}\!\left(q~{}\|~{}p_{\theta}\right)=\frac{1}{\gamma} \nabla_{\theta}{\mathcal{C}}_{\gamma}(q,p_{\theta})=-\int q({\mathbf{x}})\bigg {(}\frac{p_{\theta}({\mathbf{x}})}{\left\|p_{\theta}\right\|_{1+\gamma}}\bigg{ )}^{\gamma}\Big{(}\nabla_{\theta}\log p_{\theta}({\mathbf{x}})-\mathbb{E}_{ \tilde{p}_{\theta}({\mathbf{x}})}[\nabla_{\theta}\log p_{\theta}({\mathbf{x}}) ]\Big{)}d{\mathbf{x}}$ (163) $$

This completes the proof. Intuitively, the second term inside the integral in Eq. 163 ensures *unbiasedness* of the gradients. Therefore, the scalar coefficient $\gamma$ controls the weighting on the likelihood gradient and can be set accordingly to ignore or model outliers when modeling the data distribution.

Figure: Algorithm 3 Training (t-Flow)

## Appendix B Extension to Flows

Here, we discuss an extension of our framework to flow matching models (Albergo et al., ; Lipman et al., ) with a Student-t base distribution. More specifically, we define a *straight-line* flow of the form,

$$ ${\mathbf{x}}_{t}=t{\mathbf{x}}_{1}+(1-t)\frac{\epsilon}{\sqrt{\kappa}},\quad \epsilon\sim{\mathcal{N}}(0,{\bm{I}}_{d}),\kappa\sim\chi^{2}(\nu)/\nu$ (164) $$

where ${\mathbf{x}}_{1}\sim p({\mathbf{x}}_{1})$. Intuitively, at a given time $t$, the flow defined in Eqn. 164 linearly interpolates between data and Student-t noise. Following Albergo et al. (), the conditional vector field which induces this interpolant can be specified as (proof in App. A.8)

$$ $\frac{d{\mathbf{x}}_{t}}{dt}={\bm{b}}({\mathbf{x}}_{t},t)=\frac{{\mathbf{x}}_{ t}-\mathbb{E}\big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}}{t}.$ (165) $$

We estimate $\mathbb{E}\big{[}\frac{\epsilon}{\sqrt{\kappa}}|{\mathbf{x}}_{t}\big{]}$ by minimizing the objective

$$ ${\mathcal{L}}(\theta)\;\;=\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_ {0})}\mathbb{E}_{t\sim{\mathcal{U}}[0,1]}\mathbb{E}_{\epsilon\sim{\mathcal{N}} (0,{\bm{I}}_{d})}\mathbb{E}_{\kappa\sim\chi^{2}(\nu)/\nu}\Big{[}\Big{\|} \epsilon_{\bm{\theta}}(t{\mathbf{x}}_{0}+(1-t)\frac{\epsilon}{\sqrt{\kappa}},t )-\frac{\epsilon}{\sqrt{\kappa}}\Big{\|}^{2}_{2}\Big{]}.$ (166) $$

We refer to this flow setup as *t-Flow*. To generate samples from our model, we simulate the ODE in Eq. 165 using Heun’s solver. Figure 5 illustrates the ease of transitioning from a Gaussian flow to t-Flow. Similar to Gaussian diffusion, transitioning to t-Flow requires very few lines of code change, making our method readily compatible with existing implementations of flow models.

## Appendix C Experimental Setup

### C.1 Unconditional Modeling

#### C.1.1 HRRR Dataset

We adopt the High-Resolution Rapid Refresh (HRRR) (Dowell et al., ) dataset, which is an operational archive of the US km-scale forecasting model. Among multiple dynamical variables in the dataset that exhibit heavy-tailed behavior, based on their dynamic range, we only consider the *Vertically Integrated Liquid* (VIL) and *Vertical Wind Velocity* at level 20 (denoted as w20) channels. How to cope with the especially non-Gaussian nature of such physical variables on km-scales, represents an entire subfield of climate model subgrid-scale parameterization (e.g., Guo et al. ()). We only use data for the years $2019-2020$ for training (17.4k samples) and the data for 2021 (8.7k samples) for testing; data before 2019 are avoided owing to non-stationaries associated with periodic version changes of the HRRR. Lastly, while the HRRR dataset spans the entire US, for simplicity, we work with regional crops of size $128\times 128$ (corresponding to $384\times 384$ km over the Central US). Unless specified otherwise, we perform z-score normalization using precomputed statistics as a preprocessing step. We do not perform any additional data augmentation.

Figure: Figure 6: Inverse CDF Normalization. Using Inverse CDF Normalization (INC) can help reduce channel dynamic range during training while providing accurate denormalization. (Top Panel) INC applied to the Vertically Integrated Liquid (VIL) channel in the HRRR dataset. (Bottom Panel) INC applied to the Vertical Wind Velocity (w20) channel in the HRRR dataset.
Refer to caption: x2.png

#### C.1.2 Baselines

Baseline 1: EDM. For standard Gaussian diffusion models, we use the recently proposed EDM (Karras et al., ) model, which shows strong empirical performance on various image synthesis benchmarks and has also been employed in recent work in weather forecasting and downscaling (Pathak et al., ; Mardani et al., ). To summarize, EDM employs the following denoising loss during training,

$$ ${\mathcal{L}}(\theta)\;\;\propto\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({ \mathbf{x}}_{0})}\mathbb{E}_{\sigma}\mathbb{E}_{{\mathbf{n}}\sim{\mathcal{N}}( 0,\sigma^{2}{\bm{I}}_{d})}\big{[}\lambda(\sigma)\|{\bm{D}}_{\bm{\theta}}({ \mathbf{x}}_{0}+{\mathbf{n}},\sigma)-{\mathbf{x}}_{0}\|^{2}_{2}\big{]}$ (167) $$

where the noise levels $\sigma$ are usually sampled from a LogNormal distribution, $p(\sigma)=\text{LogNormal}(\pi_{\text{mean}},\pi_{\text{std}}^{2})$

Baseline 2: EDM + Inverse CDF Normalization (INC). It is commonplace to perform z-score normalization as a data pre-processing step during training. However, since heavy-tailed channels usually exhibit a high dynamic range, using z-score normalization for such channels cannot fully compensate for this range, especially when working with diverse channels in downstream tasks in weather modeling. An alternative could be to use a more *stronger* normalization scheme like Inverse CDF normalization, which essentially involves the following key steps:

- 1.
Compute channel-wise 1-d histograms of the training data.
- 2.
Compute channel-wise empirical CDFs from the 1-d histograms computed in Step 1.
- 3.
Use the empirical CDFs from Step 2 to compute the CDF at each spatial location.
- 4.
For each spatial location with a CDF value $p$, replace its value by the value obtained by applying the Inverse CDF operation under the standard Normal distribution.

Fig. 6 illustrates the effect of performing normalization under this scheme. As can be observed, using such a normalization scheme can greatly reduce the dynamic range of a given channel while offering reliable denormalization. Moreover, since our normalization scheme only affects data preprocessing, we leave the standard EDM model parameters unchanged for this baseline.

Baseline 3: EDM + Per Channel-Preconditioning (PCP). Another alternative to account for extreme values in the data (or high dynamic range) could be to instead add more heavy-tailed noise during training. This can be controlled by modulating the $\pi_{\text{mean}}$ and $\pi_{\text{std}}$ parameters based on the dynamic range of the channel under consideration. Recall that these parameters control the domain of noise levels $\sigma$ used during EDM model training. In this work, we use a simple heuristic to modulate these parameters based on the normalized channel dynamic range (denoted as $d$). More specifically, we set,

$$ $\pi_{\text{mean}}=-1.2+\alpha*\text{RBF}(d,1.0,\beta)$ (168) $$

where RBF denotes the Radial Basis Function kernel with radius=1.0, parameter $\beta$ and a magnitude scaling factor $\alpha$. We keep $\pi_{\text{std}}=1.2$ fixed for all channels. Intuitively, this implies that a higher normalized dynamic range (near 1.0) corresponds to sampling the noise levels $\sigma$ from a more heavy-tailed distribution during training. This is natural since a signal with a larger magnitude would need more amount of noise to convert it to noise during the forward diffusion process. In this work, we set $\alpha=3.0$, $\beta=2.0$, which yields $\pi_{\text{mean}}^{\text{vil}}=1.8$ and $\pi_{\text{mean}}^{\text{w20}}=0.453$ for the VIL and w20 channels, respectively. It is worth noting that, unlike the previous baseline, we use z-score normalization as a preprocessing step for this baseline. Lastly, we keep other EDM parameters during training and sampling fixed.

Baseline 4. Gaussian Flow. Since we also extend our framework to flow matching models (Albergo et al., ; Lipman et al., ), we also compare with a linear one-sided interpolant with a Gaussian base distribution. More specifically,

$$ ${\mathbf{x}}_{t}=t{\mathbf{x}}_{1}+(1-t)\epsilon,\quad\epsilon\sim{\mathcal{N} }(0,{\bm{I}}_{d})$ (169) $$

Similar to t-Flow (Section B), we train the Gaussian flow with the following objective,

$$ ${\mathcal{L}}(\theta)\;\;=\;\;\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_ {0})}\mathbb{E}_{t\sim{\mathcal{U}}[0,1]}\mathbb{E}_{\epsilon\sim{\mathcal{N}} (0,{\bm{I}}_{d})}\Big{[}\Big{\|}\epsilon_{\bm{\theta}}(t{\mathbf{x}}_{0}+(1-t) \epsilon,t)-\epsilon\Big{\|}^{2}_{2}\Big{]}.$ (170) $$

#### C.1.3 Evaluation

Here, we describe our scoring protocol used in Tables 2 and 3 in more detail.

Kurtosis Ratio (KR). Intuitively, sample kurtosis characterizes the heavy-tailed behavior of a distribution and represents the fourth-order moment. Higher kurtosis represents greater deviations from the central tendency, such as from outliers in the data. In this work, given samples from the underlying train/test set, we generate 20k samples from our model. We then flatten all the samples and compute empirical kurtosis for both the underlying samples from the train/test set (denoted as $k_{\text{data}}$) and our model (denoted as $k_{\text{sim}}$). The Kurtosis ratio is then computed as,

$$ $\text{KR}=\Big{|}1-\frac{k_{\text{sim}}}{k_{\text{data}}}\Big{|}$ (171) $$

Lower values of this ratio imply a better estimation of the underlying sample kurtosis.

Skewness Ratio (SR). Intuitively, sample skewness represents the asymmetry of a tailed distribution and represents the third-order moment. In this work, given samples from the underlying train/test set, we generate 20k samples from our model. We then flatten all the samples and compute empirical skewness for both the underlying samples from the train/test set (denoted as $s_{\text{data}}$) and our model (denoted as $s_{\text{sim}}$). The Skewness ratio is then computed as,

$$ $\text{SR}=\Big{|}1-\frac{s_{\text{sim}}}{s_{\text{data}}}\Big{|}$ (172) $$

Lower values of this ratio imply a better estimation of the underlying sample skewness.

Kolmogorov-Smirnov 2-Sample Test (KS). The KS (Massey, ) statistic measures the maximum difference between the CDFs of two distributions. For heavy-tailed distributions, evaluating the KS statistic at the tails could provide a useful measure of the efficacy of our model in estimating tails reliably. To evaluate the KS statistic at the tails, similar to prior metrics, we generate 20k samples from our model. We then flatten all samples in the generated and train/test sets, followed by retaining samples lying above the 99.9th percentile (quantifying right tails/extreme region) or below the 0.1th percentile (quantifying the left tail/extreme region). Lastly, we compute the KS statistic between the retained samples from the generated and the train/test sets individually for each tail and average the KS statistic values for both tails to obtain an average KS score. The final score estimates how well the model might capture both tails. As an exception,. for the VIL channel, we report KS scores only for the right tail due to the absence of a left tail in the underlying samples for this channel (see Fig. 6 (first column) for 1-d intensity histograms for this channel). Lower values of the KS statistic imply better density assignment at the tails by the model.

Histogram Comparisons. As a qualitative metric, comparing 1-d intensity histograms between the generated and the original samples from the train/test set can serve as a reliable proxy to assess the tail estimation capabilities of all models.

**Table 5: Comparison between design choices and specific hyperparameters between EDM (Karras et al., ) (+ related baselines) and t-EDM (Ours, Section 3.5) for unconditional modeling (Section 4.1). INC: Inverse CDF Normalization baselines, PCP: Per-Channel Preconditioning baseline, VIL: Vertically Integrated Liquid channel in the HRRR dataset, w20: Vertical Wind Velocity at level 20 channel in the HRRR dataset, NFE: Number of Function Evaluations**
|  | Parameters | EDM (+INC, +PCP) | t-EDM | Flow/t-Flow |
| --- | --- | --- | --- | --- |
| Preconditioner | $c_{\text{skip}}$ | $\sigma_{\text{data}}^{2}/\bigg{(}\sigma^{2}+\sigma_{\text{data}}^{2}\bigg{)}$ | $\sigma_{\text{data}}^{2}/\bigg{(}\frac{\nu}{\nu-2}\sigma^{2}+\sigma_{\text{<br>data}}^{2}\bigg{)}$ | 0 |
| $c_{\text{out}}$ | $\sigma\cdot\sigma_{\text{data}}^{2}\big{/}\sqrt{\sigma^{2}+\sigma_{\text{data}<br>}^{2}}$ | $\sqrt{\frac{\nu}{\nu-2}}\sigma\cdot\sigma_{\text{data}}^{2}\big{/}\sqrt{\frac{<br>\nu}{\nu-2}\sigma^{2}+\sigma_{\text{data}}^{2}}$ | 1 |  |
| $c_{\text{in}}$ | $1\big{/}\sqrt{\sigma^{2}+\sigma_{\text{data}}^{2}}$ | $1\big{/}\sqrt{\frac{\nu}{\nu-2}\sigma^{2}+\sigma_{\text{data}}^{2}}$ | 1 |  |
| $c_{\text{noise}}$ | $\frac{1}{4}\log{\sigma}$ | $\frac{1}{4}\log{\sigma}$ | $\sigma$ |  |
| Training | $\sigma$ | $\log{\sigma}\sim{\mathcal{N}}(\pi_{\text{mean}},\pi_{\text{std}}^{2})$ | $\log{\sigma}\sim{\mathcal{N}}(\pi_{\text{mean}},\pi_{\text{std}}^{2})$ | $\sigma=1-t$, $t\sim{\mathcal{U}}(0,1)$ |
| $\mu_{t}$ | 1 | 1 | t |  |
| $\lambda(\sigma)$ | $1/c_{\text{out}}^{2}(\sigma)$ | $1/c_{\text{out}}^{2}(\sigma,\nu)$ | 1 |  |
| Loss | Eq. 167 | Eq. 12 | t-Flow - Eq. 166<br>Gaussian Flow - Eq. 170 |  |
| Sampling | Solver | Heun’s (2nd order) | Heun’s (2nd order) | Heun’s (2nd order) |
| ODE | Eq. 11, ${\mathbf{x}}_{T}\sim{\mathcal{N}}(0,{\bm{I}}_{d})$ | Eq. 11, ${\mathbf{x}}_{T}\sim t_{d}(0,{\bm{I}}_{d},\nu)$ | t-Flow: Eq. 165, ${\mathbf{x}}_{0}\sim t_{d}(0,{\bm{I}}_{d},\nu)$<br>Flow: Eq. 165, ${\mathbf{x}}_{0}\sim{\mathcal{N}}(0,{\bm{I}}_{d})$ |  |
| Discretization | $\big{(}\sigma_{\text{max}}^{\frac{1}{\rho}}+\frac{i}{N-1}\big{[}\sigma_{\text{<br>min}}^{\frac{1}{\rho}}-\sigma_{\text{max}}^{\frac{1}{\rho}}\big{]}\big{)}^{<br>\frac{1}{\rho}}$ | $\big{(}\sigma_{\text{max}}^{\frac{1}{\rho}}+\frac{i}{N-1}\big{[}\sigma_{\text{<br>min}}^{\frac{1}{\rho}}-\sigma_{\text{max}}^{\frac{1}{\rho}}\big{]}\big{)}^{<br>\frac{1}{\rho}}$ | $\big{(}\sigma_{\text{max}}^{\frac{1}{\rho}}+\frac{i}{N-1}\big{[}\sigma_{\text{<br>min}}^{\frac{1}{\rho}}-\sigma_{\text{max}}^{\frac{1}{\rho}}\big{]}\big{)}^{<br>\frac{1}{\rho}}$ |  |
| Scaling: $\mu_{t}$ | 1 | 1 | t |  |
| Schedule: $\sigma_{t}$ | t | t | 1-t |  |
| Hyperparameters | $\sigma_{\text{data}}$ | 1.0 | 1.0 | N/A |
| $\nu$ | $\infty$ | VIL: $\{3,5,7\}$<br>w20: $\{3,5,7\}$ | Flow: $\infty$<br>t-Flow ($\downarrow$)<br>VIL=$\{3,5,7\}$<br>w20=$\{5,7,9\}$ |  |
| $\pi_{\text{mean}}$, $\pi_{\text{std}}$ | EDM: -1.2, 1.2<br>EDM (+INC) : -1.2, 1.2<br>EDM (+PCP) ($\downarrow$)<br>VIL: 1.8, 1.2<br>w20: 0.453, 1.2 | -1.2, 1.2 | N/A |  |
| $\sigma_{\text{max}}$, $\sigma_{\text{min}}$ | 80, 0.002 | 80, 0.002 | 1.0, 0.01 |  |
| NFE | 18 | 18 | 25 |  |
| $\rho$ | 7 | 7 | 7 |  |

#### C.1.4 Denoiser Architecture

We use the DDPM++ architecture from (Karras et al., ; Song et al., ). We set the base channel multiplier to 32 and the per-resolution channel multiplier to [1,2,2,4,4] with self-attention at resolution 16. The rest of the hyperparameters remain unchanged from Karras et al. (), which results in a model size of around 12M parameters.

#### C.1.5 Training

We adopt the same training hyperparameters from Karras et al. () for training all models. Model training is distributed across 4 DGX nodes, each with 8 A100 GPUs, with a total batch size of 512. We train all models for a maximum budget of 60Mimg and select the best-performing model in terms of qualitative 1-D histogram comparisons.

#### C.1.6 Sampling

For the EDM and related baselines (INC, PCP), we use the ODE solver presented in Karras et al. (). For the t-EDM models, as presented in Section 3.5, our sampler is the same as EDM with the only difference in the sampling of initial latents from a Student-t distribution instead (See Fig. 2 (Right)).
For Flow and t-Flow, we numerically simulate the ODE in Eq. 165 using the 2nd order Heun’s solver with the timestep discretization proposed in Karras et al. (). For evaluation, we generate 20k samples from each model.

We summarize our experimental setup in more detail for unconditional modeling in Table 5.

#### C.1.7 Extended Results on Unconditional Modeling

Sample Visualization. We visualize samples generated from the t-EDM and t-Flow models for the VIL and w20 channels in Figs. 7-10

Visualization of 1-d histograms. Similar to Fig. 3, we present additional results on histogram comparisons between different baselines and our proposed methods for the VIL and w20 channels in Figs. 11 and 12.

### C.2 Conditional Modeling

**Table 6: Parameters in the HRRR dataset used for conditional modeling tasks.**
| Parameter | Model Levels | Height Levels (m) |
| --- | --- | --- |
| Zonal Wind (u) | 1,2,3,4,5,6,7,8,9,10,11,13,15,20 | 10 |
| Meridonal Wind (v) | 1,2,3,4,5,6,7,8,9,10,11,13,15,20 | 10 |
| Geopotential Height (z) | 1,2,3,4,5,6,7,8,9,10,11,13,15,20 | None |
| Humidity (q) | 1,2,3,4,5,6,7,8,9,10,11,13,15,20 | None |
| Pressure (p) | 1,2,3,4,5,6,7,8,9,10,11,13,15,20 | None |
| Temperature (t) | 1,2,3,4,5,6,7,8,9,10,11,13,15,20 | 2 |
| Radar Reflectivity (refc) | N/A | Integrated |

#### C.2.1 HRRR Dataset for Conditional Modeling

Similar to unconditional modeling (See App. C.1.1), we use the HRRR dataset for conditional modeling at the 128 x 128 resolution. More specifically, for a lead time of 1hr, we sample (input, output) pairs at time $\tau$ and $\tau+1$, respectively. For the input, at time $\tau$, we use a state vector consisting of a combination of 86 atmospheric channels (including the channel to be predicted at time $\tau$), which are summarized in Table 6. For the output, at time $\tau+1$, we use either the Vertically Integrated Liquid (VIL) or Vertical Wind Velocity at level 20 (w20) channels, depending on the prediction task. Unless specified otherwise, we perform z-score normalization using precomputed statistics as a preprocessing step without any additional data augmentation.

#### C.2.2 Baselines

We adopt the standard EDM (Karras et al., ) for conditional modeling as our baseline.

#### C.2.3 Denoiser Architecture

We use the DDPM++ architecture from (Karras et al., ; Song et al., ). We set the base channel multiplier to 32 and the per-resolution channel multiplier to [1,2,2,4,4] with self-attention at resolution 16. Additionally, our noisy state ${\mathbf{x}}$ is channel-wise concatenated with an 86-channel conditioning signal, increasing the total number of input channels in the denoiser to 87. The number of output channels remains 1 since we are predicting only a single VIL/w20 channel. However, the increase in the number of parameters is minimal since only the first convolutional layer in the denoiser is affected. Therefore, our denoiser is around 12M parameters. The rest of the hyperparameters remain unchanged from Karras et al. ().

#### C.2.4 Training

We adopt the same training hyperparameters from Karras et al. () for training all conditional models. Model training is distributed across 4 DGX nodes, each with 8 A100 GPUs, with a total batch size of 512. We train all models for a maximum budget of 60Mimg.

#### C.2.5 Sampling

For both EDM and t-EDM models, we use the ODE solver presented in Karras et al. (). For the t-EDM models, as presented in Section 3.5, our sampler is the same as EDM with the only difference in the sampling of initial latents from a Student-t distribution instead (See Fig. 2 (Right)). For a given input conditioning state, we generate an ensemble of predictions of size 16 by randomly initializing our ODE solver with different random seeds. All other sampling parameters remain unchanged from our unconditional modeling setup (see App. C.1.6).

#### C.2.6 Evaluation

Root Mean Square Error (RMSE). is a standard evaluation metric used to measure the difference between the predicted values and the true values Chai & Draxler (). In the context of our problem, let ${\bm{x}}$ be the true target and $\hat{{\bm{x}}}$ be the predicted value. The RMSE is defined as:

$$ $\text{RMSE}=\sqrt{\mathbb{E}\left[\|{\bm{x}}-\hat{{\bm{x}}}\|^{2}\right]}.$ $$

This metric captures the average magnitude of the residuals, i.e., the difference between the predicted and true values. A lower RMSE indicates better model performance, as it suggests the predicted values are closer to the true values on average. RMSE is sensitive to large errors, making it an ideal choice for evaluating models where minimizing large deviations is critical.

Continuous Ranked Probability Score (CRPS). is a measure used to evaluate probabilistic predictions Wilks (). It compares the entire predicted distribution $F(\hat{{\bm{x}}})$ with the observed data point ${\bm{x}}$. For a probabilistic forecast with cumulative distribution function (CDF) $F$, and the true value ${\bm{x}}$, the CRPS can be formulated as follows:

$$ $\text{CRPS}(F,{\bm{x}})=\int_{-\infty}^{\infty}\left(F(y)-\mathbb{I}(y\geq{\bm {x}})\right)^{2}\,dy,$ $$

where $\mathbb{I}(\cdot)$ is the indicator function. Unlike RMSE, CRPS provides a more comprehensive evaluation of both the location and spread of the predicted distribution. A lower CRPS indicates a better match between the forecast distribution and the observed data. It is especially useful for probabilistic models that output a distribution rather than a single-point prediction.

Spread-Skill Ratio (SSR). is used to assess over/under-dispersion in probabilistic forecasts. Spread measures the uncertainty in the ensemble forecasts and can be represented by computing the standard deviation of the ensemble members. Skill represents the accuracy of the mean of the ensemble forecasts and can be represented by computing the RMSE between the ensemble mean and the observations.

Scoring Criterion. Since CRPS and SSR metrics are based on predicting ensemble forecasts for a given input state, we predict an ensemble of size 16 for 4000 samples from the VIL/w20 test set. We then enumerate window sizes of 16 x 16 across the spatial resolution of the generated sample (128 x 128). Since the VIL channel is quite sparse, we filter out windows with a maximum value of less than a threshold (1.0 for VIL) and compute the CRPS, SSR, and RMSE metrics for all remaining windows. As an additional caveat, we note that while it is common to roll out trajectories for weather forecasting, in this work, we only predict the target at the immediate next time step.

#### C.2.7 Extended Results on Conditional Modeling

Sample Visualization. We visualize samples generated from the t-EDM and t-Flow models for the VIL and w20 channels in Figs. 8 and 9

### C.3 Toy Experiments

**Table 7: Comparison between design choices and specific hyperparameters between EDM (Karras et al., ) and t-EDM (Ours, Section 3.5) for the Toy dataset analysis in Fig. 1. NFE: Number of Function Evaluations**
|  | Parameters | EDM | t-EDM |
| --- | --- | --- | --- |
| Preconditioner | $c_{\text{skip}}$ | $\sigma_{\text{data}}^{2}/\bigg{(}\sigma^{2}+\sigma_{\text{data}}^{2}\bigg{)}$ | $\sigma_{\text{data}}^{2}/\bigg{(}\frac{\nu}{\nu-2}\sigma^{2}+\sigma_{\text{<br>data}}^{2}\bigg{)}$ |
| $c_{\text{out}}$ | $\sigma\cdot\sigma_{\text{data}}^{2}\big{/}\sqrt{\sigma^{2}+\sigma_{\text{data}<br>}^{2}}$ | $\sqrt{\frac{\nu}{\nu-2}}\sigma\cdot\sigma_{\text{data}}^{2}\big{/}\sqrt{\frac{<br>\nu}{\nu-2}\sigma^{2}+\sigma_{\text{data}}^{2}}$ |  |
| $c_{\text{in}}$ | $1\big{/}\sqrt{\sigma^{2}+\sigma_{\text{data}}^{2}}$ | $1\big{/}\sqrt{\frac{\nu}{\nu-2}\sigma^{2}+\sigma_{\text{data}}^{2}}$ |  |
| $c_{\text{noise}}$ | $\frac{1}{4}\log{\sigma}$ | $\frac{1}{4}\log{\sigma}$ |  |
| Training | $\sigma$ | $\log{\sigma}\sim{\mathcal{N}}(\pi_{\text{mean}},\pi_{\text{std}}^{2})$ | $\log{\sigma}\sim{\mathcal{N}}(\pi_{\text{mean}},\pi_{\text{std}}^{2})$ |
| $s(t)$ | 1 | 1 |  |
| $\lambda(\sigma)$ | $1/c_{\text{out}}^{2}(\sigma)$ | $1/c_{\text{out}}^{2}(\sigma,\nu)$ |  |
| Loss | Eq. 2 in Karras et al. () | Eq. 12 |  |
| Sampling | Solver | Heun’s (2nd order) | Heun’s (2nd order) |
| ODE | Eq. 11, ${\mathbf{x}}_{T}\sim{\mathcal{N}}(0,{\bm{I}}_{d})$ | Eq. 11, ${\mathbf{x}}_{T}\sim t_{d}(0,{\bm{I}}_{d},\nu)$ |  |
| Discretization | $\big{(}\sigma_{\text{max}}^{\frac{1}{\rho}}+\frac{i}{N-1}\big{[}\sigma_{\text{<br>min}}^{\frac{1}{\rho}}-\sigma_{\text{max}}^{\frac{1}{\rho}}\big{]}\big{)}^{<br>\frac{1}{\rho}}$ | $\big{(}\sigma_{\text{max}}^{\frac{1}{\rho}}+\frac{i}{N-1}\big{[}\sigma_{\text{<br>min}}^{\frac{1}{\rho}}-\sigma_{\text{max}}^{\frac{1}{\rho}}\big{]}\big{)}^{<br>\frac{1}{\rho}}$ |  |
| Scaling: $s(t)$ | 1 | 1 |  |
| Schedule: $\sigma(t)$ | t | t |  |
| Hyperparameters | $\sigma_{\text{data}}$ | 1.0 | 1.0 |
| $\nu$ | $\infty$ | ${\mathbf{x}}_{1}=20$, ${\mathbf{x}}_{2}\in\{4,7,10\}$ |  |
| $\pi_{\text{mean}}$, $\pi_{\text{std}}$ | -1.2, 1.2 | -1.2, 1.2 |  |
| $\sigma_{\text{max}}$, $\sigma_{\text{min}}$ | 80, 0.002 | 80, 0.002 |  |
| NFE | 18 | 18 |  |
| $\rho$ | 7 | 7 |  |

Dataset. For the toy illustration in Fig. 1, we work with the Neals Funnel dataset (Neal, ), which is commonly used in the MCMC literature (Brooks et al., ) due to its challenging geometry. The underlying generative process for Neal’s funnel can be specified as follows:

$$ $p({\mathbf{x}}_{1},{\mathbf{x}}_{2})={\mathcal{N}}({\mathbf{x}}_{1};0,3){ \mathcal{N}}({\mathbf{x}}_{2};0,\exp{{\mathbf{x}}_{1}/2}).$ (173) $$

For training, we randomly generate 1M samples from the generative process in Eq. 173 and perform z-score normalization as a pre-processing step.

Baselines and Models. For the standard Gaussian diffusion model baseline (2nd column in Fig. 1), we use EDM with standard hyperparameters as presented in Karras et al. (). Consequently, for heavy-tailed diffusion models (columns 3-5 in Fig. 1), we use the *t-EDM* instantiation of our framework as presented in Section 3.5. Since the hyperparameter $\nu$ is key in our framework, we tune $\nu$ for each individual dimension in our toy experiments. We fix $\nu$ to 20 for the ${\mathbf{x}}_{1}$ dimension and vary it between $\nu\in\{4,7,10\}$ to illustrate controllable tail estimation along the ${\mathbf{x}}_{2}$ dimension.

Denoiser Architecture. For modeling the underlying denoiser ${\bm{D}}_{\theta}({\mathbf{x}},\sigma)$, we use a simple MLP for all toy models. At the input, we concatenate the 2-dimensional noisy state vector ${\mathbf{x}}$ with the noise level $\sigma$. We use two hidden layers of size 64, followed by a linear output layer. This results in around 8.5k parameters for the denoiser. We share the same denoiser architecture across all toy models.

Training. We optimize the training objective in Eq. 12 for both t-EDM and EDM (See Fig. 2 (Left)). Our training hyperparameters are the same as proposed in Karras et al. (). We train all toy models for a fixed duration of 30M samples and choose the last checkpoint for evaluation.

Sampling. For the EDM baseline, we use the ODE solver presented in Karras et al. (). For the t-EDM models, as presented in Section 3.5, our sampler is the same as EDM with the only difference in the sampling of initial latents from a Student-t distribution instead (See Fig. 2 (Right)). For visualization of the generated samples in Fig. 1, we generate 1M samples for each model.

Overall, our experimental setup for the toy dataset analysis in Fig. 1 is summarized in Table 7.

## Appendix D Optimal Noise Schedule Design

In this section we discuss a strategy for choosing the parameter $\sigma_{\text{max}}$ (denoted by $\sigma$ in this section for notational brevity) in a more principled manner as compared to EDM (Karras et al., ). More specifically, our approach involves directly estimating $\sigma$ from the empirically observed samples which circumvents the need to rely on ad-hoc choices of this parameter which can affect downstream sampler performance.

The main idea behind our approach is minimizing the statistical *mutual information* between datapoints from the underlying data distribution, ${\mathbf{x}}_{0}\sim p_{\text{data}}$, and their noisy counterparts ${\mathbf{x}}_{\sigma}\sim p({\mathbf{x}}_{\sigma})$. While a trivial (and non-practical) way to achieve this objective could be to set a large enough $\sigma$ i.e. $\sigma\rightarrow\infty$, we instead minimize the mutual information $I({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})$ while ensuring the magnitude of $\sigma$ to be as small as possible. Formally, our objective can be defined as,

$$ $\min_{\sigma^{2}}\sigma^{2}\quad\text{subject to}\quad I({\mathbf{x}}_{0},{ \mathbf{x}}_{\sigma})=0$ (174) $$

As we will discuss later, minimizing this constrained objective provides a more principled way to obtain $\sigma$ from the underlying data statistics for a specific level of mutual information desired by the user. Next, we first simplify the form of $I({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})$, followed by a discussion on the estimation of $\sigma$ in the context of EDM and t-EDM. We also extend to the case of non-i.i.d noise.

Simplification of $I({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})$. The mutual information $I({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})$ can be stated and simplified as follows,

$$ $\displaystyle I({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})$ $\displaystyle=D_{\mathrm{KL}}\!\left(p({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma}) ~{}\|~{}p({\mathbf{x}}_{0})p({\mathbf{x}}_{\sigma})\right)$ (175) $\displaystyle=\int p({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})\log\frac{p({ \mathbf{x}}_{0},{\mathbf{x}}_{\sigma})}{p({\mathbf{x}}_{0})p({\mathbf{x}}_{ \sigma})}d{\mathbf{x}}_{0}d{\mathbf{x}}_{\sigma}$ (176) $\displaystyle=\int p({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})\log\frac{p({ \mathbf{x}}_{\sigma}|{\mathbf{x}}_{0})}{p({\mathbf{x}}_{\sigma})}d{\mathbf{x}} _{0}d{\mathbf{x}}_{\sigma}$ (177) $\displaystyle=\int p({\mathbf{x}}_{\sigma}|{\mathbf{x}}_{0})p({\mathbf{x}}_{0} )\log\frac{p({\mathbf{x}}_{\sigma}|{\mathbf{x}}_{0})}{p({\mathbf{x}}_{\sigma}) }d{\mathbf{x}}_{0}d{\mathbf{x}}_{\sigma}$ (178) $$

$$ $\displaystyle=\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\Bigg{[} \int p({\mathbf{x}}_{\sigma}|{\mathbf{x}}_{0})\log\frac{p({\mathbf{x}}_{\sigma }|{\mathbf{x}}_{0})}{p({\mathbf{x}}_{\sigma})}d{\mathbf{x}}_{\sigma}\Bigg{]}$ (179) $\displaystyle=\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\Big{[}D_{ \mathrm{KL}}\!\left(p({\mathbf{x}}_{\sigma}|{\mathbf{x}}_{0})~{}\|~{}p({ \mathbf{x}}_{\sigma})\right)\Big{]}$ (180) $$

### D.1 Design for EDM

Given the simplification in Eqn. 180, the optimization problem in Eqn. 174 reduces to the following,

$$ $\min_{\sigma^{2}}\sigma^{2}\quad\text{subject to}\quad\mathbb{E}_{{\mathbf{x}} _{0}\sim p({\mathbf{x}}_{0})}\Big{[}D_{\mathrm{KL}}\!\left(p({\mathbf{x}}_{ \sigma}|{\mathbf{x}}_{0})~{}\|~{}p({\mathbf{x}}_{\sigma})\right)\Big{]}=0$ (181) $$

Since at $\sigma_{\text{max}}$ we expect the marginal distribution $p({\mathbf{x}}_{\sigma})$ to converge to the generative prior (i.e. completely destory the structure of the data), we approximate $p({\mathbf{x}}_{\sigma})\approx{\mathcal{N}}(0,\sigma^{2}{\bm{I}}_{d})$.
With this simplification, the Lagrangian for the optimization problem in Eqn. 174 can be specified as,

$$ $\sigma^{2}_{*}=\arg\min_{\sigma^{2}}\sigma^{2}+\lambda\mathbb{E}_{{\mathbf{x}} _{0}\sim p({\mathbf{x}}_{0})}\Big{[}D_{\mathrm{KL}}\!\left({\mathcal{N}}({ \mathbf{x}}_{\sigma};{\mathbf{x}}_{0},\sigma^{2})~{}\|~{}{\mathcal{N}}({ \mathbf{x}}_{\sigma};0,\sigma^{2})\right)\Big{]}$ (182) $$

Setting the gradient w.r.t $\sigma^{2}$ = 0,

$$ $1-\frac{\lambda}{\sigma^{4}}\mathbb{E}_{{\mathbf{x}}_{0}}[{\mathbf{x}}_{0}^{ \top}{\mathbf{x}}_{0}]=0$ (183) $$

which implies,

$$ $\sigma^{2}=\sqrt{\lambda\mathbb{E}_{{\mathbf{x}}_{0}}[{\mathbf{x}}_{0}^{\top}{ \mathbf{x}}_{0}]}$ (184) $$

For an empirical dataset,

$$ $\sigma^{2}=\sqrt{\frac{\lambda}{N}\sum_{i=1}^{N}\|{\mathbf{x}}_{i}\|_{2}^{2}}$ (185) $$

This allows us to choose a $\sigma_{\text{max}}$ from the underlying data statistics during training or sampling. It is worth noting that the choice of the multiplier $\lambda$ impacts the magnitude of $\sigma_{\text{max}}$. However, this parameter can be chosen in a principled manner. At $\sigma_{\text{max}}^{2}=\sigma^{2}_{*}$, the estimate of the mutual information is given by:

$$ $I_{*}({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})=\frac{1}{\sigma_{*}^{2}}\mathbb{ E}_{{\mathbf{x}}_{0}}[{\mathbf{x}}_{0}^{\top}{\mathbf{x}}_{0}]=\sqrt{\frac{ \mathbb{E}_{{\mathbf{x}}_{0}}[{\mathbf{x}}_{0}^{\top}{\mathbf{x}}_{0}]}{ \lambda}}$ (186) $$

which implies,

$$ $\lambda=\frac{\mathbb{E}_{{\mathbf{x}}_{0}}[{\mathbf{x}}_{0}^{\top}{\mathbf{x} }_{0}]}{I_{*}({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})^{2}}$ (187) $$

The above result provides a way to choose $\lambda$. Given the dataset stastistics, $\mathbb{E}_{{\mathbf{x}}_{0}}[{\mathbf{x}}_{0}^{\top}{\mathbf{x}}_{0}]$, the user can specify an acceptable level of mutual information $I({\mathbf{x}}_{0},{\mathbf{x}}_{\sigma})$ to compute the correponding $\lambda$, which can then be used to find the corresponding minimum $\sigma_{\text{max}}$ required to achieve that level of mutual information. Next, we extend this analysis to t-EDM.

### D.2 Extension to t-EDM

In the case of t-EDM, we pose the optimization problem as follows,

$$ $\displaystyle\sigma^{*}=\arg\min_{\sigma^{2}}\sigma^{2}+\lambda\mathbb{E}_{{ \mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}\Big{[}D_{\gamma}\!\left(t_{d}({ \mathbf{x}}_{0},\sigma^{2}{\bm{I}}_{d},\nu)~{}\|~{}t_{d}(0,\sigma^{2}{\bm{I}}_ {d},\nu)\right)\Big{]}$ (188) $$

where $D_{\gamma}\!\left(q~{}\|~{}p\right)$ is the Gamma-power Divergence between two distributions q and p. From the definition of the $\gamma$-Power Divergence in Eqn. 7, we have,

$$ $D_{\gamma}\!\left(t_{d}({\mathbf{x}}_{0},\sigma^{2}{\bm{I}}_{d},\nu)~{}\|~{}t_ {d}(0,\sigma^{2}{\bm{I}}_{d},\nu)\right)=\underbrace{-\frac{1}{\nu\gamma}C_{ \nu,d}^{\frac{\gamma}{1+\gamma}}(1+\frac{d}{\nu-2})^{-\frac{\gamma}{1+\gamma}} }_{=f(\nu,d)}(\sigma^{2})^{-\frac{\gamma d}{2(1+\gamma)}-1}\mathbb{E}_{{ \mathbf{x}}_{0}}[{\mathbf{x}}_{0}^{\top}{\mathbf{x}}_{0}]$ (189) $$

where $C_{\nu,d}=\frac{\Gamma(\frac{\nu+d}{2})}{\Gamma(\frac{\nu}{2})(\nu\pi)^{\frac{
d}{2}}}$ and $\gamma=-\frac{2}{\nu+d}$
Solving the optimization problem yields the following optimal $\sigma$,

$$ $\sigma_{*}^{2}=\Big{[}\lambda f(\nu,d)\Big{(}\frac{\nu-2}{\nu-2+d}\Big{)} \mathbb{E}[{\mathbf{x}}_{0}^{\top}{\mathbf{x}}_{0}]\Big{]}^{\frac{\nu-2+d}{2( \nu-2)+d}}$ (190) $$

For an empirical dataset, we have the following simplification,

$$ $\sigma_{*}^{2}=\Big{[}\lambda f(\nu,d)\Big{(}\frac{\nu-2}{\nu-2+d}\Big{)}\frac {1}{N}\sum_{i}\|{\mathbf{x}}_{i}\|_{2}^{2}\Big{]}^{\frac{\nu-2+d}{2(\nu-2)+d}}$ (191) $$

### D.3 Extension to Correlated Gaussian Noise

We now extend our formulation for optimal noise schedule design to the case of correlated noise in the diffusion perturbation kernel. This is useful, especially for scientific applications where the data energy is distributed quite non-uniformly across the (data) spectrum. Let $\mathbf{R}=\mathbb{E}_{\mathbf{x_{0}}\sim p(\mathbf{x_{0}})}[\mathbf{x_{0}}
\mathbf{x_{0}}^{\top}]\in\mathbb{R}^{d\times d}$ denote the data correlation matrix. Let also consider the perturbation kernel $\mathcal{N}(0,\bm{\Sigma})$ for the postitve-defintie covaraince matrix $\bm{\Sigma}\in\mathbb{R}^{d\times d}$. Following the steps in equation 182, the Lagrangian for noise covariance estimation can be formulated as follows:

$$ $\displaystyle\min_{\bm{\Sigma}}$ $\displaystyle\;\text{trace}(\bm{\Sigma})+\lambda\mathbb{E}_{{\mathbf{x}}_{0} \sim p({\mathbf{x}}_{0})}\Big{[}D_{\mathrm{KL}}\!\left({\mathcal{N}}({\mathbf{ x}}_{0},{\bm{\Sigma}})~{}\|~{}{\mathcal{N}}(0,{\bm{\Sigma}})\right)\Big{]}$ (192) $\displaystyle\min_{\bm{\Sigma}}$ $\displaystyle\;\text{trace}(\bm{\Sigma})+\lambda\mathbb{E}_{{\mathbf{x}}_{0} \sim p({\mathbf{x}}_{0})}\Big{[}{\mathbf{x}}_{0}^{\top}{\bm{\Sigma}}^{-1}{ \mathbf{x}}_{0}\Big{]}$ (193) $\displaystyle\min_{\bm{\Sigma}}$ $\displaystyle\;\text{trace}(\bm{\Sigma})+\lambda\mathbb{E}_{{\mathbf{x}}_{0} \sim p({\mathbf{x}}_{0})}\Big{[}\text{trace}({\bm{\Sigma}}^{-1}{\mathbf{x}}_{0 }{\mathbf{x}}_{0}^{\top})\Big{]}$ (194) $\displaystyle\min_{\bm{\Sigma}}$ $\displaystyle\;\text{trace}(\bm{\Sigma})+\lambda\,\text{trace}({\bm{\Sigma}}^{ -1}\mathbb{E}_{{\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0})}[{\mathbf{x}}_{0}{ \mathbf{x}}_{0}^{\top}])$ (195) $\displaystyle\min_{\bm{\Sigma}}$ $\displaystyle\;\text{trace}(\bm{\Sigma})+\lambda\,\text{trace}(\bm{\Sigma}^{-1 }\mathbf{R})$ (196) $$

It can be shown that the optimal solution to this minimization problem is given by,
$\bm{\Sigma}^{*}=\sqrt{\lambda}\mathbf{R}^{1/2}$, where $\mathbf{R}^{1/2}$ denotes the matrix square root of $\mathbf{R}$. This implies that the noise energy must be distributed along the singular vectors of the correlaiton matrix, where the energy is proportional to the noise singular values. We include the proof below.

Proof. We define the objective:

$$ $f(\bm{\Sigma})=\text{trace}(\bm{\Sigma})+\lambda\,\text{trace}(\bm{\Sigma}^{-1 }\mathbf{R}).$ (197) $$

We compute the gradient of $f(\bm{\Sigma})$ with respect to $\bm{\Sigma}$:

$$ $\nabla_{\bm{\Sigma}}f(\bm{\Sigma})=\frac{\partial}{\partial\bm{\Sigma}}\text{ trace}(\bm{\Sigma})+\lambda\frac{\partial}{\partial\bm{\Sigma}}\text{trace}( \bm{\Sigma}^{-1}\mathbf{R}).$ (198) $$

The gradient of the first term is straightforward:

$$ $\frac{\partial}{\partial\bm{\Sigma}}\text{trace}(\bm{\Sigma})=\mathbf{I}.$ (199) $$

For the second term, using the matrix calculus identity, the gradient is:

$$ $\frac{\partial}{\partial\bm{\Sigma}}\left(\lambda\,\text{trace}(\bm{\Sigma}^{- 1}\mathbf{R})\right)=-\lambda\bm{\Sigma}^{-1}\mathbf{R}\bm{\Sigma}^{-1}.$ (200) $$

Combining these results, the total gradient is:

$$ $\nabla_{\bm{\Sigma}}f(\bm{\Sigma})=\mathbf{I}-\lambda\bm{\Sigma}^{-1}\mathbf{R }\bm{\Sigma}^{-1}.$ (201) $$

Setting the gradient to zero to find the critical point:

$$ $\mathbf{I}-\lambda\bm{\Sigma}^{-1}\mathbf{R}\bm{\Sigma}^{-1}=0.$ (202) $$

$$ $\bm{\Sigma}^{-1}\mathbf{R}\bm{\Sigma}^{-1}=\frac{1}{\lambda}\mathbf{I}.$ (203) $$

which implies,

$$ $\mathbf{R}=\lambda\bm{\Sigma}^{2}.$ (204) $$

$$ $\bm{\Sigma}=\sqrt{\lambda}\mathbf{R}^{1/2}.$ (205) $$

which completes the proof.

## Appendix E Log-Likelihood for t-EDM

Here, we present a method to estimate the log-likelihood for the generated samples using the ODE solver for t-EDM (see Section 3.5). Our analysis is based on the likelihood computation in continuous-time diffusion models as discussed in Song et al. () (Appendix D.2). More specifically, given a probability-flow ODE,

$$ $\frac{d{\mathbf{x}}_{t}}{dt}={\bm{f}}_{\bm{\theta}}({\mathbf{x}}_{t},t),$ (206) $$

with a vector field ${\bm{f}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$, Song et al. () propose to estimate the log-likelihood of the model as follows,

$$ $\log p({\mathbf{x}}_{0})=\log p({\mathbf{x}}_{T})+\int_{0}^{T}\nabla\cdot{\bm{ f}}_{\bm{\theta}}({\mathbf{x}}_{t},t)dt.$ (207) $$

The divergence of the vector field ${\bm{f}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ is further estimated using the Skilling-Hutchinson trace estimator (Skilling, ; Hutchinson, ) as follows,

$$ $\nabla\cdot{\bm{f}}_{\bm{\theta}}({\mathbf{x}}_{t},t)=\mathbb{E}_{p(\bm{ \epsilon})}[\bm{\epsilon}^{\top}\nabla{\bm{f}}_{\bm{\theta}}({\mathbf{x}}_{t}, t)\bm{\epsilon}]$ (208) $$

where usually, $\bm{\epsilon}\sim{\mathcal{N}}(\bm{0},{\bm{I}}_{d})$. For t-EDM ODE, this estimate can be further simplified as follows,

$$ $\displaystyle\nabla\cdot{\bm{f}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ $\displaystyle=\mathbb{E}_{p(\bm{\epsilon})}[\bm{\epsilon}^{\top}\nabla{\bm{f}} _{\bm{\theta}}({\mathbf{x}}_{t},t)\bm{\epsilon}]$ (209) $\displaystyle=\mathbb{E}_{p(\bm{\epsilon})}\Big{[}\bm{\epsilon}^{\top}\nabla \Big{(}\frac{{\mathbf{x}}_{t}-{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{t} \Big{)}\bm{\epsilon}\Big{]}$ (210) $\displaystyle=\mathbb{E}_{p(\bm{\epsilon})}\Big{[}\bm{\epsilon}^{\top}\Big{(} \frac{{\bm{I}}_{d}-\nabla_{{\mathbf{x}}_{t}}{\bm{D}}_{\bm{\theta}}({\mathbf{x} }_{t},t)}{t}\Big{)}\bm{\epsilon}\Big{]}$ (211) $\displaystyle=\frac{1}{t}\mathbb{E}_{p(\bm{\epsilon})}\Big{[}\bm{\epsilon}^{ \top}\bm{\epsilon}-\bm{\epsilon}^{\top}\nabla_{{\mathbf{x}}_{t}}{\bm{D}}_{\bm{ \theta}}({\mathbf{x}}_{t},t)\bm{\epsilon}\Big{]}$ (212) $\displaystyle=\frac{1}{t}\mathbb{E}_{p(\bm{\epsilon})}\Big{[}\bm{\epsilon}^{ \top}\bm{\epsilon}-\bm{\epsilon}^{\top}\nabla_{{\mathbf{x}}_{t}}{\bm{D}}_{\bm{ \theta}}({\mathbf{x}}_{t},t)\bm{\epsilon}\Big{]}$ (213) $\displaystyle=\frac{1}{t}\Big{[}d-\mathbb{E}_{p(\bm{\epsilon})}\big{(}\bm{ \epsilon}^{\top}\nabla_{{\mathbf{x}}_{t}}{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{ t},t)\bm{\epsilon}\big{)}\Big{]}$ (214) $$

where $d$ is the data dimensionality. Thus, the log-likelihood can be specified as,

$$ $\log p({\mathbf{x}}_{0})=\log p({\mathbf{x}}_{T})+\int_{0}^{T}\frac{1}{t}\Big{ [}d-\mathbb{E}_{p(\bm{\epsilon})}\big{(}\bm{\epsilon}^{\top}\nabla{\bm{D}}_{ \bm{\theta}}({\mathbf{x}}_{t},t)\bm{\epsilon}\big{)}\Big{]}dt.$ (215) $$

When $\epsilon\sim{\mathcal{N}}(0,\sigma^{2}{\bm{I}}_{d})$, the above result can be re-formulated as,

$$ $\log p({\mathbf{x}}_{0})=\log p({\mathbf{x}}_{T})+\int_{0}^{T}\frac{1}{t}\Big{ [}d-\frac{1}{\sigma^{2}}\mathbb{E}_{p(\bm{\epsilon})}\big{(}\bm{\epsilon}^{ \top}\nabla{\bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\bm{\epsilon}\big{)}\Big{ ]}dt.$ (216) $$

Moreover, using the first-order taylor series expansion

$$ $\displaystyle D_{\theta}({\mathbf{x}}+\epsilon)=D_{\theta}({\mathbf{x}})+ \nabla D_{\theta}\epsilon+{\mathcal{O}}(\|\epsilon\|^{2})$ (217) $$

For a sufficiently small $\sigma$, higher-order terms in ${\mathcal{O}}(\|\epsilon\|^{2})$ can be ignored since $\mathbb{E}[\|\epsilon^{2}\|]=\sigma^{2}d$. Therefore,

$$ $\displaystyle D_{\theta}({\mathbf{x}}+\epsilon)$ $\displaystyle\approx D_{\theta}({\mathbf{x}})+\nabla D_{\theta}\epsilon$ (218) $\displaystyle\epsilon^{\top}\nabla D_{\theta}\epsilon$ $\displaystyle\approx\epsilon^{\top}[D_{\theta}({\mathbf{x}}+\epsilon)-D_{ \theta}({\mathbf{x}})]$ (219) $\displaystyle\mathbb{E}_{\epsilon}[\epsilon^{\top}\nabla D_{\theta}\epsilon]$ $\displaystyle\approx\mathbb{E}_{\epsilon}[\epsilon^{\top}(D_{\theta}({\mathbf{ x}}+\epsilon)-D_{\theta}({\mathbf{x}}))]$ (220) $\displaystyle\mathbb{E}_{\epsilon}[\epsilon^{\top}\nabla D_{\theta}\epsilon]$ $\displaystyle\approx\mathbb{E}_{\epsilon}[\epsilon^{\top}D_{\theta}({\mathbf{x }}+\epsilon)]$ (221) $$

Therefore, the log-likelihood expression can be further simplified as,

$$ $\displaystyle\log p({\mathbf{x}}_{0})$ $\displaystyle=\log p({\mathbf{x}}_{T})+\int_{0}^{T}\frac{1}{t}\Big{[}d-\frac{1 }{\sigma^{2}}\mathbb{E}_{p(\bm{\epsilon})}\big{(}\bm{\epsilon}^{\top}\nabla{ \bm{D}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\bm{\epsilon}\big{)}\Big{]}dt$ (222) $\displaystyle\log p({\mathbf{x}}_{0})$ $\displaystyle=\log p({\mathbf{x}}_{T})+\int_{0}^{T}\frac{1}{t}\Big{[}d-\frac{1 }{\sigma^{2}}\mathbb{E}_{p(\bm{\epsilon})}\big{(}\bm{\epsilon}^{\top}D_{\theta }({\mathbf{x}}_{t}+\bm{\epsilon},t)\big{)}\Big{]}dt$ (223) $$

The advantage of this simplification is that we don’t need to rely on expensive jacobian-vector products in Eq. 215. However, since the denoiser now depends on $\epsilon$, monte-carlo approximation of the expectation in the above equation could be computationally expensive for many samples $\epsilon$

## Appendix F Discussion and Limitations

### F.1 Related Work

Connections with Denoising Score Matching. For the perturbation kernel $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=t_{d}(\mu_{t}{\mathbf{x}}_{0},\sigma_{t}^
{2}{\bm{I}}_{d},\nu)$, the denoising score matching (Vincent, ; Song et al., ) loss, ${\mathcal{L}}_{\text{DSM}}$, can be formulated as,

$$ ${\mathcal{L}}_{\text{DSM}}(\theta)\;\;\propto\;\;\mathbb{E}_{{\mathbf{x}}_{0} \sim p({\mathbf{x}}_{0})}\mathbb{E}_{t}\mathbb{E}_{\epsilon\sim{\mathcal{N}}(0 ,{\bm{I}}_{d})}\mathbb{E}_{\kappa\sim\chi^{2}(\nu)/\nu}\Big{[}\lambda({\mathbf {x}}_{t},\nu,t)\Big{\|}{\bm{D}}_{\bm{\theta}}(\mu_{t}{\mathbf{x}}_{0}+\sigma_{ t}\frac{\epsilon}{\sqrt{\kappa}},\sigma_{t})-{\mathbf{x}}_{0}\Big{\|}^{2}_{2} \Big{]}$ (224) $$

with the scaling factor $\lambda({\mathbf{x}}_{t},\nu,t)=[(\nu+d)/(\nu+d_{1})]^{2}$ where $d_{1}=(1/\sigma_{t}^{2})\|{\mathbf{x}}_{t}-\mu_{t}{\mathbf{x}}_{0}\|_{2}^{2}$ (proof in App. A.9). Therefore, the denoising score matching loss in our framework is equivalent to the simplified training objective in Eq. 9 scaled by a data-dependent coefficient. However, in this work, we do not explore this loss formulation and leave further exploration to future work.

Prior work in Heavy-Tailed Generative Modeling. The idea of exploring heavy-tailed priors for modeling heavy-tailed distributions has been explored in several works in the past. More specifically, Jaini et al. () argue that a Lipschitz flow map cannot change the tails of the base distribution significantly. Consequently, they use a heavy-tailed prior (modeled using a Student-t distribution) as the base distribution to learn Tail Adaptive flows (TAFs), which can model the tails more accurately. In this work, we make similar observations where standard diffusion models fail to accurately model the tails of real-world distributions. Consequently, Laszkiewicz et al. () assess the *tailedness* of each marginal dimension and set the prior accordingly. On a similar note, we note that learning the tail parameter $\nu$ spatially and across channels can provide greater modeling flexibility for downstream tasks and will be an important direction for future work on this problem. More recently, Kim et al. () introduce heavy-tailed VAEs (Kingma & Welling, ; Rezende & Mohamed, ) based on minimizing $\gamma$-power divergences (Eguchi, ). This is perhaps the closest connection of our method with prior work since we rely on $\gamma$-power divergences to minimize the divergence between heavy-tailed forward and reverse diffusion posteriors. However, VAEs often have scalability issues and tend to produce blurry artifacts (Dosovitskiy & Brox, ; Pandey et al., ). On the other hand, we work with diffusion models, which are known to scale well to large-scale modeling applications (Pathak et al., ; Mardani et al., ; Esser et al., ; Podell et al., ).
In another line of work, Yoon et al. () presents a framework for modeling heavy-tailed distributions using $\alpha$-stable Levy processes while Shariatian et al. () simplify the framework proposed in Yoon et al. () and instantiate it for more practical diffusion models like DDPM. In contrast, our work deals with Student-t noise, which in general (with the exceptions of Cauchy and the Gaussian distribution) is not $\alpha$-stable and, therefore, a distinct category of diffusion models for modeling heavy-tailed distributions. Moreover, prior works like Yoon et al. (); Shariatian et al. () rely on empirical evidence from light-tailed variants of small-scale datasets like CIFAR-10 (Krizhevsky, ) and their efficacy on actual large-scale scientific datasets like weather datasets remains to be seen.

Prior work in Diffusion Models. Our work is a direct extension of standard diffusion models in the literature (Karras et al., ; Ho et al., ; Song et al., ). Moreover, since it only requires a few lines of code change to transition from standard diffusion models to our framework, our work is directly compatible with popular families of latent diffusion models (Pandey et al., ; Rombach et al., ) and augmented diffusion models (Dockhorn et al., ; Pandey & Mandt, ; Singhal et al., ). Our work is also related to prior work in diffusion models on a more theoretical level. More specifically, PFGM++ (Xu et al., ) is a unique type of generative flow model inspired by electrostatic theory. It treats d-dimensional data as electrical charges in a $D+d$-dimensional space, where the electric field lines define a bijection between a heavy-tailed prior and the data distribution. $D$ is a hyperparameter controlling the shape of the electric fields that define the generative mapping. In essence, their method can be seen as utilizing a perturbation kernel:

$$ $\displaystyle p({\mathbf{x}}_{t}|{\mathbf{x}}_{0})\propto(||{\mathbf{x}}_{t}-{ \mathbf{x}}_{0}||_{2}^{2}+\sigma_{t}^{2}D))^{-\frac{D+d}{2}}=t_{d}({\mathbf{x} }_{0},{\sigma_{t}^{2}}{\bm{I}}_{d},D)$ $$

When setting $\nu=D$, the perturbation kernel becomes equivalent to that of t-EDM, indicating the Student-t perturbation kernel can be interpreted from another physical perspective — that of electrostatic fields and charges. The authors demonstrated that using an intermediate value for $D$ (or $\nu$) leads to improved robustness compared to diffusion models (where $D\to\infty$), due to the heavy-tailed perturbation kernel.

### F.2 Limitations

While our proposed framework works well for modeling heavy-tailed data, it is not without its limitations. Firstly, while the parameter $\nu$ offers controllability for tail estimation using diffusion models, it also increases the tuning budget by introducing an extra hyperparameter. Moreover, for diverse data channels, tuning $\nu$ per channel could be key to good estimation at the tails. This could result in a combinatorial explosion with manual tuning. Therefore, learning $\nu$ directly from the data could be an important direction for the practical deployment of heavy-tailed diffusion models. Secondly, our evaluation protocol relies primarily on comparing the statistical properties of samples obtained by flattening the generated or train/test set samples. One disadvantage of this approach is that our current evaluation metrics ignore the structure of the generated samples. In general, developing metrics like FID (Heusel et al., ) to assess the perceptual quality of synthetic data for scientific domains like weather analysis remains an important future direction.

Figure: Figure 7: Random samples generated from t-EDM (Top Panel) and EDM (Bottom Panel) for the Vertically Integrated Liquid (VIL) channel. KS: Kolmogorov-Smirnov 2-sample statistic. Samples have been scaled logarithmically for better visualization
Refer to caption: x3.png

Figure: Figure 8: Random samples generated from t-Flow (Top Panel) and Gaussian Flow (Bottom Panel) for the Vertically Integrated Liquid (VIL) channel. KS: Kolmogorov-Smirnov 2-sample statistic. Samples have been scaled logarithmically for better visualization
Refer to caption: x4.png

Figure: Figure 9: Random samples generated from t-EDM (Top Panel) and EDM (Bottom Panel) for the Vertical Wind Velocity (w20) channel. KS: Kolmogorov-Smirnov 2-sample statistic
Refer to caption: x5.png

Figure: Figure 10: Random samples generated from t-Flow (Top Panel) and Gaussian Flow (Bottom Panel) for the Vertical Wind Velocity (w20) channel. KS: Kolmogorov-Smirnov 2-sample statistic
Refer to caption: x6.png

Figure: Figure 11: 1-d Histogram Comparisons between samples from the generated and the Train/Test set for the Vertically Integrated Liquid (VIL, see Top Pandel) and Vertical Wind Velocity (w20, see Bottom Panel) channels using t-EDM (with varying $\nu$).
Refer to caption: x7.png

Figure: Figure 12: 1-d Histogram Comparisons between samples from the generated and the Train/Test set for the Vertically Integrated Liquid (VIL, see Top Pandel) and Vertical Wind Velocity (w20, see Bottom Panel) channels using t-Flow (with varying $\nu$).
Refer to caption: x8.png

**Figure 8: Qualitative visualization of samples generated from our conditional modeling for predicting the next state for the Vertically Integrated Liquid (VIL) channel. The ensemble mean represents the mean of ensemble predictions (16 in our case). Columns 2-3 represent two samples from the ensemble. The last column visualizes an animation of all ensemble members (Best viewed in a dedicated PDF reader). For each sample, the rows correspond to predictions from EDM, t-EDM ($\nu=3$), and t-EDM ($\nu=5$) from top to bottom, respectively. Samples have been scaled logarithmically for better visualization**
| Ground Truth | Ensemble Mean | Prediction 1 | Prediction 2 | Ensemble (gif)<br>\animategraphics[loop, autoplay, width=0.2]1figures/nfp_vil/edm/sample_0/ensemble/vil_pred_0_015 |
| --- | --- | --- | --- | --- |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_vil/tedm_nu=3/sample_0/ensemble/vil_pred_0_015 |  |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_vil/tedm_nu=5/sample_0/ensemble/vil_pred_0_015 |  |
| Ground Truth | Ensemble Mean | Prediction 1 | Prediction 2 | Ensemble (gif)<br>\animategraphics[loop, autoplay, width=0.2]1figures/nfp_vil/edm/sample_1/ensemble/vil_pred_1_015 |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_vil/tedm_nu=3/sample_1/ensemble/vil_pred_1_015 |  |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_vil/tedm_nu=5/sample_1/ensemble/vil_pred_1_015 |  |

**Figure 9: Qualitative visualization of samples generated from our conditional modeling for predicting the next state for the Vertical Wind Velocity (w20) channel. The ensemble mean represents the mean of ensemble predictions (16 in our case). Columns 2-3 represent two samples from the ensemble. The last column visualizes an animation of all ensemble members (Best viewed in a dedicated PDF reader). For each sample, the rows correspond to predictions from EDM, t-EDM ($\nu=3$), and t-EDM ($\nu=5$) from top to bottom, respectively.**
| Ground Truth | Ensemble Mean | Prediction 1 | Prediction 2 | Ensemble (gif)<br>\animategraphics[loop, autoplay, width=0.2]1figures/nfp_w20/edm/sample_0/ensemble/w20_pred_0_015 |
| --- | --- | --- | --- | --- |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_w20/tedm_nu=3/sample_0/ensemble/w20_pred_0_015 |  |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_w20/tedm_nu=5/sample_0/ensemble/w20_pred_0_015 |  |
| Ground Truth | Ensemble Mean | Prediction 1 | Prediction 2 | Ensemble (gif)<br>\animategraphics[loop, autoplay, width=0.2]1figures/nfp_w20/edm/sample_1/ensemble/w20_pred_1_015 |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_w20/tedm_nu=3/sample_1/ensemble/w20_pred_1_015 |  |
|  |  |  | \animategraphics[loop, autoplay, width=0.2]1figures/nfp_w20/tedm_nu=5/sample_1/ensemble/w20_pred_1_015 |  |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Questions
* Since the preconditioning for EDM helps a good amount for unconditional generation (see Table 2), is not there some preconditioning for EDM in a similar spirit to the unconditional case that could help for the conditional generation case?
* As mentioned in strengths the t-EDM has a free parameter $\nu$ that can be tuned to predict more accurately the heavy-tail behavior for different datasets. However, the INC preconditioning for the EDM could be given a parameter by reparameterizing into a normal with variance $\sigma^2$ instead of a standard normal. (More precisely, this means that in step 4 in Baseline 2 in Appendix C.1.2. we would replace 'standard Normal distribution' to centered normal distribution with variance $\sigma^2.$) This change in variance is qualitatively different from a change in decay (which is what happens when we vary $\nu$) so this approach may not be theoretically justified, but it may work in practice. The PCP preconditioning already has parameters. Is there any claim for INC/PCP that can be made about whether tuning this parameters helps or not in capturing different type of heavy-tail distributions?

In summary, I am wondering if the performance of EDM can be increased only by preconditioning to reach that of t-EDM or if t-EDM does provide a significant difference not achievable by preconditioning.
soundness: 3
strengths: The t-EDM provides a straightforward and mathematically justified extension for diffusion models to heavy-tailed data. This is to be compared, for instance, to Yoon et al (2023) who also provide a framework for heavy-tailed diffusion models (using Levy processes), but whose implementation requires heavier mathematics and to make the model work in practice they make theoretically unjustified changes to the model. 

The authors do a good job presenting the concepts as natural extensions of what was already known in the field. 

The experiments make the case for t-EDM being useful for weather modeling, an application that requires capturing tails and extreme events. Although the present reviewer is not knowledgable in this application, it is clear that having an extension of EDM to heavy-tails is an important contribution to scientists working on weather forecasting and several other natural science applications. It is especially useful that this implementation requires only minimal changes over previous EDM implementation, so it is quick to implement for practitioners.

One advantage of the t-EDM is that it has a free parameter $\nu$ that can be tuned to model different datasets (see results in Table 4), providing a flexible framework useful for practitioners. 

Yoon, E. B., Park, K., Kim, S., & Lim, S. (2023). Score-based Generative Models with Lévy Processes. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, & S. Levine (Eds.), Advances in Neural Information Processing Systems (Vol. 36, pp. 40694–40707). Curran Associates, Inc.
summary: The authors provide a framework, termed t-EDM, for diffusion models where the added noise is heavy-tailed, following a Student-t distribution. They build upon the EDM introduced by Karras et al 2022. Interestingly, starting with the EDM framework, replacing the Gaussian noise by the Student-t distribution and changing the KL divergence in the objective function to $\gamma-$Power Divergence for a specific $\gamma$ leads to same loss function as that of the EDM, except that the noisy state is now drawn from a Student-t distribution. This is of practical interest since their framework can be obtained by only performing minimal changes to previous implementations of EDM. It is of theoretical interest that since the Student-t distribution with $\nu$ degrees of freedom converges to a Gaussian distribution as $\nu$ goes to infinity, t-EDM is a generalization of the EDM. The authors provide experiments on weather prediction showing that t-EDM outperforms EDM for both unconditional and conditional generation for certain metrics.
weaknesses: The experiments for unconditional generation in Table 2 show that the EDM (using preconditioning) outperforms the t-EDM for the KR and SR metric in the VIL channel with test data. This hinders the claim the authors make that standard diffusion models, even with preconditioning, fail to capture the heavy-tailed distributions (line 043) and that t-EDM outperforms standard diffusion models. Qualitatively, the plots in Figure 3 do show that t-EDM outperforms EDM with preconditioning in capturing the heavy-tailed behavior. However, I believe a justification of why the KR and SR metrics do not reflect this should be included. As addressed in Appendix F.2, this metrics rely on flattening the data. This may be one reason, but the metric KS also relies on flattening the data but does show that t-EDM outperforms EDM. In any case, I believe a more throrough justification or explanation is needed here.

Also see questions regarding preconditioning for EDM vs t-EDM. If these questions are successfully addressed I will raise my score.

Smaller comments: 
* in line 198 there is a typo, the equal is in the subscript and it should not.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: - Can the authors comment on the draw backs of using student t noise discussed in "Heavy-tailed denoising score matching" and why this makes sense in the context proposed?
- proposition 2 to me is not well defined. The construction of the stochastic process with student t increments. Can the authors discuss how such a process is constructed in more detail?
- In the proof of proposition 2 the authors derive the last term converging to Brownian motion, not a student process, can they please comment on this discrepancy.
- It seems unintuitive to define a diffusion model via its finite increments, only to show it converges to a continuous time diffusion and later make use of this perspective. Can the authors comment on why the did not simply define the noising process as a student-t diffusion in the first place and later discretise this?
soundness: 3
strengths: - The paper introduces a new way of training a score based model with heavy tailed noise, based on the student t distribution, and shows clear benefits in training.
- The experimental results are reasonably convincing.
summary: This paper adds to the line of work investigating training diffusion models using heavy-tailed noise. Specifically the paper investigates the setting where the increments of noise added in a discrete diffusion model are drawn from a student T distribution.

New methods are developed for training the diffusion model based on power divergences.
weaknesses: - The paper fails to discuss previous work, for example "Heavy-tailed denoising score matching" covering the subject matter.
- The paper lacks experiments on synthetic data which is generated as heavy tailed. While not needed, I believe it would help make the case that this method definitively helps with modeling heavy tails by demonstrating it on definitively heavy tailed data.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Figure 1 - I liked this! I thought it could be improved by adding axis labels including a description of what the colour map shows (log ground truth density?), calling the original density “ground truth” ("original hists” in the legend was slightly cryptic),  would add axis labels. 

I was a bit confused by the notation around the degree of freedom parameter v. Initially in figure 1 it is a vector, then in equation 2 it is a scalar, but then in the equations on line 155 we switch back to a vector with t_{2d}() being a composition of d independent 2-dimensional multivariate student-t distributions (although this isn’t explained). The issue around clarity here bleeds into the discussion of the divergences where then the divergence parameter \gamma would depend on the dimension. Balancing notational clarity and correctness is tricky here — perhaps say that "we’ll consider v to be constant across dimensions for the mathematical development, but the generalisation is simple" somewhere early on to clarify these issues?

There is a lack of clarity in my mind around the noising process and the approximations made when arriving at the denoising equation (equation 5). Wouldn’t the true denoising process be non-Markov? If you consider the student-t as a Gaussian scale mixture, I believe the latent scale parameter in your noising construction is effectively shared across all time-steps (a consequence of the equations on line 155). So, when generating x_t the whole history of previous x_{t+1:T} will inform us about the latent scale and so p(x_t | x_{t+1:T}) \ne  p(x_t | x_{t+1}). Have I misunderstood? Moreover, dropping the data dependent coefficient (v+d_1)/(v+d) is rather glossed over. In the Gaussian setting, dropping this can be somewhat justified by 1. it being an exact choice in the limit of small step sizes, and 2. at least it’s exact when the true distribution of x_0 is Gaussian. In the current case neither of these statements hold I believe. I felt that this should be discussed as the suboptimality of this choice in this case is likely far larger than in the usual case.  Finally, could you clarify how the continuous case fits in with this? I presume, both of the approximations mentioned above are made, and then we take the continuum limit?

Connecting the specific divergence used to a parameters in the model / denoising process (i.e. \gamma = 2/(v+d)) might be limiting e.g. the objective now can’t be used to learn the degree of freedom, the optimisation objective (8) cannot be used to compare between models that use different settings for v. Also, the correspondence to the traditional objective only then holds as q becomes Gaussian (ideally you’d like to recover the original ELBO without having to limit the model back to the DDPM). I’d have liked to have seen more discussion of this point. 

In table 1, it might be useful to add p(x_T) for the two cases (by the way do you mention somewhere that q(x_T) is a product of student-ts and that you set p(x_T) = q(x_T))?

Why not consider surface fields, arguably where modelling extremes have most impact? 

Have you considered a baseline that trains the EDM and then uses a base distribution for generation which is zero-mean unit-variance student-t, sweeping the degree of freedom parameter to find a suitable one? 

For the HRRR experiments, I presume that you are training off HRRR analysis (rather than, say forecasts) this should be mentioned (I didn’t see it in the main text or in the appendix). Also, is the time interval between inputs mentioned? Is it 1hr like in StormCast? 

The spatial fields shown in the appendix were useful. For the forecasts results shown in figure 8 and 9, it would be useful to show the persistence forecast.

For the unconditional sampling, it looks like on the test set t-EDM outperforms the baselines 1/3 of the time (and only on the KS metric). 

Is the KS metric to be more trusted? Why weren’t the other metrics applied to the forecasting case too? 

It would be interesting to see how t-Flow behaved during auto-regressive rollouts. I could imagine that the under-dispersion of EDM might actually help maintain stability in these cases, whereas heavier tails could lead to greater instability.
soundness: 3
strengths: I really liked the problem the paper was looking into — better characterisation of extremes in diffusion and flow based generative models — especially in the context of environmental forecasting.  I think that this is an important practical problem. I also liked the use of a simple example to analyse how well diffusion models capture tail behaviour. (In fact I would have liked to have seen more investigation to toy examples of this sort to pin down the failure models of diffusion more generally.) 

 I thought the writing was very clear and struck a good balance between communicating technical content and the high level idea. 

The technical contribution is interesting. I liked the fact that the authors provide both the discrete-time version and the continuous-time version of the new approach. The use of the power-divergence for producing a more tractable training objective was creative. 

HRRR was a sensible domain to test the new method on. A range of sensible metrics that are used by practitioners e.g. CRPS for the forecasting experiments. I liked the simple modifications to EDM using the inverse CDF transform and per-channel preconditioning.
summary: Diffusion models are widely used, but there has been little work studying how well they capture extreme events in the tail of distributions. This is especially important in environmental prediction where extreme events are often the ones that we care most about. This paper shows that standard approaches can fail to capture the extremes accurately. It develops a new form of diffusion by leveraging the multivariate Student-t distribution to form new noising and denoising processes in discrete and continuous time. It develops a method to fit these new processes using a power-divergence which avoids the intractability of the ELBO.  It then applies the new methods to modelling a pair of variables from HRRR, NOAA’s high-resolution local area system for the continental US. This product has recently been a focus of AI diffusion emulators e.g. in NVIDIA’s StormCast.
weaknesses: I would have liked to have seen a bit more clarity on the formulation of the denoising process (the sub-optimalities here seem significantly more severe than in the Gaussian case, see questions below for more detail). The use of the power-divergence for training was creative, but connecting the parameters of the model to the divergence parameters also seemed sub-optimal as it does not allow the framework to learn the degree of freedom parameter in the student-t distribution. 

The experimentation was a little light in that only two variables were considered in the conditional and unconditional settings, there were no experiments that considered forecasting multiple variables at one time, or autoregressive rollouts.

Setting the degree of freedom parameter via a parameter sweep is a bit painful. It’s a shame that it can’t be learned directly, especially if it is going to be different for different dimensions of x  (locations / variables / pressure levels etc.).

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1.  Can the authors discuss potential applications in other domains where rare events or extreme values are important, such as finance, seismology, or network traffic analysis. This would help demonstrate the broader impact of the method beyond just weather forecasting.
2. Do you think it possible to extend your framework into other heavy-tailed distributions beyond the Student-t, such as the Cauchy or Lévy distributions, or even unifying all these classes of perturbation kernels? Are there any theoretical or practical challenges you foresee in such extensions?
soundness: 4
strengths: 1. The use of multivariate Student-t distributions for diffusion models is innovative, allowing the model to better capture rare or extreme events with controllable tail behavior through a hyperparameter. It is a very brilliant idea. 
2. The overall presentation of this paper is very clear and understandable. 
3. As an extension of previous EDM and Rectified Flow, the authors established a new framework for Student-T distribution and flow, which is able to capture the rare events more effectively. 
4. The authors provide theoretical backing for their approach, linking it to robust statistical estimators and demonstrating its foundation in existing heavy-tailed distribution literature.
5. Choosing the weather dataset to verify the effectiveness is a great idea, and is novel.
summary: The paper presents a novel approach to adapting diffusion models for heavy-tail estimation by leveraging multivariate Student-t distributions. This method addresses the limitations of traditional Gaussian-based diffusion models in capturing rare or extreme events, which is particularly relevant for applications such as weather forecasting. The authors introduce extensions of existing diffusion and flow models, termed t-EDM and t-Flow, which utilize a Student-t prior to enable controllable tail generation through a single scalar hyperparameter. The framework is compatible with standard Gaussian diffusion models and requires minimal code changes. Empirical validation on high-resolution datasets, such as the HRRR weather dataset, demonstrates that the proposed t-EDM outperforms existing diffusion models in accurately modeling heavy-tailed distributions.
weaknesses: I like this paper a lot and I believe there are no major weakness. The only thing is that, there are some duplicated references in the reference list.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See weakness
soundness: 2
strengths: The approach appears interesting and addresses a relevant problem. Also, the methodology can be straightforwardly implemented using existing DM frameworks.
summary: The paper proposes a novel generative model that, in a sense, mirrors traditional diffusion models (DM) by substituting Gaussian noise in the noising process with noise following a Student’s t-distribution.

More precisely, the authors introduce a sequence of joint conditional distributions $q_{t-\Delta t, t}(x_{t-\Delta t, t} \,|\, x_0)$ over $\mathbb{R}^{2d}$, which are Student distributions with mean $[\mu_{t - \Delta t} x_0, \mu_t x_0]$, covariance matrix $\Sigma_t$, and degrees of freedom parameter $\nu > 1$. Leveraging properties of conditional Student distributions, they demonstrate that the marginals of $q_{t-\Delta t, t}(x_{t-\Delta t, t} \,|\, x_0)$ are also Student-distributed, as are the conditional distributions. Given a family of transitions 
$p_{t-\Delta t \, |\, t}$, the authors fit this family by minimizing a modified DM loss, where the standard Kullback-Leibler divergence is replaced with a power divergence. They then incorporate their chosen conditionals $q_{t-\Delta t, t}(x_{t-\Delta t, t} \,|\, x_0)$. Finally, the developed methodology is benchmarked against standard EDM and several of its variations for generating dynamical variables in the High-Resolution Rapid Refresh dataset.
weaknesses: In my opinion, there are several key points the authors should address.

First, a major concern is the soundness of the proposed method. In particular, it is unclear to me that the sequence of conditional distributions $ q $ introduced by the authors genuinely defines a Markov noising process. As such, the loss function they employ does not appear to serve as a meaningful metric. More precisely, it remains uncertain why this loss function would serve as a good proxy for measuring discrepancies between the model marginal $p_0^{\theta}$ and the data distribution. The authors acknowledge that the loss function is not a proper ELBO and claim that the limiting diffusion model is a special case of their method, but further justification of the loss function is needed.

My second concern relates to the results and mathematical formalism of the paper. The mathematical treatment seems quite weak. For instance, the proof in the first section of the appendix contains several typos and appears incorrect in its current form. The reference to Ding (2016) would be sufficient, as it already includes the result. I also found Proposition 2 challenging to interpret, and the accompanying proof does not clarify matters. For example, it is unclear what $f$ and $g$ represent, whether they are any functions satisfying the proposition's conditions, and whether such functions exist. Similarly, the term $ dS_t \sim t_d(0, dt, \nu + d)$ lacks clarity. I am also unclear on the main conclusions to be drawn from Proposition 2.

My third concern is related to the experiments and the metrics used. First, there is a large gap in performance when using $\nu = 3$ versus $\nu = 5$, raising questions about the robustness and reliability of the method. This gap suggests that the chosen metrics (KR, KS, SR) might be sensitive only to the tails of the target distributions. Additional metrics such as precision, recall, and density coverage could provide a more balanced assessment. Moreover, the claims regarding t-EDM's performance appear somewhat overstated, given the tables and the likely margin of error. In my opinion, the experiments should include toy examples for a deeper comparison with existing diffusion models.

Finally, the authors should provide a comprehensive literature review on extensions of diffusion models, acknowledging previous relevant works. The current literature review is too limited.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: - How do we tune the parameter $\nu$ in practice? It looks from Table 3 like the choice matters a lot, so it might be nice to say something about how one should choose a suitable value of $\nu$ or know that one has found a good value.
- How confident are you that the 1-dimensional summary statistics you are looking in Table 2 (and maybe Table 3 also) at appropriate? From looking at Figure 3, t-EDM seems clearly qualitatively superior to the other methods. The fact that this is not reflected in the summary statistics in Table 2 makes me suspect the summary statistics are not capturing the most relevant aspects of the distributions.
soundness: 4
strengths: My overall feeling is that this is a very technically solid paper which takes a relatively simple core idea - exchanging the Gaussian distributions in diffusion models for Student-t distributions in order to have a heavy-tailed sampling distribution - and executes on it methodically and very well. The presentation and flow of the argument is clear throughout, and all the obvious questions (parameterization of the reverse process, tractable training objective, continuous-time sampling with SDEs and ODEs, empirical performance) are dealt with comprehensively in my opinion.

The problem of sampling from heavy-tailed distributions is one of significant interest to the field, and I'm not aware of other work providing such a nice and integrated framework for heavy-tailed diffusion modeling, so I consider the contribution of an appropriate size for the conference.

Particular highlights were the connection of the diffusion modeling loss to the $\gamma$-Power Divergence, showing that this choice of divergence makes the loss tractable and reduces to a denoising $L^2$ loss similar to that of DDPM, and the discussion in Section 5 on various theoretical intuitions and insights the authors have concerning their framework.
summary: This paper proposes a method for adapting denoising diffusion models to sample from heavy tailed distributions. The core idea is to replace the Gaussian prior and transition distributions from standard DDPM with Student-t distributions. The authors show how to construct a heavy-tailed diffusion model by defining the relevant noising process in terms of the Student-t distribution, calculating the corresponding reverse noising process, showing how it can be appropriately parameterized in terms of a prediction of $x_0$, and deriving a training objective defined in terms of the $\gamma$-Power Divergence by analogy to the standard KL objective used to train diffusion models. They then demonstrate the empirical success of their method on distributions from the HRRR dataset, showing that it outperforms Gaussian-based diffusion models in most settings they study.
weaknesses: I do not think this paper has any major weaknesses. The addition I would be most keen to see is a more thorough discussion of the results in Tables 2 and 3 (empirical results in the unconditional setting). For example, it would appear from Table that in the test setting for VIL, their method actually underperforms compared to EDM (+PCP) and compared to EDM (+INC) in the other setting, and the authors do not discuss this in the text as far as I can tell. (Though, Figure 3 tells a somewhat different story, and I can't work out how to reconcile the two - see question below.) The authors also don't really discuss the results in Table 3 in the main body, and a short commend about what we should take from those results might be appropriate.

Secondly, I also feel that many of the derivations in the work are somewhat routine once one has the core idea of using Student-t distributions to parameterize the forward process, and the novelty of the work is therefore not outstanding. Nevertheless, the problem is of sufficient interest and the execution of the authors is sufficiently good that I believe this still clearly merits acceptance to the conference.

Meta-review: This paper introduces t-EDM, a diffusion model framework that enhances the modeling of heavy-tailed distributions, particularly relevant for applications in extreme-weather prediction.  Building upon the existing EDM framework, t-EDM replaces the Gaussian noise typically used in diffusion models with a Student-t distribution, allowing for controllable tail generation via a single hyperparameter.  This change can be implemented with minimal modifications to existing EDM implementations. The paper also introduces a new training method based on power divergence to fit the model. Empirical evaluations on the HRRR weather dataset show that t-EDM outperforms standard Gaussian-based diffusion models, including EDM, in capturing and generating extreme values, making it a valuable contribution to the field of generative modeling, especially for data with heavy-tailed characteristics.

In general the reviewers found the paper to be well-written, while providing a mathematically justified extension to diffusion models. The numerical methods back up the claims with some caveats raised by the reviewers (preconditioned versus non-preconditioned diffusion models). The implementation of the methods is quite straight forward, the mathematical ration is clear. Even though some of the more rigorous mathematical justifications are a bit wobbly, the numerical experiments are deemed sufficient. As such I recommend acceptance.

Decision: Accept (Poster)