一种基于表征注入与内化的图像生成模型训练方法与系统

本技术涉及人工智能和计算机视觉，特别是涉及一种基于表征注入与内化的图像生成模型训练方法与系统。

背景技术：

1、近年来，扩散模型已成为图像生成领域的重要技术路线，其中，以transformer作为去噪网络主体的扩散transformer（diffusion transformer，简称dit）具有较好的模型扩展能力和生成性能。然而，为获得较高质量的生成结果，此类模型通常需要进行大量训练迭代，并消耗较高的计算资源，导致训练周期较长、训练成本较高，训练效率成为制约其进一步应用的重要因素之一。特别是在医学图像生成等训练数据获取成本较高、模型训练资源受限或对生成质量要求较高的应用场景中，提高生成模型的训练效率并充分利用已有视觉表征，同样具有实际应用价值。

2、为提高扩散transformer的训练效率，现有技术提出了利用预训练视觉编码器辅助模型学习的方法。例如，表征对齐方法（representation alignment，简称repa）将扩散transformer的中间隐藏状态通过可训练投影头映射至预训练视觉编码器的表征空间，并通过表征相似性约束，使扩散transformer的中间特征与预训练视觉编码器从图像中提取的特征进行对齐，从而利用预训练视觉编码器所具备的语义表征能力辅助扩散transformer学习具有语义结构的内部表征。

3、然而，在上述表征对齐方法中，预训练视觉编码器产生的表征主要作为辅助损失的监督目标，用于约束扩散transformer的内部表征，并未直接参与其去噪计算过程。因此，预训练视觉表征对去噪过程的作用主要通过间接的对齐监督实现，其所包含的语义结构信息未被直接用于扩散transformer的内部去噪计算。

4、因此，如何在保证图像生成模型去噪过程有效进行以及推理阶段正常使用的前提下，更充分地利用预训练视觉编码器提供的语义表征，使其能够更直接地作用于扩散transformer的训练及去噪过程，并进一步提高模型的训练效率和生成质量，仍是现有图像生成模型训练技术中有待解决的问题。

技术实现思路

1、基于此，本技术的发明目的是针对上述技术问题，提供一种基于表征注入与内化的图像生成模型训练方法与系统，使预训练视觉编码器提取的视觉表征经过表征空间适配后直接参与扩散transformer的内部去噪计算，并通过表征内化机制将外部视觉表征所带来的增益内化至模型自身的表征能力中，从而提高图像生成模型的训练效率和生成质量。

2、为实现上述发明目的，本技术第一方面提供一种基于表征注入与内化的图像生成模型训练方法，包括：

3、获取包括干净图像且用于训练图像生成模型的训练样本，根据所述干净图像构造扩散训练的含噪输入；

4、将所述干净图像输入预训练视觉编码器，通过所述预训练视觉编码器提取所述干净图像的视觉表征；对所述视觉表征进行表征空间适配，得到与所述图像生成模型的生成网络自身计算得到的原生中间表征的表征空间相匹配的目标视觉表征；

5、在训练第一阶段，将所述目标视觉表征注入所述生成网络的网络模块内部，以所述目标视觉表征替代对应的原生中间表征参与去噪计算；

6、在训练第二阶段，停止所述替代，恢复所述生成网络对所述原生中间表征的计算，并施加内化约束，以减小所述生成网络计算得到的所述原生中间表征与目标表征之间的差异；其中，所述目标表征为在训练第二阶段对所述干净图像提取视觉表征并进行所述表征空间适配而得到的表征；

7、在完成训练后，移除所述预训练视觉编码器，得到包括训练后的所述生成网络的图像生成模型。

8、优选地，所述根据所述干净图像构造扩散训练的含噪输入，包括：

9、获取随机噪声以及时间变量，对所述干净图像与所述随机噪声进行加权组合，得到与所述时间变量对应的含噪输入样本：

10、

11、其中，为所述干净图像，为服从标准正态分布的随机噪声，为所述时间变量，且，，；

12、所述生成网络执行去噪预测的基础训练目标为最小化基础去噪损失：

13、

14、其中，为所述生成网络，表示生成网络的所有参数，为与所述训练样本对应的条件信息，为与所述含噪输入对应的预测目标，所述预测目标根据所述干净图像以及所述随机噪声确定，表示求期望。

15、优选地，所述生成网络为基于transformer架构的扩散生成网络，所述基于transformer架构的扩散生成网络包括若干个网络模块且其中至少一个为自注意力模块；

16、所述将所述目标视觉表征注入所述生成网络的网络模块内部，包括：将所述目标视觉表征注入选定的所述自注意力模块内部的中间表征计算过程；在所述训练第一阶段，所述若干个网络模块之间传递的隐藏状态由所述生成网络基于所述含噪输入原生计算得到。

17、优选地，所述视觉表征包括所述预训练视觉编码器产生的键表征和值表征；

