GoKawiilRoad to Linum v3 · Issue 02 previously: data filtering
TL;DR 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.
3.6× faster to train, at 4× the pixels Linum v2 (ours, previous)* · 256 × 256 2.0B latent-space DiT + VAE 256 latent tokens * image-only checkpoint JiT-DDT (ours, new) · 512 × 512 2.5B active pixel-space DiT 320 pixel tokens = 64 encoder + 256 decoder GPU-hours 0 3.6 × fewer 0 samples seen 0M 4.2 × fewer 0M
Research release 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 replay input x Encoder encoder μ = [?, ?] σ = [?, ?] μ, σ z ∈ ℝᵏ sample z Decoder decoder output x̂ ‖x − x̂‖² + β · KL(q‖N) loss ‹ ready ›
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;
... continue reading