PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Nanye Ma, Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, Saining Xie

arXiv:2401.08740v2cs.CVcs.LG

TL;DR

Standard diffusion models restrict how data connects to Gaussian noise, while stochastic interpolants offer greater flexibility. SiT systematically studies these design choices and surpasses DiT on ImageNet benchmarks with the same model structure, parameters, and GFLOPs.

  • Problem

    Standard diffusion models use a restricted connection between data and Gaussian noise, motivating study of more flexible noise-data connections and sampling choices.

  • Method

    SiT uses an interpolant framework to systematically vary time discretization, objectives, interpolants, model prediction, and deterministic or stochastic sampling.

  • Results

    2.06 and 2.62 FID-50K scores on conditional ImageNet 256×256 and 512×512, respectively, surpassing DiT without modifying model structure or hyperparameters.

  • Takeaways & Limitations

    SiT provides a framework for evaluating transport-based generative-model design choices, with diffusion coefficients tunable separately from learning.

  • Takeaways & Limitations

    Score models generally perform worse than velocity models, and standard diffusion sampling can incur bias from mismatch between practical noise sampling and the process endpoint.

Abstract

from arXiv · show

We present Scalable Interpolant Transformers (SiT), a family of generative models built on the backbone of Diffusion Transformers (DiT). The interpolant framework, which allows for connecting two distributions in a more flexible way than standard diffusion models, makes possible a modular study of various design choices impacting generative models built on dynamical transport: learning in discrete or continuous time, the objective function, the interpolant that connects the distributions, and deterministic or stochastic sampling. By carefully introducing the above ingredients, SiT surpasses DiT uniformly across model sizes on the conditional ImageNet 256x256 and 512x512 benchmark using the exact same model structure, number of parameters, and GFLOPs. By exploring various diffusion coefficients, which can be tuned separately from learning, SiT achieves an FID-50K score of 2.06 and 2.62, respectively.

1 Introduction

SiT uses the interpolant framework to systematically examine generative-model design choices beyond standard diffusion, while preserving the DiT backbone for fair comparison. Combining these choices yields faster convergence and better ImageNet generation than DiT under matched model structure and compute.

  • Design space: SiT systematically varies time discretization, model prediction, interpolant, and sampler to identify sources of performance gain.The framework transitions from a typical denoising diffusion model through orthogonal design changes and evaluates each move.
  • Motivation: Standard diffusion restricts the noise–data connection, whereas stochastic interpolants provide a more flexible connection that SiT evaluates for large-scale image generation.The paper frames this flexibility as varying how the data distribution connects to a Gaussian.
  • Sampling: SiT separates inference-time sampling choices from learning, including deterministic or stochastic sampling and tunable diffusion coefficients.The diffusion coefficient can be adapted after learning because it does not affect the learned velocity or score fields.
  • Results: SiT converges much faster than DiT across model sizes when FID-50K is tracked over training iterations with 250-step Euler–Maruyama sampling.The comparison is reported for both DiT and SiT across model sizes.
  • Results: SiT surpasses DiT on 256×256 and 512×512 ImageNet, achieving FID-50K scores of 2.06 and 2.62, respectively, without modifying model structure or hyperparameters.The comparison uses the same model structure, parameter count, and GFLOPs.

2 SiT: Scalable Interpolant Transformers

