Training Text-to-Image Models 3.6× Faster
7 hours ago
- Linum v2's attention context window was a bottleneck, with a 720p, 5-second clip using 110K tokens, making pruning the context window key for acceleration.
- Most generative image/video models are Latent Diffusion Models (LDMs) with separate VAE and DiT components, but pixel-space models like JiT offer a unified alternative.
- JiT struggles with fine-grained details despite achieving 32x32 token reduction by integrating compression into the DiT.
- JiT-DDT is a novel encoder-decoder architecture that recovers details and trains 3.6x more efficiently than Linum v2, generating 4x more pixels.
- The VAE compression wall caps at 16x16 token reduction; aggressive compression is unlocked by pushing compression into the DiT.
- JiT-DDT uses a unified model with a single-stream attention over image and text tokens, and a wider, shallower transformer.
- The DDT splits the model into an encoder for low-frequency structure and a decoder for high-frequency detail, with the encoder predicting a downsampled image.
- Compared to Linum v2, JiT-DDT generates more faithful art styles and realistic lighting but still struggles with realistic human faces.
- Future work includes training longer, rebalancing datasets, DPO post-training, or scaling up the model to address remaining issues.
- The findings are shared as a research artifact to encourage more efficient training methods in the community.