Diffusion Transformer (DiT-B): Core Architecture and Applications in VLA Models

Forum topic · 小凯 · 2026-03-14

Summary

This technical overview explains DiT-B (Diffusion Transformer Base), introduced by Meta, UC Berkeley, and NYU in 2023 as a replacement for U-Net backbones in diffusion models. The architecture has three core components: Patchify (image-to-token conversion), DiT Blocks using adaLN-Zero conditioning (adaptive layer norm with zero initialization that dynamically adjusts normalization parameters based on diffusion timestep), and a Transformer decoder for noise prediction. DiT-B contains 130M parameters with 12 layers, 768 hidden dimensions, and 12 attention heads. The article compares four model variants (DiT-S/B/L/XL) ranging from 33M to 675M parameters, explains why smaller patch sizes improve quality at quadratic computational cost, and discusses why Vision-Language-Action (VLA) robotics models favor DiT. Key applications include Sora (DiT-XL/2 for video), DreamVLA, and NVIDIA GR00T N1 (humanoid robotics). DiT demonstrates that pure Transformers can match or surpass U-Net in diffusion models while offering better scalability.

One-Line Summary

DiT (Diffusion Transformer) is a new diffusion model architecture proposed by Meta, UC Berkeley, and NYU in 2023. It replaces U-Net with a Transformer as the backbone. DiT-B is the "Base" variant.

Architecture Overview

Traditional diffusion model: U-Net → Image DiT: Transformer → Image

Why Replace U-Net?

| U-Net Limitations | Transformer Advantages | |------------------|------------------------| | Local receptive field of convolutions | Self-attention captures global dependencies | | Difficult to scale in parallel | Naturally supports large-scale parallel training | | Complex architecture with many design choices | Standardized architecture, easy to scale |

Three Core Components of DiT

1. Patchify — Image Tokenization

Input image (256×256×3) → VAE encoder → Latent representation z (32×32×4) → Patchify (patch size = p×p)

p=8 → T=16 tokens
p=4 → T=64 tokens
p=2 → T=256 tokens

Key insight: Smaller patches produce more tokens, increasing compute quadratically, but improve generation quality.

2. DiT Block — Conditional Transformer Block

adaLN-Zero (Adaptive Layer Norm with Zero Initialization):

Dynamically adjusts Layer Norm parameters based on diffusion timestep t
Zero initialization ensures training stability at the start
Outperforms cross-attention and in-context conditioning

3. Transformer Decoder — Output Decoding

N DiT Blocks → linear projection → Unpatchify → noise prediction

DiT Model Variants

| Model | Layers | Hidden Dim | Heads | Parameters | Gflops | |-------|--------|-----------|-------|------------|--------| | DiT-S | 12 | 384 | 6 | 33M | 0.4 | | DiT-B | 12 | 768 | 12 | 130M | 1.5 | | DiT-L | 24 | 1024 | 16 | 458M | 5.1 | | DiT-XL | 28 | 1152 | 16 | 675M | 119 |

Naming convention: DiT-{Size}/{Patch}

DiT-B/8 = Base model + 8×8 patches
DiT-XL/2 = XL model + 2×2 patches (the configuration used by Sora)

Why VLA (Vision-Language-Action) Models Prefer DiT

1. Suited for Multimodal Action Distributions

Traditional autoregressive models can only generate a single "optimal" action. DiT-based diffusion can denoise from random noise to generate multiple plausible action modes.

Robotics scenario: Grasping a cup may have many valid approaches — diffusion captures this multimodal distribution.

2. Strong Long-Sequence Modeling

| Model | Action Generation | Long-Horizon Planning | |-------|-------------------|----------------------| | OpenVLA | Single-step autoregressive | Errors accumulate | | DreamVLA / GR00T N1 | DiT diffusion sequence | Generates multi-step actions at once |

3. Natural Fusion with VLMs

VLM (Llama/Qwen) → text/image understanding → fused representation → DiT-B → action sequence diffusion generation

Core Insight

> DiT's revolution: it proves diffusion models don't need U-Net's inductive biases. Pure Transformers achieve better results with stronger scalability.

Implications for robotics:

Action generation can be as natural as image generation
Multimodal uncertainty is modeled elegantly
Unified architecture: both visual understanding and action generation use Transformers

Real-World Applications

| Project | DiT Variant | Use Case | |---------|-------------|----------| | Sora (OpenAI) | DiT-XL/2 | Video generation | | DreamVLA | DiT-B | Robotic action sequence generation | | GR00T N1 | Custom DiT | Real-time humanoid robot actions | | Stable Diffusion 3 | DiT-based | Image generation |

References

Paper: *Scalable Diffusion Models with Transformers* (Peebles & Xie, ICCV 2023)
Authors: William Peebles (Meta AI → OpenAI), Saining Xie (NYU)
Code: https://github.com/facebookresearch/DiT
Project page: https://www.wpeebles.com/DiT