PaperlayerView exact revision on arXiv

Source-linked AI summary

TerDiT: Ternary Diffusion Models with Transformers

Xudong Lu, Aojun Zhou, Ziyi Lin, Qi Liu, Yuhui Xu, Renrui Zhang, Xue Yang, Junchi Yan, Peng Gao, Hongsheng Li

arXiv:2405.14854v2cs.CVcs.LG

TL;DR

Large-scale DiT models are expensive to deploy, and existing quantization research provides limited support for extremely low-bit DiTs. TerDiT uses ternary quantization-aware training with DiT-specific RMS normalization and efficient 2-bit deployment, demonstrating competitive generation quality while scaling from 600M to 4.2B parameters.

  • Problem

    Extremely low-bit quantization of large-scale DiT models remains insufficiently explored, while deploying their excessive parameters is costly.

    Source: §I · p. 2, B. Model Quantization · p. 4

  • Method

    TerDiT applies weight-only ternary quantization-aware training, adds RMS Norm after ternary adaLN linear layers, and deploys models with 2-bit CUDA kernels.

    Source: §I · p. 2

  • Results

    TerDiT trains ternary DiT models from scratch from 600M to 4.2B parameters across 256×256 to 512×512 resolutions while achieving competitive generation results against full-precision and baseline quantized models.

    Source: §I · p. 2, §V · p. 10

  • Takeaways & Limitations

    The results demonstrate the feasibility of training and efficiently deploying large extremely low-bit DiT models.

    Source: §I · p. 2, §V · p. 10

Abstract

from arXiv · show

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.

I. INTRODUCTION

Large DiT models offer strong image-generation quality but are costly to deploy, while existing low-bit quantization work leaves extremely low-bit DiT models insufficiently addressed. TerDiT introduces ternary QAT and deployment techniques that scale across large DiT models while retaining competitive generation quality.

  • 2-bit Q-DiT and Q-Diffusion both fail to generate images, illustrating the difficulty of post-training quantization at extremely low bit widths.The comparison is shown in Fig. 2 and concerns 2-bit weight quantization.Source: §I · p. 2, Figure 2 · p. 2
  • TerDiT replaces direct adaLN weight ternarization with RMS Norm after ternary adaLN linear layers to improve convergence and model performance.Direct ternarization produces large dimension-wise scale and shift values, slowing convergence and degrading performance.Source: §I · p. 2
  • Ternary DiT models scale from 600M to 4.2B parameters and from 256×256 to 512×512 resolution, with 2-bit CUDA deployment reducing checkpoint size by over 10× and inference memory by about 8×.A 4.2B DiT can run with less than 2GB GPU memory at 256×256 resolution.Source: §I · p. 2
  • On the ImageNet benchmark for image generation, TerDiT delivers competitive evaluation results compared with full-precision models and baseline quantization methods, showcasing its effectiveness.Source: §I · p. 2
  • TerDiT introduces the first quantization-aware training and efficient deployment scheme for extremely low-bit DiT models.The work focuses on ternary DiT models and positions this as the first study of extremely low-bit DiT quantization.Source: §I · p. 2

II. RELATED WORKS

Related work spans diffusion-model foundations, diffusion-model quantization, ternary weight networks, and the TerDiT framework. These strands motivate applying QAT and efficient ternary deployment to large-scale DiT models.

A. Diffusion Models

  • Diffusion models reverse a diffusion process and support high-quality generation across images, text-to-image synthesis, audio, video, and 3D modalities.Source: A. Diffusion Models · p. 2

B. Quantization of Diffusion Models

  • PTQ methods improve diffusion-model efficiency but often perform suboptimally at low bit widths, while existing QAT methods remain limited to small U-Net-based models.This leaves QAT for large DiT models with extremely low bit widths underexplored.Source: B. Quantization of Diffusion Models · pp. 2–3

C. Ternary Weight Networks

  • Ternary weight networks quantize weights alone or weights and activations, offering memory and computation efficiency and supporting ternary large-language-model training and deployment.TerDiT extends these developments to ternary DiT models.Source: C. Ternary Weight Networks · p. 3

III. TERDIT

  • TerDiT combines weight-only QAT, DiT-specific structure improvements, and ternary deployment for large-scale DiT models.The framework covers quantization functions, QAT, model-structure improvement, and deployment.Source: §III · p. 3

A. Preliminary: Diffusion Transformer Models

DiT replaces U-Net denoising with transformer blocks operating on latent patches, while adaLN injects conditioning into those blocks. TerDiT applies weight-only ternarization and QAT to large DiTs, but naive ternary training suffers from slow convergence and unstable activations.

A. Preliminary: Diffusion Transformer Models

  • DiT patches spatial inputs into token sequences and uses transformer blocks, adaLN conditioning, and a final linear decoder to predict noise and covariance.The architecture can be trained similarly to U-Net-based diffusion models.Source: A. Preliminary: Diffusion Transformer… · pp. 3–4
  • AdaLN contains an MLP representing approximately 10% to 20% of model parameters, making its ternarization behavior important for TerDiT.Directly ternarizing this module produces undesirable training results.Source: A. Preliminary: Diffusion Transformer… · p. 4

