PTQ4DiT: Post-training Quantization for Diffusion Transformers

Authors: junyi  wu
haoxuan  wang
yuzhang  shang
mubarak  shah
yan  yan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 03:14:05

Abstract: The recent introduction of Diffusion Transformers (DiTs) has demonstrated exceptional capabilities in image generation by using a different backbone architecture, departing from traditional U-Nets and embracing the scalable nature of transformers. Despite their advanced capabilities, the wide deployment of DiTs, particularly for real-time applications, is currently hampered by considerable computational demands at the inference stage. Post-training Quantization (PTQ) has emerged as a fast and data-efficient solution that can significantly reduce computation and memory footprint by using low-bit weights and activations. However, its applicability to DiTs has not yet been explored and faces non-trivial difficulties due to the unique design of DiTs. In this paper, we propose PTQ4DiT, a specifically designed PTQ method for DiTs. We discover two primary quantization challenges inherent in DiTs, notably the presence of salient channels with extreme magnitudes and the temporal variability in distributions of salient activation over multiple timesteps. To tackle these challenges, we propose Channel-wise Salience Balancing (CSB) and Spearmen's $\rho$-guided Salience Calibration (SSC). CSB leverages the complementarity property of channel magnitudes to redistribute the extremes, alleviating quantization errors for both activations and weights. SSC extends this approach by dynamically adjusting the balanced salience to capture the temporal variations in activation. Additionally, to eliminate extra computational costs caused by PTQ4DiT during inference, we design an offline re-parameterization strategy for DiTs. Experiments demonstrate that our PTQ4DiT successfully quantizes DiTs to 8-bit precision (W8A8) while preserving comparable generation ability and further enables effective quantization to 4-bit weight precision (W4A8) for the first time.

