SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers

Flow and Diffusion

In recent years a family of flexible generative model based on transforming pure noise \(\varepsilon \sim \mathcal{N}(0, \mathbf{I})\) into data \(x_* \sim p(x)\) has emerged. This transformation can be described by a simple time-dependent process $$ x_t = \alpha_t x_* + \sigma_t \varepsilon $$ with t defined on \([0, T]\), \(\alpha_t, \sigma_t\) being time-dependent functions and chosen such that \(x_0 \sim p(x)\), \(x_T \sim \mathcal{N}(0, \mathbf{I})\). At each \(t\), \(x_t\) has a conditional density \(p_t(x | x_*) = \mathcal{N}(\alpha_t x_*, \sigma_t^2\mathbf{I})\), and our goal is to estimate the marginal density \(p_t(x) = \int p_t(x | x_*) p(x) \mathrm{d}x \). In most cases, the marginal density \(p_t(x)\) is in tractable. Some previous methods focus on maximizing the likelihood \(\log p_t(x)\), whereas modern approaches take advantage of differential equations with the corresponding marginal density \(p_t(x)\) to directly estmate data samples from \(p(x)\).

Diffusion-Based Models. Diffusion-Based Models is the most commonly used framework for this transformation. The \(\alpha_t\) and \(\sigma_t\) are set indrectly by a forward-time stochastic differential equation (SDE) with \(\mathcal{N}(0, \mathbf{I})\) as equilibrium distribution. $$ dX_t = f(X_t, t) \mathrm{d}t + g(t) \mathrm{d} W_t $$ where \(W_t\) is a standard Brownian motion.
In practice, the model samples the process by learning the gradient of the likelihood \(\nabla \log p_t(x)\) (score) with a generative model \(s_\theta(x_t, t)\) under the score matching objective $$ \mathcal{L}_s(\theta) = \int \mathbb{E}[\Vert \sigma_t s_\theta(x_t, t) + \varepsilon \Vert^2] \mathrm{d}t $$ Inference is done by solving either a probability flow ODE $$\mathrm{d}X_t = [f(X_t, t) - \frac{1}{2} g^2(t) \nabla \log p_t(x)] \mathrm{d}t $$ or a reverse-time SDE $$ \mathrm{d}X_t = [f(X_t, t) - g^2(t) \nabla \log p_t(x)] \mathrm{d}t + g(t) \mathrm{d} \bar{W}_t $$ where \(\bar{W}_t\) is a standard Brownian motion with reversed time. Integrating both equations from \(t = T\) to \(t=0\) with initial condition of a Gaussian noise will push to a data sample approximating \(p(x)\).

Stochastic Interpolant and Flow-Based Models. Stochastic Interpolant and other Flow-Based models are the recent additions to this family, where the \(\alpha_t\) and \(\sigma_t\) are restricted on time interval \([0,1]\) with \(\alpha_0 = \sigma_1 = 1\) and \(\alpha_1 = \sigma_0 = 0\), so that \(x_t\) exactly interpolate between \(x_*\) and \(\varepsilon\). We note that this gives more flexibility in the choice of the interpolating functions, as they are no longer subject to a forward SDE.
Furthermore, these models use a simpler probability flow ODE for inference \[\begin{align*} \mathrm{d}X_t &= \underbrace{[f(X_t, t) - \frac{1}{2} g^2(t) \nabla \log p_t(x)]}_{\text{directly learn this}} \mathrm{d}t \\ \implies \mathrm{d}X_t &= v(X_t, t) \mathrm{d}t \end{align*}\] where the velocity \(v(X_t, t)\) is estimated by the flow matching objective $$ \mathcal{L}_v(\theta) = \int \mathbb{E}[\Vert v(X_t, t) - \dot \alpha_t x_* - \dot \sigma_t \varepsilon \Vert^2] \mathrm{d}t $$ Intuitively, this can be viewed as predicting the velocity of a particle starting moving from some \(\varepsilon\) at time \(t\).

We summarize the components of the above models in the following table:

Scalable Interpolant Transformers

From the above table, we summarize the design space into four components.

• Timespace: discrete or continuous time interval;
• Model Prediction: the objectives of \(\mathcal{L}_s\) or \(\mathcal{L}_v\);
• Interpolant: the choices of \(\alpha_t\) and \(\sigma_t\);
• Sampler: ODE or SDE.

To investigate such performance improvement, we gradually transition from a DiT model, a typical denoising diffusion model (discrete, denoising, variance preserving, and SDE) to our SiT model via a series of orthogonal steps in the design space. As we progress, we carefully evaluate how each move away from the diffusion model impacts the performance in the following sections.

To conduct such study, we use the same backbone architecture of DiT-B and training hyperparameters for all models. We also maintain number of parameters, GFLOPs, and training schedule (400K steps) to be identical for all models. All the numbers presented in tables are FID-50K score evaluated with respect to ImageNet256 training set, and produced by a 250 steps Heun ODE solver without otherwise specified. For solving the SDE, we used an Euler-Maruyama integrator.

