Hitting the VAE compression wall
Unlocking aggressive compression with a unified model
Extending JiT for text-to-image models
Recovering finegrained details in pixel-space
Linum v2 was bottlenecked by the enormous size of its attention context window. A 720p, 5 second clip cost a whopping 110K tokens. To put that in perspective, LLMs see samples with fewer than 8K tokens for 97% of their pretraining . Attention is quadratic in cost, so the biggest lever we have to accelerate model training is pruning the context window down.
Most generative image and video systems are Latent Diffusion Models (LDMs). They split compression and generation into independently trained modules: the Variational Autoencoder (VAE) and the DiT (Diffusion Transformer). Recently, pixel-space models like the JiT have shown to be a promising alternative. It reduces two models into one and allows the diffusion model to construct a latent space specifically for generation, rather than rely on one built for reconstruction.
When trained on our (image, caption) dataset, the JiT seems to struggle to produce finegrained details. We propose a novel encoder-decoder architecture (JiT-DDT) that recovers this detail and trains much more efficiently than its LDM counterpart. Against our Linum v2 baseline, the JiT-DDT trains a text-to-image model with 3.6× fewer GPU-hours, even though it generates images with 4× the pixels.
320 pixel tokens = 64 encoder + 256 decoder
JiT-DDT 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.
Hitting the VAE compression wall
Almost all generative image and video models are Latent Diffusion Models (LDMs). These have two key components, a Variational Auto Encoder (VAE) for compression and a Diffusion Transformer (DiT) for generation.
Operating in raw pixels is too expensive (especially for video), so we first need to find a way to reduce RGB pixels into a smaller amount of tokens for the DiT. This is where the VAE comes in. It's trained for compression and reconstruction. Specifically, it pushes our pixel-space samples through a probabilistic encoder, spits out -dimensional tokens, and then pushes these latent tokens through a probabilistic decoder to land back in pixel-space.
The VAE is trained to compress and reconstruct
When building a LDM, you train the VAE separately and then freeze it (i.e. no gradient flow from the DiT into the VAE). This way the latent space stays static throughout the course of DiT training. You run the VAE's encoder to embed your data, train the DiT to traverse the VAE's latent space, and then transform the DiT-generated latent tokens into pixel space using the VAE's decoder.
The VAE is trained once and frozen; the DiT learns to move through its latent space
We want to eke out as much token compression as possible from the VAE, so that we can curb the cost of attention in our DiT. But if you take a survey of the popular open source text-to-image models like FLUX, Ideogram, and Z-Image, you'll notice that they all cap out at 16×16 token reduction. This aligns with our experiments on Image-Video VAEs from a few years ago . Unfortunately, it seems like there is an empirical ceiling on the amount of compression we can get out of a standard CNN VAE without degrading the reconstructions.
Unlocking aggressive compression with a unified model
Last fall, Tianhong Li and Kaiming He published a paper ( JiT ) that achieves 32×32 token reduction by throwing away the VAE altogether and pushing the compression task into the DiT itself.
Patchify: 4×4-pixel patches → 48-dim tokens → linear bottleneck to 12
each patch is tokenized independently: 16 pixels × 3 channels become one 48-dim token
a linear layer W ∈ ℝ 12×48 projects 48 dims down to 12
This approach to reducing token counts isn't particularly new. It was invented for vision transformers ( ViT ) half a decade ago, and it's pretty commonly paired with a VAE to further condense token sequences before they enter the DiT. In Linum v2, our VAE gave us 8×8 (h×w) compression and 16-dimensional latents. At the base of the DiT, we applied 2×2 patchification to get 16×16 token compression and 64-dimensional latents. We used it in Linum v2 and so do models like FLUX.
So, why hasn't anyone tried this before? This feels like a free lunch. You get a (potentially) lossless way to cut down attention cost, and it's bone-dead simple.
In early 2025, papers like VA-VAE demonstrated that DiTs struggle to learn from high dimensional inputs. There are small hacks like using an external model as a regularizer during VAE training (e.g. DINOv3) that (likely) enabled models like FLUX-2 to make the leap from 64 latent dimensions to 128 latent dimensions for their DiT. But, these strategies just kick the can down the road on a clear learnability problem within the DiT. Aggressive patchification explicitly pushes information into the channel dimension, so it triggers this instability. But as it turns out, this is not intrinsic to the architecture. Rather, it's downstream of the v-prediction, v-loss flow matching objective that everyone's been using to train diffusion models these past few years.
A quick refresher on flow matching
In old school 2022-era denoising diffusion ( DDPM ), we iteratively noise a sample and train a neural network to remove the noise. This way at inference time we can use our neural network to transform Gaussian noise into a sample from our data distribution over a sequence of steps. This formulation has a host of issues (e.g. oversaturation in generation , unstable learning , distillation collapse ), so in the intervening years the field has shifted away from it towards flow matching.
In flow matching , we construct a straight line path between every sample in our data distribution and a sample of Gaussian noise: The path between noise and samples does not have to be straight. But in practice, we all do it.
At , we recover . At , we get , where . We follow the DDPM convention throughout this post: is data, is noise. Some flow matching papers run the other way, with as noise and as data. The two formulations are equivalent. Then we train a network to approximate the velocity along that path:
We call this v-prediction, v-loss because the neural network is explicitly predicting velocity and it's trained on the MSE between its velocity prediction and the ground-truth, conditional velocity field.
V-prediction and the curse of dimensionality
If you're training a flow matching model you don't necessarily need to train your neural network to predict and regress velocity. The three terms are linearly re-arrangeable; so you can mix and match , , and across prediction and regression targets:
Three targets, each linear in the other two
Rearrange one identity to fill each off-diagonal cell
In JiT, Li and He revisited the v-prediction, v-loss decision that the field's been making since the inception of flow matching. They took a toy distribution (points on a spiral) and then projected these points from 2D to different high dimensional spaces of increasing size. For each of these spaces, they trained flow matching models with x-prediction, -prediction, and velocity-prediction; and found that the x-prediction was the only model type to accurately generate samples from the spiral distribution at large dimensions. DiTs have been struggling to learn from high-dimensional inputs because of the curse of dimensionality.
Velocity is . When we do v-prediction, our neural network has to implicitly learn the signal ( ) and noise ( ). Noise is a random Gaussian that will cover the entire -dimensional space. So, as we scale the problem of fitting noise (within the velocity term) becomes exponentially harder. This is why aggressive patchification failed in the past and why LDMs have been struggling to learn from high-dimensional VAE latents. As we grow the channel dimension, we end up in the degenerate case where our DiT is struggling to learn high dimensional Gaussian noise.
By switching to x-prediction, we can try to side-step the curse of dimensionality. If we believe that images and videos naturally lie on a low dimensional manifold, we should be able to have our models predict effectively even with high .
Noise fills the whole D-dimensional ball; images sit on a thin sliver of it
In high-dimensional space (D = 512), noise is truly random. It spreads across the entire space, is incompressible, and cannot be described by any smaller number of dimensions (left). Images are intrinsically low-dimensional, so even in a high-dimensional space they cluster in a small subspace (right).
By predicting v, the model has to learn both the noise ε and the structure. Noise is the harder of the two, and the bigger D gets, the more of the model's capacity goes to fitting it. By switching to x-prediction, the model can spend its full capacity on the low-dimensional signal, even when D becomes large.
Empirically this works, if you do x-prediction, v-loss. The network predicts . We convert that prediction to a velocity and take the MSE against the true velocity. Since , this is just x-loss scaled by : the same objective, weighted toward small (low noise, nearly clean images). In turn, this unlocks our ability to apply aggressive patchification, blow up the channel dimension, and push the compression problem into the DiT.
Extending JiT for text-to-image models
When we read about JiT, we were really excited to give it a go, since it was explicitly able to achieve 32×32 token reduction. But, we'd be remiss to say this is the only way to achieve this level of compression. Or, that everyone agrees that this is the best way to achieve this amount of compression.
LTX has been able to do this in their video models by altering their VAE's decoder to make it an explicit denoiser (i.e. they finetune the VAE decoder with a flow matching objective). More recently, Minimax H3 has achieved 32×32 compression in their VAE by swapping out the standard ~80-150M parameter CNN Decoder with a 2B parameter transformer (roughly the size of our entire Linum v2 model). And on toy benchmarks like ImageNet, LDMs still out-perform pixel space models by a smidge. Nevertheless, we think we can overcome some of the limitations present in the original JiT paper and match LDMs' performance in generative image and video.
Ideologically, we believe simple tends to beat complex when it comes to training neural networks at scale. Papers like E2E-VAE from last fall have shown that allowing your DiT to backpropagate (smartly) into the VAE can improve generation results and accelerate convergence dramatically.
To us, it makes logical sense that if we can specifically tailor the "latent space" for generation rather than rely on one built for compression, we can get better results. And we get the added benefit of having one cohesive model, rather than two disjoint ones.
We also think that ImageNet benchmarks on JiT understate its potential. The JiT might be able to achieve better compression than an equivalent VAE, by leaning on the scaffolding provided by the text prompts.
Text-to-image baselines
When we pretrained Linum v2, we relied on a VAE + patchification stack that afforded 16×16 token reduction. So, we trained on ~600M samples at 256px resolution before introducing 180p video and scaling up to 512px resolution.
For our JiT baseline, we wanted to get a sense of the output quality with the same image-latent-token budget. That meant we trained on 512px images with 32×32 token reduction.
Our JiT setup
By moving from LDM to pixel-space, we transitioned from v-prediction, v-loss to x-prediction, v-loss. But, we also made a slew of other tweaks to the network:
One single-stream DiT does the compressing and the generating
Single stream backbone Instead of alternating blocks of self-attention (image/video) and cross-attention (text-to-image/video), we concatenate visual tokens and text tokens into a single stream that goes through the DiT. This increases the attention sequence in every block and increases the FLOPs per token, but should allow for significantly more expressive relationships between text and image tokens. v2 block: self-attention, then cross-attention v3 block: one self-attention over image and text
Wider instead of deeper Our old model was a 40-layer transformer with 2048 hidden size. Here, we switch to a 23-layer transformer with a wider 2944 hidden size. Wider networks have become standard in recent DiT architectures (e.g. Z-Image), so we adopted the same.
Perceptual losses When you train a VAE, you use perceptual losses like LPIPS and adversarial loss via a GAN to push the reconstructions towards what humans like. MSE on its own gives you a blurry mess. Now that we don't have a VAE decoder, we need the JiT itself to leverage these losses to generate stuff humans like. We still use LPIPS , but instead of a GAN we use a P-DINO loss . Both are only applied when . perceptual losses · both towers frozen · on x̂₀ vs x₀
SiLU to SwiGLU We swap standard SiLU non-linear activations with gated SiLUs (i.e. SwiGLU). SwiGLU FFN · 2,944 → 7,936 → 2,944 · elementwise multiply
Muon optimizer Moonshot's Kimi models proved that the Muon optimizer works really well at scale . As they recommend, all the 2D matrices in our network (e.g. q/k/v matrices for attention, FFN weights) are optimized with Muon , while layers at the input/output of the network (e.g. patchification, output head) and scales/biases (e.g. AdaLN) are still optimized by AdamW.
PixelREPA auxiliary loss It's become pretty common to accelerate the convergence of your DiT by having an earlier layer in the network (e.g. layer 8 of a 23-layer transformer) align to the embedding of in an auxiliary model (e.g. DINOv3). This technique is referred to as REPresentation Alignment (REPA). We'll dig into this (and the limitations) later in the blog, so hold on for that. But for now, plain REPA did not work well for the JiT. Instead, we adopted PixelREPA which masks out x% of tokens in our visual token hidden state, pushes it through a shallow transformer, and then applies the typical cosine-distance loss between all visual tokens (including the masked tokens) and the auxiliary representation from DINOv3. DINOv3 uses 16×16 patches. We need the token count between the DINO representation and our hidden state to match, so we downsample the images before they go through DINO. For example, if we're doing a 32×32 patchification on 512×512 images, we will have 256 tokens. We downsample the image to 256×256 before passing it through DINO's 16×16 patchification to also get 256 tokens. PixelREPA · tapped after block 8 · x₀ downsized to 256px so DINOv3 gives 256 tokens
Sigmoid attention gating Now that we're moving from a cross-attention to a single-stream DiT architecture, we may be at a higher risk of attention sinks . We adopt sigmoid attention gating to neutralize this issue. attention gate · 23 gates per block (1 for each attention head) · elementwise multiply
Qwen text embeddings Instead of T5-XXL text embeddings, we use hidden states from a more modern decoder-only LLM, Qwen3.5-4B. One downside to using a LLM is that it's unclear what hidden state to take as your embedding. Most modern LLMs use some sort of alternating sequence of sparse/linear attention and full-attention. We take the hidden states calculated after full-attention blocks, concatenate them together, and have the DiT learn a transform to combine these representations into a single text condition. Recent Ideogram and FLUX models are more aggressive here, using larger LLMs and aggregating information across all hidden states. Given the size of our DiT, it seemed like overkill to go down that path. text conditioning · three hidden states, one learned projection
At train time, we get to pick the distribution from which t is sampled. Empirically, there's a small band of values at high t where the structure of the image is determined. This is the hardest part of the trajectory for the model to learn, so we skew timesteps accordingly.
Note that the σ term is tied to pixel count: σ = 1 is for 256×256, σ = 2 is for 512×512. The intuition is that we need to shift more aggressively at higher pixel counts. There is more redundant information within the image, so we need to noise more. This is the exact schedule from JiT.
The MSE is x-prediction with the velocity weighting (x-prediction, v-loss), clamped at t = 0.1 so the weight caps out at 100×. LPIPS and P-DINO are perceptual losses on the predicted image. PixelREPA aligns the block-8 hidden state with DINOv3 features of the clean image (downsampled 2× to match the token grid of h₈).
Recovering finegrained details in pixel-space
One of the biggest limitations that folks have observed about JiTs is that they struggle to generate the finegrained details. Our baselines corroborate this. If we want to really get our pixel space models to sing, we need to fix this.
DDT: Decoupled Diffusion Transformer
LDMs face the same issue, but to a much lesser extent.
In early 2025, Shuai Wang and team tackled this problem directly with their DDT (Decoupled Diffusion Transformer), scoring SOTA on ImageNet gFID at the time. They observed that —
In each denoising step, diffusion transformers encode the noisy inputs to extract the lower-frequency semantic component and then decode the higher frequency with identical modules. This scheme creates an inherent optimization dilemma: encoding low-frequency semantics necessitates reducing high-frequency components, creating tension between semantic encoding and high-frequency decoding. DDT abstract
DDT splits the model into two components, a "conditional encoder" for low-frequency structure and a "velocity decoder" for high-frequency detail. They give the encoder most of the layers, since the most difficult portion of the probability path to master is the transition from random noise to basic structure. When you train a diffusion model, you have to sample timesteps between and to turn your clean samples into noise-interpolated . Since Stable Diffusion 3 , it's been widely known that there seems to be a small band of high- values where the overall structure of the image is determined. We oversample from this part of the distribution, to accelerate convergence. This finding was the inspiration for the DDT to allocate the bulk of its parameters to the encoder.
DDT splits one DiT into a large structure encoder and a shallow detail decoder
Within the encoder they apply REPresentation Alignment (REPA) , an auxiliary loss that accelerates training by aligning the hidden states of an early layer of the model to the DINO representation of the clean image ( ). If you keep REPA loss active throughout all of training in standard DiTs, it actually hurts overall FID .
The DDT avoids this problem by giving the decoder the noised image ( ) so it can extract the finegrained details that REPA might otherwise destroy. Moreover, the DDT frees up the decoder to focus solely on details by creating an