论文状态:已完成

Q-DiT: Accurate Post-Training Quantization for Diffusion Transformers

发表:2024/06/25
原文链接PDF 下载
价格:0.100000
价格:0.100000
已有 3 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文提出了一种新方法Q-DiT,用于精确量化Diffusion Transformers(DiTs),针对其在权重和激活中的空间及时间方差问题。通过自动量化粒度分配和样本级动态激活量化,Q-DiT显著降低模型计算成本,同时在ImageNet上实现了高保真度的图像和视频生成,建立了新基准。

摘要

Recent advancements in diffusion models, particularly the architectural transformation from UNet-based models to Diffusion Transformers (DiTs), significantly improve the quality and scalability of image and video generation. However, despite their impressive capabilities, the substantial computational costs of these large-scale models pose significant challenges for real-world deployment. Post-Training Quantization (PTQ) emerges as a promising solution, enabling model compression and accelerated inference for pretrained models, without the costly retraining. However, research on DiT quantization remains sparse, and existing PTQ frameworks, primarily designed for traditional diffusion models, tend to suffer from biased quantization, leading to notable performance degradation. In this work, we identify that DiTs typically exhibit significant spatial variance in both weights and activations, along with temporal variance in activations. To address these issues, we propose Q-DiT, a novel approach that seamlessly integrates two key techniques: automatic quantization granularity allocation to handle the significant variance of weights and activations across input channels, and sample-wise dynamic activation quantization to adaptively capture activation changes across both timesteps and samples. Extensive experiments conducted on ImageNet and VBench demonstrate the effectiveness of the proposed Q-DiT. Specifically, when quantizing DiT-XL/2 to W6A8 on ImageNet (256×256256 \times 256), Q-DiT achieves a remarkable reduction in FID by 1.09 compared to the baseline. Under the more challenging W4A8 setting, it maintains high fidelity in image and video generation, establishing a new benchmark for efficient, high-quality quantization in DiTs. Code is available at \href{https://github.com/Juanerx/Q-DiT}{https://github.com/Juanerx/Q-DiT}.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

Q-DiT: Diffusion Transformers 的精确训练后量化 (Q-DiT: Accurate Post-Training Quantization for Diffusion Transformers)

1.2. 作者

Lei Chen, Yuan Meng, Chen Tang, Xinzhu Ma, Jingyan Jiang, Xin Wang, Zhi Wang, Wenwu Zhu。 作者主要来自清华大学 (Tsinghua University)、香港中文大学多媒体实验室 (MMLab, CUHK) 和深圳技术大学 (Shenzhen Technology University)。

1.3. 发表期刊/会议

预印本 (Preprint)。论文发布于 arXiv,发布时间为 2024-06-25T07:57:27.000Z。arXiv 是一个开放获取的预印本服务器,允许研究人员在正式同行评审和出版之前分享他们的工作。尽管是预印本,但它在学术界具有重要的影响力,常用于快速传播最新研究成果。

1.4. 发表年份

2024年。

1.5. 摘要

这篇论文研究了扩散模型 (diffusion models) 的后训练量化 (Post-Training Quantization, PTQ) 问题,特别是针对架构已从基于 UNet 模型转向 Diffusion Transformers (DiTs) 的最新进展。DiTs 在图像和视频生成质量及可扩展性方面取得了显著提升,但其高计算成本对实际部署构成了挑战。PTQ 是一种无需昂贵再训练即可实现模型压缩和加速推理的解决方案。然而,针对 DiT 的量化研究仍很稀疏,且现有主要为传统扩散模型设计的 PTQ 框架存在量化偏差,导致性能显著下降。

本文识别出 DiTs 在权重 (weights) 和激活 (activations) 中通常表现出显著的空间方差 (spatial variance),并且在激活中还存在时间方差 (temporal variance)。为解决这些问题,论文提出了 Q-DiT,一种无缝集成两种关键技术的新方法:

  1. 自动量化粒度分配 (automatic quantization granularity allocation):用于处理权重和激活在输入通道 (input channels) 之间存在的显著方差。

  2. 样本级动态激活量化 (sample-wise dynamic activation quantization):用于自适应地捕捉激活在不同时间步 (timesteps) 和样本 (samples) 之间的变化。

    ImageNetVBench 上进行的广泛实验证明了所提出的 Q-DiT 的有效性。具体而言,在 ImageNet (256×256256 \times 256) 上将 DiT-XL/2 量化到 W6A8 (权重6比特,激活8比特) 时,Q-DiT 相较于基线模型实现了 FID 1.09的显著降低。在更具挑战性的 W4A8 设置下,它依然在图像和视频生成中保持了高保真度,为 DiTs 的高效高质量量化建立了新基准。

1.6. 原文链接

原文链接: https://arxiv.org/abs/2406.17343v2 PDF 链接: https://arxiv.org/pdf/2406.17343v2.pdf 发布状态: 预印本

2. 整体概括

2.1. 研究背景与动机

2.1.1. 研究背景

扩散模型 (diffusion models) 已经成为图像和视频生成等多种任务的强大基础模型。随着其发展,核心架构从传统的 UNet [28] 转向了 Diffusion Transformers (DiTs) [27]。DiTs,例如 Stable Diffusion 3 [8] 和 Sora [4],在复杂的生成任务中展现出卓越的性能和可扩展性。这种架构的转变,通过将卷积操作替换为自注意力 (self-attention) 机制,使得模型能够更好地处理长距离依赖关系,并实现更好的可扩展性。

2.1.2. 核心问题与动机

尽管 DiTs 取得了巨大成功,但其固有的高计算成本是一个显著的限制。生成过程通常需要大量的迭代去噪步骤,这使得实时或大规模应用在计算上非常密集,导致推理延迟高。为了解决这一问题,模型量化 (model quantization) 成为一个有前景的方向,其中 后训练量化 (Post-Training Quantization, PTQ) 尤为吸引人,因为它无需耗时的模型再训练即可实现模型压缩和加速推理。

然而,针对 Transformer 架构的扩散模型 (DiTs) 的量化研究仍然有限。现有针对扩散模型的量化方法 [12, 13, 30] 主要集中在 UNet 架构上,并且大多依赖于基于重建 (reconstruction-based) 的方法。这些方法在应用于大型 DiTs 时面临可扩展性挑战 [21, 25],且往往会导致量化偏差 (biased quantization) 和显著的性能下降。

2.1.3. 论文的切入点与创新思路

本文的切入点在于深入探索 DiT 模型的独特特性,并针对这些特性设计定制化的量化方法。论文识别出 DiT 量化中的两个关键问题:

  1. 权重和激活在输入通道上存在显著方差 (significant variance of weights and activations across input channels)

  2. 激活在不同时间步上存在显著分布漂移 (significant distribution shift of activations across timesteps)

    针对这些挑战,论文提出了 Q-DiT,一个包含细粒度组量化 (fine-grained group quantization)动态激活量化 (dynamic activation quantization) 策略的框架,以期在 DiT 模型上实现高精度量化。

2.2. 核心贡献/主要发现

本文的主要贡献概括如下:

  • 提出了 Q-DiT 方法: 引入了一种针对 DiTs 的精确 PTQ 方法。该方法采用细粒度组量化 (fine-grained group quantization) 来有效管理权重和激活在输入通道上的方差,并采用样本级动态激活量化 (sample-wise dynamic activation quantization) 来适应不同时间步和样本之间的激活变化。
  • 优化了量化组大小分配: 识别出默认组大小配置的次优性,并提出了一种进化搜索策略 (evolutionary search strategy) 来优化组大小分配,从而提高量化过程的效率和效果。
  • 卓越的实验性能:ImageNetVBench 上进行了广泛实验,结果表明 Q-DiTW6A8 配置下实现了近乎无损的压缩,并在 W4A8 这一更具挑战性的设置下,在图像和视频生成中保持了最小的性能下降,展示了其优越的性能,并为 DiTs 的高效高质量量化设立了新基准。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 扩散模型 (Diffusion Models)

概念定义: 扩散模型是一类生成模型,它通过一个被称为“扩散过程”的逐步噪声添加过程,将数据(如图像)逐渐转化为纯噪声。然后,模型学习一个反向的“去噪过程”,从噪声中逐步恢复出清晰的数据。这个去噪过程通常是一个迭代过程,模型在每个时间步预测并移除噪声,直到生成高质量的数据。 在本文中的重要性: DiTs 是扩散模型的一种先进架构,因此理解扩散模型的基本工作原理是理解本文的基础。扩散模型的迭代去噪特性也导致了激活在不同时间步上的动态变化,这是本文量化方法需要解决的核心问题之一。

3.1.2. UNet

概念定义: UNet [28] 是一种经典的编码器-解码器 (encoder-decoder) 架构,最初为生物医学图像分割而设计。它以其 U 形结构而闻名,包含一个下采样路径(编码器)和一个上采样路径(解码器),并通过跳跃连接 (skip connections) 将编码器中的特征图直接传递给解码器中相应层,以保留空间信息。 在本文中的重要性: UNet 是传统扩散模型中广泛使用的架构。本文指出,现有许多针对扩散模型的量化方法是为 UNet 设计的,而 DiTs 的架构差异使得这些方法不再适用。

3.1.3. Diffusion Transformers (DiTs)

概念定义: Diffusion Transformers (DiTs) [27] 是一种将 Transformer 架构应用于扩散模型的新范式。它用 Transformer 块取代了传统 UNet 架构中的卷积层和下采样/上采样操作,从而在处理图像和视频生成任务时,能够更好地捕捉长距离依赖关系,并展现出更好的可扩展性和生成质量。 在本文中的重要性: DiTs 是本文研究的核心对象。论文指出 DiTs 具有独特的量化挑战,例如显著的空间和时间方差,这要求设计专门的量化策略。

3.1.4. 模型量化 (Model Quantization)

概念定义: 模型量化是一种模型压缩技术,旨在通过将模型中的权重 (weights) 和/或激活 (activations) 从高精度浮点数(如 FP32FP16)转换为低精度整数(如 INT8INT4)来减小模型大小、降低内存占用并加速推理。 在本文中的重要性: 量化是本文提出的 Q-DiT 方法的最终目标,即通过量化 DiT 模型来解决其高计算成本问题。

3.1.4.1. 量化感知训练 (Quantization-Aware Training, QAT)

概念定义: QAT [2, 5, 9] 是一种在模型训练过程中集成量化过程的方法。它通过在前向传播中模拟量化效应,并在反向传播中使用 Straight-Through Estimator (STE) [1] 来近似梯度,从而在训练期间同时优化量化器参数和模型参数。 特点: QAT 通常能恢复量化造成的性能下降,但需要访问原始训练数据集并进行资源密集型的再训练。

3.1.4.2. 训练后量化 (Post-Training Quantization, PTQ)

概念定义: PTQ [21, 25] 是一种在模型已经训练好之后进行量化的方法。它不需要模型再训练,而是利用少量校准数据集 (calibration dataset) 来调整权重和激活的量化参数。 特点: PTQ 效率更高、更实用,尤其适用于大型模型。然而,当进行低比特量化时,PTQ 可能导致显著的性能下降。本文的 Q-DiT 就是一种 PTQ 方法。

3.1.5. 均匀量化 (Uniform Quantization)

概念定义: 均匀量化是一种常用的量化方法,它将浮点数值范围划分为等间距的区间,并将每个区间内的浮点值映射到离散的整数值。 在本文中的重要性: 本文的 Q-DiT 方法采用了均匀量化,因为它对硬件友好 [11, 19]。

3.1.5.1. 缩放因子 (Scaling Factor, ss) 和零点 (Zero Point, ZZ)

概念定义: 在均匀量化中,缩放因子 (s) 决定了浮点值和整数值之间的映射比例,而 零点 (Z) 则表示浮点数 0 在量化后的整数表示。它们共同定义了量化范围和精度。 在本文中的重要性: 这些参数的校准是量化过程的关键,决定了量化误差的大小。本文的动态激活量化就是实时计算这些参数。

3.1.6. 评估指标

3.1.6.1. Fréchet Inception Distance (FID)

概念定义: FID [14] 是一种用于评估生成模型(如 GANs 和扩散模型)生成图像质量和多样性的指标。它通过比较真实图像和生成图像在预训练 Inception-v3 模型中间层特征空间中的分布来计算。具体来说,它假设这些特征服从多元高斯分布,并计算两个高斯分布之间的 Fréchet 距离。FID 值越低,表示生成图像的质量越高,多样性越好,与真实图像分布越接近。 数学公式: FID(x,g)=μxμg22+Tr(Σx+Σg2(ΣxΣg)1/2) \mathrm{FID}(x, g) = ||\mu_x - \mu_g||^2_2 + \mathrm{Tr}(\Sigma_x + \Sigma_g - 2(\Sigma_x \Sigma_g)^{1/2}) 符号解释:

  • xx: 真实图像的特征向量集合。
  • gg: 生成图像的特征向量集合。
  • μx\mu_x: 真实图像特征的协方差矩阵的均值。
  • μg\mu_g: 生成图像特征的协方差矩阵的均值。
  • Σx\Sigma_x: 真实图像特征的协方差矩阵。
  • Σg\Sigma_g: 生成图像特征的协方差矩阵。
  • 22||\cdot||^2_2: 欧几里得距离的平方。
  • Tr()\mathrm{Tr}(\cdot): 矩阵的迹 (trace)。

3.1.6.2. 空间 FID (spatial FID, sFID)

概念定义: sFID [29] 是 FID 的一个变体,旨在更好地评估局部图像质量和一致性。它通过在图像的不同空间尺度和位置计算 FID 来实现,从而捕捉到生成图像在局部细节上的表现。 数学公式: sFID 没有一个统一的、像 FID 那样的标准数学公式,因为它是一个计算多个局部 FID 的聚合指标。它通常涉及以下步骤:

  1. 将图像分割成多个重叠或不重叠的局部块。
  2. 对每个局部块计算其对应的 Inception 特征。
  3. 对这些局部特征计算 FID
  4. 最终的 sFID 可以是这些局部 FID 的平均值或加权平均值。 符号解释:
  • 其基本计算单元与 FID 相同,但应用于图像的局部区域。

3.1.6.3. Inception Score (IS)

概念定义: IS 是一种用于评估生成图像质量和多样性的指标。它利用预训练的 Inception-v3 模型,计算生成图像在分类预测上的特性。高 IS 值表示生成图像:

  1. 清晰可辨 (High Quality):图像内容具有高置信度的单一类别预测(低熵)。
  2. 多样性高 (High Diversity):所有生成图像的类别预测分布均匀(高熵)。 IS 值越高越好。 数学公式: IS(G)=exp(ExG[DKL(p(yx)p(y))]) \mathrm{IS}(G) = \exp \left( E_{x \sim G} \left[ D_{KL}(p(y|x) || p(y)) \right] \right) 符号解释:
  • GG: 生成模型。
  • xGx \sim G: 从生成模型 GG 中采样得到的图像。
  • p(yx)p(y|x): 图像 xxInception 模型中预测的类别条件概率分布。
  • p(y): 所有生成图像的边缘类别概率分布(即 ExG[p(yx)]E_{x \sim G}[p(y|x)])。
  • DKL()D_{KL}(\cdot || \cdot): Kullback-Leibler (KL) 散度,衡量两个概率分布之间的差异。
  • ExG[]E_{x \sim G}[\cdot]: 对从生成模型 GG 采样的所有图像的期望。

3.1.6.4. 精度 (Precision)

概念定义: 在生成模型的评估中,Precision 通常衡量生成图像的“逼真度”或“真实现象覆盖度”。它关注生成图像中有多少是属于真实数据分布的。例如,如果生成了许多非常逼真的图像,但这些图像只覆盖了真实数据分布的一小部分,那么 Precision 会很高,但 Recall 可能会低。 数学公式: Precision 在生成模型评估中没有一个统一的、像 FID 那样的标准数学公式,其具体计算方式可能因论文或工具而异。通常,它涉及到在特征空间中计算生成样本与真实样本之间的距离,以判断生成样本是否落在真实样本的“流形”附近。一种常见的近似计算方式是通过计算生成样本与真实样本的最近邻距离来实现。 符号解释:

  • 其具体计算依赖于所使用的评估库和方法,例如 torchmetricsclean-fid。通常,较高的 Precision 值表示生成样本更“真实”。

3.1.6.5. Fréchet 视频距离 (Fréchet Video Distance, FVD)

概念定义: FVD 类似于 FID,但专门用于评估视频生成模型的质量。它通过比较真实视频和生成视频在预训练视频分类模型(通常是 I3D 或类似模型)中间层特征空间中的分布来计算。它同样假设特征服从多元高斯分布,并计算两个高斯分布之间的 Fréchet 距离。FVD 值越低,表示生成视频的质量越高,时间一致性越好,多样性越好,与真实视频分布越接近。 数学公式: FVD(x,g)=μxμg22+Tr(Σx+Σg2(ΣxΣg)1/2) \mathrm{FVD}(x, g) = ||\mu_x - \mu_g||^2_2 + \mathrm{Tr}(\Sigma_x + \Sigma_g - 2(\Sigma_x \Sigma_g)^{1/2}) 符号解释:

  • xx: 真实视频的特征向量集合。
  • gg: 生成视频的特征向量集合。
  • μx\mu_x: 真实视频特征的均值。
  • μg\mu_g: 生成视频特征的均值。
  • Σx\Sigma_x: 真实视频特征的协方差矩阵。
  • Σg\Sigma_g: 生成视频特征的协方差矩阵。
  • 22||\cdot||^2_2: 欧几里得距离的平方。
  • Tr()\mathrm{Tr}(\cdot): 矩阵的迹 (trace)。

3.2. 前人工作

3.2.1. 模型量化 (Model Quantization)

  • QAT [2, 5, 9]: 将量化过程整合到模型微调 (fine-tuning) 阶段,通过 Straight-Through Estimator (STE) [1] 同时优化量化器参数和模型参数。优点是能恢复量化导致的性能下降,缺点是资源密集,需要原始训练数据集。
  • PTQ [21, 25]: 利用少量校准数据集调整量化参数,无需再训练。高效实用,但低比特量化可能导致显著性能下降。
  • 基于重建的方法 (Reconstruction-based methods) [21, 25]: 旨在通过最小化每层或每个块的重建误差来减少性能下降。在 CNN 中表现良好,但难以扩展到大型模型。

3.2.2. Transformer 的量化 (Quantization of Transformers)

这方面研究广泛,涵盖了 Vision Transformers (ViTs)Large Language Models (LLMs)

  • PTQ4ViT [39]: 提出 twin uniform quantizer 来处理 post-softmaxpost-GELU 激活的特殊分布。
  • RepQ-ViT [22]: 使用 scale reparameterization 来减少激活的量化误差。
  • 针对 LLM 的量化:
    • 仅权重 (Weight-only) 量化:GPTQ [10] 基于近似二阶信息将权重压缩到 4 比特。AWQ [23] 提出 activation-aware weight quantization 来减少显著权重的量化误差。
    • 权重-激活 (Weight-activation) 量化: LLM.int8() [7] 通过将异常值 (outliers) 保留在 FP16 中并进行混合精度计算来减少异常值的影响。Outlier Suppression [35] 通过使用 non-scaling LayerNorm 来减少激活的量化误差。 局限性: 这些技术可能无法直接应用于 DiTs,因为它们缺乏对扩散模型特性的考虑。

3.2.3. 扩散模型的量化 (Quantization of Diffusion Models)

扩散模型因需要大量采样步骤而推理速度慢。

  • PTQ4DM [30] 和 Q-diffusion [20]: 发现激活在不同去噪步骤间的方差,并采用基于重建的方法进行量化。
  • PTQD [13]: 发现量化噪声与模型输出之间的相关性,并提出 variance schedule calibration 来纠正不相关部分。
  • TDQ [31]: 利用 MLP 层估计每个时间步的激活量化参数。
  • TMPQ-DM [33]: 结合时间步序列长度缩减和量化精度选择来降低总成本。
  • PTQ4DiT [36]: 针对 Diffusion TransformersPTQ 方法,通过 Channel-wise Salience Balancing (CSB)Spearman's ρ-guided Salience Calibration (SSC) 解决极端通道幅度和时间激活可变性问题,实现 W4A8 量化。 局限性: 这些方法无法同时处理 Transformer 架构的特性和去噪过程中激活的动态变化,导致性能显著下降。

3.3. 技术演进

扩散模型架构从 UNetDiT 的演进,代表了处理复杂生成任务能力的提升,但同时也带来了计算成本的增加。量化技术也从最初的 QAT 转向更高效的 PTQ,并且从简单的张量级 (tensor-wise) 量化发展到通道级 (channel-wise) 或组级 (group-wise) 量化,以更好地适应数据分布的异质性。对于 Transformer 模型,量化研究已探索了处理激活异常值、权重敏感性等问题。对于扩散模型,研究开始关注激活随时间步变化的动态性。本文的 Q-DiT 正是在这一技术脉络中,结合了 DiT 的架构特点和扩散模型的动态特性,提出了更精细的量化策略。

3.4. 差异化分析

Q-DiT 与现有方法的核心区别和创新点在于:

  1. 针对 DiT 的特异性问题: 现有方法多为 UNet 或通用 Transformer 设计,未能充分解决 DiT 中权重和激活的显著空间方差以及激活在时间步和样本之间的动态变化。Q-DiT 明确识别并针对这些问题设计方案。
  2. 细粒度组量化: Q-DiT 引入了自动量化粒度分配 (automatic quantization granularity allocation),通过进化搜索优化组大小,解决了传统 channel-wise 量化效率低和简单 tensor-wise 量化精度差的问题。这与 LLM 量化中常见的 group quantization [23, 40] 类似,但 Q-DiT 进一步发现了组大小选择的非单调性,并通过进化搜索进行优化。
  3. 样本级动态激活量化: 不同于一些扩散模型量化方法(如 TDQ [31] 预测时间步量化参数,或 PTQ4DM [30] 采用静态时间步量化),Q-DiT 实现了样本级 (sample-wise) 动态激活量化。这意味着量化参数不仅随时间步变化,还随每个样本的实际激活分布实时调整,从而更精确地捕捉 DiT 激活的动态性,显著降低量化误差。
  4. 整体框架协同: Q-DiT 将上述两种技术无缝集成,协同解决 DiT 的多重量化挑战,而非孤立地处理某个问题。这使得它在保持高效性的同时,实现了更高的量化精度。

4. 方法论

本节将详细阐述 Q-DiT 的方法论,包括其观察到的 DiT 量化特性、预备知识以及提出的两个核心技术:自动量化粒度分配样本级动态激活量化

4.1. DiT 量化的观察 (Observations of DiT Quantization)

作者发现直接将最近的基于 UNet 的量化方法应用于 DiTs 会导致显著的性能下降。为了理解其根本原因,作者探索了 DiT 模型的独特特性,特别是它们在权重和激活分布方面与基于 UNet 架构的不同之处。

4.1.1. 观察 1:权重和激活在输入通道上存在显著方差

具体现象: 如原文 Figure 2a 所示,DiT 中权重和激活在输入通道 (input channels) 上的方差远大于输出通道 (output channels) 上的方差。这种显著的方差会严重影响量化效果,因为常见的量化方法通常沿输出通道应用逐通道量化 (channel-wise quantization) [13, 30]。此外,原文 Figure 2b 也显示,激活中存在异常值 (outliers),且这些异常值集中在特定通道中。 影响: 如果继续使用张量级 (tensor-wise) 量化(即整个张量使用一组量化参数),这些异常值将显著影响量化参数的校准,导致对非异常值的量化产生巨大的误差。这使得传统的量化方法对 DiT 模型的适应性较差。

下图(原文 Figure 2)展示了 DiT-XL/2 模型在不同层中的权重和激活分布。红色峰值表示较高的值,蓝色区域表示较低的值。可以明显看出输入通道上的方差更大,且存在异常值。

Figure 2. Distributions of weights and activations in different layers of DiT-XL/2. The red peaks indicate higher values, while the blue areas represent lower values.
该图像是图表,展示了DiT-XL/2中不同层次的权重和激活分布。图中的红色峰值表示较高的值,而蓝色区域则表示较低的值。这对于理解模型的特征分布和量化过程至关重要。

Figure 2. Distributions of weights and activations in different layers of DiT-XL/2. The red peaks indicate higher values, while the blue areas represent lower values.

4.1.2. 观察 2:激活在不同时间步上的显著分布漂移

具体现象: 如原文 Figure 3 和 Figure 4 所示,在去噪过程中,DiT 模型中激活的分布在不同时间步 (timesteps) 之间会发生显著变化。这种时间上的漂移还表现出在不同样本 (samples) 之间也存在显著的可变性(详细实验结果在补充材料中)。 影响: 这意味着如果在特定时间步校准的量化参数,可能无法在所有时间步上很好地泛化,从而导致跨时间步的量化误差增加。

下图(原文 Figure 3)是一个箱线图,展示了 DiT-XL/2 模型在生成一张 ImageNet 256×256256 \times 256 图像时,激活值在不同时间步(从 50 到 0)的分布。可以看出,随着时间步的变化,激活值的分布范围和中位数都有明显变化。

Figure 3. Box plot showing the distribution of activation values across various timesteps (from 50 to 0) for the DiT-XL/2 model when generating one image from ImageNet at \(2 5 6 \\times 2 5 6\) resolution..
该图像是一个箱型图,展示了DiT-XL/2模型在生成一张256 imes 256分辨率的ImageNet图像时,50到0各个时间步的激活值分布情况。

Figure 3. Box plot showing the distribution of activation values across various timesteps (from 50 to 0) for the DiT-XL/2 model when generating one image from ImageNet at 256×2562 5 6 \times 2 5 6 resolution..

下图(原文 Figure 4)展示了 DiT-XL/2 模型在生成一张 ImageNet 256×256256 \times 256 图像时,MLP 和注意力层中激活的标准差在 50 个时间步内的变化情况。这进一步证实了激活分布在时间维度上的不稳定性。

Figure 4. Standard deviation of activations in MLP and attention layers across different blocks over 50 timesteps for DiT-XL/2 when generating one image from ImageNet at \(2 5 6 \\times 2 5 6\) resolution.
该图像是图表,展示了在生成 ImageNet 上一幅图片时,DiT-XL/2 各个块在 50 个时间步骤内的 MLP 和注意力层激活的标准差变化。图中分别标注了 Blok.0 和 Blok.7 的注意力层,以及 Blok.14 和 Blok.26 的 MLP 层的标准差曲线。

Figure 4. Standard deviation of activations in MLP and attention layers across different blocks over 50 timesteps for DiT-XL/2 when generating one image from ImageNet at 256×2562 5 6 \times 2 5 6 resolution.

4.2. 预备知识:均匀量化 (Preliminary: Uniform Quantization)

本文采用均匀量化来量化权重和激活,因为它对硬件更加友好 [11, 19]。

量化函数: 均匀量化函数 QQ 将输入的浮点张量 x\mathbf{x} 量化为 bb 比特整数张量 x^\hat{\mathbf{x}},其表达式如下: x^=Q(x;b)=s(clip(xs+Z,0,2b1)Z) \hat { \mathbf { x } } = Q ( \mathbf { x } ; b ) = s \cdot ( \mathrm { c l i p } ( \lfloor \frac { \mathbf { x } } { s } \rceil + Z , 0 , 2 ^ { b } - 1 ) - Z ) 符号解释:

  • x\mathbf{x}: 输入的浮点张量。
  • bb: 量化比特数。
  • x^\hat{\mathbf{x}}: 量化后的整数张量。
  • ss: 缩放因子 (scaling factor),定义为 s=max(x)min(x)2b1s = \frac{\max(\mathbf{x}) - \min(\mathbf{x})}{2^b - 1}。它将浮点范围映射到 2b12^b - 1 个量化级别。
  • ZZ: 零点 (zero point),定义为 Z=min(x)sZ = - \lfloor \frac{\min(\mathbf{x})}{s} \rceil。它确保浮点数 0 映射到量化后的整数范围内的某个特定值。
  • \lfloor \cdot \rceil: 四舍五入操作。
  • clip(,min_val,max_val)\mathrm{clip}(\cdot, \text{min\_val}, \text{max\_val}): 裁剪操作,将值限制在 [min_val,max_val][\text{min\_val}, \text{max\_val}] 范围内。在这里,它将量化后的整数值裁剪到 [0,2b1][0, 2^b - 1] 的有效比特范围内。

4.3. 方法:Q-DiT

如原文 Figure 1 所示,Q-DiT 包含了两个新颖的设计:自动量化粒度分配 (Automatic Quantization Granularity Allocation)动态激活量化 (Dynamic Activation Quantization)

下图(原文 Figure 1)是一个示意图,展示了论文中提出的 Q-DiT 量化方法,包括离线的自动量化粒度分配和在线的样本动态激活量化模块。

该图像是一个示意图,展示了论文中提出的Q-DiT量化方法,包括离线的自动量化粒度分配和在线的样本动态激活量化模块,右上角给出了样本动态量化的计算公式 \(s_{i,t} = \\frac{\\max(x_{i,t}) - \\min(x_{i,t})}{2^b - 1}\),\(Z_{i,t} = \\left\\lfloor \\frac{\\min(x_{i,t})}{s_{i,t}} \\right\\rfloor\)。
该图像是一个示意图,展示了论文中提出的Q-DiT量化方法,包括离线的自动量化粒度分配和在线的样本动态激活量化模块,右上角给出了样本动态量化的计算公式 si,t=max(xi,t)min(xi,t)2b1s_{i,t} = \frac{\max(x_{i,t}) - \min(x_{i,t})}{2^b - 1}Zi,t=min(xi,t)si,tZ_{i,t} = \left\lfloor \frac{\min(x_{i,t})}{s_{i,t}} \right\rfloor

t

4.3.1. 自动量化粒度分配 (Automatic Quantization Granularity Allocation)

4.3.1.1. 基本解决方案与挑战 (Base Solution and Challenges)

针对观察 1 中指出的输入通道级方差问题,一个直接的解决方案是应用输入通道级量化,即为每个通道使用不同的量化参数。然而,这种方法在硬件部署时会损害计算效率,因为它需要重复的中间重新缩放 (intermediate rescaling) [3, 38],从而无法充分利用低精度计算的优势。

4.3.1.2. 细粒度组量化 (Fine-grained Group Quantization)

为了在输入通道级量化和张量级量化之间取得折衷,本文采用了组量化 (group quantization) [23, 40]。如原文 Figure 1 所示,权重和激活矩阵被分成多个组,然后对每个组分别进行量化。

矩阵乘法中的组量化: 考虑线性层中的矩阵乘法 Y=XW\mathbf{Y} = \mathbf{X}\mathbf{W},其中 XRn×din\mathbf{X} \in \mathbb{R}^{n \times d_{\mathrm{in}}} 是输入激活,WRdin×dout\mathbf{W} \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{out}}} 是权重矩阵。量化后每个输出元素 Y^i,j\hat{\mathbf{Y}}_{i,j} 可以通过以下方式获得: Y^i,j=k=0dinX^i,kW^k,j =u=0din/gll1v=0gll1QuX(Xi,ugll+v)QuW(Wugll+v,j) \begin{array}{l} { \displaystyle { \hat { \mathbf { Y } } _ { i , j } } = \sum _ { k = 0 } ^ { d _ { \mathrm { i n } } } { \hat { \mathbf { X } } _ { i , k } \hat { \mathbf { W } } _ { k , j } } } \\ { \displaystyle ~ = \sum _ { u = 0 } ^ { d _ { \mathrm { i n } } / g _ { l l } - 1 } \sum _ { v = 0 } ^ { g _ { l l } - 1 } Q _ { u } ^ { \mathbf { X } } ( \mathbf { X } _ { i , u g _ { l l } + v } ) Q _ { u } ^ { \mathbf { W } } ( \mathbf { W } _ { u g _ { l l } + v , j } ) } \end{array} 符号解释:

  • Y\mathbf{Y}: 原始的输出矩阵。

  • Y^i,j\hat{\mathbf{Y}}_{i,j}: 量化后的输出矩阵中的第 (i, j) 个元素。

  • X\mathbf{X}: 输入激活矩阵。

  • W\mathbf{W}: 权重矩阵。

  • X^i,k\hat{\mathbf{X}}_{i,k}: 量化后的输入激活矩阵中的第 (i, k) 个元素。

  • W^k,j\hat{\mathbf{W}}_{k,j}: 量化后的权重矩阵中的第 (k, j) 个元素。

  • dind_{\mathrm{in}}: 输入维度。

  • doutd_{\mathrm{out}}: 输出维度。

  • gllg_{ll}: 组大小 (group size),表示每个量化组包含的输入通道数量。

  • QuX()Q_u^{\mathbf{X}}(\cdot)QuW()Q_u^{\mathbf{W}}(\cdot): 分别是对第 uu 个组的激活和权重进行量化操作的函数,它们使用各自组的量化参数。

  • uu: 组的索引,范围从 0din/gll1d_{\mathrm{in}}/g_{ll} - 1

  • vv: 组内元素的索引,范围从 0gll1g_{ll} - 1

    通过对每个组独立量化,可以更精细地处理通道间的方差,同时相较于纯粹的逐通道量化,保留了一定的计算并行性。

