Abstract:As scaling laws in generative AI push performance, they also simultaneously concentrate the development of these models among actors with large computational resources. With a focus on text-to-image (T2I) generative models, we aim to address this bottleneck by demonstrating very low-cost training of large-scale T2I diffusion transformer models. As the computational cost of transformers increases with the number of patches in each image, we propose to randomly mask up to 75% of the image patches during training. We propose a deferred masking strategy that preprocesses all patches using a patch-mixer before masking, thus significantly reducing the performance degradation with masking, making it superior to model downscaling in reducing computational cost. We also incorporate the latest improvements in transformer architecture, such as the use of mixture-of-experts layers, to improve performance and further identify the critical benefit of using synthetic images in micro-budget training. Finally, using only 37M publicly available real and synthetic images, we train a 1.16 billion parameter sparse transformer with only \$1,890 economical cost and achieve a 12.7 FID in zero-shot generation on the COCO dataset. Notably, our model achieves competitive FID and high-quality generations while incurring 118$\times$ lower cost than stable diffusion models and 14$\times$ lower cost than the current state-of-the-art approach that costs \$28,400. We aim to release our end-to-end training pipeline to further democratize the training of large-scale diffusion models on micro-budgets.

What problem does this paper attempt to address?

The paper aims to address the high cost issue in training generative artificial intelligence models, particularly text-to-image (T2I) models, which has led to the development of such models being concentrated in research institutions or companies with substantial computational resources. Specifically, the research team aims to demonstrate how to train large-scale T2I diffusion transformer models at extremely low costs. To achieve this goal, the authors propose a method of randomly masking up to 75% of image patches during training to reduce computational costs. Since the computational cost of transformer models increases proportionally with the number of patches in each image, this method can significantly reduce the required computational resources. To mitigate the performance degradation caused by the high masking ratio, the authors propose a delayed masking strategy, where a lightweight patch-mixer preprocesses all patches before masking. This approach allows the unmasked patches to retain semantic information about the entire image, thereby improving the reliability of model training without incurring additional computational costs. Additionally, the paper explores recent improvements in transformer architecture, such as using mixture-of-experts layers to further enhance performance, and highlights the key benefits of using synthetic images in micro-budget training. Through these techniques and methods, the authors successfully trained a sparse transformer model with 1.16 billion parameters on a dataset containing 37 million publicly available real and synthetic images, at a computational cost of only $1,890. The model achieved a Fréchet inception distance (FID) score of 12.7 on the COCO dataset, demonstrating its competitiveness in zero-shot generation tasks. Notably, the cost of this method is only 1/118 of that of the Stable Diffusion model, reducing the cost by 14 times compared to the current state-of-the-art method (which requires $28,400). In summary, the paper effectively addresses the high cost of training T2I diffusion models by proposing an innovative delayed patch masking strategy combined with the latest transformer architecture improvements, significantly lowering the technical barrier to developing such large-scale models.

Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget

Diffscaler: Enhancing the Generative Prowess of Diffusion Transformers

EDT: An Efficient Diffusion Transformer Framework Inspired by Human-like Sketching

PixArt-$α$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis

Fast Training of Diffusion Models with Masked Transformers

BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion

TerDiT: Ternary Diffusion Models with Transformers

Dynamic Diffusion Transformer

On the Scalability of Diffusion-based Text-to-Image Generation

Towards Precise Scaling Laws for Video Diffusion Transformers

Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models

Efficient Scaling of Diffusion Transformers for Text-to-Image Generation

Scaling Laws For Diffusion Transformers

Adaptive Training Meets Progressive Scaling: Elevating Efficiency in Diffusion Models

SnapFusion: Text-to-Image Diffusion Model on Mobile Devices within Two Seconds

Make a Cheap Scaling: A Self-Cascade Diffusion Model for Higher-Resolution Adaptation

Masked Diffusion Models Are Fast Distribution Learners

Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models

DiffiT: Diffusion Vision Transformers for Image Generation

Exploring Vision Transformers as Diffusion Learners

MaxFusion: Plug&Play Multi-Modal Generation in Text-to-Image Diffusion Models