For $1,890, you can train a decent 1.2 billion parameter diffusion model from scratch.

With only 1890 US dollars and 37 million images, you can train a pretty good diffusion model.

Currently, visual generative models are good at creating realistic visual content, however the cost and effort of training these models from scratch is still high. For example, Stable Diffusion 2.1 took 200,000 A100 GPU hours. Even if researchers use the most advanced method, it still takes more than a month to train on an 8×H100 GPU.

In addition, training large models also poses challenges to data sets. These data are basically in units of hundreds of millions, which also brings challenges to training models.

High training costs and data set requirements create insurmountable obstacles for the development of large-scale diffusion models.

Now, researchers from Sony AI and other institutions have spent only $1,890 to train a good diffusion model, a sparse transformer with 1.16 billion parameters.

• Paper address: https://arxiv.org/pdf/2407.15811

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

• project (forthcoming) : https://github.com/SonyResearch/micro_diffusion

Specifically, in this work, the author develops a low-cost end-to-end pipeline for text-to-image diffusion model, making the training cost lower than SOTA The model is more than an order of magnitude faster, while not requiring access to billions of training images or proprietary datasets.

The author considered a latent diffusion model based on visual transformer for text to image generation. The main reason is that this method is simple to design and widely used. To reduce the computational cost, the authors exploit the strong dependence of the transformer computational cost on the input sequence size (i.e., the number of patches per image).

The main goal of this article is to reduce the number of effective patches used by the transformer to process each image during the training process. This can be easily achieved by randomly masking out some tokens in the input layer of the transformer.

However, existing masking methods are unable to extend the masking rate beyond 50% without significantly reducing performance, especially at high masking rates where a large portion of the input patch is not observed at all by the diffusion transformer.

In order to alleviate the significant performance degradation caused by masking, the author proposes a deferred masking strategy, in which all patches are preprocessed by a lightweight patch-mixer and then transferred to the diffusion transformer . Patch mixers contain a fraction of the number of parameters found in diffusion transformers.

Compared to naive masking methods, masking after patch mixing allows unmasked patches to retain semantic information about the entire image and enables reliable training of diffusion transformers at very high masking rates while competing with the best existing methods. There is no additional computational cost compared to advanced masking.

The authors also demonstrate that the delayed masking strategy achieves better performance than downsizing (i.e. reducing model size) under the same computational budget. Finally, the authors incorporate recent advances in Transformer architecture, such as layer-by-layer scaling, sparse Transformer using MoE, to improve the performance of large-scale training.

The low-cost training pipeline proposed by the author reduces experimental overhead. In addition to using real images, the authors also considered combining other synthetic images in the training dataset. The combined dataset contains only 37 million images, much less data than most existing large-scale models require.

On this combined dataset, the author trained a 1.16 billion parameter sparse transformer at a cost of $1890 and achieved 12.7 FID in zero-shot generation on the COCO dataset.

It is worth noting that the model trained in this article achieves competitive FID and high-quality generation, while costing only 1/118 of the stable diffusion model and 1/1 of the current state-of-the-art method (costing $28,400) 15.

Introduction to the method

In order to greatly reduce the computational cost, patch masking requires discarding most of the input patches before inputting to the backbone transformer, so that the transformer cannot obtain the information of the masked patch. High masking rates (such as 75% masking rate) can significantly reduce the overall performance of the transformer. Even with MaskDiT, only a weak improvement over naive masking can be observed, since this method also discards most of the image patches in the input layer itself.

Delayed masking, retaining the semantic information of all patches

Since high masking rate will remove most of the valuable learning signals in the image, the author can't help but ask, is it necessary to mask in the input layer? As long as the computational cost remains constant, this is just a design choice and not a fundamental limitation. In fact, the authors discovered a significantly better masking strategy that costs almost the same as the existing MaskDiT method. Since patches come from non-overlapping image regions in the diffusion Transformer, each patch embedding does not embed any information from other patches in the image. Therefore, the authors aim to preprocess patch embeddings before masking so that unmasked patches can embed information from the entire image. They call the preprocessing module patch-mixer.

Use patch-mixer to train diffusion transformer

The author believes that patch-mixer is any neural architecture that can fuse individual patch embeddings. In the transformer model, this goal can naturally be achieved through a combination of attention and feed-forward layers. Therefore, the authors use a lightweight transformer consisting of only a few layers as patch-mixer. After the input sequence tokens are processed by patch-mixer, they mask them (Figure 2e).

