PTQ4DiT: Post-training Quantization for
Diffusion Transformers

Although generative models built upon U-Nets have made great advancements in the last few years, transformer-like architectures are increasingly attracting attention [39, 7, 59]. The recently explored Diffusion Transformers (DiTs) [37] have achieved state-of-the-art performance in image generation. Encouragingly, DiTs exhibit remarkable scalability in model size and data representation, positioning them as a promising backbone for a wide range of generative applications [4, 30, 65].

DiTs consist of $n_{B}$ blocks, each containing a Multi-Head Self-Attention (MHSA) and a Pointwise Feedforward (PF) module [49, 11, 37], both preceded by their respective adaptive Layer Norm (adaLN) [38]. We illustrate the DiT Block structure in Figure 2 (Left). These blocks sequentially process the noised latent and conditional information, which are both represented as tokens in a lower-dimensional latent space [39]. In each block, conditional input $\mathbf{c}\in\mathbb{R}^{d_{in}}$ is converted to scale and shift parameters ($\bm{\gamma},\bm{\beta}\in\mathbb{R}^{d_{in}}$), which are regressed through MLPs then injected into the noised latent $\mathbf{Z}\in\mathbb{R}^{n\times d_{in}}$ via adaLN:

where LN$(\cdot)$ is the standard Layer Norm [1]. These adaLN modules dynamically adjust the layer normalization before each MHSA and PF module, enhancing DiTs’ adaptability to varying conditions and improving the generation quality.

Despite their effectiveness, DiTs require extensive computational resources to generate high-quality images, which impedes their real-world deployment. In this paper, we devise a model quantization method for DiTs that reduces both time and memory consumption without necessitating re-training the original models, offering a robust and practical solution for enhancing the efficiency of DiTs.

Model quantization is a compression technique that improves the inference efficiency of deep learning models by transforming full-precision tensors into $b$-bit integer approximations, leading to direct computational acceleration and memory saving [33, 62, 8, 28, 19, 64]. Formally, the quantization process can be defined as:

where $\mathbf{x}$ denotes the full-precision tensor, $\lfloor\cdot\rceil$ is the round-to-nearest operator [32], and the clamp function restricts the quantized value within the range of $[0,2^{b}-1]$. Here, $\bm{\delta}$ and $\bm{\lambda}$ are quantization parameters subject to optimization. Among various quantization methods, Post-training Quantization (PTQ) is a dominant approach for large quantized models, as it circumvents the substantial resources required for model re-training [20, 52, 25, 44, 15]. PTQ employs a small calibration dataset to optimize quantization parameters, which aims to reduce the performance gap between the quantized models and their full-precision counterparts with minimal data and computational expenses.

PTQ has been effectively applied to a wide range of neural networks, including CNNs [20, 52, 25], Language Transformers [8, 57, 24, 23, 27], Vision Transformers [62, 13, 21, 29], and U-Net-based Diffusion models [44, 18, 51, 50]. Despite its demonstrated success, PTQ’s applicability to Diffusion Transformers (DiTs) remains unexplored, presenting a significant open challenge within the research community. To bridge this gap, our work delves into the unique challenges of quantizing DiTs and introduces the first PTQ method for DiTs that can fruitfully preserve their generation performance.

A linear layer $f(\cdot;\mathbf{W})$ within MHSA and PF modules typically takes a token sequence $\mathbf{X}\in\mathbb{R}^{n\times d_{in}}$ as input and performs linear transformation with its weight matrix $\mathbf{W}\in\mathbb{R}^{d_{in}\times d_{out}}$, formulated as $f(\mathbf{X};\mathbf{W})=\mathbf{X}\cdot\mathbf{W},$ where $n$ is the sequence length, and $d_{in}$ and $d_{out}$ denote the input and output dimensions, respectively. As discussed in Section 3, both the activation $\mathbf{X}$ and the weight matrix $\mathbf{W}$ exhibit salient channels that possess elements with significantly greater absolute magnitudes, which lead to large post-quantization errors.

