bench-labs/PixelModel-v5

PixelModel v5 🖼️

PixelModel v5 is a 40,013,980 parameter text-to-image diffusion transformer. It generates 256x256 images in 50 sampling steps and it trained in about 7 hours on one A100.

The architecture is identical to v4. Same DiT, same width, same depth, same VAE, same text encoder, same objective. Nothing in the model changed. The only difference is that v5 saw 36 times more data.

That was the whole experiment, and it is worth being blunt about why.

v4 was starving, not undersized

v4 trained on MS-COCO, which is about 83,000 images. At batch 256 for 80,000 steps that works out to roughly 247 passes over the same photographs. A model that sees the same 83,000 images 247 times is not learning what the world looks like, it is memorising a small album.

We spent a while looking at architecture changes first. Linear attention, deeper stacks, a higher compression autoencoder, representation alignment. All of them are real techniques and none of them address the actual problem, because the bottleneck was never the shape of the network.

v5 trains on 2,986,571 image and caption pairs, which is about 10 epochs instead of 247.

The captions matter as much as the count

The images come from CC12M. Raw CC12M alt text is close to unusable for this, because it is scraped SEO copy. A typical caption reads like "This is the best recipe I have ever tried for Cuban bread. Cuban Recipes, Bread Recipes, Cooking Recipes, Pan Cubano".

v5 uses a recaptioned version where a vision model rewrote every caption in COCO style:

A purple-lit subway train sits at a station platform.

An old red tractor with gold fenders and red wheels stands in a snowy field.

A compact kitchenette in a studio apartment features a stainless steel refrigerator, a microwave, and a sink with a glass faucet.

Those describe what is in the frame, which is what a text-to-image model needs. They also match the caption distribution that FID and CLIP score are computed against, so the training signal and the evaluation are speaking the same language.

Training

120,000 steps at batch 256. AdamW at 2e-4 with betas (0.9, 0.99), no weight decay, gradient clipping at 1.0, cosine decay to a tenth of peak after a 1,000 step warmup, EMA at 0.9999 starting after warmup, and 10% caption dropout for classifier free guidance. It held 4.6 steps per second at 44.2 GiB the whole way.

The held-out curve climbs for the first 10,000 steps or so, which looks alarming and is not. EMA at 0.9999 has a time constant of roughly 10,000 steps, so early on those weights are still mostly initialisation. The number only becomes meaningful once EMA has caught up, and from there it falls without interruption.

v4 had no validation split at all, so it shipped whichever weights the last step happened to produce. v5 holds out 4,096 pairs, evaluates every 2,000 steps, and writes the best checkpoint to a separate file. Across the whole run the held-out loss improved at every single check and never once turned upward, which is the clearest evidence that the data change worked. Our sibling model VoxelModel v1, trained on 28,415 meshes for 811 epochs, peaked at step 64,000 and got 3.08% worse by the end. Same recipe, same team, different data regime, opposite outcome.

Results

Measured on 5,000 MS-COCO val2014 images at 256 center crop, using torchmetrics FID at 2048 features and torchmetrics CLIPScore with openai/clip-vit-base-patch32. This is v4's protocol.

	FID ↓	CLIP ↑
PixelModel v4	39.54	28.04
PixelModel v5, cfg 5.0	34.06	29.27
PixelModel v5, cfg 3.0	35.03	28.59
PixelModel v5, cfg 6.0	35.18	29.28
real photographs	n/a	30.05

FID is 13.9% lower and CLIP score is 1.23 higher, at an identical parameter count. Nothing in the network changed between v4 and v5, so the entire difference is the training data.

The last row is there because a metric without a reference point is hard to read. Feeding the real COCO photographs and their own captions through the same CLIPScore gives 30.05, so v5 at 29.27 is within 0.78 of the ceiling this metric can award. That number also served as a sanity check on the harness, since torchmetrics CLIPScore is incompatible with transformers 5.x and needed patching before it would run at all.

Guidance scale was swept rather than assumed. FID bottoms out at cfg 5.0 and CLIP keeps climbing slightly to cfg 6.0, which is the usual tension between the two. Anyone reproducing this should also know that FID is heavily biased at small sample sizes. The same checkpoint at cfg 5.0 scores 74.35 on 1,000 images and 34.06 on 5,000, so sample count has to match before two FID numbers can be compared at all.

One caveat worth stating plainly. v4's cached evaluation set is not published in its repository, so the reference images here were rebuilt from COCO val2014 with a fixed seed. The protocol is the same but the exact 5,000 photographs may not be. CLIP score is unaffected by this since it compares each generated image against its own caption, but the FID comparison carries that asterisk.

Specification

trainable parameters	40,013,980
frozen VAE, sd-vae-ft-mse	83,653,863
frozen text encoder, CLIP ViT-B/32	63,165,952
combined total	186,833,795
architecture	DiT, dim 384, depth 12, heads 6, patch 2
tokens	256, over a 32x32x4 latent
objective	rectified flow, target x1 - x0, logit-normal timesteps
conditioning	cross attention on 40 CLIP text tokens, plus adaLN-zero
resolution	256x256
sampling	50 steps, cfg 5.0 for best FID, cfg 6.0 for best CLIP

A note on that combined number, because it is higher than what v4 published. v4's card reports 161,500,000 total. That figure counts the CLIP text transformer but leaves out its 49,408 by 512 token embedding table, which is another 25.3M parameters. The encoder does not run without it, so it counts. The correct total for both v4 and v5 is 186,833,795.

What it does well and what it does not

After a few hundred samples the boundary is consistent.

It is good at material and light. Food, landscapes, skies, water, snow, foliage, stone, interiors, depth of field. A bowl of ramen, a wet cobblestone street, a wheat field under storm light, a library of wooden shelves. These come back looking like photographs.

It is bad at structure. Faces and animal heads deform. Limbs merge. Wings on aircraft are stubby or fused to the fuselage. Any text it renders, on a shop sign or an airliner livery, is convincing looking gibberish. Prompts asking for one isolated object in an empty scene often return the scene without the object, so "a lighthouse on a cliff at sunset" gives a very nice sunset over a cliff and no lighthouse.

Frequency in the training data matters more than complexity. Aircraft are structurally complicated and come out fine because CC12M is full of them. Hot air balloons are simpler and fail, because they are rare.

That failure profile is what a 40M model looks like once data is no longer the constraint. It has learned appearance thoroughly and topology poorly, and topology is what capacity buys. v6 is where that gets spent.

Usage

python main.py "a bowl of ramen with a soft boiled egg" --out ramen.png --cfg 6.0 --steps 50

As with every model in this line the weights are also stored as a PNG. model.png is not a picture of the network, it is the network, with each parameter packed across the red and green channels of one pixel.

Files

file	contents
model.safetensors	combined weights, backbone and VAE and text encoder
model.png	the same backbone weights as an image
dit.py	the model
train_v5.py	training loop with the held-out split
prep_v5.py	streams the dataset, encodes latents, caches captions
main.py	inference