Same silicon, better FID: rethinking the path through noise

What it does

SiT is a family of image generative models built on the Diffusion Transformer (DiT) backbone. It replaces the standard diffusion objective with an interpolant framework, which connects noise and data distributions through flexible continuous paths rather than a fixed noising schedule. The repo ships PyTorch model definitions, pre-trained class-conditional weights for ImageNet 256×256, and both training and sampling code.

The interesting bit

The clever part is the modularity: you can vary the interpolant, the prediction target, the time discretization, and whether sampling is deterministic or stochastic—all while keeping the exact same transformer backbone, parameter count, and GFLOPs. On ImageNet 256×256, SiT outperforms DiT across every model size, and the best XL/2 checkpoint hits an FID of 2.06.

Key highlights

• Uses the identical DiT architecture; gains come from the transport formulation, not more parameters or compute
• Supports both ODE and SDE sampling, with tunable diffusion coefficients that can be adjusted after training
• Pre-trained weights available for direct sampling, ported from the original JAX/TPU training runs
• Likelihood evaluation supported via ODE-based probability flow
• Training script uses PyTorch DDP and allows resuming from checkpoints with full EMA and optimizer state restoration

Caveats

• The PyTorch port currently lacks Flash Attention, torch.compile, and AMP/bfloat16 support, so training and sampling are likely slower than they could be
• FID and other metrics are not monitored during training out of the box
• The authors note minor platform differences: sampling on TPU scores 2.15 FID versus 2.06 on GPU, and likelihood precision could be improved with better dequantization

Verdict

Researchers and practitioners already invested in the DiT ecosystem should look here—it demonstrates that the transport math matters as much as the architecture. If you are not working on conditional ImageNet generation or flow/diffusion ablations, this is mostly a well-documented reference implementation.