Building a Diffusion Transformer (DiT) from Scratch in PyTorch

Transformers have quietly taken over computer vision.

Models like OpenAI Sora and modern image generators are increasingly moving toward transformer-based diffusion architectures because transformers scale better, parallelize better, and model global interactions more naturally.

In a previous article, I explored how Flow Matching works in modern generative models:

*Understanding Flow Matching for Software Engineers*

However, we still have not answered an important question:

How do we actually model the vector field u(x,t,y)?

Note: we use the classifier-free guidance (CFG) vector field u(x,t,y) instead of unconditional u(x,t), or u(x,t,∅) to enable class/prompt guidance during generation.

In this article, we’ll build a minimal Diffusion Transformer (DiT) from scratch in PyTorch and see how diffusion models can be reformulated entirely around:

• patch embeddings
• self-attention
• token processing
• transformer conditioning

Rather than focusing on diffusion mathematics, this article focuses on the actual architecture and implementation behind DiT.

A Quick Recap: Vision Transformers (ViT)

Before understanding DiT, we need one key idea from ViT:

Images can be treated as sequences of patches.

Instead of applying convolutions over pixels, ViT splits an image into fixed-size patches and treats each patch like a token in NLP.

For example:

• image size: 32 × 32
• patch size: 8 × 8

This gives:

n_tokens = (32/8)^2=16

So the image becomes a sequence of 16 visual tokens.

This is the core idea inherited by DiT.

From ViT to DiT

A standard ViT pipeline looks like this:

image → patches → transformer → class prediction

DiT changes the objective completely. Instead of predicting a class label, the transformer predicts a denoising vector field (or noise/velocity depending on formulation).

The pipeline becomes:

xt → Patchify → Transformer→ u(x,t,y)

This single change turns a classifier into a generative model.

Building a Minimal Diffusion Transformer

Here’s the full model:

class DiffusionTransformerFlowModel(ConditionalVectorField):
  def __init__(
      self,
      img_size: int = 32,
      patch_size: int = 8,
      num_layers: int = 12,
      c: int = 1,
      dim: int = 256,
      heads: int = 4,
      final_dim: int = 10,
      n_classes: int = 11,
    ):
      super().__init__()
      # 0. Construct time_embedder and y_embedder
      self.time_embedder = FourierEncoder(dim)
      self.y_embedder = nn.Embedding(num_embeddings = n_classes, embedding_dim = dim)

      # 1. Construct patchifier
      self.patchifier = Patchifier(
          img_size=img_size,
          patch_size=patch_size,
          c_in=c,
          dim=dim
        )

      # 2. Construct DiT
      n_tokens = (img_size // patch_size) ** 2
      self.dit = DiffusionTransformer(
          depth=num_layers,
          n_tokens=n_tokens,
          dim=dim,
          heads=heads,
      )

      # 3. Construct de-patchifier
      self.depatchifier = Depatchifier(
          img_size=img_size,
          patch_size=patch_size,
          dim=dim,
          final_dim=final_dim,
          c_out=c
        )

  def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Args:
    - x: b 1 32 32
    - t: b 1 1 1
    - c: b 1 1 1
    Returns:
    - u_t^theta(x|y): b 1 32 32
    """
    # 1. Embed time and y
    t_embed = self.time_embedder(t) # b d
    y_embed = self.y_embedder(y) # b d

    # 2. Patchify
    x = self.patchifier(x) # b n d

    # 3. Pass through DiT
    x = self.dit(x, t_embed + y_embed) # b d

    # 4. Depatchify
    x = self.depatchifier(x) # b 1 32 32

    return x

At a high level, the architecture has four stages:

1. Time + class conditioning
2. Patchification
3. Transformer denoising
4. Depatchification

Let’s go through them one by one.

1. Time Embeddings

Diffusion models are conditioned on time.

The model must know how noisy the current sample is.

That’s why we embed the timestep:

self.time_embedder = FourierEncoder(dim)

The most common approach uses sinusoidal or Fourier embeddings:

γ(t)=[sin⁡(ωkt), cos⁡(ωkt)]

This converts a scalar timestep into a high-dimensional representation that the transformer can process.

class FourierEncoder(nn.Module):
    """
    Based on https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/denoising_diffusion_pytorch/karras_unet.py#L183
    """
    def __init__(self, dim: int):
        super().__init__()
        assert dim % 2 == 0
        self.half_dim = dim // 2
        self.weights = nn.Parameter(torch.randn(1, self.half_dim))

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """
        Args:
        - t: b
        Returns:
        - embeddings: b d
        """
        # Step 1: compute frequencies f_i = 2 * pi * w_i * t
        t = t.view(-1, 1) # b 1
        freqs = t * self.weights * 2 * math.pi # b hd

        # Step 2: compute sin(f_i) and cos(f_i)
        sin_embed = torch.sin(freqs) # b hd
        cos_embed = torch.cos(freqs) # b hd

        # Step 3: Concatenate and return
        return torch.cat([sin_embed, cos_embed], dim=-1) * math.sqrt(2) # b d

Without timestep conditioning, the model would have no idea whether it should:

• remove heavy noise
• refine fine details
• generate structure
• sharpen textures

Time embeddings are therefore fundamental to diffusion architectures.

2. Class Conditioning

Next, we condition the model on labels:

self.y_embedder = nn.Embedding(
    num_embeddings=n_classes,
    embedding_dim=dim
)

This converts class IDs into dense vectors.

Later, we combine time and class information:

t_embed + y_embed

This gives the transformer a single conditioning vector.

Conceptually:

• timestep embedding → “how noisy?”
• class embedding → “generate what?”

This is one of the simplest forms of conditional generation.

More advanced DiT architectures use:

• AdaLN-Zero
• classifier-free guidance
• cross attention

But additive conditioning is a clean minimal starting point.

3. Patchification

This is where DiT inherits directly from ViT.

self.patchifier = Patchifier(
    img_size=img_size,
    patch_size=patch_size,
    c_in=c,
    dim=dim
)

Instead of processing pixels directly, we split the image into patches.

For a 32 × 32 image with patch size 8:

n_tokens = (img_size/patch_size)^2

which gives 16 tokens.

Each patch is flattened and projected into the transformer embedding dimension.

The image is no longer viewed as a spatial grid.

It becomes:

a sequence of visual tokens.

This is the conceptual leap from CNNs to transformers.

4. The Transformer Backbone

The transformer is the core denoiser:

self.dit = DiffusionTransformer(
    depth=num_layers,
    n_tokens=n_tokens,
    dim=dim,
    heads=heads,
)

Internally, the DiffusionTransformer is built from a stack of DiffusionTransformerLayer modules, each containing self-attention and feed-forward sublayers.

class DiffusionTransformerLayer(nn.Module):
  def __init__(
      self,
      dim: int,
      heads: int,
  ):
    """
    Args:
    - n_tokens: sequence length (for sake of positional embeddings)
    - dim: dimension of hidden layers
    - heads: number of attention heads
    """
    super().__init__()

    # Normalization
    self.norm1 = nn.RMSNorm(dim, elementwise_affine=False)
    self.norm2 = nn.RMSNorm(dim, elementwise_affine=False)
    self.ada_ln = nn.Sequential(
        nn.RMSNorm(dim, elementwise_affine=False),
        nn.Linear(dim, dim * 6)
    )

    # Initialize conditioning to zero - stabilizes residual connection!
    nn.init.zeros_(self.ada_ln[1].weight)
    nn.init.zeros_(self.ada_ln[1].bias)

    # Attention
    self.attn = MHA(dim, heads)

    # Feedforward
    self.ff = MLP([dim, 4 * dim, dim])

  def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
    """
    Args:
    - x: b n d
    - c: b d
    Returns:
    - x: b n d
    """
    # Compute conditioning gating, scaling, and bias
    c = rearrange(self.ada_ln(c), 'b d -> b 1 d') # b 1 d
    attn_scale, attn_bias, attn_gate, ff_scale, ff_bias, ff_gate = c.chunk(6, dim=-1)

    # Attention + FF
    x = x + attn_gate * self.attn(
      modulate(self.norm1(x), attn_scale, attn_bias)
    )
    x = x + ff_gate * self.ff(
      modulate(self.norm2(x), ff_scale, ff_bias)
    )
    return x

class DiffusionTransformer(nn.Module):
  def __init__(
      self,
      depth: int,
      n_tokens: int,
      dim: int,
      **layer_kwargs,
  ):
    """
    Args:
    - n_tokens: sequence length (for sake of positional embeddings)
    - dim: dimension of hidden layers
    - heads: number of attention heads
    - depth: number of layers
    """
    super().__init__()
    self.layers = nn.ModuleList([])
    for _ in range(depth):
      self.layers.append(DiffusionTransformerLayer(dim=dim, **layer_kwargs))

    # Positional encodings
    self.pos_encodings = nn.Parameter(torch.randn(n_tokens, dim))

  def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
    """
    Args:
    - x: b n d
    - c: b d
    Returns:
    - x: b n d
    """
    x = x + self.pos_encodings.unsqueeze(0)
    for layer in self.layers:
      x = layer(x, c)
    return x

Note: the conditioning vector c is shared across all transformer layers, while each block maintains its own parameters.

This design allows every layer to remain globally aware of the same timestep and class-conditioning information throughout the network. At the same time, different transformer blocks can learn progressively more abstract denoising behaviors.

5. Depatchification

After the diffusion transformer, the output tensor has shape

(b, n, d)

where:

• b = batch size
• n = number of patch tokens
• d = transformer embedding dimension

However, the target vector field u(x,t,y) must have the same spatial shape as the input image x (since the flow dynamics update the image through element-wise addition):

(b, c, h, w)

Therefore, after token processing, we need to reconstruct the original image layout from the patch representations.

This is the role of the Depatchifier.

self.depatchifier = Depatchifier(
    img_size=img_size,
    patch_size=patch_size,
    dim=dim,
    final_dim=final_dim,
    c_out=c
)

class Depatchifier(nn.Module):
  def __init__(self, img_size: int, patch_size: int, dim: int, final_dim: int, c_out: int):
      super().__init__()
      self.patch_size = patch_size
      assert img_size % patch_size == 0, "Image size must be divisible by patch size"
      h = w = img_size // patch_size

      self.net = nn.Sequential(
          # Norm + MLP
          nn.RMSNorm(dim, elementwise_affine=False),
          MLP([dim, 4*dim, final_dim * patch_size ** 2]),

          # Depatchify
          Rearrange("b (h w) (f ph pw) -> b f (h ph) (w pw)", h=h, w=w, f=final_dim, ph=patch_size, pw=patch_size),

          # Final convolution
          nn.Conv2d(final_dim, c_out, kernel_size=3, padding=1)
      )

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Args:
    - x: b n d
    Returns:
    - x: b 1 32 32
    """
    return self.net(x)

This converts:

tokens → image_tokens

The Forward Pass

Now let’s examine the actual forward pipeline.

def forward(self, x, t, y):
    # 1. Embed time and labels
    t_embed = self.time_embedder(t)
    y_embed = self.y_embedder(y)
    # 2. Patchify image
    x = self.patchifier(x)
    # 3. Transformer denoising
    x = self.dit(x, t_embed + y_embed)
    # 4. Reconstruct image
    x = self.depatchifier(x)
    return x

The entire model can be summarized as:

xt → tokens → attention → vector field u

Training Objective

It is L2 loss between the predicted vector field uθ​(x,t,y) and the reference conditional vector field u_ref (x,z,t), where:

t : time (noise level)

y : label / conditioning prompt

x : noisy sample along the conditional path, such as x = αz + σϵ, z ∼data distribution , ϵ ∼ N(0, I)

x = alpha_t * z + beta_t * torch.randn_like(z)

uθ​(x,t,y) : model-predicted conditional vector field

u_ref (x,z,t) : ground-truth conditional vector field. For Gaussian paths, this vector field u_ref (x,z,t) has a closed-form solution:

ut_ref = (dt_alpha_t - dt_beta_t / beta_t * alpha_t) * z + dt_beta_t / beta_t * x

Loss from trainer:

def get_train_loss(self, batch_size: int) -> torch.Tensor:
        # Step 1: Sample z,y from p_data
        z, y = self.path.p_data.sample(batch_size) # b ..., b

        # Step 2: Set each label to 10 (i.e., null) with probability eta
        xi = torch.rand(y.shape[0]).to(y.device)
        y[xi < self.eta] = self.null_label

        # Step 3: Sample t and x
        t = torch.rand(batch_size).to(z) * (1 - self.eps) # b
        x = self.path.sample_conditional_path(z,t) # b ...

        # Step 4: Regress and output loss
        ut_theta = self.model(x,t,y) # b ...
        ut_ref = self.path.conditional_vector_field(x,z,t) # b ...
        return torch.square(ut_theta - ut_ref).mean()

# supported class
class GaussianConditionalProbabilityPath(ConditionalProbabilityPath):
    # ...
    def sample_conditional_path(self, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Samples from the conditional distribution p_t(x|z)
        Args:
            - z: b ...
            - t: b
        Returns:
            - x: b ...
        """
        alpha_t = self.rearrange_scalar(self.alpha(t)) # (b 1 1 1)
        beta_t = self.rearrange_scalar(self.beta(t)) # (b 1 1 1)
        return alpha_t * z + beta_t * torch.randn_like(z)

    def conditional_vector_field(self, x: torch.Tensor, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates the conditional vector field u_t(x|z)
        Args:
            - x: b c h w
            - z: b c h w
            - t: b
        Returns:
            - conditional_vector_field: conditional vector field (num_samples, c, h, w)
        """
        alpha_t = self.rearrange_scalar(self.alpha(t)) # b
        beta_t = self.rearrange_scalar(self.beta(t)) # b
        dt_alpha_t = self.rearrange_scalar(self.alpha.dt(t)) # b
        dt_beta_t = self.rearrange_scalar(self.beta.dt(t)) # b

        return (dt_alpha_t - dt_beta_t / beta_t * alpha_t) * z + dt_beta_t / beta_t * x

Inference

Once the model predicts the conditional vector field u(x,t,y), we can perform guided generation using CFG dynamics.

During sampling, we combine conditional and unconditional predictions:

u_cfg = u(x,t,∅) + s(u(x,t,y) − u(x,t,∅)) = (1-s)u(x,t,∅) + su(x,t,y)

where s is the guidance scale.

class CFGVectorFieldODE(ODE):
    def __init__(self, net: ConditionalVectorField, null_label: int, guidance_scale: float = 1.0):
        self.net = net
        self.guidance_scale = guidance_scale
        self.null_label = null_label

    def drift_coefficient(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Args:
        - x: b ...
        - t: b
        - y: b
        """
        guided_vector_field = self.net(x, t, y)
        unguided_y = torch.ones_like(y) * self.null_label
        unguided_vector_field = self.net(x, t, unguided_y)
        return (1 - self.guidance_scale) * unguided_vector_field + self.guidance_scale * guided_vector_field

We then evolve the dynamics using ODE or SDE (see my previous article for details):

class EulerSimulator(Simulator):
    def __init__(self, ode: ODE):
        self.ode = ode

    def step(self, xt: torch.Tensor, t: torch.Tensor, h: torch.Tensor, **kwargs):
        h = h.view([-1] + [1] * (len(xt.shape) - 1))
        return xt + self.ode.drift_coefficient(xt, t, **kwargs) * h

class EulerMaruyamaSimulator(Simulator):
    def __init__(self, sde: SDE):
        self.sde = sde

    def step(self, xt: torch.Tensor, t: torch.Tensor, h: torch.Tensor, **kwargs):
        h = h.view([-1] + [1] * (len(xt.shape) - 1))
        return xt + self.sde.drift_coefficient(xt, t, **kwargs) * h + self.sde.diffusion_coefficient(xt, t, **kwargs) * torch.sqrt(h) * torch.randn_like(xt)

Starting from pure noise, the guided vector field progressively transports the sample toward the desired conditional data distribution.

Source: https://maverickframe.com/blog/diffusion-models/

Final Thoughts

Vision Transformers changed computer vision by turning images into tokens.

Diffusion Transformers push the idea further:

generation itself becomes token processing.

The architecture is surprisingly elegant:

• patchify
• condition
• attend
• reconstruct

Despite the simplicity of this formulation, DiT scales remarkably well and has become a foundation for modern generative systems.

More importantly, DiT demonstrates a broader trend in deep learning:

architectures originally designed for language are increasingly becoming general-purpose computation frameworks for generation.

What’s Next: From DiT to Latent Diffusion

References

The complete implementation used in this article can be found in the accompanying notebook:

IAP Diffusion Labs — Complete DiT Implementation Notebook