静态量化的困境

TASQ：扩散模型量化为什么不该所有步骤用同一种精度？

作者

快速导读：TASQ的关键洞察是：存储精度和运行精度可以解耦。它始终保存一份完整的8比特权重，但通过一个可学习的Temporal-Spatial LSB Mask，在每个去噪阶段和每一层动态决定实际使用几个低位。噪声大的阶段保留更多比特，噪声小的阶段大胆丢弃低位。这样既不用存多份权重，又能让计算量跟随实际需求变化。

扩散模型在图像生成中取得了卓越效果，但其迭代去噪过程计算量巨大。现有量化方法通常为所有权重分配统一的精度，忽略了不同去噪步骤对权重量化误差的敏感度差异。TASQ 观察到，FFN 层在噪声较大的早期步骤最敏感，而注意力层在去噪中期敏感度达到峰值，因此静态精度分配要么在敏感阶段精度不足导致质量下降，要么在非敏感阶段浪费计算资源。

TASQ 的核心思想是将权重的存储精度与运行精度解耦：始终保存一份最高精度的共享权重缓冲区，但通过一个可学习的时间-空间 LSB 掩码，在每个去噪阶段和每层动态决定实际使用多少最低有效位。配合最远阶段优先训练策略和时间精度引擎硬件设计，TASQ 在 PixArt-Σ、SANA-1.6B 和 SDXL-Turbo 上实现了质量无损或更优的前提下，将比特操作数降低 25–50%。

英文题目：TASQ: Temporal-Adaptive Bit Sparsification Quantization for Diffusion Models

论文出处：arXiv 每日论文精选 · arXiv:2608.03057

原始论文：PDF / 论文页面

这篇论文解决什么问题？

扩散模型通过逐步去噪生成图像，通常需要 20 到 100 个推理步骤。每一步都涉及大量矩阵乘法，权重精度直接影响计算和存储开销。主流量化方法如 SVDQuant 和 EfficientDM 采用静态精度分配，即所有去噪步骤使用相同的权重位宽。这种设计隐含假设各步骤对量化误差的敏感度相同，但实际并非如此。

论文通过分析发现，扩散模型不同层在不同去噪阶段表现出显著的时间敏感性差异。具体而言，FFN 层的敏感度在噪声最大的初始阶段最高，随着去噪推进逐渐降低；而注意力层的敏感度在去噪中期达到峰值。这意味着静态精度必须覆盖最敏感的阶段和层，导致其他阶段被迫使用超出实际需求的精度。

此前虽有 AdaTSQ 等工作探索了每步激活位宽的自适应调整，但权重精度始终固定。TCAQ-DM 调整了每步的激活范围，同样未涉及权重精度。BSQ、CSQ 等比特级混合精度方法学习了空间上的精度分配，但未扩展到时间维度。TASQ 首次将时间自适应引入权重量化，填补了这一空白。

核心创新

  • TASQ的关键洞察是：存储精度和运行精度可以解耦。它始终保存一份完整的8比特权重，但通过一个可学习的Temporal-Spatial LSB Mask，在每个去噪阶段和每一层动态决定实际使用几个低位。噪声大的阶段保留更多比特，噪声小的阶段大胆丢弃低位。这样既不用存多份权重，又能让计算量跟随实际需求变化。

方法概览

