随机微分方程发展综述（能量函数指导的随机微分方程及在图到图翻译任务的应用）

近日中国人民大学高瓴人工智能学院师生有23篇论文被CCF A类国际学术会议NeurIPS 2022录用。神经信息处理系统大会(Neural Information Processing Systems，简称NeurIPS) 是机器学习和计算神经科学领域的顶级国际会议。近期学院将推出【NeurIPS 2022论文解读专题】，详细分享被录用的优秀论文，供师生们交流学习。本篇解读论文信息如下：

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

作者：赵敏，鲍凡，李崇轩，朱军

通讯作者：李崇轩，朱军

一、概述

作为一种新的深度生成模型，扩散模型在图像生成领域取得了SOTA的效果，并且逐渐在诸多应用领域展现出强大的性能，如视频生成、分子图建模等。本文提出了一种一般的基于能量函数的扩散模型生成指导方法EGSDE。该方法旨在通过事先定义的能量函数对预训练好的SDE（扩散模型）进行指导从而实现可控生成。现有指导生成方法如classifier-guidance等可以理解为EGSDE的特例。理论上，本文从product of experts的角度解释了EGSDE的采样分布。实验上，本文聚焦于无监督的图到图翻译任务，解决了过去SDE方法没有用上源域训练数据的问题，在若干数据集上取得了最好的“可信-真实”权衡结果。原则上，通过定义不同的能量函数，EGSDE有潜力处理更广泛的可控生成任务如逆向分子设计等。

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

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

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

视频地址:

https://www.bilibili.com/video/BV1xg411h7qj/spm_id_from=333.999.0.0

二、动机

无监督的图到图翻译的任务是在没有成对数据的情况下学习源域到目标域的映射，且希望翻译后的图像一方面要去掉原图的域特有特征（例如在猫到狗任务中的鼻子、胡须）使其更“真实”，另一方面可以保留原图的域无关特征（例如在猫到狗任务中的姿势、色调）使其更“可信”。过去的几年里，基于GAN的模型成为这个领域主流的方法。随着扩散模型的发展，逐渐开始涌现一些基于扩散模型的图到图翻译方法并且取得了SOTA的结果。这些方法往往是在目标域训练一个SDE（扩散模型）保证“真实”，在测试的阶段用上原图的信息实现“可信”，但是并没有用上源域的训练数据。所以本文希望同时利用源域和目标域的数据训练一个能量函数，通过能量函数对预训练好的目标域SDE（扩散模型）进行指导生成，从而实现更“可信”且“真实”的图到图翻译（如图1所示）。

图1 EGSDE的框架图

三、方法

本方法的核心要义是如何根据任务设计合适的能量函数。在本文关注的图像翻译任务中，根据其目的，我们将能量函数拆成两项对数势函数，分别负责生成图像的“真实”和“可信”。

进一步，我们将负责“真实”的对数势函数定义为原图和生成的图片经过域独特的特征提取器后之间的相似度（余弦相似度），负责“可信”的则定义为两者经过域无关的特征提取器后之间的负相似度（负距离的平方）。

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

实现上，如图2所示，我们用分类器除去最后一层的剩下部分和低通滤波器分别当做域独特的特征提取器和域无关的特征提取器。至此，如图1所示，能量函数和SDE已知，从一个带噪声的原图出发，我们可以通过SDE solver对EGSDE进行采样实现图到图的翻译。除此之外，将依赖的二分类器拓展到多分类器，我们的方法也很容易拓展到多域翻译任务。

图2 特征提取器的实现

理论上，我们展示了对EGSDE的离散化采样近似等价于从如下product of experts形式的分布采样:

这为EGSDE中每个部分的作用提供了解释，其中SDE和对应于负责“真实”的专家，对应于负责“可信”的专家。

四、实验

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

图3 在AFHQ数据集的量化结果对比

在图4b我们展示了一些可视化结果。在图4a，我们在Male→ Female上展示了关于能量函数权重的消融实验，越大，生成的图片越像女性、越“真实”， 越大，生成的图片保留原图更多的信息、越“可信”，验证了能量函数的作用。

五、在其他可控生成任务的潜在价值

原则上，通过定义不同的能量函数，EGSDE有潜力处理更广泛的可控生成任务。以逆向分子设计任务为例，在材料科学和药物研发领域中往往希望生成满足某些性质的分子。考虑到3D分子构象的几何对称性, 我们提出了适用于分子的等变能量函数指导的随机微分方程EEGSDE。实验上，通过设计合适的能量函数，我们在生成满足特定量子性质和分子结构的分子任务中取得了SOTA的结果。此外，通过对能量函数的线性叠加，EEGSDE可以生成具有多目标特性的分子。论文地址：https://arxiv.org/abs/2209.15408。

图4 (a)关于能量函数的消融实验；

（b）EGSDE的可视化效果。