B. Model Quantization

  • TerDiT replaces linear layers in attention, feedforward, and MLP components with ternary layers using absmean normalization, rounding, and clamping to {−1, 0, +1}.A learnable α scales the quantized values to {−α, 0, +α}.Source: B. Model Quantization · p. 4
  • TerDiT quantizes weights but not activations, training ternary DiTs from scratch with STE while retaining full-precision parameters for updates.Ternary weights are computed in the forward pass and their gradients update the full-precision parameters during backpropagation.Source: B. Model Quantization · p. 4
  • Naive ternary training converges very slowly because ternary linear layers can produce large activation values that destabilize optimization.The paper therefore introduces a QAT-specific structural improvement in the following subsection.Source: B. Model Quantization · p. 4, Figure 4 · p. 4

C. QAT-specific Model Structure Improvement

Ternary linear layers can create unusually large activations, especially in adaLN, so TerDiT inserts RMS Norm after the adaLN MLP. This modification scales activations toward full-precision ranges and is associated with more stable, faster training and improved evaluation results.

C. QAT-specific Model Structure Improvement

  • Ternary linear layers produce much larger activations than full-precision layers, while RMS Norm scales them to a comparable reasonable range.The analysis uses a 512×1024 input and a ternary layer matching the Large-DiT adaLN dimensions.Source: C. QAT-specific Model Structure… · pp. 4–5
  • Table I compares TerDiT with full-precision diffusion models on ImageNet, but the supplied caption does not report the table’s cell values or winning configurations.Only the table scope is supported here.Source: Table I · p. 5
  • TerDiT adds RMS Norm after the adaLN MLP because ternarization can create large dimension-wise scale and shift values that harm training.Standard DiT blocks normalize attention and feedforward layers but do not otherwise normalize adaLN.Source: C. QAT-specific Model Structure… · p. 5
  • The RMS-Normalized adaLN modification leads to faster convergence, lower training loss, and better quantitative and qualitative evaluation results.The paper reports that trained activation distributions with and without RMS Norm show similar distribution features to the earlier analysis.Source: C. QAT-specific Model Structure… · p. 5
  • Unlike methods that normalize every quantized linear layer, TerDiT adds fewer norms within adaLN and reports faster training and better evaluation scores.The comparison is made against BitNet-style normalization and related quantized-layer approaches.Source: C. QAT-specific Model Structure… · p. 5

D. Deployment Scheme

TerDiT deploys ternary DiT models using packed 2-bit representations and on-the-fly unpacking during inference. Qualitatively, TerDiT-4.2B matches the visual quality of full-precision DiT models at 256×256.

Deployment implementation

  • Packed ternary weights are stored as INT8 values, with four ternary numbers per INT8 and on-the-fly unpacking to FP values during inference.The unpacking operation slows inference, although the scheme targets future hardware-supported speedups.Source: D. Deployment Scheme · pp. 5–6

Qualitative evaluation

  • TerDiT-4.2B generates 256×256 images with the same visual quality as DiT-XL/2 and Large-DiT-4.2B.The comparison uses class labels [355, 980, 328, 33, 928, 862, 330, 944] and cfg=4.Source: Figure 5 · p. 6

IV. EXPERIMENTS

The experiments evaluate ternary DiT models against full-precision models across image quality, deployment efficiency, and qualitative generation. TerDiT remains competitive with full-precision DiTs across 256×256 and 512×512 settings while spanning 600M to 4.2B parameters.

IV. EXPERIMENTS

  • The experiments cover comparisons with full-precision and quantized models, deployment efficiency, normalization, ablations, high-cfg generation, and additional qualitative results.Models contain 600M or 4.2B parameters.Source: §IV · p. 6

A. Comparison with Full-precision Models

  • The 4.2B model is identified as the actual parameter count of the repository's provided 3B Large-DiT model.This clarifies the model-size labeling used in the comparison.Source: A. Comparison with Full-precision Models · p. 6
  • TerDiT and full-precision diffusion models are evaluated on ImageNet using FID, sFID, Inception Score, Precision, Recall, and training-image counts.The study trains 600M and 4.2B ternary models at 256×256 and continues the 4.2B model at 512×512.Source: A. Comparison with Full-precision Models · p. 6
  • TerDiT-4.2B-G outperforms DiT-XL/2-G across all reported aspects at 512×512 and surpasses Large-DiT-4.2B-G in sFID and Precision.At 256×256 with cfg=1.5, it achieves similar results to DiT-XL/2 and slightly degrades relative to two full-precision DiT-structured models.Source: A. Comparison with Full-precision Models · pp. 6–7
  • TerDiT-4.2B, DiT-XL/2, and Large-DiT-4.2B show no significant visual difference in the 256×256 qualitative comparison.The comparison is presented in Fig. 5.Source: A. Comparison with Full-precision Models · p. 7