18、所述表征空间适配通过可训练的映射网络进行，包括：通过所述映射网络中可训练的键投影头将所述键表征映射至所述生成网络中所选定的自注意力模块的键空间，得到映射后键表征；通过所述映射网络中可训练的值投影头将所述值表征映射至所述所选自注意力模块的值空间，得到映射后值表征；所述映射后键表征和所述映射后值表征构成所述目标视觉表征；

19、所述将所述目标视觉表征注入所述生成网络的网络模块内部，以所述目标视觉表征替代对应的原生中间表征参与去噪计算，包括：以所述映射后键表征和所述映射后值表征分别替代所选定自注意力模块原生计算得到的原生键表征和原生值表征参与注意力计算，并保留由含噪输入相关内部特征原生计算得到的查询表征、输出投影以及残差连接。

20、优选地，所述键表征和所述值表征提取自所述预训练视觉编码器的最后一个自注意力层；所述键投影头和所述值投影头分别采用多层感知机实现。

21、优选地，在所述训练第一阶段，选定的所述自注意力模块的注意力输出为：

22、

23、其中，为由所述含噪输入相关内部特征原生计算得到的查询表征，为所述映射后键表征，为所述映射后值表征，为所述查询表征和所述映射后键表征的特征维度；所述选定的自注意力模块的输出为所述注意力输出经所述输出投影后，与经所述残差连接传递的所述含噪输入相关内部特征相加得到。

24、优选地，在所述训练第二阶段，继续将所述干净图像输入所述预训练视觉编码器，并经所述键投影头和所述值投影头分别得到作为内化目标的所述映射后键表征和所述映射后值表征；

25、所述内化约束包括内化损失：

26、

27、其中，和分别为所述训练第二阶段中所述所选自注意力模块原生计算得到的原生键表征和原生值表征，和分别为所述映射后键表征和所述映射后值表征，表示停止梯度操作，表示相应表征中的元素个数，下标表示frobenius范数；

28、在所述训练第二阶段，将所述内化损失与基础去噪损失加权求和得到总损失，并对所述生成网络进行联合优化：

29、

30、其中，为用于平衡所述基础去噪损失与所述内化损失的权重系数。

31、优选地，所述训练第一阶段为自训练开始至达到预设训练条件为止的训练阶段，所述预设训练条件包括训练迭代步数达到预设步数阈值，和/或训练迭代步数占预设总训练步数的比例达到预设比例阈值；在达到所述预设训练条件后，停止所述目标视觉表征对所述原生中间表征的替代，并直接进入所述训练第二阶段。

32、优选地，所述方法还包括：

33、对所述生成网络的中间隐藏状态施加表征对齐约束，所述表征对齐约束用于将通过可训练的对齐投影头映射后的所述生成网络的中间隐藏状态，与所述预训练视觉编码器从所述干净图像提取的表征在所述预训练视觉编码器的表征空间中对齐；

34、其中，所述目标视觉表征注入的所述网络模块，沿所述生成网络的前向传播方向，位于所述表征对齐约束所作用的网络层之前。

35、为实现发明目的，本技术第二方面提供一种基于表征注入与内化的图像生成模型训练系统，应用上述技术方案所述的一种基于表征注入与内化的图像生成模型训练方法，所述系统包括：

36、数据获取模块，用于获取包括干净图像且用于训练图像生成模型的训练样本，根据所述干净图像构造扩散训练的含噪输入；

37、视觉表征提取模块，用于将所述干净图像输入预训练视觉编码器，通过所述预训练视觉编码器提取所述干净图像的视觉表征；

38、表征映射模块，用于对所述视觉表征进行表征空间适配，得到与所述生成网络的网络模块内部的原生中间表征的表征空间相匹配的目标视觉表征；

39、表征注入模块，用于在训练第一阶段，将所述目标视觉表征注入所述生成网络的网络模块内部，以所述目标视觉表征替代对应的原生中间表征参与去噪计算；

40、表征内化模块，用于在训练第二阶段，停止所述目标视觉表征对所述原生中间表征的替代，恢复所述生成网络对所述原生中间表征的计算，并施加内化约束，以减小所述生成网络计算得到的所述原生中间表征与目标表征之间的差异，所述目标表征为在所述训练第二阶段对所述干净图像提取视觉表征并进行所述表征空间适配而得到的表征；

41、模型输出模块，用于在完成训练后，移除所述预训练视觉编码器，输出包括所述生成网络的图像生成模型。

42、本发明与现有技术相比，其有益效果是：

43、本发明通过将预训练视觉编码器提取的视觉表征经表征映射模块进行表征空间适配后，注入生成网络的网络模块内部，替代对应的原生中间表征参与去噪计算，使预训练视觉表征中的语义结构信息直接作用于去噪过程，有助于加快具有语义结构的内部表征的形成，从而提高图像生成模型的训练效率和收敛性能，并改善生成图像的语义一致性和视觉细节表现；本发明注入的目标视觉表征仅在训练第一阶段作为临时表征支架，通过训练第二阶段的内化约束将其增益内化至生成网络自身；训练完成后移除预训练视觉编码器及表征映射模块，所得图像生成模型在推理阶段仅依靠自身网络完成图像生成，不增加额外的推理计算开销。