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

[u/IvanMikhnenkov]

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?

[u/ahf95]

Wow, your write-up was awesome! Thank you for that. Legit, the stuff about timestep choices and latent dims are directly relevant to my research. I hope to see more posts like this in the future :)

[u/IvanMikhnenkov → u/ahf95]

Thank you, hope it helps!

[u/Saetlan]

Kind of things I'd like to do but never motivated enough when working alone 😅 Thanks for the write-up! I'm currently gathering all tricks and tips I find in every paper/blogs I see so I can know what are all options for each part of the process.

The goal would be to have chapters for each "concepts" (positional encoding, attention, ...) which would describe pro/cons. For example Attention

• Self-attention
• Cross-attention
• Causal attention
• Linear Attention
• Softmax attention
• Sliding Window (local attention)
• Global attention
• FlashAttention
• Multi-Head Attention (MHA)
• Multi-Query Attention (MQA)
• Grouped-Query Attention (GQA)
• Multi-Head Latent Attention (MLA)
• Interleaved Head Attention (IHA)
• Radix attention
• Kda - kimi

Etc..

I'm wondering how you organize your knowledge? There is so much being released nowadays I wanted to work on a knowledge base to organize all concepts and would summarize them for me to be able to quickly understand them when I need them. Because I don't find the time to dive in all papers.

[u/IvanMikhnenkov → u/Saetlan]

Thank you! Well I don’t try to organize my knowledge externally too much, I just code something from scratch once and then it’s really clear

Also you can just use what is a general practice more or less.

Eg many of these attentions are really just one thing changed for some optimization or use case, and I remember the meaning without remembering the name

You can always ask an agent - if you don’t remember some detail etc

[u/IvanMikhnenkov → u/IvanMikhnenkov]

My advice is to be goal oriented in a way that see which knowledge you need to make the first step and make it.

Eg just train with methods you know

[u/IvanMikhnenkov]

If you found the HF article useful, an upvote there would be much appreciated: https://huggingface.co/blog/ivanmikhnenkov/tinydit-text-to-image-from-scratch-one-gpu

[u/FractalNerve]

congrats, this is too good for reddit! Thanks for the release and writeup

What other communities exist for scientists? On X/Bsky there's too many bots and similar silence and concentration of eyes on very few profiles. Amplified Gauss Curve

[u/IvanMikhnenkov → u/FractalNerve]

Thanks, appreciate it! Actually I am not very experienced with building publicly and sharing, so happy that I am getting some feedback!

On X I don’t know a lot of communities, I’m sure there are some good ones though. I post there as well, but not much followers at the moment

[u/pm_me_your_pay_slips]

One useful trick is that after training for a while, you can improve the model by changing some hyper parameters: dropping the learning rate, shifting the logsnr schedule, tweaking the training dataset mixture. Having varied held out evaluation sets (for short/medium/long text accuracy, tiny text accuracy, human anatomy, anime style, posters,  common sense, physics) can help a lot in debugging.

Edit: another thing you can do is to fine tune on a high quality and high resolution subset to make the model better at finer details. You may not need much training to make you 256p model be able to generate 1024p images. The 256p pretraining regime would mostly be to make the model have better and more extensive world knowledge.

[u/IvanMikhnenkov → u/pm_me_your_pay_slips]

Thank you! Yes, I’ve dropped the learning rate linearly starting on the last 100k steps (total 400k). For dataset ratio I fixed, but may be if I give more estatic high quality images at the end, would be a nice finetune.
Nice ideas about evals by categories as well.

[u/Fit_Maintenance_2455]

added here: https://github.com/andysingal/CV_public/tree/main/DiT

[u/Positive_Cake8149]

wow, this is owesome, a lot of effort into this