一种扩散模型的在线策略蒸馏方法及装置

本发明涉及知识蒸馏，尤其涉及一种扩散模型的在线策略蒸馏方法及装置。

背景技术：

1、扩散模型已成为当前计算机视觉领域主流的图像生成技术，展现出了卓越的生成质量和多样性。然而，扩散模型依赖于数百至数千步的迭代去噪采样过程，导致其推理速度极其缓慢且计算资源消耗巨大。为了将高性能的扩散模型部署到计算能力受限的终端设备，模型压缩技术显得尤为关键。其中，知识蒸馏作为一种主流手段，被广泛应用于将大规模教师扩散模型的生成能力迁移至参数量较小的轻量级学生扩散模型中。

2、在现有的扩散模型知识蒸馏方案中，最常规的实现方式是基于模型输出层面的静态直接拟合。具体而言，现有技术通常会让学生模型在相同的时间步和输入噪声条件下，直接模仿教师模型的预测结果，通过计算两者特征空间输出结果之间的均方误差来进行知识传递。由于默认教师模型在任何时刻和任何位置的指导都具有同等价值，现有方法通常会对所有的扩散时间步以及整个图像的空间位置，施加均匀的全局监督信号。

3、然而，当教师与学生扩散模型在参数量或数据分布上存在显著差异时，现有静态均匀蒸馏方法面临着因分布偏移导致的蒸馏失效这一核心技术问题。现有方法忽视了模型在训练阶段的采样状态与实际推理阶段的采样状态存在严重的不一致现象。在这种静态拟合下，训练时的微小误差会在学生模型独立的迭代去噪阶段被急剧放大，导致概率流偏离正确的生成轨迹。因此，亟需一种在线策略蒸馏机制，使学生模型在动态生成的轨迹中自适应地选择高价值特征进行学习，以消除训练与推理阶段的分布偏移，从而实现高效且高质的知识迁移。

技术实现思路

1、本发明提供一种扩散模型的在线策略蒸馏方法及装置，用以解决现有技术中教师与学生模型因存在差异而产生概率流分布偏差，以及采用全局均匀监督导致蒸馏失效和训练效率低下的缺陷，实现轻量级学生模型高效且精准的生成能力迁移。

2、本发明提供一种扩散模型的在线策略蒸馏方法，所述方法应用于基于教师扩散模型对学生扩散模型进行知识蒸馏的训练过程中，包括：

3、基于所述教师扩散模型在不同引导条件下输出的第一教师预测特征，对所述学生扩散模型在对应条件下的预测结果进行联合优化，获得初始化的学生扩散模型；

4、控制所述初始化的学生扩散模型基于多条件组合引导采样方式进行迭代去噪，生成用于知识蒸馏的采样轨迹；

5、针对所述采样轨迹在迭代过程中的不同多维特征位置，获取所述教师扩散模型针对当前位置输出的第二教师预测特征，并结合所述第二教师预测特征对生成任务的影响程度指标以及所述学生扩散模型的当前拟合能力指标，确定目标动态加权权重；

6、基于所述目标动态加权权重，对所述第二教师预测特征与所述学生扩散模型当前输出特征之间的差异进行加权以得到蒸馏损失，并根据所述蒸馏损失更新所述学生扩散模型的网络参数。

7、根据本发明提供的扩散模型的在线策略蒸馏方法，所述第一教师预测特征包括真实条件概率流和真实基础生成概率流；

8、所述基于所述教师扩散模型在不同引导条件下输出的第一教师预测特征，对所述学生扩散模型在对应条件下的预测结果进行联合优化，获得初始化的学生扩散模型的步骤，包括：

9、获取包含目标图像以及与所述目标图像对应的目标条件信息的训练样本，并对所述目标图像进行加噪处理以获得目标扩散时间步下的加噪输入特征；

10、将所述加噪输入特征和所述目标条件信息输入所述教师扩散模型，获得对应的真实条件概率流；并将所述加噪输入特征和空条件信息输入所述教师扩散模型，获得对应的真实基础生成概率流；

11、分别计算所述学生扩散模型在相同输入下输出的条件预测得分与所述真实条件概率流之间的第一损失，以及所述学生扩散模型输出的无条件预测得分与所述真实基础生成概率流之间的第二损失；

12、联合所述第一损失和所述第二损失更新所述学生扩散模型的初始网络参数，获得初始化的学生扩散模型。