Fortunately, large values do not coincide in the same channels of activation and weight, so these extremes do not amplify each other, as observed in Figure 3. This property suggests the feasibility of complementarily redistributing the large magnitudes in salient channels between activation and weight, thereby alleviating quantization difficulties for both. Inspired by previous works on large model compression [54, 45, 61, 23], we propose Channel-wise Salience Balancing (CSB), which employs diagonal Salience Balancing Matrices $\mathbf{B^{X}}$ and $\mathbf{B^{W}}$ to adjust the channel-wise distribution of activation and weight, as expressed by:

To address the quantization difficulties, we need to achieve balanced distributions in $\mathbf{\widetilde{X}}$ and $\mathbf{\widetilde{W}}$, which requires $\mathbf{B^{X}}$ and $\mathbf{B^{W}}$ to capture the characteristics of salient channels. Considering that the quantization error is significantly influenced by the range of distributions [33, 57, 26], we measure the salience $s$ of an activation or weight channel as the maximal absolute value among its elements:

Here, $j$ is the channel index. Consequently, the balanced salience $\widetilde{s}$, representing the equilibrium between activation and weight channels, can be quantified using the geometric mean. Specifically, for the $j$-th channel, the balanced salience is calculated as follows:

Building on these concepts, we proceed to construct the Salience Balancing Matrices, which modulate the salience of activations and weights with the guidance of $\widetilde{s}$:

Following these, the balancing transformation defined by Eq. (3) will result in a complementary redistribution of channel salience between activations and weights. Specifically, for each channel $j$, we have $s(\mathbf{\widetilde{X}}_{j})=s(\mathbf{\widetilde{W}}_{j})=\widetilde{s}(\mathbf{X}_{j},\mathbf{W}_{j})$, thereby alleviating the quantization difficulties, as demonstrated by the reduction in overall channel salience:

Here, we characterize the overall salience $s_{o}$ of activations or weights using the maximum salience across channels, e.g., $s_{o}(\mathbf{X})=\text{max}(s(\mathbf{X}_{1}),s(\mathbf{X}_{2}),\ldots,s(\mathbf{X}_{d_{in}}))$, which reflects the distribution range of elements that are quantized collectively under certain granularity.

Diffusion Transformers (DiTs) utilize an iterative denoising process for image sampling [37]. Under this sequential paradigm, the linear layer $f$ receives inputs from an activation sequence $\mathbf{X}^{(1:T)}=(\mathbf{X}^{(1)},\mathbf{X}^{(2)},\ldots,\mathbf{X}^{(T)})$, which encompasses $T$ timesteps. Targeting a certain timestep $t$, the salience of all activation and weight channels can be evaluated using Eq. (4):

While $\mathbf{s}(\mathbf{W})$ remains consistent, we find that $\{\mathbf{s}(\mathbf{X}^{(t)})\}_{t=1}^{T}$ exhibits significant temporal variations during the process of transforming purely random noise into high-quality images, as demonstrated in Figure 4. These fluctuations diminish the effectiveness of our CSB since quantization errors can be exacerbated by the biased estimation of activation salience among timesteps, resulting in degraded generation quality of the quantized models.

To accurately gauge the activation channel salience under multi-timestep scenarios, we propose Spearman’s $\rho$-guided Salience Calibration (SSC). This offers a comprehensive evaluation of activation salience, with enhanced focus allocated to the timesteps where the complementarity property is more significant, facilitating effective salience balancing between activation and weight channels. Essentially, the lower the correlation between activation salience $\mathbf{s}(\mathbf{X}^{(t)})$ and weight salience $\mathbf{s}(\mathbf{W})$, the greater reduction effect in overall channel salience (Eq. (8)). The intuition of SSC is visualized in Figure 2 (Right). Mathematically, we formulate the Spearman’s $\rho$-calibrated Temporal Salience $\mathbf{s}_{\rho}$ by selectively aggregating the activation salience along timesteps:

where weighting factors $\{\eta_{t}\}_{t=1}^{T}$ are derived from a normalized exponential form of inverse Spearman’s $\rho$ statistic [47, 55, 56]:

Here, $\rho(\cdot,\cdot)$ computes the correlation between two sequences, and $\eta_{t}$ serves as the weighting factor for activation salience at timestep $t$. In this method, $\eta_{t}$ inversely reflects the correlation coefficient $\rho(\mathbf{s}(\mathbf{X}^{(t)}),\mathbf{s}(\mathbf{W}))$, thereby prioritizing timesteps where there is a higher degree of complementarity in salience between activations and weights. Subsequently, we utilize $\mathbf{s}_{\rho}$ for activation salience in Eqs. (5), (6), and (7), yielding refined Salience Balancing Matrices, denoted as $\mathbf{B}^{\mathbf{X}}_{\rho}$ and $\mathbf{B}^{\mathbf{W}}_{\rho}$. By applying SSC, we calibrate the activation salience within CSB to strategically account for the temporal variations during the denoising process. Appendix B presents the full Algorithm for PTQ4DiT.

Before quantization, we estimate $\mathbf{B_{\rho}^{X}}$ and $\mathbf{B_{\rho}^{W}}$ on a small calibration dataset generated from multiple timesteps. Then, we incorporate these matrices into the linear layers within MHSA and PF modules [37] to alleviate the quantization difficulty. Given that $\mathbf{B_{\rho}^{X}}$ and $\mathbf{B_{\rho}^{W}}$ are mutual inverses, this incorporation maintains mathematical equivalence to the original linear layer $f$:

The proof is provided in Appendix C. Furthermore, we design a re-parameterization scheme for DiTs, allowing for obtaining $\mathbf{\widetilde{X}}$ and $\mathbf{\widetilde{W}}$ without extra computational burden during inference. Specifically, we update the weight matrix of linear layer $f$ to $\mathbf{\widetilde{W}}$ offline and seamlessly integrate $\mathbf{B_{\rho}^{X}}$ into the preceding linear transformation operations. This integration includes adaptations to adaLN [38, 37] and matrix multiplications within attention mechanisms [49]. Appendix A discusses these adaptations.

Post-adaLN. For linear layers following the adaLN module, we integrate $\mathbf{B_{\rho}^{X}}$ by adjusting the scale and shift parameters ($\bm{\gamma},\bm{\beta}\in\mathbb{R}^{d_{in}}$) within adaLN:

Equivalently, we fuse $\mathbf{B_{\rho}^{X}}$ into the MLPs responsible for regressing these parameters, thus avoiding additional computation overhead at inference time. Detailed derivations are provided in Appendix D.

Post-Matrix-Multiplication. For linear layers after matrix multiplication, the effect of PTQ4DiT can be realized by directly absorbing the Salience Balancing Matrices into the preceding de-quantization functions associated with the matrix multiplication [12, 53, 61].

Our experimental setup is similar to the original study of Diffusion Transformers (DiTs) [37]. We evaluate PTQ4DiT on the ImageNet dataset [41], using pre-trained class-conditional DiT-XL/2 models [37] at image resolutions of 256$\times$256 and 512$\times$512. The DDPM solver [17] with 250 sampling steps is employed for the generation process. To further assess the robustness of our method, we conduct additional experiments with reduced sampling steps of 100 and 50.

For fair benchmarking, all methods utilize uniform quantizers for all activations and weights, with channel-wise quantization for weights and tensor-wise for activations, unless specified otherwise. To construct the calibration set, we uniformly select 25 timesteps for 256-resolution experiments and 10 timesteps for 512-resolution experiments, generating 32 samples at each selected timestep. The optimization of quantization parameters follows the implementation from Q-Diffusion [18]. Our code is based on PyTorch [36], and all experiments are conducted on NVIDIA RTX A6000 GPUs.