Main paper: 
 introduction: 
 Diffusion models have spearheaded recent breakthroughs in generation tasks [59, 7] . In the past, these models were based on convolutional U-Nets [40] as their backbone architectures [46, 17, 9, 39] . However, recent work [2, 60, 30] has revealed that the U-Net inductive bias is not essential for the success of diffusion models and even limits their scalability. Among this trend, Diffusion Transformers (DiTs) [37] have demonstrated exceptional capabilities in image generation by using a different backbone architecture. Different from U-Nets that carefully design downsampling and upsampling blocks with skip-connections, DiTs are constructed by repeatedly and sequentially stacking transformer blocks [49] . This architectural choice inherits the scaling property of transformers [5, 48, 58, 31] , facilitating more flexible parameter expansion for enhanced performance. With their versatility and scalability, DiTs have been successfully integrated into advanced frameworks like Sora [4] , demonstrating their potential as a leading architecture for future generative models [14, 6, 30, 65] .Figure 1: (Left) Illustration of salient channels in activation and weight. Note that salient activation channels exhibit variations over different timesteps (e.g., t = t 1 , t 2 , t 3 .), posing non-trivial quantization challenges. To mitigate the overall quantization difficulty, our method leverages the complementarity (activation and weight channels do not have extreme magnitude simultaneously) to redistribute channel salience between weights and activations across various timesteps. (Right) Quantization performance on W8A8 and W4A8, employing FID (lower is better) and IS (higher is better) metrics on ImageNet 256×256 [41] . The circle size indicates the model size.Nonetheless, the widespread adoption of Diffusion Transformers is currently constrained by their massive amount of parameters and computational complexity. DiTs consist of a large number of repeated transformer blocks and employ a lengthy iterative image sampling process, demanding high computational costs during inference. For instance, generating a 512×512 resolution image using DiTs can take more than 20 seconds and 10 5 Gflops on an NVIDIA RTX A6000 GPU. This substantial requirement makes them unacceptable or impractical for real-time applications, especially considering the potential for increased model sizes and feature resolutions.Model quantization [33, 32, 28] is a prominent technique for accelerating deep learning models because of its high compression rate and significant reduction in inference time. This technique transforms model weights and activations into low-bit formats, which directly reduces the computational burden and memory usage. Among various methods, Post-training Quantization (PTQ) stands out as a leading approach since it circumvents the need to re-train the original model [62, 44, 18, 63, 22] . Practically, PTQ requires only a small dataset for fast calibration, thus is highly suitable for quantizing DiTs, whose re-training process involves extensive data and computational resources [14, 6] .However, quantizing DiTs in a post-training manner is non-trivial due to the complex distribution patterns in weights and activations. We discover two major challenges that impede the effective quantization of DiTs: ❶ The emergence of salient channels, channels with extreme magnitudes, in both weights and activations of linear layers within DiT blocks. When low-bit representations are used for these salient channels, pronounced errors compared to the full-precision (FP) counterparts are observed, incurring fundamental difficulty for quantization. ❷ The extreme magnitudes within salient activation channels significantly vary as the inference proceeds across multiple timesteps. This dynamic behavior further complicates the quantization of salient channels, as quantization strategies optimized for one timestep may fail to generalize to other timesteps. Such inconsistency, especially in salient channels that dominate the activation signals, can result in significant deviations from the full-precision distribution, leading to degradation in the generation ability of quantized models.Targeting these two challenges, we propose a novel Post-training Quantization method specifically for Diffusion Transformers, termed PTQ4DiT. To address the quantization difficulty associated with salient channels, we propose Channel-wise Salience Balancing (CSB). CSB capitalizes on an interesting observation of the salient channels that extreme values do not coincide in the same channel of activation and weight within the same layer, as shown in Figure 1 (Left). Leveraging this complementarity property, CSB facilitates the redistribution of extreme magnitudes between activations and weights to minimize the overall channel salience. Concretely, we introduce Salience Balancing Matrices, derived from the statistical properties of activation and weight distributions, to channel-wise transform both activations and weights. This transformation achieves equilibrium in their salient channels, effectively mitigating the quantization difficulty of the balanced distributions.Recognizing the variability in activations over different timesteps, we further extend the concept of channel salience along the temporal dimension and propose Spearman's ρ-guided Salience Calibration (SSC). This method refines the Salience Balancing Matrices to comprehensively evaluate activation salience over timesteps, with more emphasis on timesteps where the complementarity between salient activation and weight channels is more significant. Furthermore, we design a re-parameterization scheme that can offline absorb these Salience Balancing Matrices into adjacent layers, thus avoiding additional computation overhead at the inference stage.While the performance of mainstream PTQ methods degrades on DiTs, our PTQ4DiT achieves comparable performance to the FP counterpart with 8-bit weight and activation (W8A8). In addition, PTQ4DiT can generate high-quality images with further reduced weight precision at 4-bit (W4A8). To the best of our knowledge, PTQ4DiT is the first method for effective DiT quantization. 
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 c ∈ R din is converted to scale and shift parameters (γ, β ∈ R din ), which are regressed through MLPs then injected into the noised latent Z ∈ R n×din via adaLN:EQUATIONwhere LN(•) 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:EQUATIONwhere x denotes the full-precision tensor, ⌊•⌉ is the round-to-nearest operator [32] , and the clamp function restricts the quantized value within the range of [0, 2 b -1]. Here, δ and λ 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.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: See Appendix F. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
diffusion transformer quantization challenges: 
 Diffusion Transformers (DiTs) diverge from conventional generative or discriminative models [39, 11] through their unique design. Specifically, DiTs are constructed with a series of large transformer blocks and operate under a multi-timestep paradigm to progressively transform pure noise into images.Our analysis reveals complex distribution patterns and temporal dynamics in the inference process of DiTs, identifying two primary challenges that prevent effective DiT quantization.❶ Pronounced Quantization Error in Salient Channels. The first challenge lies in systematic quantization errors in DiT's linear layers. As shown in Figure 3 , activation and weight channels with significantly high absolute values are prone to substantial errors after quantization. We term these as salient channels, characterized by extreme values that greatly exceed the typical range of magnitudes. Upon uniform quantization (Eq. ( 2)), it is often necessary to truncate these extreme values in order to maintain the precision of the broader set of standard channels. This compromise can result in notable deviations from the original full-precision distribution as the sampling process proceeds, especially given DiT's layered architecture and repetitive inference paradigm.❷ Temporal Variation in Salient Activation. Another challenge of DiT quantization arises from temporal variations in the magnitudes of salient activation channels. Rather than static inputs, DiTs operate across a sequence of timesteps to generate high-quality images from random noise. Consequently, activation distributions can vary drastically within the inference process, which is particularly evident in salient channels that dominate the signal. Figure 4 demonstrates that the distribution of maximal absolute values in activation channels exhibits significant variations over different timesteps. This temporal variability introduces a non-trivial difficulty to quantization optimization: Quantization parameters effective for salient activation channels at one timestep may not be suitable at other timesteps. Such discrepancies can exacerbate quantization errors, cumulatively impairing the generation quality. Therefore, for accurate quantization, it is imperative to capture the evolving trait of salient channels throughout the entire denoising procedure.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] • The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