4.3.1.3. 量化组选择的非单调性 (Non-monotonicity in Quantization Group Selection)

理想情况下,减小组大小(即增加组的数量)应该能减少量化误差,从而提高模型性能。然而,原文 Table 1 揭示了一个重要的观察:组大小减小并不总是能带来更好的结果。例如,当组大小从 128 减小到 96 时,FID 增加了约 11.8%,从 17.87 增加到 19.97,表明生成图像质量下降。这表明 DiT 量化与 LLMViT 量化不同,存在非单调性。这意味着存在一个最优的组大小配置,可以在相同的平均组大小下实现更好的量化效果,或者在更大的平均组大小下实现相同的量化结果。此外,模型中不同层的敏感度 (sensitivity) 不同,通过为不同层分配不同的组大小,可以实现高性能和高效率的平衡。

4.3.1.4. 自动组分配 (Automatic Group Allocation)

挑战: 自动分配组大小的主要挑战在于如何确定每个层的组大小与扩散模型最终生成性能之间的相关性。传统的混合精度量化工作 [21, 34] 侧重于识别每层的敏感度指标(如量化层与全精度层之间的 MSE),然后将其转化为整数线性规划 (ILP) 问题进行优化。然而,作者发现 MSE 的减小并不一定对应于 DiT 量化中性能下降的减少,这表明传统方法可能无效。