B. Extremely Low-bit Quantization Baselines

TerDiT baselines show that extremely low-bit PTQ and alternative QAT approaches struggle, while RMS-normalized adaLN improves ternary DiT training and generation. The ternary models also substantially reduce deployment memory and checkpoint size, though current hardware and software limitations slow inference.

B. Extremely Low-bit Quantization Baselines

  • FID 4.34 for TerDiT versus 6.60 for BitNet, while BitNet training runs at 0.9× the speed and inference at 1.15× the latency.EfficientDM with 2-bit quantization fails to generate normal images, and Q-DiT and Q-Diffusion likewise fail under 2-bit weight quantization.Source: B. Extremely Low-bit Quantization… · p. 7

C. Deployment Efficiency Comparison

  • TerDiT-4.2B uses less checkpoint and memory storage than Large-DiT-4.2B and even less memory than DiT-XL/2.The comparison evaluates checkpoint size, memory usage, and 250-step inference time on one A100-80G GPU.Source: C. Deployment Efficiency Comparison · p. 7
  • Current ternary DiT inference is slower than full-precision models because open-source deployment frameworks and efficient ternary hardware support are unavailable.The paper relates future computational benefits to continued software-hardware co-design for ternary networks.Source: C. Deployment Efficiency Comparison · p. 7

D. Ablation Study: RMS Normalized AdaLN Module

  • Adding RMS Norm after the adaLN MLP produces faster convergence and lower FID scores for both 600M and 4.2B ternary DiT models.The ablation records training loss and FID-50k on ImageNet 256×256 with cfg=1.5.Source: D. Ablation Study: RMS Normalized… · pp. 7–8, Figure 7 · p. 8, Figure 8 · p. 8
  • RMS-normalized adaLN improves qualitative results for both model sizes, while larger ternary DiTs show stronger learning ability and better training outcomes.The qualitative comparison includes models trained for 400k steps and fully trained models.Source: D. Ablation Study: RMS Normalized… · p. 8
  • RMS Norm limits the adaLN scale-MLP activation range during inference relative to the corresponding unnormalized ternary model.The analysis uses TerDiT-4.2B at the first sampling step in the second ternary DiT block and compares against the full-precision model.Source: D. Ablation Study: RMS Normalized… · p. 8, Figure 9 · p. 8

E. Ablation Study: Learning Rate Reduction

The learning-rate reduction ablation evaluates continued training of TerDiT-600M after 1,550k steps. Reducing the rate improves evaluation results and highlights its importance during later training, with qualitative comparisons shown in Fig. 11.

E. Ablation Study: Learning Rate Reduction

F. Image Quality with High CFG Values

TerDiT remains robust when classifier-free guidance is increased from cfg=4 to cfg=10. Although slight distortion appears, quantization does not severely impair generation quality at the higher guidance value.

F. Image Quality with High CFG Values

G. More Qualitative TerDiT Results

TerDiT produces good-quality images at both 256×256 and 512×512 resolutions. The discussion reports competitive generation results while identifying training and evaluation-scope limitations.

G. More Qualitative TerDiT Results

V. DISCUSSIONS AND FUTURE WORKS

  • Competitive ImageNet generation results against full-precision models and quantization baselines support training large ternary DiTs from scratch.Source: §V · p. 10
  • Ternary DiT training is less stable and more time-consuming than full-precision training, increasing training-related carbon dioxide emissions.Source: §V · p. 10
  • The study does not completely compare ternary quantization with INT8 or FP16 because those formats fall outside its extremely low-bit scope.Source: §V · p. 10
  • Ternary quantization could reduce memory usage by up to 16×, but DiT-specific hardware acceleration remains future work.Source: §V · p. 11

VI. BIOGRAPHY SECTION

The biography section presents the authors’ academic and professional affiliations, spanning universities, research laboratories, and industry research. Their stated interests cover machine learning, vision, model efficiency, multimodal modeling, and related areas.

VI. BIOGRAPHY SECTION

  • Several authors are affiliated with The Chinese University of Hong Kong and work on deep learning, computer vision, multimodal learning, or large language models.Source: §VI · p. 13
  • Qi Liu’s interests span formal verification, large language model reasoning, and efficient machine learning.Source: §VI · p. 13
  • The biographies also include researchers at Shanghai Artificial Intelligence Laboratory and Salesforce AI Research working on AIGCs, vision-language models, language models, and efficient deployment.Source: §VI · pp. 13–14
  • Junchi Yan is a professor and deputy director at Shanghai Jiao Tong University, with research interests in machine learning and vision.Source: §VI · p. 14
  • Hongsheng Li is an associate professor at The Chinese University of Hong Kong and is co-affiliated with its Multimedia Laboratory.Source: §VI · p. 14
Loading 2405.14854v2…