TerDiT: Ternary Diffusion Models with Transformers

• The technical contribution of this paper appears to be merely the introduction of an extra RMSNorm. The CUDA kernel deployment is from previous work, and the model architecture follows the previous Large-DiT. Therefore, I rise concerns about the technical contribution of this paper.

• The models architecture evaluated in this work are somewhat limited. For example, previous work on diffusion model quantization included various diffusion architectures, such as LDM, StableDiffusion, etc. However, this paper, which focus on the DiT quantization, only considers the Large-DiT architecture. Considering the diversity of diffusion transformers, such as VDT[1], evaluating only one architecture weakens the contribution of this paper. Additionally, this also raises the question of whether the large activation values problems is a common phenomenon for diffusion transformers?

• In Table 2, the authors provide a comparison of the efficiency, and it can be seen that TerDiT achieves lower checkpoint size and GPU memory usage. However, the proposed method does not significantly increase the inference time compared to the full-precision model. I have noted the authors' explanation for this, but in practical applications, inference speed is a very important factor we care about. The proposed method does not truly improve inference speed, which weakens the significance of the quantization of DiT.

• Why adding RMSNorm can alleviate large activation values? The paper only provides empirical results; a more in-depth theoretical or visualization analysis is necessary for understanding the working mechanism.

• In Table 1, why are the performance of TerDiT-600M-G is only given under the CFG settings?

[1] VDT: General-purpose Video Diffusion Transformers via Mask Modeling. ICLR24

• The evaluation on ImageNet is incomplete. Specifically, the authors consider two resolutions (256 and 512, corresponding to the original DiT) and two model sizes (600M param XL/2 and a sized-up 4.2B parameter version), but their application of the two models is inconsistent across experiments, e.g., the 600M model is only used for 256x256 for CFG in Table 1.
• Table 1 results are decently convincing but not great, e.g., L337 TerDiT-600M loses substantially to the original DiT-XL/2 (almost double FID, almost half Inception Score) and at 512 resolution TerDiT-4.2-G is not much better than the much smaller DiT-XL/2 despite being much larger.
• For all the trashing of PTQ said in L072, there is inadequate comparison to PTQ methods. Yes, PTQ methods like Q-DiT [3] and Q-Diffusion [4] fail at 2-bit precision (really anything below 4-bits), this is not a novel finding, especially for ViT-like architectures [5]. The trade-off between QAT and PTQ has always been that the former allows lower bit precision at the cost of much higher computational resources.
• No activation quantization. The 1.58-bit LLM paper [1] tried for 8-bit activations in addition to Ternary weights. Moreover, even if PTQ methods are limited to ~4 bit weights, these methods generally attempt activation quantization to at least 8-bits. This makes a comparison crucial as it could even the odds in terms of inference cost (e.g., BitOPs, latency, etc.) compared to this QAT method.
• There is at least one factual inaccuracy in the writing at lines 203-205: "AdaLN is an important component in the DiT model (Peebles & Xie, 2023) and has been proven more effective than cross-attention and in-context conditioning methods." this is absolutely not true and must be corrected. While the original DiT model removes Cross-Attention for AdaLN, it is very limited and restricted to ImageNet generation. T2I DiTs like the PixArt [6, 7] and Hunyuan [8] re-introduce cross-attention while SD3/3.5 and Flux [9] concat the text and image sequences together which is similar.
• L260 - 263: I find this line of reasoning to not be borne out by the experimental results and ablations, e.g., Figs 5/6. In Fig 5 the loss of the DiT without RMS Norm is higher initially but quickly drops to be comparable. Fig 6 shows RMS Norm achieves faster FID reduction, however, FID itself is not entirely a trustworthy metric on its own [10].

1. The paper lacks a clear motivation for applying a ternary quantizer to DiT. Given effective precedents for binary DM and 2-bit DM, I believe evidence is needed to demonstrate that a ternary DM can significantly outperform binary DM in accuracy or achieve substantial efficiency gains over 2-bit DM in practical use.

2. The issue of increased activations does not necessarily seem to be an inherent problem of ternarization. For example, better initialization of α in Equation 2 might straightforwardly address the phenomena shown in Figure 3.

3. The comparison with EfficientDM appears unfair. EfficientDM’s QAT training is based on LoRA, which has a natural disadvantage in precision. For a more compelling comparison with QAT methods, alternatives like LSQ and Q-DM [1] (including noise-estimating mimicking) would be more appropriate.

4. Table 1 compares only parameter counts but does not list the actual size, memory, or OPs.

5. The phrase "significantly reduced" in the caption of Table 2 lacks factual basis, as practical acceleration gains from weight-only binary quantizers are generally limited. The extent to which a carefully designed ternary quantizer could deliver such gains remains uncertain.

6. There is a lack of comparative data for using RMSNorm across the full model. In the Remark section at the end of Section 3.3, the manuscript states, "Comparative experiments are introduced in Sec.4.2," but I could not find any supporting experiments anywhere in the manuscript.

[1] Li, et al. Q-dm: An efficient low-bit quantized diffusion model.

1. The method lacks in novelty and the main contribution of the paper is in utilizing the ideas in [1] for replacing the attention and feedforward blocks and [2] for introducing RMSNorm for stable training process.
2. The results in Figure 6 are surprising since the training losses with and without RMSNorm converge to similar values but the FID performance shows a drastic difference. Could the authors comment further on this.
3. In Figure 6, could the authors also please provide the results with classifier free guidance if it was already not performed.
4. Ln 303: The authors may have missed out to point the different between the work and [3]. Could the authors please highlight the major differences with [1].

[1] Shuming Ma, Hongyu Wang, Lingxiao Ma, Lei Wang, Wenhui Wang, Shaohan Huang, Li Dong, Ruiping Wang, Jilong Xue, and Furu Wei. The era of 1-bit llms: All large language models are in 1.58 bits. arXiv preprint arXiv:2402.17764, 2024b

[2] Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F. and Rodriguez, A., 2023. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.