SiT builds generative models around flexible interpolants that connect data and Gaussian noise, while allowing learning objectives, time discretization, and sampling choices to be studied separately. The framework supports deterministic ODE and stochastic SDE generation, with diffusion coefficients decoupled from the interpolating path.

  • Interpolating process: SiT models define xt = αt x∗ + σt ε, explicitly interpolating between data at t = 0 and Gaussian noise at t = 1.Interpolants use α0 = σ1 = 1 and α1 = σ0 = 0, unlike score-based diffusion paths that approach Gaussian equilibrium only as t → ∞.
  • Sampling: The same interpolating process can be sampled either with a probability-flow ODE or a reverse-time SDE.The ODE uses a velocity field, while the SDE combines velocity, score, and a time-dependent diffusion coefficient.
  • Design choices: The stochastic-interpolant framework decouples αt and σt from the forward SDE, exposing more design flexibility than score-based diffusion models.The paper uses this flexibility to vary interpolants and diffusion coefficients independently, while noting numerical singularities near t = 0 for some score formulations.
  • Learning objective: SiT typically learns the velocity field and derives the score from it when stochastic sampling requires a score estimate.Velocity and score are linearly related, so only one of the two fields needs to be estimated in practice.
  • Time discretization: Continuous-time learning decouples the training objective from the sampling time grid, enabling sampling discretizations to be selected after learning.This provides flexibility in sampling efficiency and performance compared with DDPM’s coupled time grid.

3 Experiments

The experiments progressively replace DiT’s discrete score-based diffusion choices with continuous-time interpolant designs, velocity prediction, and alternative samplers. Performance depends on these choices, with SDE sampling and tuned diffusion coefficients offering benefits under suitable computational budgets.

  • Transition setup: The transition experiments compare discrete versus continuous time, score versus velocity parameterizations, and SBDM-VP versus Linear or GVP interpolants.SiT-B models use 400K training steps, with fixed solver and sampling-budget choices for comparison.
  • Discrete- to continuous-time: Continuous-time training provides a marginal performance increase and decouples sampling discretization from the training procedure.This decoupling permits different discretizations for diffusion coefficients selected after training.
  • Model parameterization: Significant performance increases are observed for weighted score and velocity objectives compared with the score objective.The weighting compensates for the score objective’s vanishing gradient near the data, which may explain part of the gain.
  • Choices of interpolant: Changing from SBDM-VP to GVP or Linear reduces path length and ODE trajectory curvature, which can reduce time-discretization errors.SBDM-VP also has a t = 0 singularity that can make learning near the data difficult; GVP and Linear do not exhibit this issue.
  • Deterministic vs stochastic sampling: SDE sampling improves performance over ODE sampling with the same trained velocity model, although ODE converges faster at low NFE while SDE reaches lower final FID with larger budgets.The comparison uses the probability-flow ODE and reverse-time SDE, whose performance differs across computation budgets.
  • Tunable diffusion coefficient: The optimal diffusion coefficient depends on both the model prediction and interpolant, and its influence can vary across model sizes.For Linear interpolants, the cost-regularized coefficient addresses endpoint integration difficulty, while score models generally perform worse than velocity models.
  • Classifier-free guidance: Classifier-free guidance produces significant performance gains for SiT velocity models and preserves the performance advantage over DiT under matched settings.The guidance combines conditional and null-token velocity predictions during sampling.

4 Related Work

Related work spans transformer-based diffusion architectures, diffusion training and sampling, interpolant and flow-matching methods, coupling and stochastic bridges, and latent-space generative modeling. SiT differs by studying these design choices together at large-scale image-generation benchmarks and by including SDE sampling.

  • Transformers: Transformer-based diffusion work builds on DiT by modifying architectures with internally masked prediction layers, choices described as orthogonal to SiT.The transformer architecture has also been applied across vision, language, chemistry, active matter, and biology.
  • Training and Sampling in Diffusions: Prior diffusion research includes denoising methods, DDPM and SBDM sampling algorithms, and improved diffusion ODE studies of parameterizations and paths.Improved Diffusion ODE focuses on lower-dimensional experiments, likelihood benchmarks, and does not consider SDE sampling.
  • Interpolants and flow matching: Interpolant and flow-matching studies examine velocity parameterizations, manifold generalizations, and KL-divergence trade-offs between SDE and ODE sampling.The cited interpolant work shows that the objectives studied here control KL for SDEs but not for ODEs.
  • Couplings and stochastic bridges: Other approaches reduce transport cost or gradient variance through optimal couplings, flow directly from conditioning variables, or stochastic bridges between arbitrary distributions.These methods alter how noise and data are coupled or construct stochastic connections between distributions.
  • Diffusion in Latent Space: Latent-space generative modeling offers a tractable approach for high-dimensional data and has been extended to video and velocity models.A prior latent-space study reported FID-50K 4.46 for a DiT-B backbone, motivating investigation of SiT’s performance factors.

