Train and Sample Like 2026: Logit-Normal Steps, EMA, and Few-Step Flow
Logit-normal timestep sampling, EMA weights and rectified-flow reflow in PyTorch, measured on a tiny DiT, plus where latents and distillation go next.
PyTorch
adaLN-Zero and Classifier-Free Guidance in PyTorch
Classifier-free guidance in PyTorch: class labels, label dropout, the guidance formula, adaLN-Zero init, and measured results on a tiny DiT.
MMDiT Joint Attention in PyTorch: Two Streams, One Generator
Build MMDiT joint attention in PyTorch: condition tokens, double-stream and single-stream blocks like SD3 and FLUX, measured against adaLN on a tiny DiT.
RMSNorm, QK-Norm, SwiGLU, and 2D RoPE in a Tiny DiT
RMSNorm, QK-norm, SwiGLU and 2D RoPE in PyTorch: what each one does, the code, and a measured before/after on a tiny diffusion transformer.
torch.compile, Mixed Precision, and SDPA: Speed Up TinyDiT
torch.compile, bf16 autocast, GradScaler and SDPA added to a PyTorch trainer, with honest CPU benchmarks: what got faster, what got slower, and why.
Build a Tiny Diffusion Transformer in PyTorch (DiT from Scratch)
Build a diffusion transformer (DiT) in PyTorch from scratch: timestep embedding, adaLN, the TinyDiT model, and a CPU training run on CIFAR-10 with scores.
Flow Matching in PyTorch: Train a Model to Follow a Straight Line
Flow matching in PyTorch, step by step: the straight noise-to-data line, velocity loss, a 2D toy you can watch learn, and an Euler sampler for images.
Patch Embeddings in PyTorch: How a Transformer Reads an Image
PyTorch patch embedding explained: patchify and unpatchify images, a Linear patch embed, learned 2D position embeddings, and an image-in, image-out ViT.
Attention and the Transformer Block in PyTorch
PyTorch attention from scratch: Q, K, V, softmax(QKᵀ/√d), multi-head attention, and the pre-norm transformer block, checked against PyTorch's SDPA.
PyTorch nn.Module and the Training Loop: Layers, Loss, and Checkpoints
Build a PyTorch training loop: nn.Module, Linear, GELU, LayerNorm, residual blocks, AdamW, and a reusable Trainer that saves and resumes checkpoints.