^{Figure 2: Compress patch sequence to reduce computational cost. Since the training cost of a diffusion transformer is proportional to the sequence size (i.e. the number of patches), it is best to reduce the sequence size without degrading performance. This can be achieved by: b) using a larger patch; c) randomly masking a portion of the patch; or d) using MaskDiT, which combines naive masking with additional autoencoding objectives. The authors found that all three methods lead to significant degradation in image generation performance, especially at high masking rates. To alleviate this problem, they proposed a straightforward delayed masking strategy that masks the patch after it has been processed by the patch-mixer. Their approach is similar to naive masking in all aspects except using patch-mixer. Compared to MaskDiT, their method does not require optimization of any surrogate objective and has almost the same computational cost.}

Assuming that the mask is a binary mask m, the author uses the following loss function to train the model:

where, M_ϕ is the patch-mixer model and F_θ is the backbone transformer. Note that compared to MaskDiT, the proposed method also simplifies the overall design and does not require additional loss functions or corresponding hyperparameter tuning between the two losses during training. During inference, this method does not mask any patches.

Unmasked fine-tuning

Since an extremely high masking rate will greatly reduce the ability of the diffusion model to learn the global structure of the image and introduce a train-test distribution shift on the sequence size, the author considers doing a small amount after masked pre-training of unmasked fine-tuning. Fine-tuning can also mitigate any generation artifacts caused by using patch masking. Therefore, in previous work, it is crucial to recover performance that drops sharply due to masking, especially when using classifier-free bootstrapping in sampling. However, the authors argue that this is not entirely necessary, as their method achieves comparable performance to baseline unmasked pre-training even with masked pre-training. The authors only use this approach in large-scale training to mitigate any unknown-unknown generation artifacts due to high patch masking.

Using MoE and layer-wise scaling to improve the backbone transformer architecture

The author also leverages innovations in transformer architecture design to improve model performance under computational constraints.

They use hybrid expert layers because they increase the parameters and expressiveness of the model without significantly increasing the training cost. They use a simplified MoE layer based on expert selection routing, where each expert decides which tokens are routed to it, as it does not require any additional auxiliary loss function to balance the load among experts. They also considered layer-wise scaling, which has recently been shown to outperform typical transformers in large language models. This method linearly increases the width of the transformer block, i.e. the hidden layer dimension of the attention and feedforward layers. Therefore, deeper layers in the network are assigned more parameters than earlier layers. The authors believe that since deeper layers in a visual model tend to learn more complex features, using higher parameters in deeper layers will lead to better performance. The authors describe the overall architecture of their proposed diffusion Transformer in Figure 3.

^{Figure 3: The overall architecture of the diffusion transformer proposed in this article. The authors added a lightweight patch-mixer to the backbone transformer model, which processes all patches in the input image before they are masked. Following the current work, the authors use an attention layer to process caption embeddings, which are then used for conditioning. They use sinusoidal embedding to represent the time step. Their model only denoises unmasked patches, so the diffusion loss (Equation 3 in the paper) is only calculated for these patches. They modified the backbone transformer to use layer-wise scaling on individual layers and mixed expert layers in alternating transformer blocks.}

Experiment

The experiment uses two variants of diffusion Transformer (DiT), DiT-Tiny/2 and DiT-Xl/2.

As shown in Figure 4, the delay masking method achieves better performance in multiple metrics. Furthermore, as the masking rate increases, the performance gap widens. For example, at 75% masking rate, naive masking reduces the FID score to 16.5 (lower is better), while our method can achieve 5.03, which is closer to the FID score of 3.79 without masking.

Table 1 shows that the layer-wise scaling method has a better fitting effect in the masking training of diffusion transformer.

Compare different masking strategies. The authors first compare our method with strategies using larger patches. Increasing the patch size from 2 to 4 equates to 75% patch masking. Compared to delayed masking, other methods perform poorly, achieving only 9.38, 6.31 and 26.70 FID, Clip-FID and Clip-score respectively. In comparison, latency masking achieves 7.09, 4.10, and 28.24 FID, Clip-FID, and Clip-score respectively.

The picture below shows the comparison of delayed masking vs. model shrinking to reduce training cost. Until the masking rate reaches 75%, the authors find that delayed masking outperforms network reduction in at least two of three metrics. However, at extremely high masking rates, delayed masking tends to achieve lower performance. This may be caused by the loss of information masked at these ratios being too high.

Table 5 provides details about the model training hyperparameters. The training process is divided into two stages.

Calculate the cost. Table 2 provides a breakdown of the computational costs for each training stage, including training FLOPs and economic costs. Phase 1 and 2 training consumed 56% and 44% of the total computational cost, respectively. The total clock training time of the model on an 8×H100 GPU cluster is 2.6 days, which is equivalent to 6.6 days on an 8×A100 GPU cluster.

For more results, please refer to the original paper.

The above is the detailed content of For $1,890, you can train a decent 1.2 billion parameter diffusion model from scratch.. For more information, please follow other related articles on the PHP Chinese website!