Training a 210M text-to-image DiT from scratch on one GPU: what I measured [P]

I trained a 210M-parameter text-to-image diffusion transformer from scratch (3.5 days, one RTX PRO 6000, 4.2M images at 256²) mainly to understand the recipe end to end. Three measurements came out of it that I have not seen stated plainly elsewhere, so I'm posting those rather than the samples.

1. Learned null attention slots become the sink. Following the register-token idea, the model has 16 register tokens in the image stream plus 2 learned key/value slots appended to every cross-attention. At mid-noise in a middle block, those 2 slots receive ~90% of the cross-attention mass; the EOS token, which is the usual sink in cross-attention models, drops to ~4%; content words keep a few percent each, sharply on their objects. Register vectors grow to 4–13× the norm of image tokens by the middle blocks.

2. The flow-matching loss is a health signal, not a quality signal. It moved 0.805 → 0.754 over the whole run while held-out FID went 33.7 → 27.0, FD-DINOv2 570 → 218 and detector-based object accuracy 65% → 90%. Most of the loss at high noise is the irreducible variance of the velocity target; training and held-out loss stayed equal to the third decimal for 24 epochs.

3. The training-time timestep shift is worth more than doubling the steps. On 2,456 held-out prompts with the final weights: 20 steps with shift 2.8 → FID 27.0; 50 steps → 26.6; 8 steps → 28.4; 20 steps with no shift → 27.3 and FD-DINOv2 218 → 228. Shift 2.8 comes from the SD3/RAE rule √(32·32·32/4096) for the 32-channel FLUX.2 latent.

Setup, briefly: cross-attention DiT (896 × 16 blocks), 2D RoPE, QK-norm, SwiGLU, adaLN-single; rectified flow with logit-normal timesteps and the shift above; cosine velocity and dispersive auxiliary losses; five aspect-ratio buckets of ~256 tokens from step one; flan-t5-base frozen, long/short caption per image sampled 50/40/10 (empty). Data: Pexels 2.8M (60%), a quality-filtered 1.2M slice of FLUX-Reason-6M (25%), COCO with GPT-4V captions (15%). batch 256, 400k steps, EMA 0.9999, linear LR decay over the last quarter, torch.compile (2.4× over eager).

Everything, with the sources for each decision: https://github.com/ivanmikhnenkov/tinydit

Write-up: https://huggingface.co/blog/ivanmikhnenkov/tinydit-text-to-image-from-scratch-one-gpu

Weights: https://huggingface.co/ivanmikhnenkov/tinydit-256

Demo: https://huggingface.co/spaces/ivanmikhnenkov/tinydit

For the next phase (Flow-GRPO on this base), which reward would you start with: PickScore/HPSv2, a detector-based object reward, or something verifiable like counting?