13、根据本发明提供的扩散模型的在线策略蒸馏方法，所述控制所述初始化的学生扩散模型基于多条件组合引导采样方式进行迭代去噪，生成用于知识蒸馏的采样轨迹，包括：

14、在所述初始化的学生扩散模型的轨迹生成阶段，分别计算当前扩散时间步下对应目标条件信息的条件预测得分，以及对应空条件信息的无条件预测得分；

15、将无分类器引导策略作为所述多条件组合引导采样方式，根据预设的引导尺度参数对所述条件预测得分和所述无条件预测得分进行加权组合计算，得到组合预测得分；

16、利用所述组合预测得分进行去噪状态更新以获得下一个时间步的隐状态，并将连续多个时间步迭代生成的隐状态序列作为所述采样轨迹，以使所述学生扩散模型在训练阶段和实际推理阶段保持一致的概率流分布。

17、根据本发明提供的扩散模型的在线策略蒸馏方法，所述针对所述采样轨迹在迭代过程中的不同多维特征位置，获取所述教师扩散模型针对当前位置输出的第二教师预测特征，并结合所述第二教师预测特征对生成任务的影响程度指标以及所述学生扩散模型的当前拟合能力指标，确定目标动态加权权重，包括：

18、针对所述多维特征位置所表征的当前扩散时间步和潜空间位置，评估所述第二教师预测特征对最终生成结果的潜在收益，获得所述影响程度指标；

19、在所述当前扩散时间步和所述潜空间位置下，评估所述第二教师预测特征与所述学生扩散模型当前特征的差异是否处于可学习范围内，获得所述当前拟合能力指标；

20、根据所述影响程度指标和所述当前拟合能力指标的数值组合，计算得到所述当前扩散时间步下每个所述潜空间位置对应的时空可教性函数值，并将所述时空可教性函数值作为所述目标动态加权权重。

21、根据本发明提供的扩散模型的在线策略蒸馏方法，所述在所述当前扩散时间步和所述潜空间位置下，评估所述第二教师预测特征与所述学生扩散模型当前特征的差异是否处于可学习范围内，获得所述当前拟合能力指标，包括：

22、计算所述第二教师预测特征与所述学生扩散模型当前输出特征在对应潜空间位置下的误差梯度；

23、获取所述学生扩散模型的历史优化状态参数，并结合所述误差梯度和所述历史优化状态参数，计算所述学生扩散模型对当前位置特征差异的有效拟合概率；

24、根据所述有效拟合概率为对应的潜空间位置分配可吸收性数值，并将所述可吸收性数值作为所述当前拟合能力指标，其中所述有效拟合概率越高的位置，对应的可吸收性数值越大。

25、根据本发明提供的扩散模型的在线策略蒸馏方法，所述评估所述第二教师预测特征对最终生成结果的潜在收益，获得所述影响程度指标，包括：计算所述第二教师预测特征在当前扩散时间步下对所述最终生成结果的贡献权重；

26、获取预设的扩散模型信噪比调度参数，并根据当前扩散时间步对应的信噪比数值对所述贡献权重进行动态放缩，得到对应的潜在收益数值；

27、将所述潜在收益数值作为所述影响程度指标，其中所述潜在收益数值越大的潜空间位置，对所述最终生成结果的影响程度越高。

28、根据本发明提供的扩散模型的在线策略蒸馏方法，所述基于所述目标动态加权权重，对所述第二教师预测特征与所述学生扩散模型当前输出特征之间的差异进行加权以得到蒸馏损失，包括：

29、将所述目标动态加权权重在对应的扩散时间步和潜空间位置上进行维度映射，生成与潜空间特征图尺寸一致的动态空间加权张量；

30、计算所述第二教师预测特征与所述学生扩散模型当前输出特征在对应位置处的初始特征距离，作为所述差异；

31、利用所述动态空间加权张量对所述初始特征距离进行逐位置的加权计算，获得可教性感知蒸馏损失作为所述蒸馏损失。

32、根据本发明提供的扩散模型的在线策略蒸馏方法，所述利用所述组合预测得分进行去噪状态更新以获得下一个时间步的隐状态，并将连续多个时间步迭代生成的隐状态序列作为所述采样轨迹，包括：

33、基于预设的逆向去噪调度规则，确定从当前扩散时间步转移至所述下一个时间步的条件概率分布；

