[论文解读] A Tale of Two Flows: Cooperative Learning of Langevin Flow and Normalizing Flow Toward Energy-Based Model

Jianwen Xie, Yaxuan Zhu+2 more|arXiv (Cornell University)|May 13, 2022

Generative Adversarial Networks and Image Synthesis被引用 9

一句话总结

该论文提出 CoopFlow，一种协作学习框架，通过联合训练归一化流与短时程 Langevin 流，以在能量模型下提升生成建模性能。通过利用归一化流初始化 Langevin 动力学，并借助 MCMC 转移过程优化归一化流，CoopFlow 即使在 MCMC 采样非混合（non-mixing）的情况下，也能实现高效的图像生成、重建与插值。

ABSTRACT

This paper studies the cooperative learning of two generative flow models, in which the two models are iteratively updated based on the jointly synthesized examples. The first flow model is a normalizing flow that transforms an initial simple density to a target density by applying a sequence of invertible transformations. The second flow model is a Langevin flow that runs finite steps of gradient-based MCMC toward an energy-based model. We start from proposing a generative framework that trains an energy-based model with a normalizing flow as an amortized sampler to initialize the MCMC chains of the energy-based model. In each learning iteration, we generate synthesized examples by using a normalizing flow initialization followed by a short-run Langevin flow revision toward the current energy-based model. Then we treat the synthesized examples as fair samples from the energy-based model and update the model parameters with the maximum likelihood learning gradient, while the normalizing flow directly learns from the synthesized examples by maximizing the tractable likelihood. Under the short-run non-mixing MCMC scenario, the estimation of the energy-based model is shown to follow the perturbation of maximum likelihood, and the short-run Langevin flow and the normalizing flow form a two-flow generator that we call CoopFlow. We provide an understating of the CoopFlow algorithm by information geometry and show that it is a valid generator as it converges to a moment matching estimator. We demonstrate that the trained CoopFlow is capable of synthesizing realistic images, reconstructing images, and interpolating between images.

研究动机与目标

• 解决能量模型（EBM）中非混合、非收敛 Langevin 动力学带来的挑战，该问题会导致梯度估计偏差与模型校准性差。
• 通过将归一化流与类似流的 Langevin 动力学结合，克服因受限可逆变换导致的表达能力不足问题。
• 设计一种协作学习机制：归一化流作为 Langevin 动力学的快速初始化器，而 Langevin 流则通过时序差分学习改进归一化流。
• 证明即使在非混合 MCMC 情况下，双流生成器仍为有效生成器，因其收敛至矩匹配估计器，从而确保模型的有效性。

提出的方法

• 提出双流生成器：使用归一化流进行可 tractable 的密度建模，使用短时程 Langevin 流基于 MCMC 从能量模型中采样。
• 利用归一化流初始化 Langevin 动力学的潜在变量，从而减少实现有效采样所需的 Langevin 步数。
• 通过最大似然法训练能量模型，使用双流生成器合成的样本作为公平样本，尽管其未完全收敛。
• 通过最大化 Langevin 流生成的合成样本的似然性来更新归一化流，从而实现近似推理（amortized inference）。
• 实施联合优化方案，使归一化流学习其初始样本与 Langevin 动力学优化后样本之间的时序差分。
• 使用重建损失，通过梯度下降将观测图像反演至潜在空间，从而支持图像重建与图像修复任务。

实验结果

研究问题

• RQ1 在非混合 MCMC 的情况下，归一化流与 Langevin 流之间的协作学习框架是否能提升能量模型中的样本质量？
• RQ2 归一化流与短时程 Langevin 流的结合是否构成一个有效的生成模型，并能收敛至矩匹配估计器？
• RQ3 归一化流能否有效初始化 Langevin 动力学，以减少所需 Langevin 步数，同时保持样本质量？
• RQ4 该双流生成器在图像重建、图像修复与潜在空间插值等下游任务中的支持能力如何？
• RQ5 CoopFlow 框架的信息几何特性如何在非混合 MCMC 的情况下保证收敛性与稳定性？

主要发现

• CoopFlow 通过结合归一化流进行初始化与短时程 Langevin 流进行优化，即使在非混合 MCMC 情况下，仍能成功生成逼真图像。
• 在 CIFAR-10 上，模型实现了高质量的图像重建，经过对潜在码 200 步梯度下降后，重建图像与原始输入高度一致。
• CoopFlow 在 CelebA 上实现了有效的图像修复，无论初始条件或优化迭代次数如何，均能生成多样且连贯的结果。
• 从观测数据中推断出的图像之间的潜在空间插值，可产生平滑且有意义的过渡，表明数据流形结构良好且平滑。
• 理论分析表明，CoopFlow 收敛至矩匹配估计器，验证了其作为一致生成模型的有效性，即使 MCMC 未完全收敛。
• 协作训练机制使两个流能够迭代改进：归一化流学习更优的初始化，Langevin 流则提供更优的反馈以优化生成结果。