萤火跑模型 | 能量函数指导的图图翻译扩散模型

Chongxuan Li October 20, 2022

扩散模型作为一种新的深度生成模型，在图像生成领域取得了 SOTA 的效果，并且逐渐在诸多应用领域展现出强大的性能，如视频生成、分子图建模等。

最近来自人民大学李崇轩课题组的研究者们提出了一种基于能量函数的扩散模型生成指导方法 EGSDE。该方法旨在通过预定义的能量函数对预训练好的扩散模型进行指导，从而实现可控生成。现有指导生成方法如 classifier-guidance 等可以理解为 EGSDE 的特例。理论上，研究者从专家产物的角度解释了 EGSDE 的采样分布。

该项目利用了幻方AI深度学习训练平台的算力及加速性能，在无监督的图到图翻译任务上进行了大量实验，在若干数据集上取得了最好的“可信-真实”权衡结果。该项工作被人工智能顶级学术会议 NeurIPS 2022 收录，本文将为大家简要描述。

论文标题：EGSDE: Unpaired Image-to-Image Translation via Energy-Guided Stochastic Differential Equations

论文地址：https://arxiv.org/abs/2207.06635

模型仓库：https://github.com/ML-GSAI/EGSDE

模型介绍

1. 研究动机

无监督的图到图翻译的任务是在没有成对数据的情况下学习源域到目标域的映射，且希望翻译后的图像一方面要去掉原图的域特有特征使其更“真实”，另一方面可以保留原图的域无关特征使其更“可信”。例如将猫翻译到狗，要把原图的鼻子、胡须这种特有特征去掉，同时保留原图的姿势、色调这种无关特征。

随着扩散模型的发展，逐渐开始涌现一些基于扩散模型的图到图翻译方法并且取得了 SOTA 的结果。这些方法往往是在目标域训练一个扩散模型保证“真实”，在测试的阶段用上原图的信息实现“可信”，但是并没有用上源域的训练数据。

基于此，该工作的研究者们希望同时利用源域和目标域的数据训练一个能量函数，通过能量函数对预训练好的目标域扩散模型进行指导生成，从而实现更“可信”且“真实”的图到图翻译。整体结构如下图所示：

2. 模型方法

该方法的核心要义是如何根据任务设计合适的能量函数。在图像翻译任务中，研究者将能量函数 $\varepsilon(y,x,t)$ 拆成两项对数势函数，分别负责生成图像的“真实”和“可信”。

\begin{aligned} d_y = [f(y,t)-g(t)^2(s(y, t) - \nabla_y \varepsilon(y,x_0,t))]dt + g(t)d\bar{w} \end{aligned}

\begin{aligned} \varepsilon(y,x,t) &= \lambda_s\varepsilon_s(y, x, t) + \lambda_i\varepsilon_i(y, x, t) \\ &= \lambda_s\mathbb{E}_{q_{t | 0}(x_t | x)}\varepsilon_s(y, x_t, t) - \lambda_i\mathbb{E}_{q_{t | 0}(x_t | x)}\varepsilon_i(y, x_t, t) \end{aligned}

进一步，研究者将负责“真实”的对数势函数 $\varepsilon_s(y,x,t)$ 定义为原图和生成的图片经过域独特的特征提取器 $\mathbb{E}_s(\cdot,\cdot)$ 后之间的相似度（余弦相似度），负责“可信”的 $\varepsilon_i(y,x,t)$ 则定义为两者经过域无关的特征提取器 $\mathbb{E}_i(\cdot,\cdot)$ 后之间的负相似度（负 $L_2$ 距离的平方）。

\begin{aligned} \varepsilon_s(y, x_t, t) &=cos(\mathbb{E}_s(y,t), \mathbb{E}_s(x_t, t)) \\ \varepsilon_i(y, x_t, t) &=-||\mathbb{E}_i(y,t), \mathbb{E}_i(x_t, t)||_2^2 \end{aligned}

直观上来说，去噪过程中能量函数的减少一方面导致原图和翻译图片间的特有特征相似度变低，从而去掉原图中源域特有的特征，另一方面导致两者的域无关的特征相似度变高，从而保留原图中域无关的特征。

实现上，如上图所示，研究者用分类器除去最后一层的剩下部分和低通滤波器分别当做域独特的特征提取器 $\mathbb{E}_s(\cdot,\cdot)$ 和域无关的特征提取器 $\mathbb{E}_i(\cdot,\cdot)$ 。

至此，能量函数和扩散模型已知，从一个带噪声的原图出发，可以通过扩散模型对 EGSDE 进行采样实现图到图的翻译。除此之外，将 $\mathbb{E}_s(\cdot,\cdot)$ 依赖的二分类器拓展到多分类器，该方法也很容易拓展到多域翻译任务。

理论上，研究者展示了对 EGSDE 的离散化采样近似等价于从如下专家产物形式的分布采样:

\begin{aligned} \tilde{p}(y_t|x_0) = \frac{p_{r1}(y_t|x_0)p_{r2}(y_t|x_0)p_{f}(y_t|x_0)}{Z_t} \end{aligned}

这为 EGSDE 中每个部分的作用提供了解释，其中扩散模型和 $\varepsilon_s(y,x,t)$ 对应于负责“真实”的专家， $\varepsilon_i(y,x,t)$ 对应于负责“可信”的专家。

实验

在实验中，研究者选择 FID 作为对“真实”的衡量，L2、PSNR、SSIM 作为对“可信”的衡量，人工测量 AMT 作为同时两方面的衡量，在 AFHQ 和 CelebA-HQ 数据集验证方法的有效性。如下表所示，EGSDE 取得了最好的“可信-真实”权衡结果，且通过调整超参可以进一步改善 FID。

为了验证 DAT 的有效性，研究者们展示了一些可视化结果。

如上图所示，研究者在 “Male→Female” 任务上展示了关于能量函数权重的消融实验， $\lambda_s$ 越大，生成的图片越像女性、越“真实”， $\lambda_i$ 越大，生成的图片保留原图更多的信息、越“可信”，验证了能量函数的作用。

总结

研究者们提出了一种基于能量函数的扩散模型生成指导方法 EGSDE。该方法旨在通过预定义的能量函数对预训练好的扩散模型进行指导，从而实现可控生成。该方法实现了更“可信”且“真实”的图到图翻译，同时可以推广应用到更多的场景领域。

同时，研究者采用幻方 AI 的一系列深度学习自研套件，极大加速了模型的训练过程，获得了大量的实验结果。我们欢迎更多优秀的课题与幻方合作，一道推动 AI 技术的发展与落地。

综合体验打分如下：

研究指数：★★★★

该模型改进提升了扩散模型方法，推动了图图翻译场景下的技术发展。
开源指数：★★★★★

数据和代码都进行了完整且规整的开源。
门槛指数：★★★

数据规模，模型大小适中，适合多级多卡数据并行训练。一般单卡训练难度比较大。
通用指数：★★★★

该方法能适用于很多可控生成任务，如材料和药物领域新分子的生成。
适配指数：★★★★★

依赖简单，很容易与幻方 AI 的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力 AI 落地与价值创造，欢迎各方数据研究者与开发者们一同共建。