解决方案: 为了解决上述挑战,本文直接使用 FID 作为图像生成的评估指标,或者使用 FVD 作为视频生成的评估指标。

目标函数: 对于图像生成: L(g)=FID(R,Gg) L ( \mathbf { g } ) = \mathrm { FID } ( R , G _ { \mathbf { g } } ) 对于视频生成: L(g)=FVD(R,Gg) L ( \mathbf { g } ) = \mathrm { FVD } ( R , G _ { \mathbf { g } } ) 符号解释:

  • g={g1,g2,,gN}\mathbf{g} = \{g_1, g_2, \ldots, g_N\}: 层级组大小配置 (layer-wise group size configuration),其中 NN 是被量化层的数量。每个 gig_i 代表第 ii 层的组大小。
  • RR: 真实样本 (real samples)。
  • GgG_{\mathbf{g}}: 由使用特定组大小配置 g\mathbf{g} 量化后的模型生成的样本。
  • FID(R,Gg)\mathrm{FID}(R, G_{\mathbf{g}}): 真实样本和生成样本之间的 Fréchet Inception Distance
  • FVD(R,Gg)\mathrm{FVD}(R, G_{\mathbf{g}}): 真实视频样本和生成视频样本之间的 Fréchet Video Distance

优化问题: 本文采用进化算法 (evolutionary algorithm) 来优化以下目标函数: g=argmingL(g),s.t.B(g)Nbitops { \bf g } ^ { * } = \underset { { \bf g } } { \arg \operatorname* { m i n } } { \cal L } ( { \bf g } ) , \mathrm { s . t . } { \cal B } ( { \bf g } ) \leq N _ { b i t o p s } 符号解释:

  • g\mathbf{g}^*: 最优的层级组大小配置。

  • L(g)\mathcal{L}(\mathbf{g}): 基于 FIDFVD 的损失函数,如上所述。

  • B(g)\mathcal{B}(\mathbf{g}): 衡量比特操作 (BitOps) 的函数,表示在给定组大小配置 g\mathbf{g} 下的计算成本。

  • NbitopsN_{bitops}: 预定义的比特操作阈值,作为计算约束。

    这种方法能够更好地捕捉组大小对量化性能的细微影响,从而在效率和图像质量方面取得改进。

