TerDiT: Ternary Diffusion Models with Transformers

## Abstract

Abstract Recent developments in large-scale pre-trained text-to-image diffusion models have significantly improved the generation of high-fidelity images, particularly with the emergence of diffusion models based on transformer architecture (DiTs). Among these diffusion models, diffusion transformers have demonstrated superior image generation capabilities, boosting lower FID scores and higher scalability. However, deploying large-scale DiT models can be expensive due to their extensive parameter numbers. Although existing research has explored efficient deployment techniques for diffusion models such as model quantization, there is still little work concerning DiT-based models. To tackle this research gap, in this paper, we propose TerDiT , a quantization-aware training (QAT) and efficient deployment scheme for ternary diffusion models with transformers. We focus on the ternarization of DiT networks and scale model sizes from 600M to 4.2B. Our work contributes to the exploration of efficient deployment strategies for large-scale DiT models, demonstrating the feasibility of training extremely low-bit diffusion transformer models from scratch while maintaining competitive image generation capacities compared to full-precision models. Code will be available at https://github.com/Lucky-Lance/TerDiT .

## 1 Introduction

The advancements in large-scale pre-trained text-to-image diffusion models [, , , , ] have led to the successful generation of images characterized by both complexity and high fidelity to the input conditions. Notably, the emergence of diffusion models based on transformer architecture (DiTs) [] represents a significant stride in this research domain. Compared with other diffusion models, diffusion transformers have demonstrated the capability to achieve lower FID scores with higher computation Gflops []. Recent research highlights the remarkable image generation capabilities of diffusion transformer architecture, as demonstrated in methods like Stable Diffusion 3 [], and its impressive performance in video generation, showcased by work such as Sora(^1^11 [https://openai.com/index/sora](https://openai.com/index/sora)).

Given the impressive performance of diffusion transformer models, researchers are now increasingly delving into understanding the scaling law of these vision models [], which resembles large language models (LLMs). For instance, the Stable Diffusion 3 provides trained DiT models across a spectrum of parameter sizes, from 800 million to 8 billion. Additionally, there is speculation among researchers that Sora might boast around 3 billion parameters. Given the enormous parameter numbers, deploying these DiT models can often be costly, especially on certain end devices.

To deal with the deployment dilemma, there have already been some works on the efficient deployment of diffusion models recently [, , , ], most of which focus on model quantization. However, as far as we are concerned, there are still two main shortcomings in current research. Firstly, while much attention has been given to quantizing U-Net-based diffusion models, the exploration of quantization methods for transformer-based diffusion models remains quite limited. Secondly, most prevailing approaches in the current literature heavily rely on post-training quantization (PTQ) techniques for model quantization [, , , , ], which leads to unacceptable performance degradation, particularly with extremely low bit width (e.g., 2-bit and 1-bit). However, the extremely low-bit quantization of neural networks is important as it can significantly reduce the computation resources required for deployment [], especially for models with huge parameter sizes. During our research, we find that there is still no work considering the extremely low-bit quantization of DiT models.

To tackle these shortcomings, we propose to leverage the quantization-aware training (QAT) technique for the extremely low-bit quantization of large-scale DiT models. Low-bit QAT methods for large-scale models have been discussed in the LLM domain. Recent works observe that training large language models with extremely low-bit parameters (e.g., binary and ternary) from scratch can also lead to competitive performance [, ] compared with their full-precision counterparts. This indicates that there still exists significant precision redundancy in large-scale models, and implies the feasibility of the QAT scheme for large-scale DiT models.

In this paper, we focus primarily on ternary weight networks [] and provide TerDiT, the first quantization scheme for DiTs to our best knowledge. Our method achieves quantization-aware training (weight-only) and efficient deployment for ternary diffusion transformer models. Different from the naive quantization of linear layers in LLMs and CNNs [, ], we find that the direct weight ternarization of the adaLN module [] in DiT blocks [, ] leads to large dimension-wise scale and shift values in the normalization layer compared with full-precision models (due to weight quantization, gradient approximation), which result in slower convergence speed and poor model performance. Consequently, we propose a variant of adaLN by applying an RMS Norm [] after the ternary linear layers of the adaLN module to effectively mitigate this training issue.

With this modification, we scale the parameters of the ternary DiT model from 600M (size of DiT-XL/2 []) to 4.2B (size of Large-DiT-4.2B []), finding that models with a larger number of parameters can converge to better results. We further employ existing 2-bit CUDA kernels to deploy the trained ternary DiT models, resulting in over a tenfold reduction in the model checkpoint size and about sixfold reduction in the inference memory consumption, while achieving competitive generation quality. The contributions are summarized as follows:

- 1.
Inspired by the quantization-aware training scheme for low-bit LLMs, we study the QAT methods for ternary DiT models and introduce DiT-specific improvements for better training, which have not been explored in DiT literature.
- 2.
We scale ternary DiT models from 600M to 4.2B parameters, and further deploy the trained ternary DiT on GPUs based on existing 2-bit CUDA kernels, enabling the inference of a 4.2B DiT model with less than 3GB GPU memory.
- 3.
Competitive evaluation results compared with full-precision models on the ImageNet [] benchmark (image generation) showcase the effectiveness of our proposed TerDiT scheme.

To our best knowledge, our study is the first attempt to explore the quantization of DiT models. We focus on quantization-aware training and efficient deployment for large ternary DiT models, offering valuable insights for future research on deploying DiT models with extremely low-bit precision.

## 2 Related Works

Diffusion Models. Diffusion models have gained significant attention in recent years due to their ability to generate high-quality images and their potential for various applications. The concept of diffusion models was first introduced by [], proposing a generative model that learns to reverse a diffusion process. This work laid the foundation for subsequent research in the field. [] further extended the idea by introducing denoising diffusion probabilistic models (DDPMs), which have become a popular choice for image generation tasks. DDPMs have been applied to a wide range of domains, including unconditional image generation [], image inpainting [], and image super-resolution []. Additionally, diffusion models have been used for text-to-image synthesis, as demonstrated by the DALL-E model [] and the Imagen model []. These models showcase the ability of diffusion models to generate highly realistic and diverse images from textual descriptions. Furthermore, diffusion models have been extended to other modalities, such as audio synthesis [] and video generation [], demonstrating their versatility and potential for multimodal applications.

Quantization of Diffusion Models. The quantization of diffusion models has been studied in recent years to improve the efficiency of diffusion models. Post-training quantization (PTQ) methods, such as those presented in [, , , , ], offer advantages in terms of quantization time and data usage. However, these methods often result in suboptimal performance when applied to low-bit settings. To address this issue, [] proposes combining quantization-aware low-rank adapters (QALoRA) with PTQ methods, leading to improved evaluation results. As an alternative to PTQ, quantization-aware training (QAT) methods have been introduced specifically for low-bit diffusion model quantization [, , , ]. Despite their effectiveness, these QAT methods are currently only limited to small-sized U-Net-based diffusion models, revealing a research gap in applying QAT to large-scale DiT models. Further exploration of QAT techniques for large DiT models with extremely low bit width could potentially unlock even greater efficiency gains and enable the effective deployment of diffusion models in resource-constrained environments.

Ternary Weight Networks. Ternary weight networks [] have emerged as a memory-efficient and computation-efficient network structure, offering the potential for significant reductions in inference memory usage. When supported by specialized hardware, ternary weight networks can also deliver substantial computational acceleration. Among quantization methods, ternary weight networks have garnered notable attention, with two primary approaches being explored: weight-only quantization and weight-activation quantization. In weight-only quantization, as discussed in [], solely the weights are quantized to ternary values. On the other hand, weight-activation quantization, as presented in [, ], involves quantizing both the weights and activations to ternary values. Recent research has demonstrated the applicability of ternary weight networks to the training of large language models [], achieving results comparable to their full-precision counterparts. Building upon these advancements, our work introduces, for the first time, quantization-aware training and efficient deployment schemes specifically designed for ternary DiT models. By leveraging the benefits of ternary quantization in the context of DiT models, we aim to push the boundaries of efficiency and enable the deployment of powerful diffusion models in resource-constrained environments, opening up new possibilities for practical applications.

## 3 TerDiT

In this section, we introduce TerDiT, a framework designed to conduct weight-only quantization-aware training and efficient deployment of large-scale ternary DiT models. We first give a brief review of diffusion transformer (DiT) models in Sec. 3.1. Then building upon the previous open-sourced Large-DiT [], we illustrate the quantization function and quantization-aware training scheme in Sec. 3.2, conduct QAT-specific model structure improvement for better network training in Sec. 3.3, and introduce ternary deployment scheme in Sec. 3.4.

Figure: Figure 2: Model structure comparison between (A) Ternary DiT block, (B) Large-DiT block, and the (C) original ViT block. The Large-DiT (DiT) block adds an adaLN module to the original ViT block for condition injection. Ternary DiT block further adds an RMS Norm in the adaLN module for better ternarization-aware training.

### 3.1 Diffusion Transformer Models

Diffusion Transformer. Diffusion transformer [] (DiT) is an architecture that replaces the commonly used U-Net backbone in the diffusion models with a transformer that operates on latent patches. Similar to the Vision Transformer (ViT) architecture shown in Fig. 2 (C), DiT first pachifies the spatial inputs into a sequence of tokens, then the denoising process is carried out through a series of transformer blocks (Fig. 2 (B)). To deal with additional conditional information (e.g., noise timesteps $t$, class labels $l$, natural language inputs), DiT leverages adaptive normalization modules [] (adaLN-Zero) to insert these extra conditional
inputs to the transformer blocks. After the final transformer block, a standard linear decoder is applied to predict the final noise and covariance. The DiT models can be trained in the same way as U-Net-based diffusion models.

AdaLN Module in DiT. The main difference between DiT and traditional ViT is the need to inject conditional information for image generation. DiT employs a zero-initialized adaptive layer normalization (adaLN-Zero) module in each transformer block, as shown in the red part of Fig. 2 (B), which calculates the dimension-wise scale and shift values from the input condition $c$:

$$ $\displaystyle\text{adaLN}(c)=\text{MLP}(\text{SiLU}(c)).$ (1) $$

AdaLN is an important component in the DiT model [] and has been proven more effective than cross-attention and in-context conditioning methods. Within the DiT architecture, the adaLN module integrates an MLP layer with a substantial number of parameters, constituting approximately 10% to 20% of the model’s total parameters. Throughout the training of TerDiT, we observe that the direct weight-ternarization of this module yields undesirable training results (analyzed in Sec. 3.3).

### 3.2 Model Quantization

As illustrated in Sec. 1, there is an increasing popularity in understanding the scaling law of DiT models, which has been proven crucial for developing and optimizing LLMs. In recent explorations, Large-DiT [] successfully scales up the model parameters from 600M to 7B by incorporating the methodologies of LLaMA [, ] and DiT. The results demonstrate that parameter scaling can potentially enhance model performance and improve convergence speed for the label-conditioned ImageNet generation task. Motivated by this, we propose to further investigate the ternarization of DiT models, which can alleviate the challenges associated with deploying large-scale DiT models. In this subsection, we introduce the quantization function and quantization-aware training scheme.

Figure: Figure 3: Activation value analysis. We compare activation values passing through a ternary weight linear layer with and without RMS Norm, using a full-precision linear layer as a reference. The ternary linear layer without RMS Norm results in extremely large activation values, introducing instability in neural network training. However, when the normalization layer is applied, the activation values are scaled to a reasonable range, similar to those observed in the full-precision layer.

Quantization Function. To construct a ternary weight DiT network, we replace all the linear layers in self-attention, feedforward, and MLP of the original Large-DiT blocks with ternary linear layers, obtaining a set of ternary DiT blocks (Fig. 2 (A)). For ternary linear layers, we adopt an absmean quantization function similar to BitNet b1.58 []. First, the weight matrix is normalized by dividing each element by the average absolute value of all the elements in the matrix. After normalization, each value in the weight matrix is rounded to the nearest integer and clamped into the set $\{-1,0,+1\}$.

Referring to current popular quantization methods for LLMs [, ], we also multiply a learnable scaling parameter $\alpha$ to each ternary linear matrix after quantization, leading to the final value set as $\{-\alpha,0,+\alpha\}$. The quantization function is formulated as:

$$ $\displaystyle\widetilde{W}=\alpha\cdot\textrm{RoundClip}\left(\frac{W}{\gamma+ \epsilon},-1,1\right),$ (2) $$

where $\epsilon$ is set to a small value (e.g., $10^{-6}$) to avoid division by 0, and

$$ $\displaystyle\text{RoundClip}(x,a,b)$ $\displaystyle=\text{Clamp}(\text{round}(x),a,b)\text{ and }\gamma=\frac{1}{mn} \sum_{ij}|W_{ij}|.$ (3) $$

TerDiT is a weight-only quantization scheme and we do not quantize the activations.

Quantization-aware Training Scheme. Based on the above-designed quantization function, we train a DiT model from scratch(^2^22It is observed in [] that for ternary LLMs, the conversion or post-training quantization from trained LLMs does not help. So we also train ternary DiT models from scratch.) utilizing the straight-through estimator (STE) [], allowing gradient propagation through the undifferentiable network components. We preserve the full-precision parameters of the network throughout the training process. For each training step, ternary weights are calculated from full-precision parameters by the ternary quantization function in the forward pass, and the gradients of ternary weights are directly applied to the full-precision parameters for parameter update in the backward pass.

However, we find the convergence speed is very slow. Even after many training iterations, the loss cannot be decreased to a reasonable range. We find that this issue may arise from the trait that ternary linear layers usually cause large activation values, and propose to tackle the problem with QAT-specific model structure improvement in the following subsection.

### 3.3 QAT-specific Model Structure Improvement

Activation Analysis for Ternary Linear Layer. In a ternary linear layer, all the parameters take one value from the set {-$\alpha$, 0, +$\alpha$}. The values passing through this layer would become large activation values, which might hamper the stable training of neural networks. We conduct a pilot study to qualitatively demonstrate the impact of ternary linear weights on activation values.

We randomly initialize a ternary linear layer with the input feature dimension set to 1024 and the output feature dimension to 9216 (corresponding to the linear layer of the adaLN module in Large-DiT). The weight parameters pass through the quantization function and receive a 512$\times 1024$ sized matrix input (filled with 1). The box plot of activation distribution is shown in the center part of Fig. 3. We also calculate the activation distribution after passing the matrix through a full-precision linear layer generated with the same random seed, shown on the right part of Fig. 3. As can be seen, the ternary linear layer leads to very large activation values compared with full-precision linear layers.

The large-activation problem brought about by the ternary linear weights can be alleviated by applying a layer norm to the output of the ternary linear layer. We add an RMS Norm (similar to LLaMA) after the ternary linear layer and obtain the activation distribution (as shown in the left part of Fig. 3). In this case, the activation values are scaled to a reasonable range after passing the normalization layer and lead to more stable training behavior. The observation also aligns with [], where a layer normalization function is applied before the activation quantization for each quantized linear layer.

RMS Normalized AdaLN Module. We analyze the DiT model for QAT-specific model structure improvement based on the above insights. In a standard ViT transformer block, the layer norm is applied for every self-attention layer and feedforward layer. This is also the case with the self-attention layers and feedforward layers in the DiT block, which can help to properly scale the range of activations. However, the DiT block differs from traditional transformer blocks due to the presence of the AdaLN module, as introduced in Sec. 3.1. Notably, there is no layer normalization applied to this module. In the context of full-precision training, the absence of layer normalization does not have a significant impact. However, for ternary DiT networks, its absence can result in large dimension-wise scale and shift values in the adaLN (normalization) module, posing bad influences on model training. To mitigate this issue, we introduce an RMS Norm after the MLP layer of the adaLN module in each ternary DiT block:

$$ $\displaystyle\text{adaLN_norm}(c)=\text{RMS}(\text{MLP}(\text{SiLU}(c))),$ (4) $$

and the final model structure of TerDiT is illustrated in Fig. 2 (A). This minor modification can result in faster convergence speed and lower training loss, leading to better quantitative and qualitative evaluation results. To better showcase the effect, the actual activation distribution with/without the RMS Norm after model training is analyzed in Sec. A.1.

### 3.4 Deployment Scheme

After training the DiT model, we find that there are currently no effective open-source deployment solutions for ternary networks. In this case, we deploy the trained networks with a 2-bit implementation. To be specific, we pack the ternary linear weights to int8 values (4 ternary numbers into one int8 number) with the pack_2bit_u8() function provided by []. During the inference process of the DiT model, we call the complementary unpack_2bit_u8() function on the fly to recover packed 2-bit numbers to floating-point values, then perform subsequent calculations. The addition of the unpacking operation will slow down the inference process, but we believe that with further research into model ternarization, more hardware support for inference speedup will become available.

## 4 Experiments

In this section, a set of experiments are carried out to evaluate our proposed TerDiT. We show our main evaluation results in Sec. 4.1, carry out deployment efficiency comparison in Sec. 4.2, and illustrate the effectiveness of the RMS Normalized adaLN module in Sec. 4.3. Our DiT implementation is based on the open-sourced code of Large-DiT-ImageNet(^3^33[https://github.com/Alpha-VLLM/LLaMA2-Accessory/tree/main/Large-DiT-ImageNet](https://github.com/Alpha-VLLM/LLaMA2-Accessory/tree/main/Large-DiT-ImageNet)). We conduct experiments on ternary DiT models with 600M (size of DiT-XL/2) and 4.2B(^4^44The provided 3B model in the Large-DiT-ImageNet repository actually has 4.2B parameters.) (size of Large-DiT-4.2B) parameters respectively.

### 4.1 Main Evaluation Results

We provide quantitative and qualitative evaluation results of TerDiT in this subsection. To the best of our knowledge, there is no published work done on the quantization of diffusion transformer models at this point, so we mainly focus on the comparison with representative full-precision diffusion models in this subsection.

Remark on TerDiT Baselines. There is still no work studying the quantization of DiT models to our best knowledge. Other than comparison with full-precision models in this subsection, we have also established some baselines for comparison in other subsections. For the QAT baseline, we directly train the ternary DiT model without RMS Norm in the adaLN module in Sec. 4.3. To compare with existing PTQ [] methods, we perform 4-bit weight quantization for pre-trained models on the same set of parameters as TerDiT, and find they just fail to generate viewable images (detailed in Sec. A.2).

Experiment Setup. Following the evaluation setting of the original DiT paper [], we train 600M and 4.2B ternary DiT models on the ImageNet dataset. Due to computation resource limitations, we train and evaluate the model under 256$\times$256 resolution, but we think that the evaluation results are already quite representative. We compare TerDiT with a series of full-precision diffusion models, report FID [], sFID [], Inception Score, Precision, and Recall (50k generated images) following []. We also provide the total number of images (million) during the training
stage as in [] to offer further insights into the convergence speed of different generative models.

**Table 1: Comparison between TerDiT and a series of full-precision diffusion models on the ImageNet $256\times 256$ label-conditional generation task. For generation with classifier-free guidance, we use cfg=1.5. As can be seen, TerDiT achieves comparable results with full-precision models.**
| ImageNet 256×256 Benchmark |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Models | Images (M) | FID $\downarrow$ | sFID $\downarrow$ | Inception Score $\uparrow$ | Precision $\uparrow$ | Recall $\uparrow$ |
| BigGAN-deep [] | - | 6.95 | 7.36 | 171.40 | 0.87 | 0.28 |
| StyleGAN-XL [] | - | 2.30 | 4.02 | 265.12 | 0.78 | 0.53 |
| ADM [] | 507 | 10.94 | 6.02 | 100.98 | 0.69 | 0.63 |
| ADM-U [] | 507 | 7.49 | 5.13 | 127.49 | 0.72 | 0.63 |
| LDM-8 [] | 307 | 15.51 | - | 79.03 | 0.65 | 0.63 |
| LDM-4 [] | 213 | 10.56 | - | 103.49 | 0.71 | 0.62 |
| DiT-XL/2 (675M) [] | 1792 | 9.62 | 6.85 | 121.50 | 0.67 | 0.67 |
| \hdashline<br>TerDiT-4.2B | 604 | 9.66 | 6.75 | 117.54 | 0.66 | 0.68 |
| Classifier-free Guidance |  |  |  |  |  |  |
| ADM-G [] | 507 | 4.59 | 5.25 | 186.70 | 0.82 | 0.52 |
| ADM-G, ADM-U [] | 507 | 3.94 | 6.14 | 215.84 | 0.83 | 0.53 |
| LDM-8-G [] | 307 | 7.76 | - | 209.52 | 0.84 | 0.35 |
| LDM-4-G [] | 213 | 3.60 | - | 247.67 | 0.87 | 0.48 |
| DiT-XL/2-G (675M) [] | 1792 | 2.27 | 4.60 | 278.24 | 0.83 | 0.57 |
| Large-DiT-4.2B-G | 435 | 2.10 | 4.52 | 304.36 | 0.82 | 0.60 |
| \hdashline<br>TerDiT-600M-G | 448 | 4.34 | 4.99 | 183.49 | 0.81 | 0.54 |
| TerDiT-4.2B-G | 604 | 2.42 | 4.62 | 263.91 | 0.82 | 0.59 |

### 4.2 Deployment Efficiency Comparison

The improvement in deployment efficiency is the motivation of our proposed TerDiT scheme. In this subsection, we provide a comparison between TerDiT-600M/4.2B, DiT-XL/2, and Large-DiT-4.2B to discuss the actual deployment efficiency TerDiT can bring about. Tab. 4.1 shows the checkpoint sizes of the four DiT models. We also record the memory usage and inference time of the total diffusion sample loop (step=250) on a single A100-80G GPU.

From the table, we can see that TerDiT greatly reduces checkpoint size and memory usage. The checkpoint size and memory usage of the 4.2B ternary DiT model are significantly smaller than those of Large-DiT-4.2B, even smaller than DiT-XL/2. This brings significant advantages to deploying the model on end devices (e.g., mobile phones). We observe a slower inference speed due to the required unpack operations, but we believe that the computational advantage of ternary weight networks will be showcased through future hardware support.

### 4.3 Discussion on the RMS Normalized AdaLN Module

The main modification of TerDiT to the structure of the DiT model is the addition of an RMS Norm after the MLP in the adaLN module. In this part, we compare with the baseline ternary model to demonstrate the influence of RMS Norm on both the training process and the training outcomes.

Experiment Setup. We train ternary DiT models with 600M and 4.2B parameters on the ImageNet [] dataset in 256$\times$256 resolution. For each parameter size, we train two models, one with RMS Norm in the adaLN module and one without (our baseline). We record the loss curves during training and measure the FID-50k score (cfg=1.5) every 100k training steps.

Training Details. For a fair comparison, we train all the ternary DiT models on 8 A100-80G GPUs with batchsize set to 256. The learning rate is set to 5e-4 throughout the training process.

Results Analysis. The training loss(^6^66The actual training process of diffusion models is not as ‘smooth’ as one might assume. To better illustrate the training dynamics, we employ exponential smoothing (with a smoothing factor of 0.995) during visualization.) and evaluation scores are shown in Fig. 5 and Fig. 6 respectively. As can be seen, training with the RMS Normalized adaLN module will lead to faster convergence speed and lower FID scores. Another observation is that models with more parameters tend to achieve faster and better training compared to models with fewer parameters. This also to some extent reflects the scaling law of the ternary DiT model. Qualitative comparison results are shown in Sec. A.4.

Figure: Figure 5: Training loss comparison of ternary DiT models with/without RMS Norm in the adaLN module. We show the loss curves training both 600M (left) and 4.2B (right) DiT models. Adding the RMS Norm will lead to faster convergence speed and lower training loss.

Figure: Figure 6: FID-50k score comparison on class-conditional ImageNet 256$\times$256 generation task (cfg=1.5) with/without RMS Norm for both 600M and 4.2B ternary DiT models (100k steps to 400k steps). Training with RMS Norm will lead to lower FID scores.

## 5 Discussions and Future Works

In this paper, based on the successful low-bit training methods for large language models, we propose quantization-aware training (QAT) and efficient deployment methods for large-scale ternary DiT models. Competitive evaluation results on the ImageNet dataset (256$\times$256) demonstrate the feasibility of training a large ternary DiT from scratch while achieving results comparable to those of full-precision models. To the best of our knowledge, this work is also the first study concerning the quantization of DiT models.

While we believe this work provides valuable insights into the low-bit quantization of DiT models, it still has some limitations. Firstly, training ternary DiT is less stable and more time-consuming than full-precision networks. In our paper, although we discuss how to make the training more stable by adding norms, it still remains more time-consuming than training full-precision networks (Large-DiT-4.2B), which will lead to an increase in carbon dioxide emissions during model training in a broader context. Secondly, limited by computational resource constraints, we do not conduct ImageNet 512$\times$512 experiments, nor do we conduct experiments on the text-to-image generation task. However, we believe that the evaluation results on the ImageNet 256$\times$256 benchmark are already quite representative. The remaining tasks will be reserved for our future work. We hope that our work can reduce the deployment requirements of image generation models, and can inspire the community to join us in the future, fostering a broader advancement of this field of study.

## Appendix A Appendix / supplemental material

### A.1 Activation Distribution after Training

Here we conduct an analysis of the activation distribution during inference after model training, as a supplementary for the experiment and explanations in Sec. 3.3.

Following Sec. 4.3, we train the TerDiT-4.2B model with/without RMS Norm in the adaLN module for 50k steps. We then analyze the activation distribution of the ‘scale_mlp’ output in the adaLN module during inference, specifically at the first sampling step within the second ternary DiT block. For comparison, we also calculate the activation distribution of the original full-precision Large-DiT-4.2B model at the same layer. As shown in Fig. 7, training with RMS Norm can help limit the range of the activation values.

Figure: Figure 7: Activation value analysis. We train the TerDiT-4.2B model with/without RMS Norm in the adaLN module for 50k steps and show the activation distribution of the ‘scale_mlp’ output (one output of the adaLN module) in the second ternary DiT block during inference (at the first sampling step). The activation distribution of the original full-precision model is also provided.
Refer to caption: x6.png

### A.2 Comparison with Existing PTQ Methods

For a simple comparison with PTQ methods, we adopt the fake quant function for model weight quantization(^7^77[https://github.com/mit-han-lab/smoothquant/blob/main/smoothquant/fake_quant.py#L17](https://github.com/mit-han-lab/smoothquant/blob/main/smoothquant/fake_quant.py#L17)) provided in SmoothQuant [] and set n_bits=4. We do not carry out activation quantization. Same as TerDiT, we quantize all the linear layer weights in self-attention, feedforward, and MLP of the original DiT blocks. We then sample images with the 4-bit quantized model. In this case, both original DiT-XL/2 and Large-DiT-4.2B fail to generate images. As shown in Fig. 8. However, TerDiT can generate high-quality images with ternary precision. This highlights the significance of QAT for extremely low-bit DiT quantization.

Figure: Figure 8: DiT-XL/2 and Large-DiT-4.2B fail to generate images at 4-bit weight quantization.

### A.3 Effectiveness of Learning Rate Reduction

In Sec. 4.1, we adopt a learning rate reduction after training for certain steps for more fine-grained parameter updates. Here we provide an ablation study on this learning rate reduction.

We choose the TerDiT-600M model for convenience. Following the setting of Sec. 4.1, we train a TerDiT-600M model with RMS Normalized adaLN module and 5e-4 learning rate for 1550k steps. We then continue training this model with 1e-4/5e-4 for another 200k steps.

Quantitative Results. We evaluate FID, sFID, Inception Score, Precision, and Recall of these two models. As can be seen in Tab. 3, the reduction in learning rate will lead to better evaluation results.

**Table 3: Effectiveness of learning rate reduction. Training with a reduced learning rate after 1550k steps for the TerDiT-600M model will lead to better evaluation results.**
| ImageNet 256×256 Benchmark, Classifier-free Guidance |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Models | LR | FID $\downarrow$ | sFID $\downarrow$ | Inception Score $\uparrow$ | Precision $\uparrow$ | Recall $\uparrow$ |
| TerDiT-600M-G (cfg=1.50) | 1e-4 | 4.34 | 4.99 | 183.49 | 0.81 | 0.54 |
| TerDiT-600M-G (cfg=1.50) | 5e-4 | 6.38 | 5.00 | 147.79 | 0.76 | 0.54 |

Qualitative Results. We also provide qualitative results comparison on TerDiT-600M with/without learning rate reduction in Fig. 9.
For reference, the image generated by TerDiT-4.2B (fully trained) on the provided class label is shown in Fig. 10.
The quality comparison highlights the importance of learning rate reduction in the later stages of training.

Figure: Figure 9: TerDiT-600M, class label [207, 360, 387, 974, 88, 979, 417, 279], cfg=4.

Figure: Figure 10: TerDiT-4.2B, class label [207, 360, 387, 974, 88, 979, 417, 279], cfg=4.
Refer to caption: extracted/2405.14854v1/figures/sample_ema.png

### A.4 Qualitative Results Analysis for RMS Normalized AdaLN Module

In Sec. 4.3 we provide quantitative comparison results for models trained with/without RMS Norm in the adaLN module. Here we provide more qualitative results with models trained for 400k steps, together with the fully trained models.

Figure: Figure 11: Qualitative results comparison for 600M and 4.2B ternary DiT models with/without RMS Norm in the adaLN module. We choose class labels [972, 390, 417, 812, 118, 129, 447, 309] with cfg=4. Training the 4.2B model with the RMS Normalized adaLN module will lead to better qualitative generation results.

We compare TerDiT-600M and TerDiT-4.2B models trained with/without RMS Norm in the adaLN module. We sample images from the models trained for 400k steps in Sec. 4.3 and also show the results of the fully trained models. The comparisons are demonstrated in Fig. 11.

We can come to two conclusions:

- 1.
For both the 600M and 4.2B TerDiT model, training with the RMS Normalized adaLN module will lead to better qualitative results.
- 2.
Models with more parameters show more learning ability and can achieve better training results compared to models with fewer parameters.

These observations are consistent with the quantitative results provided in Sec. 4.3.

### A.5 More Qualitative Results

Figure: Figure 12: TerDiT-4.2B, class label [257, 300, 975, 973, 478, 388, 427, 809], cfg=4.
Refer to caption: extracted/2405.14854v1/figures/sample_ema_ternary_demo1.png

Figure: Figure 13: TerDiT-600M, class label [257, 300, 975, 973, 478, 388, 427, 809], cfg=4.
Refer to caption: extracted/2405.14854v1/figures/sample_ema_demo1_600M.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. In addition there are some small grammatical errors like
  a. pachifies in Ln 192
soundness: 3
strengths: 1. The paper is the first to utilize Ternary blocks in Diffusion transformers for quantization aware training.
2. The paper identifies a major problem in post training quantization of DiT and illustrates the drastic drop in performance that happens when about 1-2 bits of post training based quantization is used.
3. The paper identifies a major problem that happens in low bit quantization aware training using Ternary blocks, which is the high activation values and proposes a simple fix based on LLAMA to solve this problem. The authors also show that the with the utilization of RMSNorm, the performance in terms of FID drastically improves for TerDiT.
summary: The paper proposes TerDiT, a quantization aware training process that could be leverage to perform memory efficient training of DIT based architectures. The paper also introduces an efficient deployment scheme for ternary DITs. For this the authors introduce Ternary DiT blocks that replaces the conventional DiT blocks during training. The overall idea is to perform 1 bit bit quantization at different layers of the diffusion model during the training process.  The authors also noted that a naive replacement of the attention and feed forward layers in DiT leads to high activations which worsens performance. Hence they introduce an RMS normalization to alleviate the effect. The proposed changes bring in a reduction of maximum memory allocated  by a factor of 6 and checkpoint size by a factor of 14 while causing minor reduction in FID scores.
weaknesses: 1. The method lacks in novelty and the main contribution of the paper is in utilizing the ideas in [1] for replacing the attention and feedforward blocks and [2] for introducing RMSNorm for stable training process. 
2. The results in Figure 6 are surprising since the training losses with and without RMSNorm converge to similar values but the FID performance shows a drastic difference. Could the authors comment further on this.
3. In Figure 6, could the authors also please provide the results with classifier free guidance if it was already not performed.
4. Ln 303: The authors may have missed out to point the different between the work and [3]. Could the authors please highlight the major differences with [1].

[1] Shuming Ma, Hongyu Wang, Lingxiao Ma, Lei Wang, Wenhui Wang, Shaohan Huang, Li Dong, Ruiping Wang, Jilong Xue, and Furu Wei. The era of 1-bit llms: All large language models are in 1.58 bits. arXiv preprint arXiv:2402.17764, 2024b

[2] Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F. and Rodriguez, A., 2023. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Why choose a ternary quantizer for DiT rather than a 2-bit quantizer?

Would a 2-bit quantizer for DiT potentially yield better results than the ternary approach?

As the manuscript’s deployment relies on 2-bit operators, why insist on using 3-bit?
	
2. How is α in Equation 2 initialized?

3. Would using RMS Norm for the entire model achieve better results? If so, does this mean the strategy improvement in this manuscript is limited compared to BitNet b1.58?
soundness: 2
strengths: 1.	This is the first attempt to ternarize DiT, achieving generative results that can even rival those of full-precision models.

2.	Compared to BitNet b1.58, the manuscript selectively applies RMS Norm based on the structural characteristics of DiT.

3.	Comparative results across multiple models confirm the effectiveness of TerDiT.
summary: The manuscript presents the first ternary-weight DiT, applying a ternary quantizer to DiT and incorporating RMS Norm into the adaLN module, achieving high-accuracy image generation through QAT training.
weaknesses: 1. The paper lacks a clear motivation for applying a ternary quantizer to DiT. Given effective precedents for binary DM and 2-bit DM, I believe evidence is needed to demonstrate that a ternary DM can significantly outperform binary DM in accuracy or achieve substantial efficiency gains over 2-bit DM in practical use.

2. The issue of increased activations does not necessarily seem to be an inherent problem of ternarization. For example, better initialization of α in Equation 2 might straightforwardly address the phenomena shown in Figure 3.

3. The comparison with EfficientDM appears unfair. EfficientDM’s QAT training is based on LoRA, which has a natural disadvantage in precision. For a more compelling comparison with QAT methods, alternatives like LSQ and Q-DM [1] (including noise-estimating mimicking) would be more appropriate.

4. Table 1 compares only parameter counts but does not list the actual size, memory, or OPs.

5. The phrase "significantly reduced" in the caption of Table 2 lacks factual basis, as practical acceleration gains from weight-only binary quantizers are generally limited. The extent to which a carefully designed ternary quantizer could deliver such gains remains uncertain.

6. There is a lack of comparative data for using RMSNorm across the full model. In the Remark section at the end of Section 3.3, the manuscript states, "Comparative experiments are introduced in Sec.4.2," but I could not find any supporting experiments anywhere in the manuscript.

[1] Li, et al. Q-dm: An efficient low-bit quantized diffusion model.

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 4
questions: No questions, but more nitpicks and comments:
- L192: 'pachifies'
- Fig 2 - I do not see why this figure needs to be as large and detailed as it is just for some simple changes.
- Fig 3 too much whitepsace.
- Eq. 4 "adaln_norm" spells out as adaptive layer normalization normalization... change. 

References:

[1] Ma, Shuming, et al. "The era of 1-bit llms: All large language models are in 1.58 bits." arXiv preprint arXiv:2402.17764 (2024).

[2] Peebles, William, and Saining Xie. "Scalable diffusion models with transformers." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[3] Chen, Lei, et al. "Q-dit: Accurate post-training quantization for diffusion transformers." arXiv preprint arXiv:2406.17343 (2024).

[4] Li, Xiuyu, et al. "Q-diffusion: Quantizing diffusion models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[5] Frumkin, Natalia, Dibakar Gope, and Diana Marculescu. "Jumping through local minima: Quantization in the loss landscape of vision transformers." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[6] Chen, Junsong, et al. "Pixart-$\alpha $: Fast training of diffusion transformer for photorealistic text-to-image synthesis." arXiv preprint arXiv:2310.00426 (2023).

[7] Chen, Junsong, et al. "Pixart-\sigma: Weak-to-strong training of diffusion transformer for 4k text-to-image generation." arXiv preprint arXiv:2403.04692 (2024).

[8] Li, Zhimin, et al. "Hunyuan-DiT: A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding." arXiv preprint arXiv:2405.08748 (2024).

[9] Esser, Patrick, et al. "Scaling rectified flow transformers for high-resolution image synthesis." Forty-first International Conference on Machine Learning. 2024.

[10] Podell, Dustin, et al. "Sdxl: Improving latent diffusion models for high-resolution image synthesis." arXiv preprint arXiv:2307.01952 (2023).
soundness: 3
strengths: - Competent QAT method for Diffusion Transformers.
- QAT jumps to very low bit, ternary quantization.
- Experimental results, when presented, are convincing.
- The paper is well-written, easy to understand, and competently presented. 
- Visual results per Fig. 4 and otherwise are impressive.
summary: This paper extends the concept of 1.58-bit LLMs, which were introduced earlier this year [1], to Diffusion Transformers [2]. Specifically, 1.58-bits refers to ternary weights - each weight is either -1, 0 or 1. Specific contributions include the addition of RMSNorm in several components of the DiT including the Adaptive LayerNorm (AdaLN), usage of 2-bit CUDA kernels displaying model checkpoint size savings, and experimental results on ImageNet.
weaknesses: - The evaluation on ImageNet is incomplete. Specifically, the authors consider two resolutions (256 and 512, corresponding to the original DiT) and two model sizes (600M param XL/2 and a sized-up 4.2B parameter version), but their application of the two models is inconsistent across experiments, e.g., the 600M model is only used for 256x256 for CFG in Table 1. 
- Table 1 results are decently convincing but not great, e.g., L337 TerDiT-600M loses substantially to the original DiT-XL/2 (almost double FID, almost half Inception Score) and at 512 resolution TerDiT-4.2-G is not much better than the much smaller DiT-XL/2 despite being much larger. 
- For all the trashing of PTQ said in L072, there is inadequate comparison to PTQ methods. Yes, PTQ methods like Q-DiT [3] and Q-Diffusion [4] fail at 2-bit precision (really anything below 4-bits), this is not a novel finding, especially for ViT-like architectures [5]. The trade-off between QAT and PTQ has always been that the former allows lower bit precision at the cost of much higher computational resources. 
- No activation quantization. The 1.58-bit LLM paper [1] tried for 8-bit activations in addition to Ternary weights. Moreover, even if PTQ methods are limited to ~4 bit weights, these methods generally attempt activation quantization to at least 8-bits. This makes a comparison crucial as it could even the odds in terms of inference cost (e.g., BitOPs, latency, etc.) compared to this QAT method.
- There is at least one factual inaccuracy in the writing at lines 203-205: "AdaLN is an important component in the DiT model (Peebles & Xie, 2023) and has been proven more effective than cross-attention and in-context conditioning methods." this is absolutely not true and must be corrected. While the original DiT model removes Cross-Attention for AdaLN, it is very limited and restricted to ImageNet generation. T2I DiTs like the PixArt [6, 7] and Hunyuan [8] re-introduce cross-attention while SD3/3.5 and Flux [9] concat the text and image sequences together which is similar. 
- L260 - 263: I find this line of reasoning to not be borne out by the experimental results and ablations, e.g., Figs 5/6. In Fig 5 the loss of the DiT without RMS Norm is higher initially but quickly drops to be comparable. Fig 6 shows RMS Norm achieves faster FID reduction, however, FID itself is not entirely a trustworthy metric on its own [10].

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: - The proposed method focuses solely on weight quantization. Since this paper uses more computationally expensive QAT, considering activation quantization would further benefit this work.

-  It is well-known that DiT is famous for its role as a video generation backbones; how effective of this method in accelerating video generation models?

- Why use RMSNorm instead of other normalization mehotds?

- Overall, although this work claims to be the first to study extremely low-bit quantization of DiT, in my view, this work is not in-depth and solid. It merely adds an extra RMSNorm as its technical contribution. Therefore, I am tend to give a score of reject.
soundness: 2
strengths: - The proposed method achieves the ternarization of DiT model through the QAT technique, enhancing the efficiency of the DiT model. 

- The additional RMSNorm after AdaLN to address the issue of large activation values is simple and efficient.
summary: This work proposes a Quantization Aware Training (QAT) paradigm to obtain the proposed ternaryDiT models. By adding an additional RMSNorm layer after the AdaLN of the DiT model, the proposed method can effectively overcome the issue of large activation values. The authors evaluate the performance of the proposed method on image generation tasks with different resolutions and across various model sizes(600M and 4.2B).
weaknesses: - The technical contribution of this paper appears to be merely the introduction of an extra RMSNorm. The CUDA kernel deployment is from previous work, and the model architecture follows the previous Large-DiT. Therefore, I rise concerns about the technical contribution of this paper. 

- The models architecture evaluated in this work are somewhat limited. For example, previous work on diffusion model quantization included various diffusion architectures, such as LDM, StableDiffusion, etc. However, this paper, which focus on the DiT quantization, only considers the Large-DiT architecture. Considering the diversity of diffusion transformers, such as VDT[1], evaluating only one architecture weakens the contribution of this paper. Additionally, this also raises the question of whether the large activation values problems is a common phenomenon for diffusion transformers?

- In Table 2, the authors provide a comparison of the efficiency, and it can be seen that TerDiT achieves lower checkpoint size and GPU memory usage. However, the proposed method does not significantly increase the inference time compared to the full-precision model. I have noted the authors' explanation for this, but in practical applications, inference speed is a very important factor we care about. The proposed method does not truly improve inference speed, which weakens the significance of the quantization of DiT.

- Why adding RMSNorm can alleviate large activation values? The paper only provides empirical results; a more in-depth theoretical or visualization analysis is necessary for understanding the working mechanism.

-  In Table 1, why are the performance of TerDiT-600M-G is only given under the CFG settings?

> [1] VDT: General-purpose Video Diffusion Transformers via Mask Modeling. ICLR24