TerDiT: Ternary Diffusion Models with Transformers

TL;DR

TerDiT tackles the high deployment cost of diffusion transformer models by introducing a weight-only quantization-aware training framework for extremely low-bit, weight-only ternary DiTs. It stabilizes training with an RMS Norm after the adaLN module and demonstrates deployment using 2-bit CUDA kernels that pack weights into INT8, achieving significant memory and checkpoint reductions while maintaining competitive image generation on ImageNet at both $256\times256$ and $512\times512$. The approach scales TerDiT from $600\mathrm{M}$ to $4.2\mathrm{B}$ parameters and shows competitive or superior quantitative performance compared with full-precision DiTs such as DiT-XL/2 and Large-DiT-4.2B, validating a practical path toward efficient diffusion transformers on resource-constrained hardware. Code and pretrained TerDiT checkpoints are released to facilitate adoption and further research in efficient diffusion-model deployment.

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 transformer models (DiTs). Among 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 excessive 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, we propose TerDiT, the first quantization-aware training (QAT) and efficient deployment scheme for extremely low-bit diffusion transformer models. We focus on the ternarization of DiT networks, with model sizes ranging from 600M to 4.2B, and image resolution from 256$\times$256 to 512$\times$512. Our work contributes to the exploration of efficient deployment of large-scale DiT models, demonstrating the feasibility of training extremely low-bit DiT models from scratch while maintaining competitive image generation capacities compared to full-precision models. Our code and pre-trained TerDiT checkpoints have been released at https://github.com/Lucky-Lance/TerDiT.

Figures (14)

• Figure 1: Sample images (256$\times$256) generated by TerDiT with 4.2B parameters (using 2GB of GPU memory) are shown. For comparison, images generated by full-precision diffusion transformer models—DiT-XL/2 with 675M parameters (using 3GB of GPU memory) and Large-DiT-4.2B with 4.2B parameters (using 17GB of GPU memory)—are provided in Fig. \ref{['fig:sample']}.
• Figure 2: 2-bit Q-DiT (A) and 2-bit Q-Diffusion (B) quantization results.
• Figure 3: 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.
• Figure 4: 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.
• Figure 5: Qualitative results analysis (256$\times$256). We compare images generated by DiT-XL/2 (A), Large-DiT-4.2B (B), and TerDiT-4.2B (C) with class labels [355, 980, 328, 33, 928, 862, 330, 944] and cfg=4. TerDiT-4.2B generates images of the same visual quality as two full-precision DiT models.