CN118014044B - 基于多基二值化的扩散模型量化方法和装置 - Google Patents

Abstract

本发明公开了一种基于多基二值化的扩散模型量化方法和装置。该量化方法包括如下步骤：在前向传播中，使用可学习的多基二值化器，以增强扩散模型的二值化权重；在反向传播中，并行进行多基二值化器中多个基的计算；在训练过程中，通过在低秩空间中模仿全精度表示，使得二值化扩散模型模仿全精度扩散模型的表示结果。与现有技术相比较，本发明提供了显著的生产力优势，展示了在边缘硬件上部署扩散模型的巨大潜力。

Description

基于多基二值化的扩散模型量化方法和装置

技术领域

本发明涉及一种基于多基二值化的扩散模型量化方法，同时也涉及相应的扩散模型量化装置，属于神经网络量化技术领域。

背景技术

扩散模型（Diffusion Models，简写为DM）在各个领域的生成任务中展现了出色的能力。凭借优越的生成质量和多样性的优势，扩散模型已经成为最受欢迎的生成模型范式之一。扩散模型包括两个过程：前向传播和反向传播，其中前向传播又称为扩散过程；反向传播用于生成数据样本，因此又称为生成过程。

然而，扩散模型的生成过程比较缓慢，动辄多达上千步的迭代步骤会减慢推理过程，并依赖于昂贵的硬件资源，对其广泛实施构成了重大挑战。量化和二值化是流行的神经网络压缩方法，可以将神经网络的全精度参数量化为较低的比特位数（例如，1～8比特）。通过将浮点权重和激活函数转换为量化值，可以减小神经网络的模型大小，并且还可以降低计算复杂度，从而显著提高推理速度，节省内存使用量，并有效降低能耗。现有技术中，对于扩散模型的量化方法通常分为两类：后训练量化和量化感知训练。作为一种无需训练的方法，后训练量化被认为是一种更实用的解决方案，它通过搜索最佳缩放因子候选值和优化校准策略来以较低的成本获得量化模型。然而，通过后训练量化的扩散模型在生成质量上会显著降低。因此，量化感知训练出现了，以提高扩散模型量化后的准确性。由于具有充足数据和训练资源的训练/微调过程的好处，通过量化感知训练获得的扩散模型通常比后训练量化获得的扩散模型具有更高的准确性。然而，针对扩散模型权重的1比特量化（即二值化）仍然远未实现。

在申请号为202311540202.0的中国发明申请中，公开了一种扩散模型的轻量化方法，包括如下步骤：获取目标数据和预训练的老师模型；将所述目标数据输入所述老师模型，得到所述多个中间块分别输出的第一结果；将每个所述中间块的输入，输入所述中间块对应的所述块，得到每个所述块在采用不同的路径操作的情况下，输出的多个第二结果；根据每个所述中间块输出的第一结果，以及所述中间块对应的所述块输出的多个第二结果，对每个所述块进行分块搜索，得到每个所述块对应的目标块；根据每个所述目标块，生成子网；对所述子网进行重新训练，得到训练好的目标子网；根据所述目标子网，建立目标扩散模型。

发明内容

本发明所要解决的首要技术问题在于提供一种基于多基二值化的扩散模型量化方法。

本发明所要解决的另一技术问题在于提供一种基于多基二值化的扩散模型量化装置。

为实现上述发明目的，本发明采用以下的技术方案：

根据本发明实施例的第一方面，提供一种基于多基二值化的扩散模型量化方法，包括如下步骤：

在前向传播中，使用可学习的多基二值化器，以增强扩散模型的二值化权重；在反向传播中，并行实现所述多基二值化器中多个基的计算；其中，所述可学习的多基二值化器通过下式表示：

其中，表示二值量化后的近似数值，w表示量化前的原始权重，/>和/>是可学习的标量，初始值分别设为/>和/>，其中 /> 表示 /> 范数；

在训练过程中，通过在低秩空间中模仿全精度表示，使二值化扩散模型模仿全精度扩散模型的表示结果。

其中较优地，在扩散模型的噪声估计网络中，只有当特征尺度大于或等于输入尺度的一半时，使用可学习的多基二值化器。