TASQ的关键洞察是：存储精度和运行精度可以解耦。它始终保存一份完整的8比特权重，但通过一个可学习的Temporal-Spatial LSB Mask，在每个去噪阶段和每一层动态决定实际使用几个低位。噪声大的阶段保留更多比特，噪声小的阶段大胆丢弃低位。这样既不用存多份权重，又能让计算量跟随实际需求变化。

  • 时间-空间 LSB 掩码：TASQ 维护一份 8 比特的共享权重缓冲区，并为每个去噪阶段和每层学习一个连续门控参数，控制是否保留该权重的各个最低有效位。门控值经 sigmoid 函数映射到 [0,1]，训练中通过 LSB 掩码退火逐步二值化为 0 或 1，最终确定每层每阶段的运行精度。
  • 存储与运行精度解耦：所有阶段共享同一份 8 比特权重，无需为不同阶段存储多份副本。运行精度由掩码动态决定，可在 2 到 8 比特之间变化。这解决了时间自适应量化的核心存储挑战——若为每个阶段单独存储权重，内存开销将成倍增加。
  • 最远阶段优先训练：直接联合训练所有阶段的掩码会导致相邻阶段梯度高度相关，优化困难。TASQ 将去噪步骤划分为四个阶段，按时间轨迹上距离最大化的顺序更新共享权重：先训练阶段 0 和阶段 3，再训练阶段 1 和阶段 2。这种调度使每次更新面对差异最大的数据分布，稳定了训练过程。
  • LSB 掩码退火与正则化：训练初期门控初始化为 1，保留所有权重位。随着训练进行，通过逐渐降低温度参数将连续门控推向 0 或 1，同时施加稀疏正则化鼓励更多位被剪枝。退火和剪枝间隔是关键的调参超参数。
  • 时间精度引擎：为将自适应精度分配转化为实际加速，TASQ 设计了比特串行 GEMM 数据流。该引擎按位平面顺序处理权重，仅传输掩码选定的位平面，执行周期与有效精度成正比。与固定精度比特串行执行相比，无需精度切换开销。
  • 初始化策略：主实验从 SVDQuant 的 PTQ 结果初始化共享权重和低秩分支，然后进行 TASQ 自适应训练。消融实验表明，即使不使用 SVDQuant 初始化，TASQ 仍优于非时间自适应的 QAT 基线。
  • 阶段划分：基于敏感度分析的肘部法则将去噪过程划分为四个阶段。消融实验显示，SDXL-Turbo 在四阶段时性能最佳（匹配其四步去噪），PixArt-Σ 在四阶段后增益趋于平缓。
  • 训练开销：TASQ 的额外训练迭代数为 1.4k–5.0k，在 SVDQuant 基础上增加约 20–40% 的训练时间，但换来了显著的推理计算节省。

逐图理解论文

静态量化的困境

静态量化的困境
Figure 1: Temporal sensitivity and TASQ bit allocation across diffusion stages: FFN sensitivity peaks at noisy steps while attention peaks mid-trajectory, and TASQ assigns higher bits where sensitivity is high.

图 1 展示了扩散模型不同层在不同去噪阶段对量化误差的敏感度差异。FFN 层在噪声最大的早期阶段敏感度最高，注意力层在中期达到峰值。TASQ 据此为高敏感阶段分配更多比特，低敏感阶段使用更少比特，从而在不损失质量的前提下降低整体计算量。

核心创新：时间自适应精度

核心创新：时间自适应精度
Figure 2: Process of Temporal–Spatial LSB Mask Regularization and LSB Mask Annealing.

图 2 说明了时间-空间 LSB 掩码的训练流程。左侧展示掩码正则化过程，连续门控参数控制每个权重位平面的保留与否。右侧展示 LSB 掩码退火，训练中温度逐渐降低，将连续值推向 0/1 二值化，最终确定每层每阶段的离散精度分配。

训练与硬件协同设计

训练与硬件协同设计
Table 6: Measured four-tile GEMM cost on the Temporal- Precision Engine with A4. The mixed schedule incurs no precision-switching overhead.

直接联合训练所有阶段的掩码会遇到梯度耦合问题。TASQ提出Farthest-Stage-First Training，按时间轨迹上距离最远的顺序交替更新，让每次优化面对差异最大的数据分布。推理端则设计了Temporal-Precision Engine，一种比特串行数据流，只传输掩码选定的位平面，执行周期与有效精度成正比，没有任何精度切换开销。

结论与局限

结论与局限
Table 2: Quality–compute comparison. “Eff. W” denotes the average operating weight precision and “Rel. BitOPs” the compute cost normalized to static W8A8. TASQ stores one shared 8-bit weight buffer across all stages. W4 rows are compute-matched, while W8 TASQ retains W8 quality with 74–75% compute. Non-SVDQuant controls are reported in Table 4.

TASQ证明扩散模型去噪过程存在显著的阶段性精度冗余，通过时间自适应量化可以在不损失质量的前提下将计算量降低25%到50%。时间与空间自适应缺一不可，两者结合才能达到最佳效果。不过，TASQ并没有减少模型存储大小，共享权重缓冲区仍需容纳最高精度。以上就是这篇论文的核心内容，感谢收听。