ptq4dit: 
 To overcome the identified challenges, we propose Channel-wise Salience Balancing (CSB) and Spearman's ρ-guided Salience Calibration (SSC) in our PTQ4DiT in Sections 4.1 and 4.2, respectively. Subsequently, we devise a re-parameterization scheme in Section 4.3, eliminating extra computational demands of PTQ4DiT during inference while maintaining the mathematical equivalence.A linear layer f (•; W) within MHSA and PF modules typically takes a token sequence X ∈ R n×din as input and performs linear transformation with its weight matrixW ∈ R din×dout , formulated as f (X; W) = X • 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 X and the weight matrix 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 B X and B W to adjust the channel-wise distribution of activation and weight, as expressed by:EQUATIONTo address the quantization difficulties, we need to achieve balanced distributions in X and W, which requires B X and 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:EQUATIONHere, j is the channel index. Consequently, the balanced salience 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:EQUATION)Building on these concepts, we proceed to construct the Salience Balancing Matrices, which modulate the salience of activations and weights with the guidance of s:EQUATIONEQUATIONFollowing 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( X j ) = s( W j ) = s(X j , W j ), thereby alleviating the quantization difficulties, as demonstrated by the reduction in overall channel salience:EQUATIONHere, we characterize the overall salience s o of activations or weights using the maximum salience across channels, e.g., s o (X) = max(s(X 1 ), s(X 2 ), . . . , s(X din )), 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 X (1:T ) = (X (1) , X (2) , . . . , 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):s(X (t) ) = (s(X (t) 1 ), s(X (t)2 ), . . . , s(XEQUATION) While s(W) remains consistent, we find that {s(X (t) )} T t=1 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 ρ-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 s(X (t) ) and weight salience s(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 ρ-calibrated Temporal Salience s ρ by selectively aggregating the activation salience along timesteps:EQUATIONwhere weighting factors {η t } T t=1 are derived from a normalized exponential form of inverse Spearman's ρ statistic [47, 55, 56] :EQUATIONHere, ρ(•, •) computes the correlation between two sequences, and η t serves as the weighting factor for activation salience at timestep t. In this method, η t inversely reflects the correlation coefficient ρ(s(X (t) ), s(W)), thereby prioritizing timesteps where there is a higher degree of complementarity in salience between activations and weights. Subsequently, we utilize s ρ for activation salience in Eqs. ( 5), (6) , and ( 7), yielding refined Salience Balancing Matrices, denoted as B X ρ and B W ρ . 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 B X ρ and B 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 B X ρ and B W ρ are mutual inverses, this incorporation maintains mathematical equivalence to the original linear layer f :EQUATION)The proof is provided in Appendix C. Furthermore, we design a re-parameterization scheme for DiTs, allowing for obtaining X and W without extra computational burden during inference. Specifically,RepQ* (W4A8) Q-Diffusion (W4A8) PTQ4DiT (W4A8)Figure 5 : Random samples generated by PTQ4DiT and two strong baselines: RepQ* [21] and Q-Diffusion [18] , with W4A8 quantization on ImageNet 512×512 and 256×256. Our method can produce high-quality images with finer details. Appendix E presents more visualization results.we update the weight matrix of linear layer f to W offline and seamlessly integrate B 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 B X ρ by adjusting the scale and shift parameters (γ, β ∈ R din ) within adaLN:EQUATIONEquivalently, we fuse B 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] . 
