Pressure testing the Pyramid-JiT (P-JiT)

For 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.

If 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.

Today, 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.

When 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.

We've written extensively about the steps we've taken in this direction, but here's the long story short:

• 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.
• 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 on the amount of compression they can achieve. There are other ways to get more compression, but they require substantial surgery to the traditional VAE. More importantly, improving the VAE on its task of reconstruction can produce worse generations in the downstream diffusion model.
• 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.

In 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). P-JiT enables 32×32 token reduction. This implies a 4× reduction in attention context-window.

For the details on how we got there, we'd recommend you take a look at our previous blogs on JiT-DDT and P-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.

Adjacent work in REPA, RAE-v2, and Self-Flow 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.

The 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). For context, generative video models pretty much always start as generative image models. You have to learn nouns (people, places, things) before you can learn verbs (actions, movement). It's a pressure test of sorts along the way towards larger models (and video).

Training Regime

Text 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. This is a 3:1 compression rate. A 32×32 patch with RGB values is a 32×32×3 = 3,072 representation, which gets projected into 1024 dimensions.

At block 5, we have a 2-layer PixelREPA readout that aligns the model's latent space to DINOv3-L. We mask 20% of the tokens in our hidden state and push it through a 2-layer transformer to predict DINOv3-L's feature representation of the same image. 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 and P-DINO).

Every DiT block is conditioned on the timestep through a low-rank AdaLN, following Z-Image. The timestep embedding goes through one down-projection to 256-dim. The image refiner and trunk share one down-projection; the text refiner has its own. Each block then has its own up-projection to a scale and a gate, zero-initialized so every block starts as the identity.

As 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.

We 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. 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. We keep a fixed grid aligned to 512² and multiply the 1024² patch indices by ½ to align them to the 512² coordinate axes. Moreover, we start the 1024² EMA 2M samples in, so it has adjusted to the new resolution before we start the averaging process.

For post-training, we used a variant of Finite Difference Flow Optimization (FDFO) with JEV-Image and human-in-the-loop judges.

It's still early days for post-training diffusion models

To 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.

Of 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.

DPO has been around for a while, but tends to dramatically overfit to specific aesthetic choices. Flow-GRPO 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 and this June's FDFO (the best articulation of diffusion post-training so far, in our opinion).

We 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.

Model Limitations

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, 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.

This 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.

Appendix A – Training Settings

Authorship statement

We wrote all the words on this page. We used Claude Opus 5.5 to help us build the image gallery and diagrams.

The Huggingface Model Card and Github Repo were written automatically by Claude Opus 5.5.