算法 1:进化搜索算法 (Evolutionary Search Algorithm) 下面是用于自动组分配的进化算法流程: 输入:

  • Group size search space Sg\mathcal{S}_g: 组大小的搜索空间,例如 {32,64,128,192,288}\{32, 64, 128, 192, 288\}
  • number of layers LL: 模型中可量化层的数量。
  • population size NpN_p: 进化算法中的种群大小。
  • iterations NiterN_{iter}: 进化算法的迭代次数。
  • mutation probability pp: 变异操作的概率。
  • Constraint NbitopsN_{bitops}: 比特操作的计算约束阈值。

算法步骤:

  1. 初始化种群 (Initialize population) P\mathcal{P}
    • PP 被初始化为 NpN_p 个随机生成的组大小配置 gj\mathbf{g}^j,其中每个配置 gjRL\mathbf{g}^j \in \mathbb{R}^L 的元素都从 search space Sg\mathcal{S}_g 中随机选择。
  2. 初始化 TopK 候选集 STopKS_{\mathrm{TopK}}
    • STopKS_{\mathrm{TopK}} 初始化为空集。
  3. 开始进化迭代 (Start evolutionary iterations):
    • 对于 t=1,2,,Nitert = 1, 2, \ldots, N_{iter} 执行以下步骤:
      • 评估种群 (Evaluate population):
        • 对于 PP 中的每个配置 gj\mathbf{g}^j,根据公式 (4) 或 (5) 计算其 FID (或 FVD) 值。
      • 更新 TopK 候选集 (Update STopKS_{\mathrm{TopK}}):
        • 根据 FID (或 FVD) 分数对配置进行排名,并将表现最好的 KK 个配置更新到 STopKS_{\mathrm{TopK}} 中。
      • 清空当前种群 (Clear current population) P=\mathcal{P} = \emptyset
      • 交叉操作 (Crossover operation):
        • 重复以下步骤直到 PP 中的配置数量达到 Np/2N_p/2
          • 以概率 1-pSTopKS_{\mathrm{TopK}} 中选择父代进行交叉 (Crossover) 操作,生成新的子代配置 gcross=Crossover(STopK)\mathbf{g}_{cross} = \mathrm{Crossover}(S_{\mathrm{TopK}})
          • 如果新生成的 gcross\mathbf{g}_{cross} 满足比特操作约束 B(gcross)<NbitopsB(\mathbf{g}_{cross}) < N_{bitops},则将其添加到 PP 中。
      • 变异操作 (Mutation operation):
        • 重复以下步骤直到 PP 中的配置数量达到 NpN_p
          • 以概率 ppSTopKS_{\mathrm{TopK}} 中选择父代进行变异 (Mutate) 操作,生成新的子代配置 gmutate=Mutate(STopK)\mathbf{g}_{mutate} = \mathrm{Mutate}(S_{\mathrm{TopK}})
          • 如果新生成的 gmutate\mathbf{g}_{mutate} 满足比特操作约束 B(gmutate)<NbitopsB(\mathbf{g}_{mutate}) < N_{bitops},则将其添加到 PP 中。
  4. 获取最佳配置 (Get the best configuration):
    • 从最终的种群中选择具有最佳 FID (或 FVD) 分数的组大小配置 gbest\mathbf{g}^{\mathrm{best}}

    • 使用 gbest\mathbf{g}^{\mathrm{best}} 来量化模型。

      输出: 量化后的模型。