experimental settings: 
 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×256 and 512×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×256, and 5,000 for ImageNet 512×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 √ 2 quantizers specifically for DiTs.Tables 1 and 2 report the outcomes on large-scale class-conditional image generation for ImageNet 256×256 and 512×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×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×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 B X and B W are estimated based on distributions at the midpoint timestep 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 . ρ is absorbed into the weight matrix of the linear layer f . Meanwhile, B X ρ is integrated offline into the MLPs layer prior to adaLN modules for Projection1 and FC1, and into the preceding matrix multiplication operation for Projection2.The Multi-Head Self-Attention (MHSA) and Pointwise Feedforward (PF) modules are essential for processing input tokens and conditional information in DiT Blocks [37] . As depicted in Figure 7 , we incorporate our Channel-wise Salience Balancing (CSB) and Spearman's ρ-guided Salience Calibration (SSC) techniques into the linear layers within both modules. These techniques are designed to mitigate the quantization difficulties by dynamically adjusting the salience of activations and weights via Salience Balancing Matrices. Through the adjustments, CSB and SSC allow for more uniform distributions of activation and weight magnitudes across salient channels, reducing the impact of extreme values and enhancing the overall stability of the quantization process.To eliminate additional computational demands during inference, the Salience Balancing Matrices, B W ρ and B X ρ , are pre-integrated into the DiT Blocks. Specifically, we replace the weight matrix of the linear layer f with W = B W ρ W and integrate B X ρ into the preceding linear transformations. For Projection1 and FC1, B X ρ is absorbed into the MLPs before the adaLN modules, while for Projection2, it can be absorbed within the matrix multiplication [12, 53, 43] . Derivations of the integration are provided in Appendix D.This section provides a comprehensive description of the PTQ4DiT pipeline, detailed in Algorithm 1. The PTQ4DiT is designed to enhance the performance of quantized DiTs by addressing quantization challenges through sophisticated salience estimation and balancing strategies.The full algorithm mainly consists of five steps: ❶ The pipeline begins with estimating activation and weight salience for the pre-trained model using a calibration dataset. ❷ Following the estimation, we employ Spearman's ρ-guided Salience Calibration to compute correlation coefficients between activation salience and weight salience, which helps determine the weighting factors for each timestep. These factors are crucial for computing a temporally adjusted salience, which aims to minimize quantization errors that typically occur due to misalignment in salience peaks across the timesteps.❸ The Channel-wise Salience Balancing step follows, wherein Salience Balancing Matrices are constructed to redistribute the activation and weight values channel-wise. Specifically, for each 4: Estimate activation salience s(X (t) ) at each timestep t ▷ Using Eq. ( 9) 5: Estimate weight salience s(W) ▷ Using Eq. ( 9)6: ❷ Spearman's ρ-guided Salience Calibration: , W j ) for each channel j ▷ Using Eqs. ( 5), ( 14) 12: Construct refined Salience Balancing Matrices B X ρ and B W ρ ▷ Using Eqs. ( 6), ( 7)13: ❹ Re-Parameterization:14: Integrate B W ρ into the weight matrix of linear layers offline▷ By W = B W ρ W 15: Integrate B Xρ into the MLPs before adaLN modules offline ▷ Using Eqs. ( 13), ( 20)16: ❺ Quantization:17: Obtain X = XB X ρ and W = B W ρ W without extra computational demand during inference 18: Perform quantization on balanced activation X and weight W channel j, the balanced salience s ρ (X (1:T ) j , W j ) is given by:EQUATIONwhere s ρ (X(1:T ) j) is the j-th element of s ρ (X (1:T ) ). Then, we formulate the refined Salience Balancing Matrices B X ρ and B W ρ based on s ρ (XEQUATION), as detailed in Eq. ( 15). This step is pivotal in aligning the activation and weight distributions, thereby minimizing the overall quantization difficulty. ❹ In the Re-Parameterization phase, these balancing matrices are integrated into the pre-trained model, ensuring that no additional computational cost is required during inference. This integration maintains computational efficiency while retaining the benefits of our salience balancing technique. ❺ Finally, we perform quantization on the model with balanced activations and weights, setting the stage for the deployment of efficient and effective quantized DiTs in resource-constrained environments.In this section, we provide detailed proof demonstrating that our PTQ4DiT maintains mathematical equivalence to the original linear layers. This proof ensures that the balancing operation does not alter the original computational outcomes of the full-precision models.In PTQ4DiT, we introduce the Salience Balancing Matrices B X ρ and B W ρ , which are diagonal matrices intended to balance the salience across activation and weight channels. We verify the inverse relationship of B X ρ and B W ρ mentioned in Section 4.3:B X ρ • B W ρ = diag( s ρ (X (1:T ) 1 , W 1 ) s ρ (X (1:T ) 1 ) , . . . , s ρ (X (1:T ) din , W din ) s ρ (X (1:T ) din ) ) • diag( s ρ (X (1:T ) 1 , W 1 ) s(W 1 ), . . . , s ρ (X(1:T ) din , W din ) s(W din ) ) = diag( s ρ (X (1:T ) 1 , W 1 )s ρ (X(1:T ) 1 ) • s ρ (X (1:T ) 1 , W 1 ) s(W 1 ), . . . , s ρ (X(1:T ) din , W din ) s ρ (X (1:T ) din ) • s ρ (X (1:T ) din , W din ) s(W din ) ) = diag( (s ρ (X (1:T ) 1 ) • s(W 1 )) 1 2 •2 s ρ (X (1:T ) 1 ) • s(W 1 ), . . . , (s ρ (XEQUATIONwhere I denotes the identity matrix. Therefore, we can derive the mathematical equivalence:EQUATIONThis section details the integration of the Salience Balancing Matrix B X ρ into the MLPs before the adaptive Layer Norm (adaLN) modules [37] , aimed at eliminating extra computational overhead at the inference stage. Recall that the initial formulation adaLN on the input latent Z ∈ R n×din is given by:EQUATIONwhere γ, β ∈ R din are scale and shift parameters, respectively, regressed by MLPs based on the conditional input c ∈ R din :EQUATIONHere, W γ , W β are weight matrices, and b γ , b β are bias terms. In PTQ4DiT, X is obtained by applying B X ρ to the output of adaLN as follows:EQUATIONwhich echos with Eq. ( 13). To avoid additional matrix multiplications in γB X ρ and βB X ρ , we can pre-absorb B X ρ in MLPs' weights and biases offline, expressed as:EQUATIONThus, the re-parameterized MLPs can directly produce the adjusted scale and shift parameters:EQUATIONThis allows for obtaining X without extra computational burden at the inference stage.Figures 8 and 9 supplement visualization results of our PTQ4DiT on W8A8 quantization, compared with baseline PTQ methods and the full-precision (FP) counterpart, on ImageNet 512×512 and 256×256. Our method generates results that closely mirror those of the FP models, presenting finer details and richer semantic content than the baseline approaches.This work introduces a pioneering solution facilitating the broad deployment of Diffusion Transformers (DiTs) through Post-training Quantization. Our method substantially lowers computational and memory demands, thereby improving the accessibility of DiTs. Currently, our research concentrates on visual generation. For future work, we plan to extend our methodology to other generative models across various modalities, such as audio and 3D. However, there remains an inherent risk that these generative models could be utilized to produce disinformation. While our study contributes to the widespread application of DiTs, it does not address such ethical risks. We recognize the importance of developing safeguards and encourage further research into strategies that can prevent the misuse of these powerful generative technologies. NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All external assets such as datasets and code libraries used in the research are properly credited, and their licenses are respected and clearly mentioned.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: The paper does not release new assets. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: All reviewers agree on the quality of this paper, I recommend acceptance

