Pioneering 4-Bit FP Quantization for Diffusion Models: Mixup-Sign Quantization and Timestep-Aware Fine-Tuning

📄 arXiv: 2505.21591v1 📥 PDF

作者: Maosen Zhao, Pengtao Chen, Chong Yu, Yan Wen, Xudong Tan, Tao Chen

分类: cs.LG, cs.AI

发布日期: 2025-05-27


💡 一句话要点

首创扩散模型4比特浮点量化:提出混合符号量化与时间步感知微调

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 模型量化 浮点量化 低比特量化 后训练量化 时间步感知 LoRA微调

📋 核心要点

  1. 现有基于整数量化和后训练微调的方法在扩散模型4比特量化中性能不稳定,无法有效压缩模型。
  2. 提出混合符号浮点量化(MSFP)框架,结合无符号FP量化、时间步感知LoRA和去噪因子损失对齐,实现精确微调。
  3. 实验证明,该方法首次在扩散模型上实现了优于现有4比特整数量化方法的4比特浮点量化性能。

📝 摘要(中文)

模型量化通过降低权重和激活的比特宽度来提高扩散模型的内存效率和推理速度。然而,实现4比特量化仍然具有挑战性。现有的方法,主要基于整数量化和后训练量化微调,在性能上表现不一致。受到浮点(FP)量化在大型语言模型中成功的启发,我们探索了扩散模型的低比特FP量化,并发现了关键挑战:有符号FP量化无法处理非对称激活分布,微调过程中对去噪过程中的时间复杂性考虑不足,以及微调损失与量化误差之间的不一致。为了应对这些挑战,我们提出了混合符号浮点量化(MSFP)框架,首先在模型量化中引入无符号FP量化,以及时间步感知LoRA (TALoRA)和去噪因子损失对齐(DFA),以确保精确和稳定的微调。大量实验表明,我们首次在扩散模型的4比特FP量化中实现了卓越的性能,优于现有的4比特INT量化中的PTQ微调方法。

🔬 方法详解

问题定义:论文旨在解决扩散模型中4比特量化的问题,特别是浮点量化。现有的整数量化方法在低比特量化时性能下降明显,而直接应用浮点量化则面临激活分布不对称、时间步信息利用不足以及微调目标不明确等问题。这些问题导致量化后的模型性能不稳定,难以达到理想的压缩效果。

核心思路:论文的核心思路是针对扩散模型的特性,设计一种新的浮点量化方案,即混合符号浮点量化(MSFP)。该方案通过引入无符号浮点量化来处理激活分布不对称的问题,并结合时间步感知LoRA和去噪因子损失对齐来提升微调的稳定性和有效性。这样设计的目的是为了在保证模型性能的前提下,尽可能地降低模型的比特宽度。

技术框架:MSFP框架主要包含三个核心模块:1) 混合符号浮点量化(MSFP):使用无符号浮点量化处理激活分布不对称问题。2) 时间步感知LoRA(TALoRA):在LoRA微调中引入时间步信息,使模型更好地适应去噪过程中的时间复杂性。3) 去噪因子损失对齐(DFA):通过调整损失函数,使微调目标与量化误差更加一致,从而提高微调的有效性。整体流程是先对模型进行量化,然后使用TALoRA和DFA进行微调,最终得到一个高性能的4比特量化模型。

关键创新:论文最重要的技术创新点在于提出了混合符号浮点量化(MSFP)框架,该框架首次将无符号浮点量化引入到模型量化中,并结合时间步感知LoRA和去噪因子损失对齐,有效地解决了扩散模型4比特量化中的关键问题。与现有方法的本质区别在于,MSFP框架能够更好地处理激活分布不对称的问题,并能够更有效地利用时间步信息进行微调。

关键设计:在混合符号浮点量化中,论文使用了一种混合的量化策略,即对一部分层使用有符号浮点量化,对另一部分层使用无符号浮点量化,具体选择哪种量化方式取决于该层的激活分布。在时间步感知LoRA中,论文将时间步信息作为LoRA模块的输入,使LoRA模块能够根据不同的时间步调整模型的参数。在去噪因子损失对齐中,论文调整了损失函数的权重,使损失函数更加关注去噪过程中的重要时间步。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,该方法在扩散模型的4比特浮点量化中取得了显著的性能提升,优于现有的4比特整数量化方法。具体来说,在图像生成任务上,该方法能够达到与全精度模型相近的生成质量,同时将模型大小降低到原来的1/8。与现有的最佳4比特整数量化方法相比,该方法在FID指标上取得了显著的提升。

🎯 应用场景

该研究成果可应用于各种需要高效推理和低内存占用的扩散模型场景,例如移动设备上的图像生成、视频编辑和3D内容创作。通过降低模型的大小和计算复杂度,可以使这些应用在资源受限的设备上运行,并提高用户体验。此外,该技术还可以用于加速扩散模型的训练过程,从而降低训练成本和时间。

📄 摘要(原文)

Model quantization reduces the bit-width of weights and activations, improving memory efficiency and inference speed in diffusion models. However, achieving 4-bit quantization remains challenging. Existing methods, primarily based on integer quantization and post-training quantization fine-tuning, struggle with inconsistent performance. Inspired by the success of floating-point (FP) quantization in large language models, we explore low-bit FP quantization for diffusion models and identify key challenges: the failure of signed FP quantization to handle asymmetric activation distributions, the insufficient consideration of temporal complexity in the denoising process during fine-tuning, and the misalignment between fine-tuning loss and quantization error. To address these challenges, we propose the mixup-sign floating-point quantization (MSFP) framework, first introducing unsigned FP quantization in model quantization, along with timestep-aware LoRA (TALoRA) and denoising-factor loss alignment (DFA), which ensure precise and stable fine-tuning. Extensive experiments show that we are the first to achieve superior performance in 4-bit FP quantization for diffusion models, outperforming existing PTQ fine-tuning methods in 4-bit INT quantization.