其中较优地，在所述反向传播中，采用直通估计器近似函数。

其中较优地，在低秩空间中模仿全精度表示，包括如下子步骤：

采用时间步嵌入模块对全精度扩散模型进行分组，使用主成分分析将中间表示结果投影到低秩空间。

其中较优地，所述时间步嵌入模块由残差卷积和变换器块组成。

其中较优地，在所述训练过程的早期阶段，采用渐进二值化策略使二值化扩散模型稳定收敛。

其中较优地，所述渐进二值化策略包括如下子步骤：

在最初的迭代中，对扩散模型中的第个时间步嵌入模块进行量化；在随后的/>个迭代中，在第/>次迭代时将第 />个和第 />个时间步嵌入模块进行二值化，直到所有时间步嵌入模块都被二值化，其中/>表示向下取整函数，/>表示扩散模型的噪声估计网络中的时间步嵌入模块的数量。

根据本发明实施例的第二方面，提供一种基于多基二值化的扩散模型量化装置，包括处理器和存储器；其中，

所述存储器与所述处理器耦接，用于存储计算机程序，当该计算机程序被所述处理器执行时，使处理器实现上述基于多基二值化的扩散模型量化方法。

与现有技术相比较，本发明提供的基于多基二值化的扩散模型量化方法（简写为BinaryDM）在各种模型和数据集上都提供了显著的生产力优势。在相同的1比特权重下，尤其是在超低比特激活下，本扩散模型量化方法可以持续地优于像素空间扩散模型（DDIM）和潜空间扩散模型（LDM）的基线。例如，在CIFAR-10 32×32 DDIM上，本扩散模型量化方法的精度指标甚至比基线高出49.04%，使得二值化扩散模型避免了崩溃。作为扩散模型的领先二值化方法，本扩散模型量化方法在FLOP（浮点运算数）和模型大小上实现了惊人的16.0倍和27.1倍的节省，展示了在边缘硬件（例如手机、智能手表等）上部署扩散模型的巨大优势和潜力。

附图说明

图1为本发明实施例提供的扩散模型量化方法的逻辑框架图；

图2为本发明实施例中，卷积层中滤波器（通道）逐个二值化权重的比较图；

图3为本发明实施例中，BinaryDM与基线在生成图像中细节的比较图；

图4为本发明实施例中，不同蒸馏损失函数对全精度扩散模型和二值化扩散模型中每个块的输出特征的影响，以距离衡量；

图5为本发明实施例提供的扩散模型量化装置的架构示意图。

具体实施方式

下面结合附图和具体实施例对本发明的技术内容进行详细具体的说明。

如图1所示，本发明实施例提供的扩散模型量化方法主要包括两方面的技术内容：(1) 在前向传播和反向传播中引入可学习的多基二值化器（Learnable Multi-basisBinarizer，简写为LMB），用于恢复由二值化扩散模型生成的表示结果。其中，LMB应用至少两组带有可学习的标量的二进制基，以增强权重的特征提取能力。(2) 在优化过程中引入低秩表示模仿（Low-rank Representation Mimicking，简写为LRM）的解决方案来增强扩散模型的二值化感知优化能力。LRM将二值化和全精度表示投影到低秩，使得二值化扩散模型的优化能够集中在主要方向上，并在细粒度监督下稳定地减轻歧义。在此基础上，本扩散模型量化方法进一步应用逐步初始化策略到早期训练阶段，使得优化从易于收敛的位置开始。

下面，首先介绍二值化扩散模型的基线（Baseline），以便于理解本发明相对于现有技术的改进之处。在扩散模型前向传播的过程中，通过控制噪声强度的调度，对数据进行/>次高斯噪声叠加，该过程可以表示为：

其中，表示第/>步的噪声样本。反向传播旨在通过去除噪声生成样本，用学习到的分布/>来近似不可用的条件分布/>，可以表示为：

均值和方差/>可以使用现有的重新参数化技巧导出：

其中，与/>表示具有可学习参数/>的函数近似，它从/>预测/>。对于扩散模型的训练，通常采用简化的变分下界变体作为损失函数，以实现高质量样本，可以表示为：