5 Conclusion

The paper presents SiT as a framework for studying key design choices in image generation and reports that careful choices can improve performance.

  • SiT explores continuous versus discrete time, interpolants, model predictions, and diffusion coefficients.The framework examines tradeoffs among these choices and their advantages and disadvantages.
  • Careful decisions among these design choices can lead to significant performance improvements.
  • Applying SiT to a wide variety of downstream tasks is left for future work.

A Proofs

The proofs assume vector-valued variables in R^d and use the dot product notation ·.

  • All bold variables, including x and ε, are assumed to be real-valued vectors in R^d.
  • The proofs use · to denote the dot product.
  • Most proofs are derived from prior work cited as [2].

A.1 Proof of the probability flow ODE (2) with the velocity in Eq. (3).

The proof differentiates the characteristic function of the interpolant distribution and uses Fourier-transform identities to derive its transport equation and probability-flow ODE.

  • The characteristic function of p_t is defined as E[e^{i k · x_t}].Here, expectation is over x* and ε.
  • Differentiating the characteristic function yields an expression involving the conditional velocity v(x,t).The velocity is defined as E[(α̇_t x* + σ̇_t ε) | x_t = x].
  • Fourier-transform manipulation and integration by parts imply that p_t satisfies ∂_t p_t(x) + ∇_x · (v(x,t)p_t(x)) = 0.
  • Solving the transport equation by characteristics leads to the probability-flow ODE.

A.2 Proof of the SDE (4)

The proof shows that the proposed SDE preserves the interpolant marginals for any nonnegative diffusion coefficient, connecting stochastic sampling to the same transport dynamics.

  • The SDE has marginal density p_t(x) for any diffusion coefficient w_t ≥ 0.
  • Its density satisfies a Fokker–Planck equation involving the velocity, score, and Laplacian terms.The score is defined as s(x,t) = ∇ log p_t(x).
  • When w_t = 0, the Fokker–Planck equation becomes a continuity equation and the SDE becomes an ODE.
  • The conditional noise expectation is related to the score by ∇_x log p_t(x) = -σ_t^-1 E[ε | x_t = x] when σ_t ≠ 0.
  • The relation is extended to all t ∈ [0,1] through w_t s(x_t,t) = -1/2 E[ε | x_t = x], removing the singularity at t = 0.
  • The velocity and score are connected through v(x,t) = (α̇_t/α_t)x - λ_tσ_t s(x,t), with λ_t = σ̇_t - α̇_tσ_t/α_t.

A.5 Proof for the optimal wt for tightening the KL bound

The appendix derives KL-bound-based choices for the diffusion coefficient w_t and connects velocity, score, and conditional guidance formulations within the interpolant framework.

  • KL bound: The KL divergence is bounded by a time-weighted squared drift error, which can be expressed through the velocity-model loss L_t.The reverse-SDE drift is related to the velocity through the interpolant coefficients.
  • Endpoint behavior: For Linear and GVP interpolants, w_t^KL becomes zero at t = 0 and infinite at t = 1, and the KL bound is non-integrable at t = 0 unless lim_t→0 L_t = 0.The appendix identifies this endpoint behavior as a limitation of the theoretical bound.
  • Integration cost: Adding an integration-cost term produces a modified bound whose optimum depends on η, the parameter trading integration cost against error.The resulting optimized coefficient is well defined when L_t is well defined, but the integrand can still diverge at t = 0 without the same limiting-loss condition.
  • Velocity versus score: The velocity and score objectives place interpolant-dependent factors differently, so their numerical stability depends on the interpolant.For SBDM-VP, L_sλ performed better than L_v in Table 3 despite instability near t = 0; Linear and GVP have different instability sources.
  • Conditional guidance: Classifier-free guidance follows from mixing unconditional and conditional scores, with the associated velocity mixture representing a tempered conditional distribution.The paper reports a significant empirical performance boost from classifier-free guidance.