Reviews: review: paper_summary: paper_summary: This paper proposes a new PTQ method for DiTs. It develops a Channel-wise Salient Balancing method to suppress the outliers of linear layers in transformer blocks when applying activation quantization. Besides, it designs the Spearmen’s ρ-guided Salience Calibration to tackle the timestep dimension’s variety. It can improve the performance of the quantized models, comparing to directly applying the conventional quantization methods for Unet-based diffusion models, ViTs or traditional convolution networks to DiTs.
strength_weakness: strength_weakness: -Lacks of experiments. This paper does not report the W8A8 results of ImageNet 512x512. Besides, it only conducts the experiment on one DiT model. The proposed method should be validated on more architectures.
-Lack of novelty. The idea of Re-Parameterization activation range has already been used in PTQ of ViT and LLM, such as Rep-Q [1] or outlier suppression+ [2]. 
-Lack of direct evidence. It lacks direct evidence to display the effectiveness of the proposed method, such as the visualization of activation before and after CSB and SSC to verify it actually suppress the salience channel. 
-Lack of motivation. This paper proposes several modules, but there is no direct evidence to prove that these modules can solve the problems. For example, in RepQ’s, it chooses the mean value of activation range to scale the distribution, but this paper uses the equilibrium between weight and activation without explanation. Besides, it also not explains why using inverse Spear-man’s ρ statistic as the weigh.
[1]Li Z, Xiao J, Yang L, et al. Repq-vit: Scale reparameterization for post-training quantization of vision transformers. In ICCV, 2023.
[2]Wei X, Zhang Y, Li Y, et al. Outlier Suppression+: Accurate quantization of large language models by equivalent and effective shifting and scaling. In EMNLP, 2023.
questions: questions: 1. Please see the weaknesses.
2. Did the authors quantize the post-softmax layer? 
4. The SOTA DiT models can achieve a FID around 3 on ImageNet, why are the results reported in this paper not aligned with the origin paper? Why did the authors choose different settings?
limitations: limitations: N/A
paper_summary: paper_summary: This paper proposes the first post-training quantization (PTQ) method for Diffusion Transformer (DiT). It addresses the presence of salient channels with extreme magnitudes and the temporal variability in the distributions of salient activations over multiple timesteps. Experimental results demonstrate comparable performance in low-bit scenarios.
strength_weakness: strength_weakness: 1. In line 158, the author proposes an important property: large values do not coincide in the same channels of activation and weight. This property is crucial for the proposed method but is only briefly illustrated by a sketch. We believe that such an important property should be demonstrated with statistical experiments.
2. In line 170, the author attempts to use the geometric mean to balance activation and weight channels. Is this approach well-designed? Could another type of mean be used?
3. PTQ4DiT is designed for Diffusion Transformers and is tested on traditional DiT. However, many new DiT-based models, such as pixart-alpha [1], pixart-sigma [2], SD3 [3], and Lumina [4], have emerged. To verify the generality of the method, it should be tested on a broader range of DiT-based models.
4. This method appears to work only with linear layers.

