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.
翻译:近年来,大规模预训练文本到图像扩散模型的发展显著提升了高保真图像的生成质量,特别是基于Transformer架构的扩散模型(DiTs)的出现。在这些扩散模型中,扩散Transformer展现出卓越的图像生成能力,实现了更低的FID分数和更高的可扩展性。然而,大规模DiT模型因其庞大的参数量而部署成本高昂。尽管现有研究已探索了扩散模型的高效部署技术(如模型量化),但针对基于DiT的模型仍缺乏相关研究。为填补这一空白,本文提出TerDiT——一种面向Transformer三值化扩散模型的量化感知训练(QAT)与高效部署方案。我们专注于DiT网络的三值化,并将模型规模从6亿参数扩展至42亿参数。本工作推动了对大规模DiT模型高效部署策略的探索,证明了从头训练极低比特扩散Transformer模型的可行性,同时在与全精度模型的对比中保持了具有竞争力的图像生成能力。代码将在https://github.com/Lucky-Lance/TerDiT公开。