Let's first revisit the mathmatical details of diffusion model and rectified flow. If you are familiar with them, feel free to jump to Section 2. Otherwise, you could also refer to the other blogs that introduces rectified flow and diffusion models in detail.

Diffusion Model

A diffusion model[1,2] is mathematically defined by the forward SDE, which can be written as:

where $f(\cdot,t): R^{d}\rightarrow R^{d}$ is a vector-valued function known as the drift coefficient of $\bm{x}(t)$, $g(t): R\rightarrow R$ is a scalar function known as the diffusion coefficient of $\bm{x}(t)$, and $\bm{w}$ is the standard Wiener process. When sampling along the forward process, we could generate a series of samples $\{\bm{x}(t)\}_{t=0}^{T}$ indexed by a continous time variable $t\in[0, T]$, such that $\bm{x}(0) \sim p_0$ and $\bm{x}(T) \sim p_T$. $p_0$ is the target data distribution for which we have i.i.d samples for training our model, while $p_T$ is commonly chosen as a Gaussian distribution $\mathcal{N}(0, \sigma(T)^2\mathbf{I})$.

Given the forward SDE, we have the corresponding reverse SDE that starts from $p_T$ to $p_0$.

Here, $\bar{\bm{w}}$ is the standard Wiener process when time flow backwards from $T$ to $0$. $dt$ is an infinitesimal negative timestep. $p_t(\bm{x})$ denotes the probability distribution of sample $\bm{x}(t)$ at time step $t$. As a result, given the score function at time step $t$, we can sample along the reverse process and simulate to generate data from $p_0$.

The score function can be approximated by the following score matching objective:

where $\lambda(t)$ is a weighting funciton, $t$ is uniformly sampled over $[0,T]$, $\bm{x}(0)\sim p_0(\bm{x})$, $\bm{x}(t) \sim p_{0t}(\bm{x}(t)\mid\bm{x}(0))$, which denotes the transition kernel from $0$ to $t$. It has the general form:

where

where we assume a linear structured drift coefficient: $f(\bm{x}, t) = f(t)\bm{x}$.

Given the definition of the transition kernel, the score of the conditional distribution can be decomposed as:

In practice, we typically parameterize the score model $s_{\bm{\theta}}(\bm{x}(t), t)$ as $\epsilon$-prediction or $x$-prediction. Namely: $s_{\bm{\theta}}(\bm{x}(t), t) = -\frac{D_{\bm{\theta}}(\bm{x}(t),t)}{s(t)\sigma(t)}$ or $s_{\bm{\theta}}(\bm{x}(t), t) = -\frac{\bm{x}(t)-s(t)D_{\bm{\theta}}(\bm{x}(t),t)}{s(t)^2\sigma(t)^2}$.

For brevity, we consider the diffusion model defined in EDM, where $f(\bm{x}, t)=0$ and $g(t)=\sqrt{2t}$, $t\in[T_{min}, T_{max}]$, $T_{min}=0.002$ and $T_{max}=80$. As a result, we have $s(t)=1$ and $\sigma(t)=t$. The $\epsilon$-prediction and $x$-prediction parameterization can be written as: $s_{\bm{\theta}}(\bm{x}(t), t) = -\frac{D_{\bm{\theta}}(\bm{x}(t),t)}{t}$ and $s_{\bm{\theta}}(\bm{x}(t), t) = -\frac{\bm{x}(t)-D_{\bm{\theta}}(\bm{x}(t),t)}{t^2}$ respectively.

Rectified FLow

Rather than starting from SDE training to ODE sampling as in diffusion models, the rectified flow[4] proposes approximating a forward ODE with a velocity model $v_{\bm{\theta}}$. The forward ODE is defined as

where $t\in[0,1]$, $\bm{x}(1)\sim\mathcal{N}(0,\mathbf{I})$, $\bm{x}(0)\sim p_0$. The above ODE moves sample $\bm{x}(0)$ from $p_0$ to $\bm{x}(1)$ in $\mathcal{N}(0,\mathbf{I})$. To transport backwards from $\mathcal{N}(0,\mathbf{I})$ to $p_0$, it proposes to approximate an ODE that yields the same marginal distribution of $\bm{x}(t)$ as the above equation. The training objective is

When parameterizing the score model with the $\bm{\epsilon}$- or $x$-prediction, the model $D_\theta$ is dictated to predict the noise $\bm{\epsilon}$ or $\bm{x}(0)$ at time step $t$. Considering that these two parameterizations only result in different optimization weight coefficient, without loss of generality, let's focus on analyzing the weakness of the $\bm{\epsilon}$-prediction formulation. With the $\bm{\epsilon}$-prediction model parameterization, the signal is reconstructed via $\hat{\bm{x}}(0)=\bm{x}(t) - t\cdot D_{\bm{\theta}}(\bm{x}(t),t)$. This leads to the model prediction eror being magnified by a factor of $t$, which introduces excessive error during early sampling process and results in poor sample generation quality in particular when the total number of discretization step is small.