量化通过将权重和激活值离散化为低比特位，以此压缩和加速噪声估计模型。在二值化扩散模型的基线中，权重被二值化为1比特：

其中，函数将/>限制在＋1或－1之间，阈值为0，/>表示二值化权重，/>表示扩散模型的噪声估计网络，/>是浮点标量，初始化为/>（/>表示权重元素的数量）。

在目前的基础二值化过程中，权重被量化为二进制值，以节省推理时的存储和计算资源。此外，激活值也可以进行量化以便进一步节省资源。然而，在扩散模型中将权重广泛离散化为1位导致生成结果明显恶化。将每个权重元素的位宽限制为原始位宽的局部导致了由这些权重形成的滤波器或线性投影的丰富性大大降低。因此，扩散模型从输入中提取特征的能力大大受损。这种提取能力的妥协，再加上权重的数值离散化，使得在表示结果中准确保留细节面临挑战，而这对于保持扩散模型中合成数据的复杂细节和纹理至关重要。因此，目前的基础二值化过程导致了扩散模型表示质量的显著下降。

为了解决这个问题，本发明实施例中率先提出一个可学习的多基二值化器（简写为LMB），以增强扩散模型的二值化权重并合成高质量的表示结果。在前向传播中，LMB的过程定义如下：

其中，表示二值量化后的近似数值，w表示量化前的原始全精度权重，/>和/>是可学习的标量，它们的初始值分别设为/>和/>，其中 /> 表示 范数。

由LMB进行二值化的层的推理涉及多个基的计算。例如，在二值化扩散模型中，卷积的推理过程为：

其中，表示激活，/>和/>分别表示由乘法和加法指令组成的卷积。在推理过程中，由于某个LMB涉及的多个基是独立的，因此它们的计算可以并行化，从而加速了二值化扩散模型的推理过程。

在LMB的反向传播中，可学习的标量的梯度计算如下：

其中，表示求导函数。在反向传播中，本发明实施例采用直通估计器（StraightThrough Estimator）来近似/>函数。通过具有多个可学习的标量的二进制基，量化权重的表示能力得到显著增强。残差初始化使得二值化扩散模型的优化从一个最小化误差的状态开始。与二值化扩散模型的基线相比，通过LMB使权重的表示显著多样化，其中卷积层中滤波器（通道）逐个二值化权重的比较图如图2所示。LMB丰富的特征表示使得本发明实施例提供的扩散模型量化方法在整体去噪的基础之上还能拥有更完整的表达，正如图3中各种细节所展示的那样。

需要说明的是，发明人只在扩散模型的关键位置应用了LMB，同时在其他位置保持简洁的基础二值化，避免了不必要的额外开销，并在精度和效率之间实现了平衡。例如在本发明的一个实施例中，在扩散模型的噪声估计网络中，只有在特征尺度大于或等于输入尺度的一半时，即U-Net中的前六层和最后六层时，发明人才应用LMB。在其他位置，量化方式与基线保持一致，如公式(5)所示的那样。上述配置方式可以显著减少权重存储，与直接在所有位置应用LMB相比，存储量减少了41.5%。因此，LMB在几乎不增加额外负担的情况下显著改善了二值化扩散模型。

另一方面，由于极端的权重二值化函数（以及激活量化函数）的参与，直接训练二值化扩散模型在收敛方面存在困难。因此，发明人应用表示模仿来辅助二值化扩散模型的训练。在训练过程中，它会对全精度扩散模型和二值化扩散模型的中间和输出表示进行对齐。

然而，在优化过程中直接对齐二值化和全精度扩散模型的中间表示结果存在问题。首先，高维表示的细粒度对齐会导致扩散模型的优化方向模糊，尤其是当引入中间特征的模仿时。其次，与全精度扩散模型相比，二值化扩散模型中的中间特征源自离散的潜在空间，因为权重（和激活）的离散化使得直接模仿全精度扩散模型变得困难。