4.3.2. 样本级动态激活量化 (Sample-wise Dynamic Activation Quantization)

4.3.2.1. 现有方法的局限性 (Limitations of Existing Methods)

在基于 UNet 的扩散模型量化中,以往的研究 [12, 31] 要么为每个时间步的所有激活分配一组量化参数,要么设计一个多层感知机 (MLP) [31] 来根据时间步预测量化参数。然而,根据观察 2(激活在不同时间步和样本之间存在显著分布漂移),这些方法与本文的细粒度量化不兼容。具体来说,如果在每个时间步为每个组分配量化参数,将导致巨大的内存开销。例如,对于一个有 50 个时间步的采样器,内存开销可能达到全精度模型大小的 39%。

4.3.2.2. 动态量化方法 (Dynamic Quantization Approach)

LLM 优化最新工作 [37] 的启发,本文设计了一种即时 (on-the-fly) 动态量化方法用于激活。具体来说,在推理过程中,每个激活组的量化参数是根据其最小值 (min)最大值 (max) 实时计算的。

量化参数的计算: 对于给定时间步 tt 的样本 ii,激活 xi,t\mathbf{x}_{i,t} 的量化参数可以表示为: si,t=max(xi,t)min(xi,t)2b1,Zi,t=min(xi,t)si,t. \begin{array}{l} { s _ { i , t } = \frac { \operatorname* { m a x } ( \mathbf x _ { i , t } ) - \operatorname* { m i n } ( \mathbf x _ { i , t } ) } { 2 ^ { b } - 1 } , } \\ { Z _ { i , t } = - \left\lfloor \frac { \operatorname* { m i n } ( \mathbf x _ { i , t } ) } { s _ { i , t } } \right\rceil . } \end{array} 符号解释:

  • si,ts_{i,t}: 针对样本 ii 在时间步 tt 时计算的激活缩放因子。

  • Zi,tZ_{i,t}: 针对样本 ii 在时间步 tt 时计算的激活零点。

  • xi,t\mathbf{x}_{i,t}: 样本 ii 在时间步 tt 的激活张量。

  • max(xi,t)\max(\mathbf{x}_{i,t}): xi,t\mathbf{x}_{i,t} 中的最大值。

  • min(xi,t)\min(\mathbf{x}_{i,t}): xi,t\mathbf{x}_{i,t} 中的最小值。

  • bb: 激活的量化比特数。

  • 2b12^b - 1: 对于 bb 比特量化,可表示的非负整数值的最大值。

  • \lfloor \cdot \rceil: 四舍五入操作。

    效率集成: 此外,这种动态量化与 min-max 计算被整合到前置操作符 (prior operator) 中,可以受益于操作符融合 (operator fusion)。与 Transformer 块中昂贵的矩阵乘法相比,其开销变得可以忽略不计。

5. 实验设置

本节详细介绍 Q-DiT 在图像和视频生成任务上的实验设置,包括数据集、评估指标、对比基线以及其他量化参数。

5.1. 数据集

5.1.1. 图像生成 (Image Generation)

  • 数据集: ImageNet
  • 来源与特点: ImageNet 是一个大规模图像数据集,包含数百万张图像,涵盖数千个类别,常用于图像分类、目标检测和图像生成等任务的模型训练和评估。
  • 使用方式: 评估 Q-DiTImageNet 上生成 256×256256 \times 256512×512512 \times 512 分辨率图像的能力。