[1] PixArt-$\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, ICLR24.
[2] PixArt-\Sigma: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation, arXiv.
[3] Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, arXiv.
[4] Lumina-T2X: Transforming Text into Any Modality, Resolution, and Duration via Flow-based Large Diffusion Transformers, arXiv.
questions: questions: Please refer to the weaknesses part.
limitations: limitations: Yes. The authors have addressed the limitations and broader impacts.
paper_summary: paper_summary: - The paper introduces PTQ4DiT, a post-training quantization method for Diffusion Transformers. This approach can facilitate the widespread deployment of DiTs. By investigating the distribution of activation and weight of DiTs, PTQ4DiT designs the Channel-wise Balancing and timestep-aware Salience Calibration. Through these, PTQ4DiT effectively quantizes DiTs to W4A8 and reduce computational costs while maintaining image generation quality.
strength_weakness: strength_weakness: - **Effectiveness of Salience Calibration:** The methodology sections 4.1 and 4.3 of the paper  resembles existing work on quantization [1]. The primary novel contribution of this paper is the 'Salience Calibration' technique, which addresses the specific challenge of activation timestep-variance in DiTs. It is critical to verify its importance. However, the evidence supporting the effectiveness of the Salience Calibration (SSC) is currently confined to Table 3, which presents results for ImageNet with 256x256 resolution, using a W4A8 configuration. The improvement observed with SSC is only moderate when compared to the CSB method. To more effectively emphasize the paper's unique contribution, additional evidence demonstrating the effectiveness of CSB would be beneficial."
- **Overheads for Salience Calibration:** The 'Salience Calibration' technique introduces a timestep-wise correction coefficient to refine the activation $S(X(t))$ and weight $S(W)$ distributions, resulting in the adjusted binary matrices $B_{\rho}^x$ and $B_{\rho}^w$. This refinement implies that $B_{\rho}^w$ may differ for each timestep t. However, the weights are reparameterized offline prior to quantization. Introducing a time-varying $B_{\rho}^w$ could lead to the generation of distinct quantized integer (INT) weights for each timestep, which may complicate efforts to reduce the model's weight memory cost. This could involve either storing multiple sets of INT weights or repeatedly offloading weights for each timestep.