因此，在本发明的一个实施例中提出了低秩表示模仿的解决方案，即通过在低秩空间中模仿全精度表示，有效地优化二值化扩散模型。具体地说，发明人基于时间步嵌入模块（由残差卷积和变换器块组成）对全精度扩散模型进行分组，中间表示结果可以记为，其中h、w、c分别表示高、宽、通道维度。接下来，使用主成分分析（PCA）将中间表示结果投影到低秩空间。

在本发明的一个实施例中，全精度扩散模型表示的协方差矩阵为：

其中，表示前 /> 个模块的组合。

然后，通过以下方式可以得到特征向量矩阵：

其中，是/>的特征值的对角矩阵，按降序排列。发明人将由/>的前/>列特征向量组成的矩阵作为变换矩阵，记为/>，其中/>表示向上取整函数，/>表示降维的次数。

发明人使用来投影全精度扩散模型和二值化扩散模型的中间表示结果：

其中， 表示具有二值化参数/>的扩散模型中第/>层的中间表示结果，和 />分别表示具有相同形状的/> 的全精度扩散模型和二值化扩散模型的低秩表示。

随后，利用获得的低秩表示来推动二值化扩散模型学习全精度对应物。具体地说，发明人构建了全精度扩散模型和二值化扩散模型之间第个模块低秩表示之间的均方误差（MSE）损失：

总的损失由公式(4)与公式(13)共同组成：

其中，表示扩散模型的噪声估计网络中的时间步嵌入模块的数量，/>是一个超参数系数，用于平衡损失项。由于在低秩表示模仿的解决方案中获取变换矩阵/>的计算成本相当昂贵，发明人通过第一批输入计算该矩阵，并在训练过程中保持其不变。固定表示之间的映射也有助于从稳定的角度优化二值化扩散模型。

低秩表示模仿的解决方案使得二值化扩散模型能够模仿全精度扩散模型的表示结果，通过引入额外的监督来改善二值化扩散模型的优化。正如图4所示，低秩表示模仿的解决方案有效地将本地模块靠近完整精度模块。此外，通过在表示模仿之前基于完整精度表示的主要成分应用低秩投影，二值化扩散模型可以沿着清晰稳定的方向进行优化，加速扩散模型的收敛。此外，二值化扩散模型和全精度扩散模型具有完全一致的架构，使得它们之间的表示模仿十分自然。

接下来，进一步说明本发明实施例提供的扩散模型量化方法的训练策略。前已述及，该扩散模型量化方法利用可学习的多基二值化器对扩散模型的噪声估计网络进行了二值化，并利用低秩表示模仿的解决方案对目标函数进行了优化。然而，尽管在架构和优化方面取得了显著的改进，但二值化扩散模型在训练过程中的收敛仍然缓慢且难以实现。

因此，在训练过程的早期阶段，本发明实施例中采用了渐进二值化策略，使二值化扩散模型在加速并稳定收敛的同时，不会产生额外的成本。具体地说，在最初的迭代中，对扩散模型中的第个时间步嵌入模块进行量化。在随后的/>个迭代中，在第/>次迭代时将第个和第 />个时间步嵌入模块进行二值化，直到所有时间步嵌入模块都被二值化，其中/>表示向下取整函数。这个过程使得二值化扩散模型能够调整到有利的起始位置进行优化。与直接在训练开始时对整个扩散模型进行二值化相比，这种渐进二值化策略在相同的训练持续时间内展现了显著的性能优势。此外，值得注意的是，整个渐进量化过程通常仅占用了训练所需的总200K迭代中的约0.002%（5个迭代），可以被视为对二值化扩散模型进行极其高效的初始化策略。

以上对本发明实施例提供的扩散模型量化方法的具体步骤进行了详细说明。为了验证本发明实施例提供的扩散模型量化方法的实际性能，发明人在多个数据集上进行了实验，包括 CIFAR-10 32×32、LSUN-Bedrooms 256×256、LSUN-Churches 256×256、FFHQ256×256和 ImageNet 256×256，用于无条件图像生成任务和有条件图像生成任务，覆盖了像素空间扩散模型（DDIM）和潜空间扩散模型（LDM）。具体的实验结果说明如下：

1. 无条件图像生成任务