Timespace

The first move away is well-studied: we switch from a discrete-time denoising model to a continuous-time score model. Marginal performance improvement is observed.

Model Prediction

We claim that velocity model is related to score model by a time-dependent weighting function. To be specific, we discovere that \[ v(x_t, t) = \frac{\dot \alpha_t}{\alpha_t} x_t - \lambda_t\sigma_t s(x_t, t) \] with \( \lambda_t = \dot \sigma_t - \frac{\dot \alpha_t \sigma_t}{\alpha_t} \). Plug this linear relation into \(\mathcal{L}_v\), we obtain \[ \begin{align*} \mathcal{L}_{v}(\theta) &= \int_0^T \lambda_t^2 \mathbb{E}[\Vert \sigma_t s_\theta(x_t, t) + \varepsilon \Vert^2] \mathrm{d} t \\ &= \mathcal{L}_{s_\lambda}(\theta) \end{align*} \] this aligns with the observation made in , that different model predictions of diffusion models corresponding to a vanilla denosiing objective weighted by different time-dependent functions. We trained all three models and present the results below.

Interpolant

We mainly experiment with three choices of interpolants:

• VP: \(\alpha_t = e^{-\tfrac12 \int_0^t\beta_s \mathrm{d} s} \), \( \sigma_t = \sqrt{1- e^{- \int_0^t\beta_s \mathrm{d} s}}\)
• GVP: \(\alpha_t = \cos(\frac12 \pi t) \), \( \sigma_t = \sin(\frac12 \pi t)\)
• Linear: \(\alpha_t = 1 - t\), \(\sigma_t = t\)

Sampler

We introduce more flexibility in sampling the velocity model in this section. Firstly, under SBDM setting, the reverse-time SDE for velocity can be constructed in the following way: \[ dX_t = [v(X_t, t) - \frac12 g^2(t)s(X_t, t)]\mathrm{d}t + g(t) \mathrm{d} \bar{W}_t \] where we take advantage of the linear relationship between velocity and score to construct the drift term. We denote \(g(t)\) as the SBDM diffusion coefficient. Following previous sections, we can also construct such SDE for GVP and Linear interpolants given the relationship \(g^2(t) = 2\lambda_t\sigma_t\). The results are resented below

We further proposed that the diffusion coefficient \(g(t)\) can be tuned separately from the learning process. In fact, any non-negative function \(w(t)\) (not necessarily monotone) is eligible to be used as diffusion coefficient, and the reverse-time SDE can thus be generalized to \[ dX_t = [v(X_t, t) - \frac12 w(t)s(X_t, t)]\mathrm{d}t + \sqrt{w(t)} \mathrm{d} \bar{W}_t \] Apart from SBDM coefficient, we also experiment with the choices of \(w(t) = \sigma_t\) (to eliminate the singularity of score near data distribution) and \(w(t) = \sin^2(\pi t)\), as well as their effects on either velocity or score models.

Lastly, we note that the performance of ODE and SDE integrators may differ under different computation budget. As shown below, the ODE converges faster with fewer number of functions evaluations, while the SDE is capable of reaching a much lower final FID score when given a larger computational budget.

Classifier-free Guidance

In this section, we give a concise justification for adopting it on the velocity model, and then empirically show that the drastic gains in performance for DiT case carry across to SiT.

Guidance for a velocity field means that: (i) that the velocity model \(v_\theta(x_t, t, y)\) takes class labels \(y\) during training, where \(y\) is occasionally masked with a null token \(\emptyset\); and (ii) during sampling the velocity used is \(v_\theta^\zeta(x_t, t, y) = \zeta v_\theta(x_t, t, y) + (1 - \zeta) v_\theta(x_t, t, \emptyset)\) for a fixed \(\zeta > 0\). Given this observation, one can leverage the usual argument for classifier-free guidance on score-based models. For a CFG scale of \(\zeta = 1.5\), DiT-XL sees an improvement in FID from a 9.6 (non-CFG) down to 2.27 (CFG). We observed similar performance improvement with our largest SiT-XL model under identical computation budget and CFG scale. Sampled with an ODE, the FID-50K score improved from 9.4 to 2.15; with an SDE, the FID improved from 8.6 to 2.06. This shows that SiT benefits from the same training and sampling choices explored previously, and can surpass DiT's performance in each training setting, not only with respect to model size, but also with respect to sampling choices.

Conclusion

In this work, we have presented Scalable Interpolant Transformers, a simple and powerful framework for image generation tasks. Within the framework, we explored the tradeoffs between a number of key design choices: the choice of a continuous or discrete-time model, the choice of interpolant, the choice of model predcition, and the choice of samplers . We highlighted the advantages and disadvantages of each choice and demonstrated how careful decisions can lead to significant performance improvements. Many concurrent works explore similar approaches in a wide variety of downstream tasks, and we leave the application of SiT to these tasks for future works.