In contrast, RFs train $D_{\bm{\theta}}(\bm{x}(t), t)$, $t\in[0,1]$, to predict the velocity $\bm{x}(1)-\bm{x}(0)$. During sampling, we traverse backwards in time step-by-step via:

Although the model prediction is also multiplied by a coefficient $t$, the prediction error remains constrained. This is because $t$ ranges in $[0,1]$ in RFs. In addition, the RFs' training velocity objective also provides a practically better training dynamics by making the training more robust, reducing unexpected prediction error made by $D_{\bm{\theta}}$. To demonstrate this, we decompose the training objective in Equation 2 at time $t$ into a irreducible constant error and approximation error:

The irreducible constant error is the lower bound of the optimization objective at time $t$ while the approximation error determines how good the model performs. An optimal training dynamics largely benefits model convergence and pushes the approximation error to lower values. The analysis on the decomposed training objective is also applicable to diffusion models, where the optimal solution for Equation 1 is $\mathop{\mathbb{E}}[\bm{x}(0) \mid\bm{x}(t)]$. We refer readers to Figure 5 in FasterDiT[5] for more empirical evidence on this.

For diffusion model, to mitigate model prediction error during training and bring it under control during sampling, it's more reasonable to predict the expected signal $\hat{\bm{x}}(0)$ directly at large $t$. As a result, EDM proposes to predict a mixture of noise and clean image at different $t$. Specifically, it parameterizes the score model $s_\theta$ as:

where $c_{out}(t)$ and $c_{skip}(t)$ is the scalar function. The training objective becomes:

As demonstrated in Figure 1, for $c_{out}(t)$ and $c_{skip}(t)$ are chosen such that the variance of effective target equals to 1: $Var[\frac{1}{c_{out}}(\bm{x}_0 - c_{skip}\bm{x}(t)]=1$. Besides, when $t=T_{max}$, $c_{skip}(t)=0$, $D_{\bm{\theta}}$ predicts signal $\bm{x}(0)$; when $t=T_{min}$, $c_{out}(t)=0$, $D_{\bm{\theta}}$ predicts $\bm{\epsilon}$. The above design ensures the model prediction error is amplified as little as possible across all time steps. By adopting the parameterization and preconditioning, the training is more stable and robust, as evidenced by the experiments (Table 2) in EDM[3].

One mistaken concept of RF is its learned sampling trajectory is "straight", leading to a faster sampling and better generation perforamnce. However, this is wrong. RF is preferred in many scenarios primarily due to its simpler and more concise implementations. When training DMs, in particular DDPM, one should write forward and reverse sampling functions, computing the drift and diffusion coefficients and the posterior mean. The implementation details are notoriously tedious yet largely impact the final generation performance. While for RF, the forward and sampling is simple and staightfoward.

Here, we discuss one toy example to compare the learned sampling trajectories of RF and DM. Consider target data distribution $p(\bm{x}_0) \sim \mathcal{N}(\mu, \sigma^2)$, RF sampling relies the velocity function $v(\bm{x}_t, t)$, which has analytical solution in this scenario. Recall that $\bm{x}_t = (1-t)*\bm{x}_0 + t*\bm{x}_1, \bm{x}_1 \sim \mathcal{N}(0, 1), t\in[0, 1].$ Take posterior expectation with respect to $x_t$ on both sides, we have:

Here, $E[\bm{x}_t|\bm{x}_t]=\bm{x}_t$. Besides, $v(\bm{x}_t, t)=E[\bm{x}_1 - \bm{x}_0|\bm{x}_t] = E[\bm{x}_1 |\bm{x}_t] - E[\bm{x}_0 |\bm{x}_t]$. By substituting it into Eq 3, we have:

With Tweedie's formula, the clean image is $(1-t)\bm{x}_0$, thus we could compute the posterior mean $E[\bm{x}_0|\bm{x}_t]$ by:

$\nabla_{\bm{x}}\log p_t(\bm{x}_t)=-\frac{\bm{x}_t-(1-t)*\mu}{(1-t)^2*\sigma^2 + t^2}$, Therefore, we have:

By setting $\mu=-4, \sigma=1$, we could plot its sampling trajectory in Fig 2. Similarly, we could derive the score function thus the analytic solution of EDM's ODE sampling function (described above).