实验如何设计？

  • 模型与数据集：在 PixArt-Σ、SANA-1.6B 和 SDXL-Turbo 三个主流扩散模型上评估，使用 MJHQ 基准测试的图像质量指标 FID 和 ImageReward。
  • 精度设置：主要对比 W4A4 和 W8A8 两种配置。W4A4 下 TASQ 的平均运行精度为 4 比特，但各阶段各层可在 2–8 比特间变化。W8 TASQ 保持 8 比特质量同时降低计算量。
  • 基线方法：包括静态 W4A4/W8A8、SVDQuant、QAT、EfficientDM 等。所有方法使用相同的校准数据池以隔离优化方法的影响。
  • 硬件测量：在时间精度引擎上实测四瓦片 GEMM 的执行周期，对比朴素位平面循环和静态 8 比特比特串行执行。
  • 消融设计：分别消融时间自适应和空间自适应的贡献，验证两者互补性；消融 SVDQuant 初始化的必要性；消融阶段数量对性能的影响。
  • CIFAR-10 验证：在 100 步 DDIM 模型上与 PTQ4DM、Q-Diffusion、TFMQ-DM、EfficientDM 等对比，验证方法在小规模任务上的泛化性。

关键结果与论文证据

  • 质量-计算权衡（表 2，第 5 页）：TASQ W4A4 在 PixArt-Σ 上取得 MJHQ FID 16.1，优于静态 W4A4 的 17.1–19.2。SANA-1.6B 上 FID 15.9，SDXL-Turbo 上 23.6，分别改善 1.0–3.1 FID。
  • W8 计算节省（表 2，第 5 页）：W8 TASQ 保持与 W8 静态量化相同的图像质量，同时将相对 BitOPs 降至 0.74–0.75，即节省约 25% 计算量。
  • 时间与空间自适应的互补性（表 3，第 7 页）：在 PixArt-Σ W4A4 上，仅空间自适应 FID 为 17.1，仅时间自适应 FID 为 18.4，两者结合达到 16.1，证明时间与空间维度共同贡献了性能提升。
  • 无需 SVDQuant 初始化（表 4，第 7 页）：TASQ 从随机初始化训练时，在所有设置下 ImageReward 均优于 QAT，FID 在除一个设置外的所有情况下也更低，表明方法具有独立有效性。
  • 硬件加速实测（表 6，第 7 页）：时间精度引擎上 TASQ 混合精度调度相比朴素位平面循环减少执行周期 3.2–5.1 倍，相比静态 8 比特比特串行执行减少 6.1–7.5 倍。
  • 精度分配模式（图 4，第 6 页；图 5，第 6 页）：TASQ 学到的精度分配与敏感度分析一致——噪声最大的阶段 0 保留最多位平面，后续阶段精度逐渐降低。不同层之间的精度分配差异显著，FFN 和注意力层呈现不同的时间模式。
  • 阶段数量影响（图 9，第 15 页）：SDXL-Turbo 性能随阶段数从 1 增至 4 持续改善，匹配其四步去噪特性。PixArt-Σ 在四阶段达到肘部点，更多阶段增益边际递减。
  • CIFAR-10 验证（表 12，第 15 页）：在 100 步 DDIM 模型上，TASQ 在 A8 和 A4 两种激活精度下均取得最低 FID，且 IS 与最佳基线持平，证明方法在不同规模任务上均有效。

阅读时需要注意

  • 存储未减少：TASQ 的共享权重缓冲区必须容纳最敏感阶段所需的最高精度，因此模型存储大小与静态 8 比特量化相同，未实现存储压缩。
  • 仅自适应权重精度：激活精度在 TASQ 中保持固定，未利用激活在时间维度上可能存在的冗余。
  • 训练开销：需要额外的微调阶段（1.4k–5.0k 迭代），虽然换来了推理加速，但增加了部署前的准备成本。
  • 硬件依赖：相对 BitOPs 和周期节省假设使用精度可伸缩的比特串行硬件，固定 INT4/INT8 内核无法利用自适应调度带来的计算节省。
  • 超参数敏感性：LSB 掩码退火和剪枝的超参数（λM、剪枝间隔）可能需要针对新模型架构进行调整。

关联工作

  • SVDQuant：PTQ 基线方法，通过保护低秩分支减少量化误差。TASQ 主实验从 SVDQuant 初始化，在其基础上进一步进行时间自适应训练。
  • EfficientDM：基于 LoRA 的扩散模型 QAT 方法。TASQ 在类似 LoRA 微调框架上构建，但增加了时间-空间掩码学习。
  • BSQ/CSQ/MSQ：比特级混合精度方法，学习空间上的精度分配。TASQ 将这一思想扩展到时间维度，通过 LSB 掩码实现阶段级精度变化。
  • AdaTSQ：搜索每步激活比特数，但权重精度保持静态。与 TASQ 的权重自适应互补，两者可结合使用。
  • TCAQ-DM：在统一位宽下调整每步激活范围，不改变权重精度。TASQ 首次在权重端引入时间自适应。

发表评论