To comprehensively assess generated image quality, we employ four metrics: Fréchet Inception Distance (FID) [16], spatial FID (sFID) [42, 34], Inception Score (IS) [42, 3], and Precision, all computed using the ADM toolkit [10]. For all methods under evaluation, including the full-precision (FP) models, we sample 10,000 images for ImageNet 256$\times$256, and 5,000 for ImageNet 512$\times$512, consistent with conventions from prior studies [35, 44].

We present a comprehensive assessment of our PTQ4DiT against prevalent baseline methods in various settings. Our evaluation focuses on mainstream Post-training Quantization (PTQ) methods that are widely used and adaptable to DiTs, including PTQ4DM [44], Q-Diffusion [18], and PTQD [15].

We reimplement these methods to suit the unique structure of DiTs. Considering the architectural similarity between DiTs and ViTs [11], our analysis also includes RepQ-ViT [21], the state-of-the-art PTQ method initially designed for ViTs. We enhance RepQ-ViT (denoted as RepQ*) by extending the calibration set to integrate temporal dynamics and customizing its advanced channel-wise and log$\sqrt{2}$ quantizers specifically for DiTs.

Tables 1 and 2 report the outcomes on large-scale class-conditional image generation for ImageNet 256$\times$256 and 512$\times$512, respectively. Table 1 demonstrates the effectiveness of PTQ4DiT across various quantization settings and timesteps. Notably, our finding reveals that at 8-bit precision (W8A8), PTQ4DiT closely matches the generative capabilities of the

FP models, whereas most baseline methods experience significant performance losses. At the more stringent 4-bit weight precision (W4A8), all baseline methods exhibit more considerable degradation. For instance, under 250 timesteps, PTQ4DM [44] sees a drastic FID increase of 68.05. In contrast, our PTQ4DiT only incurs a slight increase of 2.56. This resilience remains evident as the number of timesteps decreases, underscoring the robustness of PTQ4DiT in resource-limited environments. Moreover, PTQ4DiT markedly outperforms mainstream methods at the higher 512$\times$512 resolution, further validating its superiority. For example, using 250 timesteps, PTQ4DiT substantially lowers FID by 41.26 and sFID by 9.83 over the second-best method, Q-Diffusion. Figure 6 depicts the efficiency-vs-efficacy trade-off on W8A8 across various timestep configurations. Our PTQ4DiT achieves comparable performance levels to FP models but with considerably reduced computational costs, offering a viable alternative for high-quality image generation. Figures 5, 8, and 9 also present randomly generated images for visual comparisons, highlighting PTQ4DiT’s ability to produce images of superior quality.

To verify the efficacy of CSB and SSC, we conduct an ablative study on the challenging W4A8 quantization. Experiments are performed on ImageNet 256$\times$256 using 250 sampling timesteps. Three method variants are considered in our ablation: (i) Baseline, which applies basic linear quantization on DiTs, (ii) Baseline + CSB, which integrates CSB in the linear layers within MHSA and PF modules, where the Salience Balancing Matrices $\mathbf{B^{X}}$ and $\mathbf{B^{W}}$ are estimated based on distributions at the midpoint timestep $\frac{T}{2}$, and (iii) Baseline + CSB + SSC, which is the complete PTQ4DiT. Results detailed in Table 3 indicate that each proposed component improves the performance, validating their effectiveness. Particularly, CSB enhances upon the Baseline by a large margin, decreasing FID by 14.37 and sFID by 2.35, suggesting its critical role in alleviating the severe quantization difficulties inherent in DiTs. Note that with the addition of CSB, our method surpasses Q-Diffusion [18], a leading PTQ method for diffusion models. Moreover, integrating SSC further boosts our PTQ4DiT towards state-of-the-art performance, facilitating high-quality image generation at W4A8 precision, as shown in Figure 5.