5.1.2. 视频生成 (Video Generation)

  • 数据集: VBench 提示词套件 [18]。
  • 来源与特点: VBench 是一个综合性的视频生成模型基准测试套件,包含一系列精心设计的提示词 (prompts),用于评估视频生成模型在 16 个不同维度上的性能,如主题一致性、整体一致性、时间风格、外观风格等。
  • 使用方式:VBench 提示词套件中的每个提示词,采样 5 个 2 秒的视频进行评估。

5.2. 评估指标

除了在 3.1.6 节中已详细解释的 FIDsFIDISPrecision (用于图像生成) 以及 FVD (用于视频生成) 外,VBench 还提供了 16 个视频生成质量维度。由于 VBench 的 16 个维度没有标准的数学公式,其评估通常基于对生成视频的特征提取和比较,或采用人类评估/自动化评估的组合,因此在此仅列举其关注点而非公式。

VBench 评估维度 (用于视频生成):

  1. Subject Consistency (主体一致性)

  2. Overall Consistency (整体一致性)

  3. Temporal Style (时间风格)

  4. Appearance Style (外观风格)

  5. Scene (场景)

  6. Spatial Relationship (空间关系)

  7. Color (颜色)

  8. Human Action (人类动作)

  9. Multiple Objects (多物体)

  10. Object Class (物体类别)

  11. Imaging Quality (成像质量)

  12. Aesthetic Quality (美学质量)

  13. Dynamic Degree (动态程度)

  14. Motion Smoothness (运动平滑度)

  15. Temporal Flickering (时间闪烁)

  16. Background Consistency (背景一致性)

    这些维度涵盖了视频生成质量的各个方面,包括内容准确性、视觉风格、动态表现和时间连贯性。

5.3. 对比基线 (Baselines)

本文将 Q-DiT 与五个强劲的基线模型进行比较:

  1. PTQ4DM [30]: 一种专门为基于 UNet 的扩散模型设计的 PTQ 方法,主要关注激活的校准。
  2. RepQ-ViT [22]: 一种为 Vision Transformers (ViTs) 量化开发的技术,旨在减少 Transformer 激活中的量化误差。
  3. TFMQ-DM [17]: 一个专门为扩散模型开发的 PTQ 框架,旨在在量化过程中保留时间特征 (temporal features)。
  4. PTQ4DiT [36]: 一种为 Diffusion Transformers (DiTs) 定制的 PTQ 方法,通过 Channel-wise Salience Balancing (CSB)Spearman's ρ-guided Salience Calibration (SSC) 解决极端通道幅度和时间激活可变性问题。
  5. G4W+P4AG4W+P4A 本文构建的一个鲁棒基线,用于视频和图像生成任务。它结合了 GPTQ [10] 进行权重 (weight) 量化,并使用 PTQ4DM 进行激活 (activation) 量化。

5.4. 其他设置

  • 基线模型: 使用预训练的 DiT-XL/2 模型,图像分辨率为 256×256256 \times 256512×512512 \times 512,并将其转换为 FP16 作为全精度 (full precision) 基线模型。
  • 采样器: 采用 DDIM 采样器 [32],采样步数为 50 或 100。
  • 分类器无条件引导 (Classifier-Free Guidance, CFG): 评估在有和无 CFG [15] 情况下的性能。cfg 参数表示 classifier-free guidance scale
  • 图像生成样本数量: 在每种设置下,ImageNet 256×256256 \times 256ImageNet 512×512512 \times 512 分别采样 10K 张图像。
  • 视频生成设置: 采用 30 步的 rectified flow schedulercfg 规模为 7.0。
  • 量化类型: Q-DiT 对权重和激活都采用非对称量化 (asymmetric quantization)。
  • 权重精确度: 使用 GPTQ [10] 进行权重量化。
  • 默认组大小: 默认组大小为 128。
  • 组大小搜索空间: 进化搜索算法的组大小搜索空间 Sg\mathcal{S}_g{32,64,128,192,288}\{32, 64, 128, 192, 288\}
  • 权重与激活组大小关系: 同一层内的权重和激活的组大小必须相同。

6. 实验结果与分析

本节将深入分析 Q-DiT 在图像和视频生成任务上的实验结果,并进行消融研究,以评估所提出组件的有效性。

6.1. 核心结果分析

6.1.1. 图像生成结果

以下是原文 Table 2 的结果,展示了 DiT-XL/2ImageNet 256×256256 \times 256512×512512 \times 512 分辨率上的量化结果。W/A 表示权重和激活的比特宽度。

以下是原文 Table 2 的结果:

Model Bit-width (W/A) Method Size (MB) Metrics
FID ↓ sFID ↓ IS ↑ Precision ↑
DiT-XL/2 256×256 (steps = 100) 16/16 FP 1349 12.40 19.11 116.68 0.6605
6/8 PTQ4DM 508 17.86 25.33 92.24 0.6077
RepQ-ViT 508 27.74 20.91 63.41 0.5600
TFMQ-DM 508 22.33 27.44 72.74 0.5869
PTQ4DiT 508 15.21 21.34 105.03 0.6440
G4W+P4A 520 16.72 24.61 100.09 0.6123
Ours 518 12.21 18.48 117.75 0.6631
4/8 PTQ4DM 339 213.66 85.11 3.26 0.0839
RepQ-ViT 339 224.14 81.24 3.25 0.0373
TFMQ-DM 339 143.47 61.09 5.61 0.0497
PTQ4DiT 339 28.90 34.56 65.73 0.4931
G4W+P4A 351 25.48 25.57 73.46 0.5392
Ours 347 15.76 19.84 98.78 0.6395
DiT-XL/2 256×256 (steps = 100 cfg = 1.5) 16/16 FP 1349 5.31 17.61 245.85 0.8077
6/8 PTQ4DM 508 8.41 25.56 196.73 0.7622
RepQ-ViT 508 10.77 18.53 163.11 0.7264
TFMQ-DM 508 8.87 23.52 194.08 0.7737
PTQ4DiT 508 5.34 18.48 209.90 0.8047
G4W+P4A 520 6.41 19.52 225.48 0.7705
Ours 518 5.32 17.40 243.95 0.8044
4/8 PTQ4DM 339 215.68 86.63 3.24 0.0741
RepQ-ViT 339 226.60 77.93 3.61 0.0337
TFMQ-DM 339 141.90 56.01 6.24 0.0439
PTQ4DiT 339 7.75 22.01 190.38 0.7292
G4W+P4A 351 7.66 20.76 193.76 0.7261
Ours 347 6.40 18.60 211.72 0.7609
DiT-XL/2 512×512 (steps = 50) 16/16 FP 1349 16.01 20.50 97.79 0.7481
6/8 PTQ4DM 508 21.22 22.01 80.07 0.7131
RepQ-ViT 508 19.67 20.11 75.78 0.7082
TFMQ-DM 508 20.99 22.35 71.08 0.6918
PTQ4DiT 508 19.42 21.94 77.35 0.7024
G4W+P4A 520 19.55 22.43 85.56 0.7158
Ours 517 16.21 20.41 96.78 0.7478
4/8 PTQ4DM 339 131.66 75.79 11.54 0.1847
RepQ-ViT 339 105.32 65.63 18.01 0.2504
TFMQ-DM 339 80.70 59.34 29.61 0.2805
PTQ4DiT 339 35.82 28.92 48.62 0.5864
G4W+P4A 351 26.58 24.14 70.24 0.6655
Ours 348 21.59 22.26 81.80 0.7076
DiT-XL/2 512×512 (steps = 50 cfg == 1.5) 16/16 FP 1349 6.27 18.45 204.47 0.8343
6/8 PTQ4DM 508 9.84 26.57 164.91 0.8215
RepQ-ViT 508 8.30 19.19 158.80 0.8153
TFMQ-DM 508 8.34 17.94 162.16 0.8262
PTQ4DiT 508 7.69 18.86 178.34 0.8121
G4W+P4A 520 7.28 19.62 185.92 0.8143
Ours 517 6.24 18.36 202.48 0.8341
4/8 PTQ4DM 339 88.45 50.80 26.79 0.3206
RepQ-ViT 339 79.69 49.76 29.46 0.3413
TFMQ-DM 339 54.61 44.27 58.77 0.4215
PTQ4DiT 339 11.69 22.86 117.34 0.7121
G4W+P4A 351 9.98 20.76 156.07 0.7840
Ours 347 7.82 19.60 174.18 0.8127