34、将所述组合预测得分作为均值预测参数代入所述条件概率分布中进行随机采样，获取所述下一个时间步的隐状态；

35、按照扩散时间步递减的顺序重复执行所述状态更新与随机采样过程，直至达到最终生成阶段，并将迭代过程生成的隐状态序列作为所述采样轨迹。

36、根据本发明提供的扩散模型的在线策略蒸馏方法，所述根据所述蒸馏损失更新所述学生扩散模型的网络参数，包括：

37、基于反向传播算法，计算所述蒸馏损失相对于所述学生扩散模型当前网络参数的梯度矩阵；

38、利用预设的优化器算法，根据所述梯度矩阵对所述学生扩散模型的网络参数进行单次迭代更新；

39、在下一次训练迭代中，控制参数更新后的所述学生扩散模型重新生成新的采样轨迹进行下一轮的在线知识迁移，重复执行该过程，直至所述学生扩散模型满足预设的收敛条件。

40、根据本发明提供的扩散模型的在线策略蒸馏方法，所述获取包含目标图像以及与所述目标图像对应的目标条件信息的训练样本，并对所述目标图像进行加噪处理以获得目标扩散时间步下的加噪输入特征的步骤，包括：获取包含所述目标图像以及对应文本描述的样本作为所述训练样本，并将所述文本描述对应的特征提取为所述目标条件信息；

41、从预设的时间步分布中随机采样获得所述目标扩散时间步，并获取所述目标扩散时间步对应的噪声方差；

42、根据所述噪声方差向所述目标图像中注入标准高斯噪声，生成所述目标扩散时间步下的加噪输入特征。

43、本发明还提供一种扩散模型的在线策略蒸馏装置，所述装置应用于基于教师扩散模型对学生扩散模型进行知识蒸馏的训练过程中，包括：

44、冷启动初始化模块，用于基于所述教师扩散模型在不同引导条件下输出的第一教师预测特征，对所述学生扩散模型在对应条件下的预测结果进行联合优化，获得初始化的学生扩散模型；

45、在线轨迹生成模块，用于控制所述初始化的学生扩散模型基于多条件组合引导采样方式进行迭代去噪，生成用于知识蒸馏的采样轨迹；

46、时空权重确定模块，用于针对所述采样轨迹在迭代过程中的不同多维特征位置，获取所述教师扩散模型针对当前位置输出的第二教师预测特征，并结合所述第二教师预测特征对生成任务的影响程度指标以及所述学生扩散模型的当前拟合能力指标，确定目标动态加权权重；

47、加权蒸馏更新模块，用于基于所述目标动态加权权重，对所述第二教师预测特征与所述学生扩散模型当前输出特征之间的差异进行加权以得到蒸馏损失，并根据所述蒸馏损失更新所述学生扩散模型的网络参数。

48、本发明还提供一种电子设备，包括存储器、处理器及存储在存储器上并可在处理器上运行的计算机程序，所述处理器执行所述计算机程序时实现如上述任一种所述扩散模型的在线策略蒸馏方法。

49、本发明还提供一种非暂态计算机可读存储介质，其上存储有计算机程序，该计算机程序被处理器执行时实现如上述任一种所述扩散模型的在线策略蒸馏方法。

50、本发明还提供一种计算机程序产品，包括计算机程序，所述计算机程序被处理器执行时实现如上述任一种所述扩散模型的在线策略蒸馏方法。

51、本发明提供的扩散模型的在线策略蒸馏方法及装置，通过基于教师扩散模型在不同引导条件下输出的第一教师预测特征对学生模型在对应条件下的预测结果进行联合优化，有效降低了教师与学生模型之间的初始概率流差异；然后通过控制初始化的学生模型基于多条件组合引导采样方式进行迭代去噪以生成采样轨迹，显著缓解了训练阶段与实际推理阶段之间的概率分布偏移，避免了误差在迭代过程中的累积放大；在此基础上，针对采样轨迹迭代过程中的不同多维特征位置，结合第二教师预测特征对生成任务的影响程度指标与学生模型的当前拟合能力指标来确定目标动态加权权重，并基于该权重对特征差异进行加权以计算蒸馏损失，从而使学生模型能够优先学习对生成结果收益大的教师知识，这样在教师模型与学生模型存在显著分布差异的场景下，有效提升了在线策略蒸馏的效率与稳定性，显著改善了轻量学生模型的最终生成质量。