Trained 210M text-to-image model from scratch on one GPU: what actually mattered

I wanted to deeper understand how modern image generation works, not by reading about it but by training one: a text-to-image model where every part that learns is written by me. The result is tinydit, a 210M-parameter diffusion transformer trained from scratch in 3.5 days on a single RTX PRO 6000, on 4.2M images at 256². It sits on two frozen pretrained pieces, the FLUX.2 autoencoder and the flan-t5-base text encoder; the transformer, the training recipe, the data pipeline and the evaluation are original. You can try it in the browser, download the weights, and read every decision with its sources in the repo.

This is the write-up of what mattered, in the order it mattered, with the numbers.

1. The data decided the outcome before training started

My first attempt at scaling, months earlier, used a recaptioned web crawl of about a million images. It made the model worse than COCO alone: the images were random internet quality and the captions did not fit them. So this time I did the unglamorous thing first: pulled twenty real rows from every candidate dataset, looked at them as a 256² model would see them, and only then chose.

The mix that trained this model:

Three details did real work. Every image has a long caption (up to 128 T5 tokens) and a short one, and training samples the long one 50% of the time, the short one 40%, and an empty caption 10% for classifier-free guidance; long-only captions hurt short prompts, short-only captions hurt composition. The synthetic set is capped at a quarter, because a model fed mostly generated images learns the generator's look. And nothing is centre-cropped to a square.

That last point deserves its own paragraph. A naive "keep only near-square images" filter would have discarded 94% of COCO and 99% of Pexels, because real photography is 4:3 and 3:2. Instead the model trains on five aspect-ratio buckets of about 256 tokens each (256×256, 288×224, 224×288, 320×208, 208×320), every batch drawn from one bucket, from the very first step. Mean crop: 3 to 5% of the long side. Position comes from 2D rotary embeddings, so the transformer needs no change per shape, and it even generalises to a 448×256 frame it never saw.

2. The architecture is the current recipe, plus registers

Nothing in the transformer is invented here, and that is the point: it is the block every current model uses, written from scratch so I understood each line. A DiT with cross-attention to the text tokens (the PixArt and SANA lineage; FLUX and SD3 use joint attention instead), 2D RoPE, QK-normalisation, SwiGLU feed-forwards, and adaLN-single modulation, which spends one shared modulation network plus a small per-block table instead of a full linear per block. That change alone freed 27% of the old parameter count for zero compute, and I spent it on depth: 896 wide, 16 blocks, 14 heads.

The one addition from recent research is registers. Sixteen learned tokens join the image sequence at block 3 and are dropped before the output, and every cross-attention block gets two learned "null" key/value slots that are always attendable. The idea comes from Vision Transformers Need Registers and the 2026 papers that brought it to diffusion transformers: give the model a place to park global state and "attend to nothing" so it stops abusing a random patch or the text encoder's end-of-sequence token for that job.

It worked, and measurably. At mid-noise in a middle block, the two null slots receive about 90% of all cross-attention mass; the end-of-sequence token, which used to be the sink, gets 4%; the content words get the remaining few percent, sharp and on their objects. The register vectors grow to 4 to 13 times the norm of an image token by the middle blocks, the signature of tokens acting as scratch space, and each image patch sends 8 to 18% of its self-attention to them.

3. The recipe: a shift, two extra losses, and compile

Rectified flow with a timestep shift of 2.8. The model predicts the velocity from noise to image, and timesteps are sampled log-normally as in SD3. The shift matters because the FLUX.2 latent has 32 channels: an image is still recognisable at a noise level that would have destroyed a 4-channel SD latent, so to get equally hard high-noise problems you need more noise. With the shift, half of all training samples are more than 74% noise, and the same shift bends the 20-step sampling schedule so that 15 steps land in that region, where the layout is decided.

Two auxiliary losses. A cosine term on the velocity direction (from LightningDiT), because at high noise the MSE is dominated by magnitude while the direction is what determines the image. And a dispersive loss on block-5 features (Wang and He, 2025) that pushes different samples' internal representations apart with no external encoder; it keeps the representation from collapsing and costs nothing.

torch.compile, for training. People assume compilation is an inference tool. It is not: Dynamo traces the forward pass, AOTAutograd derives the backward, and Inductor fuses the small elementwise operations around the matrix multiplies in both. Measured on the real step: 2.4× faster and half the memory. The constraint is static shapes, so each bucket gets its own compiled graph, five in total.

The rest. AdamW at 2e-4 with betas (0.9, 0.95), fp32 master weights under bf16 autocast, an EMA of the weights with decay 0.9999 (a 10,000-step horizon; the common 0.999 is far too short for a 400k-step run), batch 256, and a linear decay of the learning rate to zero over the last quarter of the run. 0.76 seconds per step; 400k steps in 3.5 days.

4. The loss is not the signal

Here is the curve that would have fooled me if I had trusted it:

The flow-matching loss fell from 0.805 to 0.754 over the whole run while the images went from blobs to photographs. The reason is structural. For a given noisy input, many image-and-noise pairs are consistent with it, so the best possible prediction is their average, and the spread around that average is error no model can remove. At high noise, where half the training happens, that irreducible part is most of the loss. What the model can still learn is a thin layer on top of a large floor.

So the loss is a health signal (training and held-out stayed equal to the third decimal throughout, which is how I know nothing was memorised across 24 epochs) but not a quality signal. Quality was tracked every 10k steps on a frozen set of held-out captions: FID and FD-DINOv2 against the real held-out images, an object accuracy where a COCO-trained detector has to find the requested class in "a photo of a {class}" for all 80 classes, and the human-preference models PickScore and HPSv2.1, which are the same family of model I will use as rewards in the RL phase. FID and object accuracy saturated around step 300k; FD-DINOv2 and the preference scores kept improving to the end, and the learning-rate decay in the last 100k steps bought a further 5% on FD-DINOv2.

5. What it can and cannot do

Single objects, animals, scenes, colours and simple spatial relations are reliable, in all five shapes. Long training-style captions are followed in their main clauses. And here is the failure grid, which I think belongs in every write-up like this:

Readable text, close faces, crowds, counts above three, clock hands, stacked geometry. These are the known limits of a 210M model at 256², and they are the natural targets for the next phase.

One more figure, because it explains a question I kept getting: why 20 steps are enough. The bottom row is the model's own prediction of the final image at each step of sampling; the layout is fixed after three steps and everything after that is sharpening.

Measured on 2,456 held-out prompts: 50 steps instead of 20 is worth 0.4 FID and nothing on the preference scores; removing the training-time shift costs more than doubling the step count gains.

6. Five things I would tell someone starting this run

1. Look at your data at the resolution you train at, before anything else. Twenty real rows per candidate, on one page, decided more than any architectural choice.
2. Match the timestep distribution to the latent. A 32-channel latent wants a shift near 2.8; without it the model under-trains exactly where layouts are formed.
3. Buckets from step one. The square-crop warm-up phase is a habit, not a requirement, and it throws away the parts of the frame that captions describe.
4. Do not read progress off the loss. Pick two or three metrics you can compute every hour and freeze the prompt set so every evaluation is comparable.
5. Compile the training step. It is the cheapest 2× you will ever get.

What's next

The base model is deliberately not distilled and uses plain classifier-free guidance, because the next phase is reinforcement learning with Flow-GRPO, using the same preference models tracked above as rewards, and the failure grid as the target list. The results will be a second release in the same repository.

Code and every decision with its sources: github.com/ivanmikhnenkov/tinydit · Weights: tinydit-256 · Demo: Space · Updates: x.com/ivanmikhnenkov