发明人首先在 CIFAR-10 32×32上进行了实验。如表1所示，在这种低分辨率的情况下，二值化扩散模型的基线遭遇了严重的性能衰退，而本扩散模型量化方法显著恢复了性能。在CIFAR-10 32×32 DDIM上，本扩散模型量化方法的精度指标甚至比基线高出49.04（51.22－2.18）%，使得二值化扩散模型避免了崩溃。在 W1A4 位宽下，本扩散模型量化方法在CIFAR-10 32×32上的IS指标显著超过了二值化基线，并在W1A32中从各种评价指标上远远超过了二值化扩散模型的基线。

表1 CIFAR-10上无条件图像生成任务的比较结果

2. 有条件图像生成任务

对于有条件图像生成任务，在分辨率为256×256的ImageNet数据集上的性能，重点关注LDM-4。发明人采用三种不同的采样器生成图像：DDIM、PLMS和DPM-Solver。表2中的结果突显了本扩散模型量化方法的显著有效性，在几乎所有评估指标上持续超越基线，甚至在某些情况下还优于全精度扩散模型。二值化扩散模型的基线在配置W1A32和W1A8下表现相对稳定，但在W1A4下显著下降。具体来说，当使用DPM-Solver采样器时，IS下降到了85.99，FID急剧增加至25.85。与之鲜明对比的是，本发明实施例中的二值化扩散模型保持了持续高性能，实现了156.19的IS和11.15的FID，并在大多数场景中优于基线。

表2 ImageNet 256×256 上有条件图像生成任务的量化结果

3. 消融实验结果

接下来，发明人对 LSUN-Bedrooms 256×256数据集上的 LDM-4 进行了全面的消融研究，以评估本扩散模型量化方法的有效性。

表 3 LSUN-Bedrooms 256×256数据集上的消融结果

在此，评估了发明人提出的LMB和LRM的有效性，结果如表3所示。当仅将本发明实施例提供的LMB应用于二值化扩散模型时，性能显著恢复。另一方面，仅在优化中应用LRM并没有完全解决二值化扩散模型的性能瓶颈，导致准确性显著降低。在本发明实施例提供的扩散模型量化方法中结合这两种技术可以显著提升性能。

4. 推理效率分析

对于推理效率，发明人展示了不同激活比特宽度下扩散模型量化方法的参数大小和 FLOPs。表4中的结果表明，本发明实施例提供的扩散模型量化方法在推理过程中可以实现高达27.1倍的空间节省，同时在推理过程中获得高达16.0倍的加速，充分体现了二值化计算的优势。

表4 扩散模型量化方法的推理效率

5. 训练效率分析

对于训练效率，尽管与后训练量化方法相比，本发明实施例提供的扩散模型量化方法在训练过程通常会产生更高的开销，但实际观察表明，本扩散模型量化方法在各种模型和数据集上都提供了生产力优势。正如表5所示，尽管本扩散模型量化方法的训练时间比现有的Q-Diffusion量化方法需要的校准时间要短，但本扩散模型量化方法在低位宽下获得了显著优越的生成质量。

表5 本扩散模型量化方法与Q-Diffusion量化方法

在训练时间成本方面的比较结果

基于上述基于多基二值化的扩散模型量化方法，本发明实施例进一步提供一种基于多基二值化的扩散模型量化装置。如图5所示，该扩散模型量化装置包括一个或多个处理器和存储器。其中，存储器与处理器耦接，用于存储一个或多个计算机程序，当一个或多个计算机程序被一个或多个处理器执行，使得一个或多个处理器实现如上述实施例中基于多基二值化的扩散模型量化方法。

其中，处理器用于控制该扩散模型量化装置的整体操作，以完成上述基于多基二值化的扩散模型量化方法的全部或部分步骤。该处理器可以是中央处理器（CPU）、图形处理器（GPU）、现场可编程逻辑门阵列（FPGA）、专用集成电路（ASIC）、数字信号处理（DSP）芯片等。存储器用于存储各种类型的数据以支持在该扩散模型量化装置上的操作，这些数据例如可以包括用于扩散模型量化装置操作的任何应用程序或方法的指令，以及应用程序相关的数据。存储器可以由任何类型的易失性或非易失性存储设备或者它们的组合实现，例如静态随机存取存储器（SRAM）、电可擦除可编程只读存储器（EEPROM）、可擦除可编程只读存储器（EPROM）、可编程只读存储器（PROM）、只读存储器（ROM）、磁存储器、快闪存储器等。

