{"slug": "training-a-text-to-image-model-for-21k", "title": "Training a Text-to-Image Model for $21K", "summary": "A two-person research lab trained a text-to-image model from scratch for $21,000 as a proof-of-concept for its next video model, releasing it as an experimental model. The lab's Pyramid-JiT (P-JiT) architecture is a pixel-space diffusion model that converges in 11× fewer samples than its prior Linum v2 and uses 4× more aggressive token reduction, representing a 512² image in the same number of tokens previously used for a 256² image. The pressure-test run trained on roughly 600M samples rather than the 50–100M samples used in initial experiments, with text and image each passing through two single-modality refiner DiT blocks before merging into one stream of 22 identical DiT blocks.", "body_md": "Click any picture for full resolution and its prompt\n\n## Pressure testing the Pyramid-JiT (P-JiT)\n\nFor us, this $21K training run is a proof-of-concept along the way to our next video model. We were surprised by how good the results looked given the size of the model, so we thought we'd share it with the world as an experimental release.\n\nIf you haven't followed our work previously, we'd love to take a step back and introduce ourselves. We're a 2-person research lab training our own generative video models entirely from scratch. The ultimate ambition is creative tools that democratize access to animation. You can read more about that [here](https://linum.ai/about).\n\nToday, generative video models have two major flaws: they are too slow/expensive and they lack proper controls. Creativity is an iterative process. You make, you tweak, you get better. With this double whammy, video models (as they stand today) are effectively useless for storytelling.\n\nWhen we set out to build Linum v3 earlier this year, the primary goal was to bring the cost envelope of inference down significantly. If our model is fast and cheap, it's the first step (of many) towards useful creative tools.\n\nWe've written extensively about the steps we've taken in this direction, but here's the long story short:\n\n- **The bottleneck for training and inference for video is attention, O(n²).** If we can compress more aggressively, we'll reduce the context window for attention, slashing the cost of training and inference.\n- **Typically, compression is handled by a separate specialized model like a Variational Autoencoder, but that's not necessarily optimal.** VAEs seem to have an[empirical ceiling](https://linum.ai/field-notes/jit-ddt#hitting-the-vae-compression-wall) on the amount of compression they can achieve. More importantly, improving the VAE on its task of reconstruction[can produce worse generations](https://linum.ai/field-notes/vae-reconstruction-vs-generation#better-reconstruction-better-generation) in the downstream diffusion model.\n- **If you get rid of the VAE altogether, you should be able to compress more aggressively and learn faster.** Diffusion models are a more powerful family of generative models than VAEs. We should be able to use them to compress and generate at the same time. Instead of using another model's latent space, the diffusion model should learn a latent space purpose-built for its use case. In turn, we should be able to squeeze more juice out of this compression.\n\nIn our last blog, we iterated towards a new architecture that we called the Pyramid-JiT (P-JiT). It's a pixel-space model that converges in 11× fewer samples than Linum v2 with 4× more aggressive token reduction (i.e., we represent a 512² image in the same number of tokens we used to represent a 256² image).\n\nFor the details on how we got there, we'd recommend you take a look at our previous blogs on [JiT-DDT](https://linum.ai/field-notes/jit-ddt) and [P-JiT](https://linum.ai/field-notes/pyramid-jit). But at a high level, we found a noise schedule for pixel-space models that helps the models better learn fine-grained detail. And, we added readouts along the trunk of the DiT.\n\nAdjacent work in [REPA](https://arxiv.org/pdf/2410.06940), [RAE-v2](https://arxiv.org/pdf/2605.18324), and [Self-Flow](https://arxiv.org/pdf/2603.06507) all point to weak gradients in the early blocks of the DiT. By operating in pixel space, every block of our DiT is quite literally predicting a version of the final image, so we can boost gradient signal earlier in the network by predicting low-resolution versions of the final image. This helps accelerate convergence.\n\nThe architecture did well in our initial experiments, but we wanted to see how it performed when we trained it for longer (i.e., ~600M samples instead of 50–100M samples). It's a pressure test of sorts along the way towards larger models (and video).\n\n## Training Regime\n\nText and image each pass through two single-modality \"refiner\" DiT blocks, then merge into one stream of 22 identical DiT blocks. The text representation comes from a frozen Qwen3.5-4B model, which we read at layers 7, 15 and 27 (its three full-attention layers). We concatenate and project these representations into the text refiner. The image is split into 32×32 patches and compressed into a 1024-dimensional space, before being projected up into the image refiner blocks.\n\nAt block 5, we have a 2-layer [PixelREPA](https://arxiv.org/pdf/2603.14366) readout that aligns the model's latent space to DINOv3-L. We have two linear readout heads along the trunk after blocks 10 and 16. As discussed earlier, these predict the clean image at a quarter and half resolution, respectively. The total loss is given by MSE at the readouts, PixelREPA cosine-distance, MSE at the final output head, and perceptual losses on the full resolution image ([LPIPS](https://arxiv.org/pdf/1801.03924) and [P-DINO](https://arxiv.org/pdf/2602.02493v1)).\n\nEvery DiT block is conditioned on the timestep through a low-rank AdaLN, following [Z-Image](https://arxiv.org/pdf/2511.22699). The timestep embedding goes through one down-projection to 256-dim. Each block then has its own up-projection to a scale and a gate, zero-initialized so every block starts as the identity.\n\nAs described in our last blog, we employ a curriculum over the noise schedule. We start concentrated at high noise for the first 80M samples, so the model learns basic structure. Then, we widen the noise distribution for the remaining training time so it can learn fine-grained details. See [Appendix A](#appendix-a-training-settings) for details.\n\nWe start EMA at 100M samples, so the model has already learned decent representations of structure and detail. When we switch to 1024², we start with a fresh optimizer and resume training from the 512² EMA weights. At 1024² we double the noise scale from 2 to 4 to account for the redundancy of information (see [JiT, Appendix A](https://arxiv.org/html/2511.13720#A1.SS0.SSS0.Px2)). The higher noise level and need to reproduce fine detail exactly at high resolution result in a higher overall loss, even as image quality improves rapidly. We use RoPE interpolation in the resolution transition. Moreover, we start the 1024² EMA 2M samples in, so it has adjusted to the new resolution before we start the averaging process.\n\nFor post-training, we used a variant of [Finite Difference Flow Optimization (FDFO)](https://arxiv.org/pdf/2603.12893) with [JEV-Image](https://docs.typesafe.ai/images) and human-in-the-loop judges.\n\nTo be blunt, post-training is a very young discipline within diffusion models, and we've only seen the first sparks of it working in the past few months.\n\nOf course there's been a lot of work in model and schedule distillation, but when it comes to post-training in terms of improving fundamental model capabilities (like we've seen with coding in LLMs over the past 18 months), it hasn't really come to fruition just yet for diffusion.\n\nDPO has been around for a while, but tends to dramatically overfit to specific aesthetic choices. [Flow-GRPO](https://arxiv.org/pdf/2505.05470) was one of the first articulations of bringing RLVR-style frameworks to diffusion, and that only came out last fall. Since then there have been a few papers like [Diffusion NFT](https://arxiv.org/pdf/2509.16117) and this June's [FDFO](https://arxiv.org/pdf/2603.12893) (the best articulation of diffusion post-training so far, in our opinion).\n\nWe think that this difficulty in adapting techniques comes down to the fact that we don't have access to strong likelihood estimates in diffusion models like we have in auto-regressive LLMs. We'll dive more into these concerns and our full post-training stack in the coming weeks.\n\n## Model Limitations\n\n**This model was not trained with images of rendered text (no infographics, PDFs, pages of books, PowerPoint slides), so it cannot write text or generate posters.** As we detailed in our [data filtration blog](https://linum.ai/field-notes/data-filtering-gen-video), we don't consider this data useful for our ultimate goal of animation tools, so this text-heavy data has been filtered out entirely from our dataset.\n\nThis is a small model (2B-parameter DiT, 4B text encoder), and accordingly there are a few areas where the model struggles: full body, many entities (e.g. 6 cars, crowds), multi-person interactions, and world knowledge. The model is capable of correctly generating these images on occasion. Training a larger model (e.g. 7B parameters) and further post-training would likely mitigate these issues.\n\nExample Model Failures\n\nFull body: legs deformed on the starting blocks\n\n## Appendix A – Training Settings\n\nEvery setting of the run, by stage\n\n| Setting | 512² stage | 1024² stage | \n|---|---|---|\n| Samples | 610.27M (297,985 steps) | 49.97M (24,400 steps) | \n| Global batch | 2,048 | 2,048 | \n| Data | Images at native 512px-class sizes: 368×512, 512×512, 512×384, 512×336, 640×360 | Images at native 1024px-class sizes: 864×1216, 1024×1024, 1152×864, 1248×832, 1280×720 | \n| Aspect-ratio mix | 20% each, no resizing | 20% each, no resizing | \n| Noising |  | : the noise scale doubles with the side length | \n| Timesteps | LogitNormal(0.8, 0.8) for the first 80M samples (step 39,063), then LogitNormal(−0.2, 1.0) | LogitNormal(−0.2, 1.0) throughout | \n| Objective | x-prediction, v-loss, weight clamped at t = 0.1 · readouts at weight 1 · LPIPS 0.1 and P-DINO 0.01 on the final head for t ≤ 0.7 · PixelREPA 0.1 |  | \n| Optimizer | Muon (momentum 0.95, Nesterov) on the 2D block matrices · AdamW (β = 0.9, 0.999, ε = 1e-8) on the rest · Muon updates RMS-matched to AdamW |  | \n| Learning rate | Peak 2e-4 · linear warmup over 1M samples · flat to 550M · 1 − √ decay to 2.5e-5 over the last 60M | Linear re-warm over 0.5M samples to 2.5e-5, then flat | \n| Weight decay | 0.1 (0 on biases) |  | \n| Gradient clipping | 1.0 (global norm) |  | \n| EMA | Half-life 1,954 steps (4M samples), from 100M samples (step 48,829) | Half-life 977 steps (2M samples), from 2M samples | \n| Spike-aware EMA | EMA updates pause for 250 steps when the gradient norm exceeds 10× its rolling median |  | \n| Caption dropout | 5%, for classifier-free guidance |  | \n| Hardware | 64× H100 (8 nodes) · bf16 · FSDP2 · 32 images per GPU | Same, as 4 accumulated microbatches of 8 | \n\nFor further details on the architecture and training losses, refer to the [Appendix of P-JiT](https://linum.ai/field-notes/pyramid-jit#appendix-a-pyramid-jit-details).\n\n## Authorship statement\n\nWe wrote all the words on this page. We used Claude Opus 5.5 to help us build the image gallery and diagrams.\n\nThe Huggingface Model Card and Github Repo were written automatically by Claude Opus 5.5. We pointed Claude to our internal, experiment repo and had it pull out (and clean up) the necessary code.\n\n### Get Field Notes\n\nTechnical deep dives on building generative video models from the ground up, plus updates on new releases from Linum.", "url": "https://wpnews.pro/news/training-a-text-to-image-model-for-21k", "canonical_source": "https://linum.ai/field-notes/pjit-21k", "published_at": "2026-11-03 20:50:12.052107+00:00", "updated_at": "2026-11-03 20:50:14.109466+00:00", "lang": "en", "topics": ["generative-ai", "artificial-intelligence", "ai-research", "machine-learning"], "entities": ["Pyramid-JiT", "P-JiT", "Linum v2", "Linum v3", "REPA", "RAE-v2", "Self-Flow"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/training-a-text-to-image-model-for-21k", "markdown": "https://wpnews.pro/news/training-a-text-to-image-model-for-21k.md", "text": "https://wpnews.pro/news/training-a-text-to-image-model-for-21k.txt", "jsonld": "https://wpnews.pro/news/training-a-text-to-image-model-for-21k.jsonld"}}