D Sampling with ODE and SDE

SiT supports deterministic ODE and stochastic SDE sampling with separate numerical solvers, while experiments compare their FID behavior under controlled evaluation budgets. The sampling choices expose a trade-off between quality, stability, and computational cost.

  • Samplers: The main experiments use second-order Heun integration for ODEs and first-order Euler-Maruyama integration for SDEs, with 250 function evaluations.This setup matches the DiT sampling budget used for evaluation.
  • ODE versus SDE: Table 8 reports better SDE performance across all model sizes, while ODE sampling has an advantage in the lower-NFE regime.The table presents each model's ODE and SDE FID-50K results as paired entries.
  • Model parameterization: Deterministic sampling generally favors learning a velocity model, whereas stochastic sampling requires both velocity and score quantities during integration.The paper recommends clipping the time interval near t = 0 to avoid numerical issues in stochastic sampling.
  • Guidance cost: Classifier-free guidance doubles sampling computation because each step evaluates both conditional and unconditional model outputs.The guidance mechanism is compatible with both samplers.
  • Comparison with DiT: SiT converges faster than DiT in FID across all model sizes, but the comparison is not directly apples-to-apples because SiT uses Heun and DiT uses DDIM.Both deterministic samplers use 250 NFE.
  • Heun sampler: Deterministic Heun sampling updates an initial Gaussian sample with an Euler predictor followed by an explicit trapezoidal correction.The procedure returns the final state after N fixed steps.

E Additional Implementation Details

The implementation follows DiT’s architecture and training setup while providing JAX-based solver implementations, EMA evaluation, and standardized FID measurement across SiT configurations.

  • Software: The implementation uses JAX following the DiT PyTorch codebase, with diffrax supplying the Heun solver and existing implementations supporting Euler-Maruyama.The paper provides separate deterministic and stochastic sampler algorithms.
  • Architecture: SiT uses the identical transformer architectures as DiT across Small, Base, Large, and XLarge configurations.The configurations vary in model size and compute.
  • Training: Training retains DiT’s structure and hyperparameters, including AdamW, a 1×10^-4 learning rate, batch size 256, and random horizontal flips with probability 0.5.The authors did not tune learning rates, schedules, optimizer parameters, or add gradient clipping.
  • Throughput: SiT-XL trains at approximately 6.8 iterations per second on a TPU v4-64 pod, compared with 6.4 for DiT-XL under identical settings.The reported throughput is slightly higher for SiT-XL.
  • Evaluation: All reported samples use EMA checkpoints with decay 0.9999, and sampler endpoints are tuned to improve performance and avoid numerical instability.FID is evaluated with the ADM suite on GPU-generated samples against ImageNet training images.

F Additional Visual results

The appendix presents uncurated SiT-XL samples at 512×512 and 256×256 resolution across diverse ImageNet classes, all generated with classifier-free guidance scale 4.0.

  • 512×512 samples: The 512×512 visual results show uncurated SiT-XL samples for volcano, arctic fox, loggerhead turtle, balloon, red panda, and geyser classes.Each figure specifies classifier-free guidance scale 4.0 and the corresponding ImageNet class label.
  • 256×256 samples: The 256×256 visual results show uncurated SiT-XL samples for macaw, golden retriever, ice cream, cliff, husky, and valley classes.Each figure specifies classifier-free guidance scale 4.0 and the corresponding ImageNet class label.
Loading 2401.08740v2…