在一个示例性实施例中，扩散模型量化装置具体可以由计算机芯片或实体实现，或者由具有某种功能的产品来实现，用于执行上述基于多基二值化的扩散模型量化方法，并达到如上述方法一致的技术效果。一种典型的实施例为计算机，例如可以为个人计算机、膝上型计算机、平板计算机、车载人机交互设备或可穿戴设备，以及蜂窝电话、相机电话、智能电话、个人数字助理等，或者这些设备中的任何组合。

与现有技术相比较，本发明提供的扩散模型量化方法在各种模型和数据集上都提供了显著的生产力优势。在相同的1比特权重下，尤其是在超低比特激活下，本扩散模型量化方法可以持续地优于像素空间扩散模型（DDIM）和潜空间扩散模型（LDM）的基线。例如，在CIFAR-10 32×32 DDIM上，本扩散模型量化方法的精度指标甚至比基线高出49.04%，使得二值化扩散模型避免了崩溃。作为扩散模型的领先二值化方法，本扩散模型量化方法在FLOP（浮点运算数）和模型大小上实现了惊人的16.0倍和27.1倍的节省，展示了在边缘硬件（例如手机、智能手表等）上部署扩散模型的巨大优势和潜力。

上面对本发明提供的基于多基二值化的扩散模型量化方法和装置进行了详细的说明。对本领域的一般技术人员而言，在不背离本发明实质内容的前提下对它所做的任何显而易见的改动，都将构成对本发明专利权的侵犯，将承担相应的法律责任。

Claims (8)

1.一种基于多基二值化的扩散模型量化方法，其特征在于包括如下步骤：

在前向传播中，使用可学习的多基二值化器，以增强扩散模型的二值化权重；在反向传播中，并行实现所述多基二值化器中多个基的计算；其中，所述可学习的多基二值化器通过下式表示：

其中，表示二值量化后的近似数值，w表示量化前的原始权重，/>和/>是可学习的标量，初始值分别设为/>和/>，其中 /> 表示 /> 范数；

在训练过程中，通过在低秩空间中模仿全精度表示，使二值化扩散模型模仿全精度扩散模型的表示结果。

2.如权利要求1所述的扩散模型量化方法，其特征在于：

在扩散模型的噪声估计网络中，只有当特征尺度大于或等于输入尺度的一半时，使用可学习的多基二值化器。

3.如权利要求1所述的扩散模型量化方法，其特征在于：

在所述反向传播中，采用直通估计器近似函数。

4.如权利要求1所述的扩散模型量化方法，其特征在于在低秩空间中模仿全精度表示，包括如下子步骤：

采用时间步嵌入模块对全精度扩散模型进行分组，使用主成分分析将中间表示结果投影到低秩空间。

5.如权利要求4所述的扩散模型量化方法，其特征在于：

所述时间步嵌入模块由残差卷积和变换器块组成。

6.如权利要求5所述的扩散模型量化方法，其特征在于：

在所述训练过程的早期阶段，采用渐进二值化策略使二值化扩散模型稳定收敛。

7.如权利要求6所述的扩散模型量化方法，其特征在于所述渐进二值化策略包括如下子步骤：

在最初的迭代中，对扩散模型中的第个时间步嵌入模块进行量化；在随后的/>个迭代中，在第/>次迭代时将第 />个和第 />个时间步嵌入模块进行二值化，直到所有时间步嵌入模块都被二值化，其中/>表示向下取整函数，/>表示扩散模型的噪声估计网络中的时间步嵌入模块的数量。

8.一种基于多基二值化的扩散模型量化装置，其特征在于包括处理器和存储器；其中，

所述存储器与所述处理器耦接，用于存储计算机程序，当该计算机程序被所述处理器执行时，使处理器实现权利要求1～7中任意一项所述的基于多基二值化的扩散模型量化方法。