dit.py

DiT - minimal runnable implementation
"""
DiT - minimal runnable implementation
=====================================

Educational PyTorch reference for the class-conditional Diffusion Transformer
(Peebles & Xie 2023): patchify -> N x [adaLN-Zero DiT block] -> final layer
-> unpatchify. Sizes are tiny so the whole script runs on CPU in seconds.

Pairs with: docs/tutorials/image_generation_systems_tutorial.md §4
(DiT architecture and adaLN-Zero).

What the sanity checks verify (the claims in §4.2):
    1. At init every block is the identity: gate alpha = 0, and the modulated
       norm uses (1 + gamma) so gamma = 0 leaves LN(x) intact.
    2. The identity is NOT a dead end: for a block with a loss on its output,
       dL/d(alpha) != 0 at step 0, so the zero-initialised modulation MLP
       receives gradient and the block starts learning.
    3. gamma and beta get zero gradient at step 0 — their only path to the
       loss runs through alpha = 0.
    4. patchify / unpatchify round-trip exactly.
    5. In the whole model the final layer is zero-initialised as well, so step
       0 updates only its output linear; the blocks are reached from step 1.

Run:
    python dit.py
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# Embeddings
# ---------------------------------------------------------------------------
def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
    """Sinusoidal embedding of (possibly continuous) timesteps. [B] -> [B, dim]."""
    half = dim // 2
    freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32) / half)
    args = t.float()[:, None] * freqs[None]
    return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)


def sincos_pos_embed_2d(dim: int, grid: int) -> torch.Tensor:
    """Fixed 2D sin-cos position embedding for a grid x grid patch grid. -> [grid*grid, dim]."""
    assert dim % 4 == 0
    quarter = dim // 4
    omega = 1.0 / 10000 ** (torch.arange(quarter, dtype=torch.float32) / quarter)
    ys, xs = torch.meshgrid(torch.arange(grid, dtype=torch.float32),
                            torch.arange(grid, dtype=torch.float32), indexing="ij")
    out = []
    for coord in (ys.reshape(-1), xs.reshape(-1)):  # [N]
        a = coord[:, None] * omega[None]            # [N, quarter]
        out += [torch.sin(a), torch.cos(a)]
    return torch.cat(out, dim=1)                    # [N, dim]


class TimestepEmbedder(nn.Module):
    def __init__(self, hidden: int, freq_dim: int = 256):
        super().__init__()
        self.freq_dim = freq_dim
        self.mlp = nn.Sequential(nn.Linear(freq_dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden))

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        return self.mlp(timestep_embedding(t, self.freq_dim))


class LabelEmbedder(nn.Module):
    """Class embedding with an extra "null" class for classifier-free guidance."""

    def __init__(self, num_classes: int, hidden: int):
        super().__init__()
        self.table = nn.Embedding(num_classes + 1, hidden)
        self.null_id = num_classes

    def forward(self, y: torch.Tensor, drop_mask: torch.Tensor | None = None) -> torch.Tensor:
        if drop_mask is not None:
            y = torch.where(drop_mask, torch.full_like(y, self.null_id), y)
        return self.table(y)


# ---------------------------------------------------------------------------
# adaLN-Zero block
# ---------------------------------------------------------------------------
def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    # (1 + scale), not scale: with a zero-initialised modulation MLP the
    # normalised activations pass through unchanged instead of being zeroed.
    return x * (1 + scale[:, None, :]) + shift[:, None, :]


class Attention(nn.Module):
    def __init__(self, hidden: int, heads: int):
        super().__init__()
        self.heads = heads
        self.qkv = nn.Linear(hidden, 3 * hidden)
        self.proj = nn.Linear(hidden, hidden)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, D = x.shape
        q, k, v = self.qkv(x).reshape(B, N, 3, self.heads, D // self.heads).permute(2, 0, 3, 1, 4)
        out = F.scaled_dot_product_attention(q, k, v)          # [B, H, N, d]
        return self.proj(out.transpose(1, 2).reshape(B, N, D))


class DiTBlock(nn.Module):
    """x -> x + alpha1 * Attn(mod(LN(x))) -> x + alpha2 * MLP(mod(LN(x)))."""

    def __init__(self, hidden: int, heads: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden, heads)
        self.norm2 = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6)
        self.mlp = nn.Sequential(nn.Linear(hidden, int(hidden * mlp_ratio)), nn.GELU(approximate="tanh"),
                                 nn.Linear(int(hidden * mlp_ratio), hidden))
        # One MLP produces all six modulation vectors; its last layer is
        # zero-initialised -> shift = scale = gate = 0 at step 0 (adaLN-Zero).
        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(hidden, 6 * hidden))
        nn.init.zeros_(self.ada[1].weight)
        nn.init.zeros_(self.ada[1].bias)

    def forward(self, x: torch.Tensor, c: torch.Tensor, mod: torch.Tensor | None = None) -> torch.Tensor:
        # `mod` may be passed in so a caller can inspect its gradient (see [D03]).
        if mod is None:
            mod = self.ada(c)
        shift1, scale1, gate1, shift2, scale2, gate2 = mod.chunk(6, dim=1)
        x = x + gate1[:, None, :] * self.attn(modulate(self.norm1(x), shift1, scale1))
        x = x + gate2[:, None, :] * self.mlp(modulate(self.norm2(x), shift2, scale2))
        return x


class FinalLayer(nn.Module):
    """adaLN (shift, scale only) + linear to patch pixels; both zero-initialised."""

    def __init__(self, hidden: int, patch: int, out_channels: int):
        super().__init__()
        self.norm = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden, patch * patch * out_channels)
        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(hidden, 2 * hidden))
        for lin in (self.ada[1], self.linear):
            nn.init.zeros_(lin.weight)
            nn.init.zeros_(lin.bias)

    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        shift, scale = self.ada(c).chunk(2, dim=1)
        return self.linear(modulate(self.norm(x), shift, scale))


# ---------------------------------------------------------------------------
# Full model
# ---------------------------------------------------------------------------
class DiT(nn.Module):
    def __init__(self, input_size: int = 8, patch: int = 2, in_channels: int = 4, hidden: int = 64,
                 depth: int = 4, heads: int = 4, num_classes: int = 10, learn_sigma: bool = True):
        super().__init__()
        assert input_size % patch == 0
        self.patch, self.in_channels = patch, in_channels
        self.out_channels = in_channels * 2 if learn_sigma else in_channels
        self.grid = input_size // patch
        self.x_embed = nn.Linear(patch * patch * in_channels, hidden)   # patchify == linear on flattened patches
        self.register_buffer("pos_embed", sincos_pos_embed_2d(hidden, self.grid), persistent=False)
        self.t_embed = TimestepEmbedder(hidden)
        self.y_embed = LabelEmbedder(num_classes, hidden)
        self.blocks = nn.ModuleList(DiTBlock(hidden, heads) for _ in range(depth))
        self.final = FinalLayer(hidden, patch, self.out_channels)

    def patchify(self, x: torch.Tensor) -> torch.Tensor:
        """[B, C, H, W] -> [B, N, p*p*C] with N = (H/p)*(W/p), row-major over patches."""
        B, C, H, W = x.shape
        p = self.patch
        x = x.reshape(B, C, H // p, p, W // p, p).permute(0, 2, 4, 3, 5, 1)   # [B, h, w, p, p, C]
        return x.reshape(B, (H // p) * (W // p), p * p * C)

    def unpatchify(self, tokens: torch.Tensor, channels: int) -> torch.Tensor:
        """[B, N, p*p*C] -> [B, C, H, W]; exact inverse of patchify."""
        B, N, _ = tokens.shape
        p, g = self.patch, self.grid
        x = tokens.reshape(B, g, g, p, p, channels).permute(0, 5, 1, 3, 2, 4)  # [B, C, h, p, w, p]
        return x.reshape(B, channels, g * p, g * p)

    def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor,
                drop_mask: torch.Tensor | None = None) -> torch.Tensor:
        h = self.x_embed(self.patchify(x)) + self.pos_embed[None]   # [B, N, hidden]
        c = self.t_embed(t) + self.y_embed(y, drop_mask)            # [B, hidden]
        for blk in self.blocks:
            h = blk(h, c)
        return self.unpatchify(self.final(h, c), self.out_channels)  # [B, out_C, H, W]


# ---------------------------------------------------------------------------
# Sanity checks — the §4.2 claims, executed
# ---------------------------------------------------------------------------
def main() -> None:
    torch.manual_seed(0)
    model = DiT()
    B = 3
    x = torch.randn(B, 4, 8, 8)
    t = torch.randint(0, 1000, (B,))
    y = torch.randint(0, 10, (B,))

    # 4. patchify / unpatchify round-trip is exact
    tokens = model.patchify(x)
    assert tokens.shape == (B, 16, 16)
    assert torch.equal(model.unpatchify(tokens, 4), x)
    print(f"[D01] patchify: {tuple(x.shape)} -> {tuple(tokens.shape)} -> back, exact  PASS")

    # 1. every block is the identity at init
    h = model.x_embed(tokens) + model.pos_embed[None]
    c = model.t_embed(t) + model.y_embed(y)
    for i, blk in enumerate(model.blocks):
        assert torch.equal(blk(h, c), h), f"block {i} not identity at init"
    out = model(x, t, y)
    assert out.shape == (B, 8, 8, 8) and torch.equal(out, torch.zeros_like(out))
    print(f"[D02] adaLN-Zero: {len(model.blocks)} blocks are the identity at init; final layer outputs 0  PASS")

    # 2./3. gradient at step 0 for ONE block with a loss on its output (the
    # per-block derivation in §4.2): dL/d(alpha) != 0, dL/d(gamma) = dL/d(beta) = 0.
    blk = model.blocks[0]
    mod = blk.ada(c.detach())                          # [B, 6D] = (shift1, scale1, gate1, shift2, scale2, gate2), all 0
    mod.retain_grad()
    loss = (blk(h.detach(), c, mod=mod) - torch.randn_like(h)).pow(2).mean()
    loss.backward()
    shift1, scale1, gate1, shift2, scale2, gate2 = mod.grad.chunk(6, dim=1)
    gates, others = torch.cat([gate1, gate2]), torch.cat([shift1, scale1, shift2, scale2])
    assert gates.abs().max() > 0, "gate alpha got no gradient at init"
    assert torch.equal(others, torch.zeros_like(others)), "shift/scale got gradient through alpha=0"
    print(f"[D03] one block, step 0: |dL/d(alpha)|max={gates.abs().max():.3e} (nonzero), "
          f"dL/d(gamma)=dL/d(beta)=0 exactly  PASS")

    # Whole model from init: the final layer is zero-initialised too, so at
    # step 0 only its output linear gets a nonzero gradient — nothing reaches
    # the blocks. After that update the linear is nonzero, step 1's backward
    # reaches the blocks' gates, and after step 1 they leave the identity.
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    target = torch.randn(B, 8, 8, 8)
    for step in range(2):
        opt.zero_grad()
        F.mse_loss(model(x, t, y), target).backward()
        blocks_have_grad = any(p.grad is not None and p.grad.abs().max() > 0 for p in model.blocks.parameters())
        assert blocks_have_grad == (step == 1), f"step {step}: blocks_have_grad={blocks_have_grad}"
        opt.step()
    assert not torch.equal(model.blocks[0](h, c), h)
    print("[D04] full model: step 0 updates only the final linear; step 1's gradient reaches the blocks; "
          "after it block 0 is no longer the identity  PASS")

    # CFG plumbing: dropping the label routes to the null class
    drop = torch.tensor([True, False, True])
    emb = model.y_embed(y, drop)
    assert torch.equal(emb[0], model.y_embed.table.weight[model.y_embed.null_id])
    assert torch.equal(emb[1], model.y_embed.table.weight[y[1]])
    print("[D05] label dropout -> null-class embedding (CFG training)  PASS")

    n_params = sum(p.numel() for p in model.parameters())
    print(f"\nall DiT sanity checks passed  ({n_params:,} params, hidden=64, depth=4, 8x8x4 latent, patch=2)")


if __name__ == "__main__":
    main()