[论文解读] Maximum Likelihood Training of Implicit Nonlinear Diffusion Models

Dongjun Kim, Byeonghu Na+4 more|arXiv (Cornell University)|May 27, 2022

Generative Adversarial Networks and Image Synthesis被引用 10

一句话总结

本文提出隐式非线性扩散模型（INDM），通过在潜在空间中学习归一化流，隐式建模数据空间中的非线性漂移与波动，从而训练一种数据自适应的非线性扩散过程。INDM 在 CelebA 上实现了 1.75 的 SOTA FID，采用最大似然估计（MLE）训练，其生成样本质量与训练稳定性均优于线性基线方法（如 DDPM++）。

ABSTRACT

Whereas diverse variations of diffusion models exist, extending the linear diffusion into a nonlinear diffusion process is investigated by very few works. The nonlinearity effect has been hardly understood, but intuitively, there would be promising diffusion patterns to efficiently train the generative distribution towards the data distribution. This paper introduces a data-adaptive nonlinear diffusion process for score-based diffusion models. The proposed Implicit Nonlinear Diffusion Model (INDM) learns by combining a normalizing flow and a diffusion process. Specifically, INDM implicitly constructs a nonlinear diffusion on the extit{data space} by leveraging a linear diffusion on the extit{latent space} through a flow network. This flow network is key to forming a nonlinear diffusion, as the nonlinearity depends on the flow network. This flexible nonlinearity improves the learning curve of INDM to nearly Maximum Likelihood Estimation (MLE) against the non-MLE curve of DDPM++, which turns out to be an inflexible version of INDM with the flow fixed as an identity mapping. Also, the discretization of INDM shows the sampling robustness. In experiments, INDM achieves the state-of-the-art FID of 1.75 on CelebA. We release our code at https://github.com/byeonghu-na/INDM.

研究动机与目标

• 为解决现有扩散模型依赖固定线性扩散过程所带来的表达能力受限与训练效率低下问题。
• 探究是否可通过潜在空间中的可学习归一化流，在数据空间中学习非线性前向 SDE，从而提升生成建模性能。
• 通过基于流的变换隐式建模前向 SDE 的漂移与波动系数，实现扩散模型的最大似然估计（MLE）训练。
• 通过确保非线性 SDE 的离散化在不同步长下保持稳定，提升采样鲁棒性。
• 证明：当训练得当时，非线性扩散模型可在似然与样本质量指标上超越线性扩散基线。

提出的方法

• INDM 通过将可学习的归一化流应用于潜在空间中的线性扩散过程，将结果映射回数据空间，从而在数据空间中构建非线性扩散过程。
• 归一化流具备可逆性与可微性，支持前向与反向过程的反向传播，这对训练非线性 SDE 至关重要。
• 数据空间中的前向 SDE 通过流映射隐式定义：$\mathbf{x}_t = \mathbf{f}_{\bm{\phi}}(\mathbf{z}_t)$，其中 $\mathbf{z}_t$ 遵循线性 SDE。
• 通过最小化近似真实数据梯度的去噪梯度损失来训练得分网络 $\mathbf{s}_{\bm{\theta}}$，从而实现基于变分界界的 MLE 训练。
• 采用权重函数 $\lambda(t) = g^2(t)$ 确保梯度损失对应于负对数似然的变分下界。
• 采样通过求解时间反向 SDE 完成，从先验 $\mathbf{x}_T \sim \mathcal{N}(0, I)$ 开始，反向过程基于学习到的得分网络与基于流的漂移校正推导得出。

实验结果

研究问题

• RQ1 与固定线性扩散相比，数据自适应的可学习非线性扩散过程是否能在数据空间中提升生成建模性能？
• RQ2 通过归一化流隐式建模前向 SDE 是否能比标准线性扩散更有效地训练漂移与波动系数？
• RQ3 所提方法能否实现扩散模型的最大似然估计（MLE）训练？其在似然与样本质量方面与基于得分的训练相比如何？
• RQ4 非线性 SDE 的离散化鲁棒性与 DDPM++ 等线性扩散模型相比如何？
• RQ5 通过流网络引入的非线性在 CIFAR-10 与 CelebA 等基准数据集上，对 FID 与 NLL 指标的提升程度如何？

主要发现

• INDM 在 CelebA $64 \times 64$ 上实现了 SOTA 的 FID 1.75，优于所有先前的扩散模型，包括 DDPM++ 与 NCSN++。
• 在 CIFAR-10 上，INDM 实现了 2.97 的 NLL（经校正后），且 NELBO 与 NLL 之间的差距仅为 0.02，表明接近最优的似然估计。
• INDM 在 CIFAR-10（VE）上的 FID 降低至 2.28，在 VP 设置下为 2.90，优于 DDPM++（FID 2.64）与 NCSN++（FID 2.20），在似然与样本质量上均表现更优。
• 该模型对采样离散化表现出强鲁棒性，即使在粗粒度时间步长下仍能保持高质量生成，而线性基线则表现不佳。
• INDM 实现了接近 MLE 的性能，CIFAR-10 上 NELBO 与 NLL 之间的差距仅为 0.02，表明似然优化效果显著。
• 消融实验确认：通过流网络实现的非线性扩散至关重要；若将流设为恒等映射（即类似 DDPM++），则训练曲线显著变差，FID 显著升高。