Pyramid-JiT: Low-Res Drafts, High-Res Images

TL;DRIn our last post, we developed a pixel-space encoder-decoder architecture (JiT-DDT). By throwing away the VAE, we were able to compress tokens more aggressively. In turn, this slashed the cost of attention and directly led to faster training and inference.While effective, it's unclear how to scale encoder-decoder architectures under fixed parameter budgets. So today, we propose the Pyramid-JiT (P-JiT), a novel decoder-only pixel space architecture that beats our prior models, by predicting the target image at several resolutions along the DiT trunk. It reaches Linum v2's FD-DINOv2 on 11.3× fewer training samples, and trains in 4.3× fewer GPU-hours at 4× the pixels. 0100M200M300M400M500M600M301003001000FD-DINOv2 · lower is bettertraining samples seenreal-vs-realfloor 25.4Linum v279.9JiT-DDT 39.1P-JiT 27.91.0×fewer samples to reachLinum v2's final FD-DINOv2Each point scores 30K generations at 256² against a fixed reference of 30K in-domain real images. The real-vs-real floor is the same measurement with 30K more real images. 25.4 is a rough lower bound on the best score a model could expect to get here.Linum v2 is our previous latent diffusion model. Only four Linum v2 image checkpoints survive, so its curve starts at 117M samples. Lines between checkpoints are interpolated in log-log. Research releaseP-JiT code and model weights are available under the Apache 2.0 license. We hope that by sharing our findings with the broader community, we can encourage others to also explore more efficient training methods. This should be treated as a research artifact, not a full model release. Stay tuned for more research checkpoints like this, en route to Linum v3. Where should the parameters go? Our goal is to build affordable animation tools, so that anyone can bring their story to screen. Right now the best video models in the world (e.g. Seedance) are expensive because they calculate attention on ridiculously large token sequences. For our own Linum v2 model, a short 720p, 5 second video cost 110K tokens. By comparison, a LLM trains on context sizes of 8K or less for 97% of their training regime. To make our vision a reality, we need to cut the cost of video generation by an order of magnitude. We're bottlenecked by attention, which is an operation with respect to context-window token count. If we can cut down the number of tokens, we can slash the cost of attention and achieve massive savings across training and inference. En-route to our Linum v3 video model, we've been using text-to-image as a testbed to try out ideas that can help us condense context windows.Intuitively, video models need to learn nouns (static people, places, things) before they can learn verbs (actions, motions, etc.). As a result, video models are pretty much always trained as extensions of image models. Previously, we wrote about how LDMs (Latent Diffusion Models) seem to have an empirical ceiling of 16x16 token compression unless you tweak the auto-encoder significantly. So, we pivoted towards pixel space models (JiT) to achieve 32x32 token compression, unifying compression and generation into a single streamlined model. From there, we developed a novel pixel-space architecture (JiT-DDT) that produced better results than our prior LDM while training 3.6x faster. While effective, it's unclear how to scale encoder-decoder architectures. Do you increase the encoder's size so that it learns structure better? If so, it seems wasteful to spend so many parameters predicting a low resolution image. Or, do you increase the decoder's size so that it spends more time learning the fine-grained details that really matter to humans? In doing so, you might cap the model's ability to learn the framing and structure of a scene (without which the details are meaningless).Our architecture was inspired by the Decoupled Diffusion Transformer (DDT), which found it optimal to allocate 4x more parameters to the encoder than the decoder. They induce structure learning while operating in full-resolution v-space by starving the decoder of the class condition. We enforce structure learning in x-space by simply having the encoder predict a low-resolution version of the target image . Given these significant differences, it wasn't immediately obvious that their optimal encoder-decoder size would transfer to the JiT-DDT. And in our early experiments, the 4:1 encoder:decoder parameter-ratio proposed by the DDT performed worse than 1:1 encoder:decoder parameter ratio. Language faced a similar dilemma ~10 years ago. If you dig through the archives, you'll see that the transformer was first proposed as the building block for an encoder-decoder trained on language translation. The field then split into encoder-only models (BERT-style classifiers) and decoder-only models (GPT-style LLMs). The beauty of a decoder-only model is that it's really easy to scale, since there is no explicit architectural bottleneck. So the driving question is: Can we port over the advances from encoder-decoder back to the vanilla JiT and get similarly high quality results? Revisiting the JiT In our prior work, we found that we could dramatically improve the quality of our generations by widening the noise schedule during training and adding separate DiT blocks to refine the image and text streams independently, before merging them into a single shared stream. So, we copied and pasted these innovations over to the JiT. Additionally, we widened our linear patchification's bottleneck dimension from 256 to 1024 to reduce the effective compression rate within the network from 12x to 3x.In the original JiT paper the authors found that the model scored worse on ImageNet FID as they increased the bottleneck dimension past 128. That's unintuitive. If you have a bigger bottleneck dimension, your effective compression rate goes down, so you should be able to capture detail more easily.For 32x32 patches that's 32*32*3 (RGB) = 3072 values being consolidated in a 256 dimensional token (3072/256 = 12x compression). If we use a 1024 bottleneck this becomes 3x effective compression (3072/1024). We tried ablating 256 bottlenecks vs. 1024 bottlenecks. In practice, the results were pretty much the same (which is weird in its own right). So, we stuck with the wider, more expressive bottleneck. And, we picked up the innovations from Z-Image.RMS-Norm instead of LayerNorm, tanh gates with post-norm, and truncated AdaLN. These other goodies didn't seem to help the JiT-DDT, but there's no harm porting them over. The conservative writes to the residual stream should help make training more stable. JiT (last post)2.0B active pixel-space DiT256 image tokens100M training samplesJiT-DDT (last post)2.5B active pixel-space DiT320 image tokens138M training samplesRevised JiT (new)2.18B active pixel-space DiT256 image tokens75M training samples With these updates, we're able to recover a lot of performance, but we're still missing two big deltas from the JiT-DDT: multiresolution prediction and dual patchification. Replacing the encoder with readouts In the JiT-DDT, we had the encoder predict a 64x64 version of the 512x512 input image. Our hypothesis was this would help the model draft an initial low-res answer so that the decoder could allocate its capacity to filling in details. When we zoomed out and looked at adjacent work like RAE v2 and Self-Flow, it seemed our encoder prediction served two more concrete purposes. First, it boosted the gradient signal in the early layers of the DiT that tend to adapt much more slowly. And second, by regressing a low resolution image it helped the early portions of the DiT explicitly learn structure (something that hacks like REPA tap into to accelerate learning as well). Even if we throw away the encoder, we can still access these improvements. Since our model is trained with x-prediction, every intermediate DiT block is implicitly predicting a version of our target image. So we can slap K OutputHeads along the trunk of the DiT and have each regress a version of our target image (either at full resolution or downsampled to a lower resolution, like we did in the JiT-DDT).This K-Readout formulation is so easy to implement, in large part because of our decision to train with an x-prediction objective in pixel-space. If the model was trained in latent space, we'd have to run intermediate outputs through a several hundred million parameter VAE. That's quite expensive, compared to running intermediate outputs through very small MLP OutputHeads. Variant512×512 loss: K-Readout vs 1-Readout1-Readout (Revised JiT)512×512 at block 22baseline · 0.10373-Readout128×128 at block 10256×256 at block 16512×512 at block 22−10.0% 0.09343-Readout, full resolution512×512 at block 10512×512 at block 16512×512 at block 22−8.5% 0.09482-Readout128×128 at block 10512×512 at block 22−7.4% 0.09603-Readout, down-weighted128×128 at block 10 · weight 0.25256×256 at block 16 · weight 0.5512×512 at block 22 · weight 1.0−5.1% 0.0983every run trains the same 22-block trunk on the same recipe for 75M samplesreported loss is the v-space MSE from the final DiT block that predicts a 512×512 image We tried predicting full-resolution images along the network but that performed slightly worse than predicting images in an inverted pyramid of increasing resolutions from 128² to 256² to 512². Another hunch we had was that we might want to place less weight on the low-resolution readouts so the model could allocate more features for high-resolution detail. In practice that worked worse than equal-weighted readouts; so 3-Readout became our new baseline for our subsequent experiments. 1-Readout3-Readout1-Readout3-Readout2.18B active pixel-space DiT · 256 image tokens · 75M training samples · same prompt, same seedEyes and freckling seem more aesthetic in the 3-Readout sample of red-headed-women. Moreover, the wrinkles seem natural in the woman's brow, while the 1-Readout has unnatural horizontal streaks across the entire face. Fur looks