分析:

  • W6A8 配置下的性能:

    • ImageNet 256×256256 \times 256 (100步) 设置下,全精度 (FP) 模型的 FID 为 12.40。基线方法如 PTQ4DMRepQ-ViTTFMQ-DMG4W+P4AG4W+P4A 都表现出显著的性能下降,FID 值显著高于 FP 模型。PTQ4DiT 表现相对较好,FID 为 15.21。
    • Q-DiTW6A8 设置下表现出色,FID 达到 12.21,甚至略优于全精度模型 (12.40),并且 IS 达到 117.75 (高于 FP 模型的 116.68)。这表明 Q-DiTW6A8 配置下实现了近乎无损的压缩效果。
    • cfg=1.5cfg = 1.5 的更高质量生成设置下,Q-DiTFID (5.32) 也非常接近甚至略低于 FP 模型 (5.31),IS 也非常接近。
    • ImageNet 512×512512 \times 512 (50步) 高分辨率设置下,Q-DiTFID (16.21) 也与 FP 模型 (16.01) 相当接近,远优于所有基线方法。
  • W4A8 配置下的性能:

    • 在更具挑战性的 W4A8 设置下,所有基线方法的性能急剧下降。例如,在 ImageNet 256×256256 \times 256 (100步) 设置下,PTQ4DMFID 高达 213.66,IS 降至 3.26,表明生成质量严重受损。PTQ4DiTG4W+P4AG4W+P4AW4A8 下表现相对较好,但 FID 仍有明显增加。
    • Q-DiTW4A8 下依然保持了高保真度。在 ImageNet 256×256256 \times 256 (100步) 设置下,Q-DiTFID 为 15.76,IS 为 98.78,显著优于所有基线方法。相较于 FP 模型,其 FID 仅增加了 3.36 (15.76 - 12.40)。
    • cfg=1.5cfg = 1.5W4A8 设置下,Q-DiTFID (6.40) 显著优于 PTQ4DiT (7.75) 和 G4W+P4AG4W+P4A (7.66)。
    • 在高分辨率 ImageNet 512×512512 \times 512 (50步) 的 W4A8 设置下,Q-DiTFID (21.59) 仍然远低于基线,显示了其在严格量化约束下的鲁棒性。
  • 模型大小 (Size): Q-DiTW6A8 下模型大小为 518MB,在 W4A8 下为 347MB,与大多数基线方法(508MB 和 339MB)相当,表明其压缩效率与现有方法类似,但性能更优。

  • CFG 影响: 在启用 classifier-free guidance (cfg=1.5cfg = 1.5) 时,生成质量普遍提高 (FP 模型的 FID 从 12.40 降至 5.31)。在这种情况下,Q-DiT 依然能保持与 FP 模型非常接近的性能,体现了其在不同生成模式下的适应性。

    下图(原文 Figure 5)提供了定性结果,展示了 G4W+P4AG4W+P4A(一个基线模型)和 Q-DiTW4A8 设置下于 ImageNet 256×256256 \times 256ImageNet 512×512512 \times 512 上生成的样本。

    Figure 5. Qualititive results. Samples generated by \(\\mathrm { G 4 W + P 4 A }\) (one of our baselines) and Q-DiT with W4A8 on ImageNet \(2 5 6 \\times 2 5 6\) and ImageNet \(5 1 2 \\times 5 1 2\) .For each example (a-e), the image generated by \(\\mathrm { G 4 W + P 4 A }\) shows notable artifacts and distortions. In contrast, our method produces cleaner and more realistic images, with better preservation of textures. 该图像是插图,展示了使用Q-DiT和基线方法G4W+P4A生成的图像质量对比。左侧列出G4W+P4A生成的样本,包括“城堡”、“雏菊”、“哈士奇”、“泡泡”等图像,右侧为Q-DiT生成的对应样本,展示了更清晰、真实的画面,特别是在“北极狐”的生成效果上,表现出了更好的纹理保留和细节表现。

Figure 5. Qualititive results. Samples generated by G4W+P4A\mathrm { G 4 W + P 4 A } (one of our baselines) and Q-DiT with W4A8 on ImageNet 256×2562 5 6 \times 2 5 6 and ImageNet 512×5125 1 2 \times 5 1 2 .For each example (a-e), the image generated by G4W+P4A\mathrm { G 4 W + P 4 A } shows notable artifacts and distortions. In contrast, our method produces cleaner and more realistic images, with better preservation of textures.

定性分析: 从 Figure 5 可以看出,G4W+P4AG4W+P4A 生成的图像存在明显的伪影和失真,例如 (a) 中的模糊纹理,(b) 中的不自然颜色,(c) 和 (d) 中的形状扭曲。相比之下,Q-DiT 生成的图像更清晰、更真实,纹理保留更好,更接近全精度模型的视觉效果。这进一步证实了 Q-DiT 在保持生成图像质量方面的优越性。

6.1.2. 视频生成结果

以下是原文 Table 3 的结果,展示了 Q-DiT 在视频生成任务上的表现。

以下是原文 Table 3 的结果:

Method Bit-width (W/A) VBench Metrics (Higher is Better)
Subject Consistency Overall Consistency Temporal Style Appearance Style Scene Spatial Relationship Color Human Action
FP 16/16 0.9522 0.2667 0.2507 0.2352 0.4094 0.3441 0.7864 0.8680
G4W+P4A 4/8 0.9444 0.2628 0.2489 0.2344 0.3924 0.3265 0.7657 0.8600
Ours 4/8 0.9498 0.2663 0.2511 0.2346 0.3871 0.3810 0.7947 0.8620
  <table>
  <thead>
  <tr>
  <th>Method</th>
  <th>Bit-width (W/A)</th>
  <th>Multiple Objects</th>
  <th>Object Class</th>
  <th>Imaging Quality</th>
  <th>Aesthetic Quality</th>
  <th>Dynamic Degree</th>
  <th>Motion Smoothness</th>
  <th>Temporal Flickering</th>
  <th>Background Consistency</th>
  </tr>
  </thead>
  <tbody>
  <tr>
  <td>FP</td>
  <td>16/16</td>
  <td>0.4143</td>
  <td>0.8383</td>
  <td>0.5829</td>
  <td>0.5173</td>
  <td>0.6139</td>
  <td>0.9855</td>
  <td>0.9917</td>
  <td>0.9678</td>
  </tr>
  <tr>
  <td>G4W+P4A</td>
  <td>4/8</td>
  <td>0.3540</td>
  <td>0.8225</td>
  <td>0.5730</td>
  <td>0.5018</td>
  <td>0.5639</td>
  <td>0.9849</td>
  <td>0.9895</td>
  <td>0.9651</td>
  </tr>
  <tr>
  <td><b>Ours</b></td>
  <td>4/8</td>
  <td><b>0.3904</b></td>
  <td><b>0.8475</b></td>
  <td><b>0.5812</b></td>
  <td><b>0.5160</b></td>
  <td><b>0.6167</b></td>
  <td><b>0.9859</b></td>
  <td><b>0.9915</b></td>
  <td><b>0.9687</b></td>
  </tr>
  </tbody>
  </table></div>

  </td>
</tr>

分析:

  • 在严格的 W4A8 量化设置下,Q-DiTVBench 的 16 项指标中,有 15 项表现优于基线 G4W+P4AG4W+P4A

  • Q-DiTSubject Consistency (主体一致性)、Overall Consistency (整体一致性)、Temporal Style (时间风格)、Appearance Style (外观风格)、Color (颜色)、Human Action (人类动作) 等多个方面都与全精度 (FP) 模型非常接近,甚至在 Temporal StyleColor 上略有超越。

  • 特别是在衡量视频生成关键指标如 Multiple Objects (多物体)、Object Class (物体类别)、Imaging Quality (成像质量)、Aesthetic Quality (美学质量) 和 Dynamic Degree (动态程度) 上,Q-DiT 相比 G4W+P4AG4W+P4A 有显著提升,且与 FP 模型表现相当。

  • Motion Smoothness (运动平滑度)、Temporal Flickering (时间闪烁) 和 Background Consistency (背景一致性) 等时间连贯性指标上,Q-DiT 同样保持了与 FP 模型相近的卓越性能,优于 G4W+P4AG4W+P4A

  • 唯一一项 Q-DiT 略低于 G4W+P4AG4W+P4A 的指标是 Scene,但差距非常小。

    这些结果表明,Q-DiT 在视频生成方面同样表现出色,能够有效保持视频质量和视频条件一致性,即使在极低比特的量化设置下,也能实现与全精度模型相当的性能。

6.2. 消融实验 (Ablation Studies)

为了评估 Q-DiT 中每个组件的有效性,作者在 ImageNet 256×256256 \times 256 上使用 DiT-XL/2 模型进行了消融研究。采样步数设置为 100,classifier-free guidance scale 设置为 1.5。

6.2.1. Q-DiT 各组件的增量分析

以下是原文 Table 4 的结果,展示了在 W4A8 设置下,Q-DiT 中各个独立组件的增量分析。

以下是原文 Table 4 的结果:

FID ↓ sFID ↓ IS ↑
FP (W16A16) 5.31 17.61 245.85
W4A8 RTN 225.50 88.54 2.96
+ Group size 128 13.77 27.41 146.93
+ Sample-wise Dynamic activation quantization 6.64 19.29 211.27
+ Automatic quantization granularity allocation 6.40 18.60 211.72

