Training Text-to-Image Models 3.6× Faster TL;DRLinum 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 perspe
By Coderz Club · 2026-09-17 · Tags: ai
Training Text-to-Image Models 3.6× Faster
TL;DRLinum 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. 3.6× faster to train, at 4× the pixelsLinum v2 (ours, previous)* · 256×2562.0B latent-space DiT + VAE256 latent tokens* image-only checkpointJiT-DDT (ours, new) · 512×5122.5B active pixel-space DiT320 pixel tokens = 64 encoder + 256 decoderGPU-hours03.6× fewer0samples seen0M4.2× fewer0MFor more comparisons, see Appendix Research releaseModel codeModel weightsJiT-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 reconstructreplayinput xencoderμ = [?, ?]σ = [?, ?]μ, σz ∈ ℝᵏsample zdecoderoutput x̂‖x − x̂‖²+ β · KL(q‖N)loss‹ready›gradient (purple) reaches every weight* simplified: in practice the KL term is ≈ 0, previously we trained a σ-VAE with an L1 reconstruction loss plus LPIPS and GAN losses; see our VAE post. 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 spaceTrainingInferencereplayinput x❄encoder= E(x)z(1−t)·z+ t·εINTERPOLATEzₜDiTtrainableditv̂pred‖v̂ − v‖²v = ε − zlossε ~ N(0, I)gaussianεSAMPLE εt ~ LogitNormalsample t‹ready›VAE frozen (dashed) · DiT trainable (purple) · gradient stops at the DiT · t = 0 clean image, t = 1 pure Gaussian noise 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 12replay16×16 pixels, 3 channels (RGB) eachcut into 4×4-pixel patches (16 patches)each patch is tokenized independently: 16 pixels × 3 channels become one 48-dim tokena linear layer W ∈ ℝ12×48 projects 48 dims down to 12‹ready›Illustrative. In JiT at 512px we use 32×32 patches, so a 512×512 image becomes 256 tokens, each starting at 32·32·3 = 3,072 dims; the b
TL;DRLinum 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. 3.6× faster to train, at 4× the pixelsLinum v2 (ours, previous)* · 256×2562.0B latent-space DiT + VAE256 latent tokens* image-only checkpointJiT-DDT (ours, new) · 512×5122.5B active pixel-space DiT320 pixel tokens = 64 encoder + 256 decoderGPU-hours03.6× fewer0samples seen0M4.2× fewer0MFor more comparisons, see Appendix Research releaseModel codeModel weightsJiT-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 reconstructreplayinput xencoderμ = [?, ?]σ = [?, ?]μ, σz ∈ ℝᵏsample zdecoderoutput x̂‖x − x̂‖²+ β · KL(q‖N)loss‹ready›gradient (purple) reaches every weight* simplified: in practice the KL term is ≈ 0, previously we trained a σ-VAE with an L1 reconstruction loss plus LPIPS and GAN losses; see our VAE post. 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 spaceTrainingInferencereplayinput x❄encoder= E(x)z(1−t)·z+ t·εINTERPOLATEzₜDiTtrainableditv̂pred‖v̂ − v‖²v = ε − zlossε ~ N(0, I)gaussianεSAMPLE εt ~ LogitNormalsample t‹ready›VAE frozen (dashed) · DiT trainable (purple) · gradient stops at the DiT · t = 0 clean image, t = 1 pure Gaussian noise 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 12replay16×16 pixels, 3 channels (RGB) eachcut into 4×4-pixel patches (16 patches)each patch is tokenized independently: 16 pixels × 3 channels become one 48-dim tokena linear layer W ∈ ℝ12×48 projects 48 dims down to 12‹ready›Illustrative. In JiT at 512px we use 32×32 patches, so a 512×512 image becomes 256 tokens, each starting at 32·32·3 = 3,072 dims; the b