[1] Xiao, Guangxuan, et al. "Smoothquant: Accurate and efficient post-training quantization for large language models." International Conference on Machine Learning. PMLR, 2023
questions: questions: - Is the 'Spearman’s ρ-guided Salience Calibration (SSC)' method capable of generalizing across a different timesteps and solvers? It seems that the calibration process depends on the number of timesteps. Does this imply that recalibration is necessary when employing SSC with different models, timesteps, or solvers?
- In the context of reparameterization, the paper specifically addresses the 'Post-adaLN' and 'Post-Matrix-Multiplication' scenarios. However, within the Feed-Forward Network (FFN) layers of DiTs, there are instances where two consecutive linear layers do not conform to either of these cases. It raises the question of how to approach reparameterization in such circumstances."
- Some unclear details about the quantization process:
  - Are all layers in the DiTs quantized? e.g., time embedding linear layers, attention matmuls. 
  - In Line 237, the authors state that " the optimization of quantization parameters follows the implementation of Q-Diffusion". Does this mean that the quantization process involves gradient-based optimization of scaling factors, and Adaround of zero-points? What is the cost of the PTQ process. 
- The statement “the lower the correlation between activation salience s(X(t)) and weight salience s(W), the greater the reduction effect in overall channel salience”  in Sec. 4.2 should be more clearly explained. What is the definition of "correlation salience"?
limitations: limitations: - The authors have discussed the limitations in the appendix.
paper_summary: paper_summary: This paper presents PTQ4DiT, a quantization method designed for diffusion transformers. The method focuses on addressing quantization challenges due to extreme magnitudes in salient channels and the temporal variability of activations across multiple timesteps. To combat these issues, it incorporates techniques like Channel-wise Salience Balancing (CSB) and Spearmen’s ρ-guided Salience Calibration (SSC). These strategies help in redistributing magnitudes to minimize quantization errors and adapt dynamically across different timesteps, significantly enhancing the performance of quantized DiTs without re-training the original models.
strength_weakness: strength_weakness: - The classifier-free guidance scales used for sampling are not specified, which could impact the reproducibility and evaluation of the model's performance.
- Under the W4A8 setting, the method exhibits significant degradation in performance, suggesting limitations in its effectiveness at lower bit-widths.
questions: questions: see weaknesses
limitations: limitations: This work reasonably discussed the limitations and future work

Comments: title: Thank you for your response
Sincere Thanks for Constructive Review and Support
Summary of Authors' Responses
Thanks for your response
Thank you for your response
comment: Dear Reviewer USPH,

We deeply appreciate your response and are grateful that you found our rebuttal beneficial. Thank you once more for your constructive comments.
Thank you for the clarification, which addressed most of my concerns. I keep my initial rating.
Dear Reviewers,

We would like to express our gratitude for your recognition of our work and your constructive feedback. Your insightful suggestions have significantly contributed to refining our research and enhancing its overall impact.

We particularly appreciate the time and effort you invested in the review process. We look forward to any further feedback you may have and remain committed to addressing any additional concerns. Thank you.

Sincerely,

The Authors
Dear Reviewers,

We sincerely appreciate the time and effort you have dedicated to reviewing our paper. We have carefully considered all the points raised and would like to summarize our responses as follows:

**(I) Additional experiments.** We conducted three additional experiments: (1) W8A8 ImageNet 512x512 generation, (2) PTQ4DiT on other DiT models, and (3) evaluation using 50K samples. These experiments further validate our method's effectiveness and generality.

**(II) Novelty and motivation.** We highlighted the unique innovations and contributions of our PTQ4DiT and clarified the motivation of the equilibrium and inverse ρ.

**(III) Direct evidence.** We supplemented two statistical experiments: (1) the Jaccard similarity of salient activation and weight channels, which provides direct evidence of the important property, and (2) the reduction in the percentage of salient channels when applying CSB and SSC, supporting their effectiveness.

**(IV) Methodology and implementation.** We clarified the quantization of Post-softmax layers. Moreover, we justified the use of the geometric mean and discussed the extension of PTQ4DiT to other structures besides the linear layers.

***Thank you again for your valuable feedback. We would like to address any remaining concerns for further enhancement of our work.***
Dear Reviewer jrKh,

Thank you for your thorough review and constructive suggestion. We are grateful for your acknowledgment of our efforts to address the concerns. Your expertise has significantly contributed to the enhancement of our work.
Dear Reviewer a8C3,

We sincerely appreciate your prompt response and are grateful that you found our rebuttal beneficial. Thank you once more for your valuable feedback in enhancing our submission.
After reading the rebuttal and the comments from other reviewers, I keep my initial rating.
Thank the authors for the clarification, most of my concerns are addressed. I keep my scoring as acceptance.