分析:

  • RTN 基线 (W4A8 RTN): 最基本的 round-to-nearest (四舍五入) 量化方法在 W4A8 配置下表现极差,FID 高达 225.50,IS 仅为 2.96,表明生成图像质量极低。这强调了对低比特量化需要更复杂的策略。

  • 添加组大小 128 (+ Group size 128):RTN 的基础上引入组大小为 128 的量化,性能得到显著提升,FID 降至 13.77,IS 升至 146.93。这验证了组量化在处理权重和激活空间方差方面的有效性。

  • 添加样本级动态激活量化 (+ Sample-wise Dynamic activation quantization): 进一步引入样本级动态激活量化,生成质量再次大幅提升,FID 降至 6.64,sFID 降至 19.29,IS 升至 211.27。这证明了动态激活量化在适应激活随时间步和样本变化的分布漂移方面的关键作用。

  • 添加自动量化粒度分配 (+ Automatic quantization granularity allocation): 最后,通过引入自动量化粒度分配(即进化搜索)来优化组大小,性能进一步提升,FID 降至 6.40,sFID 降至 18.60,IS 升至 211.72。这个结果已经非常接近全精度模型 (FP) 的 FID (5.31) 和 IS (245.85),表明进化搜索能够找到更优的组大小配置,从而最大限度地减少量化损失。

    这项增量分析清晰地表明,Q-DiT 的每个组件都对最终的卓越性能做出了重要贡献,其中样本级动态激活量化带来了最大的性能提升。

6.2.2. 动态激活量化方法的比较

以下是原文 Table 5 的结果,比较了不同动态激活量化方法在 W16A8 设置下的性能。TFMQ-DM 是一种时间步感知激活量化方法,而 Q-DiT 是同时时间步和样本感知的。

以下是原文 Table 5 的结果:

Method FID ↓ sFID ↓ IS ↑ Precision ↑
FP (W16A16) 5.31 17.61 245.85 0.8077
TFMQ-DM 7.74 19.23 204.56 0.7765
Ours 5.34 17.44 245.24 0.8048

分析:

  • 为了公平比较激活量化的影响,本实验将权重保持在全精度 (W16),只对激活进行 8 比特量化 (A8)。
  • 全精度模型 (FP) 的 FID 为 5.31。
  • TFMQ-DM 作为一种时间步感知 (timestep-aware) 的激活量化方法,其 FID 为 7.74,与 FP 模型存在一定差距。
  • Q-DiT (样本级动态激活量化) 的 FID 仅为 5.34,非常接近 FP 模型,并且在 sFIDISPrecision 等指标上也表现出极高的相似性。
  • 这突出表明,Q-DiT 的样本级动态激活量化机制能够更精确地捕获激活分布的细微变化,从而显著降低量化误差,在保持模型精度的同时最大限度地减少性能下降。

6.2.3. 搜索方法的比较

以下是原文 Table 6 的结果,比较了所提出的搜索方法(进化算法)与潜在的替代方案。

以下是原文 Table 6 的结果:

Search method FID ↓ sFID ↓ IS ↑ Precision ↑
Group size = 128 6.64 19.29 211.27 0.7548
ILP 6.71 19.20 205.54 0.7538
Hessian-based 7.38 19.41 197.48 0.7385
Ours 6.40 18.60 211.72 0.7609

分析:

  • 本实验在 W4A8 配置下进行,并包含了样本级动态激活量化,以评估组大小搜索策略的有效性。
  • 当所有层都使用默认的组大小 128 时,FID 为 6.64。
  • 基于整数线性规划 (ILP) [24] 的搜索方法,其 FID 为 6.71,略高于固定组大小 128,表明其可能未能找到最优配置或其敏感度指标与实际生成性能相关性不足。
  • 基于 Hessian [21] 的搜索方法,其 FID 为 7.38,表现最差,进一步证实了传统基于 MSE 或局部敏感度指标的方法不适用于 DiT 的复杂量化。
  • 本文提出的进化搜索方法 (Ours) 实现了最佳性能,FID 降至 6.40,sFID 降至 18.60,IS 升至 211.72。这显著优于所有对比搜索方法和固定组大小的基线。
  • 这表明,直接使用 FID (或 FVD) 作为优化目标,并通过进化算法进行全局搜索,能够更有效地找到最优的层级组大小配置,从而最大限度地提高 DiT 量化的性能。

7. 总结与思考

7.1. 结论总结

本文提出了 Q-DiT,一个专为 Diffusion Transformers (DiTs) 设计的新颖的训练后量化 (Post-Training Quantization, PTQ) 框架。该方法有效地解决了 DiTs 中存在的两个关键量化挑战:

  1. 权重和激活在输入通道上的显著空间方差: 通过引入自动量化粒度分配 (automatic quantization granularity allocation) 方法加以解决,该方法使用进化算法优化各层的组大小。

  2. 激活在不同时间步和样本间的动态变化: 通过实现样本级动态激活量化 (sample-wise dynamic activation quantization),在运行时自适应调整量化参数。

    广泛的实验结果验证了 Q-DiT 的卓越性能。在 ImageNetVBench 上,Q-DiTW6A8 配置下实现了近乎无损的图像和视频生成质量,甚至在更具挑战性的 W4A8 设置下,也能保持高保真度,显著优于现有基线方法。例如,在 ImageNet 256×256256 \times 256 数据集上,将模型量化到 W4A8 时,FID 仅增加了 1.09,这在低比特量化领域是一个显著的成就。

7.2. 局限性与未来工作

论文作者指出,当前 Q-DiT 方法的一个主要局限性在于其对进化算法的依赖来确定最优的组大小配置。这个过程计算成本高且耗时,增加了整体优化成本和持续时间。作者计划在未来的工作中优化这一部分。

7.3. 个人启发与批判

7.3.1. 个人启发

  1. 问题导向的创新: Q-DiT 的成功在于它不是盲目应用通用量化技术,而是首先深入分析了 DiT 模型的独特特性(空间方差、时间/样本动态性),然后针对性地提出了解决方案。这提醒我们,在进行技术创新时,深入理解具体任务和模型架构的痛点是至关重要的。
  2. 动态量化与细粒度控制的价值: 样本级动态激活量化和自动粒度分配的有效性表明,对于复杂模型和动态数据流,静态的、粗粒度的量化方法往往不足。更精细、更自适应的量化策略是未来量化研究的重要方向。
  3. 度量驱动的优化: 进化搜索直接以 FIDFVD 等最终生成质量指标作为优化目标,而非代理指标(如 MSE),这避免了代理指标与真实性能不一致的问题。对于像生成模型这样难以通过传统指标评估的任务,这种端到端 (end-to-end) 的优化思路非常有启发性。
  4. 硬件友好性与部署前景: Q-DiT 采用均匀量化,这通常对硬件加速友好。在保持高性能的同时实现 W4A8 量化,极大地降低了 DiT 模型的部署门槛,使其在资源受限的设备上运行成为可能,有望加速 DiT 在边缘计算和实时应用中的落地。

7.3.2. 批判与潜在改进

  1. 进化算法的计算效率问题: 论文明确指出了进化算法耗时和计算成本高的局限性。未来工作可以探索更高效的搜索策略,例如:
    • 基于强化学习的搜索: 将组大小分配建模为一个序列决策问题,通过强化学习智能体来学习最优的分配策略。
    • 可微分搜索: 设计一种可微分的量化框架,使得组大小参数可以通过梯度下降进行优化,从而避免离散搜索。
    • 启发式或代理模型加速: 使用轻量级代理模型来预测 FID/FVD,或者开发更精确的层敏感度指标来指导搜索,以减少每次评估的开销。
  2. 量化位宽的进一步探索: 论文主要关注 W6A8W4A8。虽然 W4A8 已经很有挑战性,但进一步探索更低位宽(如 W4A4 甚至 W2A2)的量化,将为更极致的边缘部署提供可能性。这可能需要引入混合精度量化 (mixed-precision quantization)、稀疏化 (sparsity) 或其他模型压缩技术。
  3. 硬件实现与实际加速: 论文主要通过 BitOps 间接衡量计算成本。未来的工作可以结合实际硬件平台进行部署和测试,量化实际的推理延迟和功耗,以提供更具体的加速数据。这可能需要开发定制的 CUDA 内核或与硬件厂商合作。
  4. 对更复杂 DiT 架构的泛化能力: 论文主要基于 DiT-XL/2STDiT3 进行实验。随着 DiT 架构的不断演进(例如更大规模、更多变体),Q-DiT 的方法是否能无缝泛化,并保持相同的性能优势,值得进一步研究。特定的 Transformer 块结构或激活模式可能需要微调量化策略。
  5. 校准数据集依赖: 作为 PTQ 方法,Q-DiT 仍然需要少量校准数据来确定量化参数(尽管动态量化减少了这种依赖)。探索完全无数据 (data-free) 的 PTQ 方法,或者在校准数据稀缺时如何保持性能,也是一个有价值的方向。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。